Fix Codex collaboration lifecycle handling

Preserve terminal tool-call status, discard child output before buffering, and reserve the subagent UI for actual spawn operations. Add regressions for structured output, command buffering, active labels, and orchestration rendering.

Generated-By: PostHog Code
Task-Id: 971044c5-e622-4f07-b1b5-dca9b06d040c

Author: richardsolomou

packages/agent/src/adapters/codex-app-server/codex-app-server-agent.test.ts

@@ -141,16 +141,29 @@ describe("CodexAppServerAgent", () => {
       "turn/start": { turn: { id: "turn_1", status: "inProgress" } },
     });
     const { client, sessionUpdates, extNotifications } = makeFakeClient();
+    const structuredOutputs: Array<Record<string, unknown>> = [];
     const agent = new CodexAppServerAgent(client, {
       processOptions: { binaryPath: "/bundle/codex" },
       model: "gpt-5.5",
       rpcFactory: stub.factory,
+      onStructuredOutput: async (output) => {
+        structuredOutputs.push(output);
+      },
     });
+    const schema = {
+      type: "object",
+      properties: { source: { type: "string" } },
+      required: ["source"],
+    };
 
     await agent.initialize(init);
     await agent.newSession({
       cwd: "/repo",
-      _meta: { environment: "cloud", taskRunId: "run_1" },
+      _meta: {
+        environment: "cloud",
+        jsonSchema: schema,
+        taskRunId: "run_1",
+      },
     } as unknown as NewSessionRequest);
     const promptDone = agent.prompt({
       sessionId: "thr_1",
@@ -178,6 +191,27 @@ describe("CodexAppServerAgent", () => {
       itemId: "subagent_message_1",
       delta: "subagent prose",
     });
+    stub.emit("item/reasoning/textDelta", {
+      threadId: "subagent_1",
+      turnId: "subagent_turn_1",
+      itemId: "subagent_reasoning_1",
+      delta: "subagent reasoning",
+    });
+    stub.emit("item/completed", {
+      threadId: "subagent_1",
+      turnId: "subagent_turn_1",
+      item: {
+        type: "agentMessage",
+        id: "subagent_message_1",
+        text: '{"source":"child"}',
+      },
+    });
+    stub.emit("item/commandExecution/outputDelta", {
+      threadId: "subagent_1",
+      turnId: "subagent_turn_1",
+      itemId: "shared_command_id",
+      delta: "child command output",
+    });
     stub.emit("thread/tokenUsage/updated", {
       threadId: "subagent_1",
       tokenUsage: {
@@ -221,6 +255,26 @@ describe("CodexAppServerAgent", () => {
       itemId: "message_1",
       delta: "parent response",
     });
+    stub.emit("item/completed", {
+      threadId: "thr_1",
+      turnId: "turn_1",
+      item: {
+        type: "commandExecution",
+        id: "shared_command_id",
+        command: "echo parent",
+        status: "completed",
+        aggregatedOutput: null,
+      },
+    });
+    stub.emit("item/completed", {
+      threadId: "thr_1",
+      turnId: "turn_1",
+      item: {
+        type: "agentMessage",
+        id: "message_1",
+        text: '{"source":"parent"}',
+      },
+    });
     stub.emit("turn/completed", {
       threadId: "thr_1",
       turn: { id: "turn_1", status: "completed" },
@@ -231,6 +285,9 @@ describe("CodexAppServerAgent", () => {
     expect(serializedUpdates).toContain("spawn_agent");
     expect(serializedUpdates).toContain("parent response");
     expect(serializedUpdates).not.toContain("subagent prose");
+    expect(serializedUpdates).not.toContain("subagent reasoning");
+    expect(serializedUpdates).not.toContain("child command output");
+    expect(structuredOutputs).toEqual([{ source: "parent" }]);
     expect(
       extNotifications.filter(
         (notification) => notification.method === "_posthog/turn_complete",

packages/agent/src/adapters/codex-app-server/codex-app-server-agent.ts

@@ -937,10 +937,12 @@ export class CodexAppServerAgent extends BaseAcpAgent {
   }
 
   private handleNotification(method: string, params: unknown): void {
-    const mappedParams = this.withBufferedCommandOutput(method, params);
-    const notificationThreadId = readNotificationThreadId(mappedParams);
+    const notificationThreadId = readNotificationThreadId(params);
     const isMainThread =
       !notificationThreadId || notificationThreadId === this.threadId;
+    const mappedParams = isMainThread
+      ? this.withBufferedCommandOutput(method, params)
+      : params;
 
     if (this.sessionId && !this.session.cancelled) {
       if (isMainThread) {

packages/agent/src/adapters/codex-app-server/mapping.test.ts

@@ -268,7 +268,7 @@ describe("mapAppServerNotification", () => {
     });
   });
 
-  it("keeps a completed spawn tool call active while its subagent is running", () => {
+  it("keeps a completed spawn tool call terminal while its subagent is running", () => {
     const result = mapAppServerNotification(
       "s-1",
       APP_SERVER_NOTIFICATIONS.ITEM_COMPLETED,
@@ -293,7 +293,7 @@ describe("mapAppServerNotification", () => {
       update: {
         sessionUpdate: "tool_call_update",
         toolCallId: "spawn-1",
-        status: "in_progress",
+        status: "completed",
       },
     });
   });

packages/agent/src/adapters/codex-app-server/mapping.ts

@@ -549,7 +549,7 @@ function mapItem(
     update: {
       sessionUpdate: "tool_call_update",
       toolCallId: item.id,
-      status: mapToolStatus(item),
+      status: mapStatus(item.status),
       ...(content ? { content } : {}),
     },
   };
@@ -595,20 +595,6 @@ function mapStatus(
   return "in_progress";
 }
 
-function mapToolStatus(
-  item: AppServerItem,
-): "completed" | "failed" | "in_progress" {
-  if (
-    item.type === "collabAgentToolCall" &&
-    Object.values(item.agentsStates ?? {}).some(
-      (agent) => agent?.status === "pendingInit" || agent?.status === "running",
-    )
-  ) {
-    return "in_progress";
-  }
-  return mapStatus(item.status);
-}
-
 function readItem(params: unknown): AppServerItem | null {
   if (params && typeof params === "object" && "item" in params) {
     const item = (params as Record<string, unknown>).item;

packages/ui/src/features/sessions/components/chat-thread/ToolGroup.test.tsx

@@ -1,12 +1,18 @@
+import { ServiceProvider } from "@posthog/di/react";
 import { posthogToolMeta } from "@posthog/shared";
 import type { ConversationItem } from "@posthog/ui/features/sessions/components/buildConversationItems";
 import { Theme } from "@radix-ui/themes";
 import { render, screen } from "@testing-library/react";
+import { Container } from "inversify";
 import { describe, expect, it } from "vitest";
 import { ToolGroup } from "./ToolGroup";
 
 function subagentItem(
   id: string,
+  options: {
+    status?: "completed" | "in_progress";
+    turnComplete?: boolean;
+  } = {},
 ): Extract<ConversationItem, { type: "session_update" }> {
   return {
     type: "session_update",
@@ -16,26 +22,53 @@ function subagentItem(
       toolCallId: id,
       title: "Subagent",
       kind: "other",
-      status: "completed",
+      status: options.status ?? "completed",
       _meta: posthogToolMeta({ toolName: "spawn_agent" }),
     },
     turnContext: {
       toolCalls: new Map(),
       childItems: new Map(),
       turnCancelled: false,
-      turnComplete: true,
+      turnComplete: options.turnComplete ?? true,
     },
   } as Extract<ConversationItem, { type: "session_update" }>;
 }
 
 describe("ToolGroup", () => {
   it("labels Codex spawn batches as subagents", () => {
     render(
-      <Theme>
-        <ToolGroup tools={[subagentItem("spawn-1"), subagentItem("spawn-2")]} />
-      </Theme>,
+      <ServiceProvider container={new Container()}>
+        <Theme>
+          <ToolGroup
+            tools={[subagentItem("spawn-1"), subagentItem("spawn-2")]}
+          />
+        </Theme>
+      </ServiceProvider>,
     );
 
     expect(screen.getByText("Used Subagents")).toBeInTheDocument();
   });
+
+  it("labels unresolved Codex spawn batches as active subagents", () => {
+    render(
+      <ServiceProvider container={new Container()}>
+        <Theme>
+          <ToolGroup
+            tools={[
+              subagentItem("spawn-1", {
+                status: "in_progress",
+                turnComplete: false,
+              }),
+              subagentItem("spawn-2", {
+                status: "in_progress",
+                turnComplete: false,
+              }),
+            ]}
+          />
+        </Theme>
+      </ServiceProvider>,
+    );
+
+    expect(screen.getByText("Using Subagents")).toBeInTheDocument();
+  });
 });

packages/ui/src/features/sessions/components/new-thread/conversationThreadConfig.ts

@@ -14,6 +14,7 @@ import {
   Wrench,
 } from "@phosphor-icons/react";
 import type { CodeToolKind } from "@posthog/ui/features/sessions/types";
+import { SUBAGENT_SPAWN_TOOL_NAMES } from "../session-update/collaborationTools";
 
 /**
  * Single source of truth for the modernized conversation thread's tuneable
@@ -57,17 +58,7 @@ export const grouping = {
   /**
    * Tool names that create a subagent. Counted separately in the chip summary.
    */
-  subagentToolNames: new Set<string>(["Task", "Agent", "spawn_agent"]),
-  /** Collaboration tools rendered with the dedicated subagent view. */
-  collaborationToolNames: new Set<string>([
-    "Task",
-    "Agent",
-    "spawn_agent",
-    "send_input",
-    "resume_agent",
-    "wait_agent",
-    "close_agent",
-  ]),
+  subagentToolNames: SUBAGENT_SPAWN_TOOL_NAMES,
   /**
    * MCP-app tool calls are excluded from collapsing so their iframes stay
    * mounted (the `keepMounted` contract). Flip to false to fold them in.

packages/ui/src/features/sessions/components/session-update/ToolCallBlock.test.tsx

@@ -115,4 +115,17 @@ describe("ToolCallBlock codex routing", () => {
     expect(screen.getByText("Run tests")).toBeInTheDocument();
     expect(screen.getByText("pnpm test")).toBeInTheDocument();
   });
+
+  it("renders Codex orchestration calls as operations rather than subagents", () => {
+    renderBlock({
+      toolCallId: "tc-wait-agent",
+      title: "Wait for subagents",
+      kind: "other",
+      status: "completed",
+      _meta: posthogToolMeta({ toolName: "wait_agent" }),
+    });
+
+    expect(screen.getByText("Wait for subagents")).toBeInTheDocument();
+    expect(screen.queryByText(/^Subagent/)).not.toBeInTheDocument();
+  });
 });

packages/ui/src/features/sessions/components/session-update/ToolCallBlock.tsx

@@ -16,7 +16,7 @@ import type { ToolCall } from "@posthog/ui/features/sessions/types";
 import { Box } from "@radix-ui/themes";
 import type { ConversationItem, TurnContext } from "../buildConversationItems";
 import { useChatThreadChrome } from "../chat-thread/chatThreadChrome";
-import { grouping } from "../new-thread/conversationThreadConfig";
+import { isSubagentSpawnTool } from "./collaborationTools";
 import {
   MCP_TOOL_BLOCK_COMPONENT,
   type McpToolBlockComponent,
@@ -48,7 +48,7 @@ export function ToolCallBlock({
 
   const props = { toolCall, turnCancelled, turnComplete };
 
-  if (isSubagentTool(toolName)) {
+  if (isSubagentSpawnTool(toolName)) {
     const subagentChildItems = childItems ?? [];
     const turnContext: TurnContext = {
       toolCalls: buildChildToolCallsMap(subagentChildItems),
@@ -109,10 +109,6 @@ export function ToolCallBlock({
   return <Box>{content}</Box>;
 }
 
-function isSubagentTool(toolName: string | undefined): boolean {
-  return grouping.collaborationToolNames.has(toolName ?? "");
-}
-
 function buildChildToolCallsMap(
   childItems: ConversationItem[],
 ): Map<string, ToolCall> {

packages/ui/src/features/sessions/components/session-update/collaborationTools.ts

@@ -0,0 +1,9 @@
+export const SUBAGENT_SPAWN_TOOL_NAMES = new Set<string>([
+  "Task",
+  "Agent",
+  "spawn_agent",
+]);
+
+export function isSubagentSpawnTool(toolName: string | undefined): boolean {
+  return SUBAGENT_SPAWN_TOOL_NAMES.has(toolName ?? "");
+}