maybe fix codex subturns

Julius Marminge committed 8 months ago

Commit ef843d41 parent 654ca0d4 files · +334 −16Browse files

Showing the changes to apps/server/src/codexAppServerManager.ts (1 of 4 files). Show all

Unified · Split

apps/server/src/codexAppServerManager.ts+94−10

   pending: Map<PendingRequestKey, PendingRequest>;
   pendingApprovals: Map<ApprovalRequestId, PendingApprovalRequest>;
   pendingUserInputs: Map<ApprovalRequestId, PendingUserInputRequest>;
+  collabReceiverTurns: Map<string, TurnId>;
   nextRequestId: number;
   stopping: boolean;
 }
         pending: new Map(),
         pendingApprovals: new Map(),
         pendingUserInputs: new Map(),
+        collabReceiverTurns: new Map(),
         nextRequestId: 1,
         stopping: false,
       };
   async sendTurn(input: CodexAppServerSendTurnInput): Promise<ProviderTurnStartResult> {
     const context = this.requireSession(input.threadId);
+    context.collabReceiverTurns.clear();
     const turnInput: Array<
       { type: "text"; text: string; text_elements: [] } | { type: "image"; url: string }
     context: CodexSessionContext,
     notification: JsonRpcNotification,
   ): void {
-    const route = this.readRouteFields(notification.params);
+    const rawRoute = this.readRouteFields(notification.params);
+    this.rememberCollabReceiverTurns(context, notification.params, rawRoute.turnId);
+    const childParentTurnId = this.readChildParentTurnId(context, notification.params);
+    const isChildConversation = childParentTurnId !== undefined;
+    if (
+      isChildConversation &&
+      this.shouldSuppressChildConversationNotification(notification.method)
+    ) {
+      return;
+    }
     const textDelta =
       notification.method === "item/agentMessage/delta"
         ? this.readString(notification.params, "delta")
       threadId: context.session.threadId,
       createdAt: new Date().toISOString(),
       method: notification.method,
-      turnId: route.turnId,
+      ...((childParentTurnId ?? rawRoute.turnId)
-      itemId: route.itemId,
+        ? { turnId: childParentTurnId ?? rawRoute.turnId }
+        : {}),
+      ...(rawRoute.itemId ? { itemId: rawRoute.itemId } : {}),
       textDelta,
       payload: notification.params,
     });
     }
     if (notification.method === "turn/started") {
+      if (isChildConversation) {
+        return;
+      }
       const turnId = toTurnId(this.readString(this.readObject(notification.params)?.turn, "id"));
       this.updateSession(context, {
         status: "running",
     }
     if (notification.method === "turn/completed") {
+      if (isChildConversation) {
+        return;
+      }
+      context.collabReceiverTurns.clear();
       const turn = this.readObject(notification.params, "turn");
       const status = this.readString(turn, "status");
       const errorMessage = this.readString(this.readObject(turn, "error"), "message");
     }
     if (notification.method === "error") {
+      if (isChildConversation) {
+        return;
+      }
       const message = this.readString(this.readObject(notification.params)?.error, "message");
       const willRetry = this.readBoolean(notification.params, "willRetry");
   }
   private handleServerRequest(context: CodexSessionContext, request: JsonRpcRequest): void {
-    const route = this.readRouteFields(request.params);
+    const rawRoute = this.readRouteFields(request.params);
+    const childParentTurnId = this.readChildParentTurnId(context, request.params);
+    const effectiveTurnId = childParentTurnId ?? rawRoute.turnId;
     const requestKind = this.requestKindForMethod(request.method);
     let requestId: ApprovalRequestId | undefined;
     if (requestKind) {
               : "item/fileChange/requestApproval",
         requestKind,
         threadId: context.session.threadId,
-        ...(route.turnId ? { turnId: route.turnId } : {}),
+        ...(effectiveTurnId ? { turnId: effectiveTurnId } : {}),
-        ...(route.itemId ? { itemId: route.itemId } : {}),
+        ...(rawRoute.itemId ? { itemId: rawRoute.itemId } : {}),
       };
       context.pendingApprovals.set(requestId, pendingRequest);
     }
         requestId,
         jsonRpcId: request.id,
         threadId: context.session.threadId,
-        ...(route.turnId ? { turnId: route.turnId } : {}),
+        ...(effectiveTurnId ? { turnId: effectiveTurnId } : {}),
-        ...(route.itemId ? { itemId: route.itemId } : {}),
+        ...(rawRoute.itemId ? { itemId: rawRoute.itemId } : {}),
       });
     }
       threadId: context.session.threadId,
       createdAt: new Date().toISOString(),
       method: request.method,
-      turnId: route.turnId,
+      ...(effectiveTurnId ? { turnId: effectiveTurnId } : {}),
-      itemId: route.itemId,
+      ...(rawRoute.itemId ? { itemId: rawRoute.itemId } : {}),
       requestId,
       requestKind,
       payload: request.params,
     return route;
   }
+  private readProviderConversationId(params: unknown): string | undefined {
+    return (
+      this.readString(params, "threadId") ??
+      this.readString(this.readObject(params, "thread"), "id") ??
+      this.readString(params, "conversationId")
+    );
+  }
+  private readChildParentTurnId(context: CodexSessionContext, params: unknown): TurnId | undefined {
+    const providerConversationId = this.readProviderConversationId(params);
+    if (!providerConversationId) {
+      return undefined;
+    }
+    return context.collabReceiverTurns.get(providerConversationId);
+  }
+  private rememberCollabReceiverTurns(
+    context: CodexSessionContext,
+    params: unknown,
+    parentTurnId: TurnId | undefined,
+  ): void {
+    if (!parentTurnId) {
+      return;
+    }
+    const payload = this.readObject(params);
+    const item = this.readObject(payload, "item") ?? payload;
+    const itemType = this.readString(item, "type") ?? this.readString(item, "kind");
+    if (itemType !== "collabAgentToolCall") {
+      return;
+    }
+    const receiverThreadIds =
+      this.readArray(item, "receiverThreadIds")
+        ?.map((value) => (typeof value === "string" ? value : null))
+        .filter((value): value is string => value !== null) ?? [];
+    for (const receiverThreadId of receiverThreadIds) {
+      context.collabReceiverTurns.set(receiverThreadId, parentTurnId);
+    }
+  }
+  private shouldSuppressChildConversationNotification(method: string): boolean {
+    return (
+      method === "thread/started" ||
+      method === "thread/status/changed" ||
+      method === "thread/archived" ||
+      method === "thread/unarchived" ||
+      method === "thread/closed" ||
+      method === "thread/compacted" ||
+      method === "thread/name/updated" ||
+      method === "thread/tokenUsage/updated" ||
+      method === "turn/started" ||
+      method === "turn/completed" ||
+      method === "turn/aborted" ||
+      method === "turn/plan/updated" ||
+      method === "item/plan/delta"
+    );
+  }
   private readObject(value: unknown, key?: string): Record<string, unknown> | undefined {
     const target =
       key === undefined