fix: Map collab agent tool call events (#223)

Matches replay behavior.

Closes #222

Author: benbrandt

src/CodexAcpServer.ts

@@ -10,7 +10,6 @@ import {ACPSessionConnection, type AcpClientConnection, type UpdateSessionEvent}
 import type {InputModality, ReasoningEffort} from "./app-server";
 import type {
     Account,
-    CollabAgentToolCallStatus,
     Model,
     ReasoningEffortOption,
     Thread,
@@ -45,6 +44,7 @@ import {
     LEGACY_SET_SESSION_MODEL_METHOD,
 } from "./AcpExtensions";
 import {
+    createCollabAgentToolCallUpdate,
     createCommandExecutionCompleteUpdate,
     createCommandExecutionUpdate,
     createDynamicToolCallUpdate,
@@ -907,7 +907,7 @@ export class CodexAcpServer {
             case "dynamicToolCall":
                 return [await createDynamicToolCallUpdate(item)];
             case "collabAgentToolCall":
-                return [this.createCollabAgentToolCallUpdate(item)];
+                return [createCollabAgentToolCallUpdate(item)];
             case "webSearch":
                 return [this.createWebSearchUpdate(item)];
             case "imageView":
@@ -947,25 +947,6 @@ export class CodexAcpServer {
         }));
     }
 
-    private createCollabAgentToolCallUpdate(
-        item: ThreadItem & { type: "collabAgentToolCall" }
-    ): UpdateSessionEvent {
-        return {
-            sessionUpdate: "tool_call",
-            toolCallId: item.id,
-            kind: "other",
-            title: `collab.${item.tool}`,
-            status: this.toAcpToolCallStatus(item.status),
-            rawInput: {
-                prompt: item.prompt,
-                senderThreadId: item.senderThreadId,
-                receiverThreadIds: item.receiverThreadIds,
-                agentsStates: item.agentsStates,
-                status: item.status,
-            },
-        };
-    }
-
     private createWebSearchUpdate(
         item: ThreadItem & { type: "webSearch" }
     ): UpdateSessionEvent {
@@ -1017,17 +998,6 @@ export class CodexAcpServer {
         };
     }
 
-    private toAcpToolCallStatus(status: CollabAgentToolCallStatus): "in_progress" | "completed" | "failed" {
-        switch (status) {
-            case "inProgress":
-                return "in_progress";
-            case "completed":
-                return "completed";
-            case "failed":
-                return "failed";
-        }
-    }
-
     private userInputToContentBlocks(input: UserInput): acp.ContentBlock[] {
         switch (input.type) {
             case "text":

src/CodexEventHandler.ts

@@ -33,6 +33,8 @@ import type { McpStartupCompleteEvent } from "./app-server";
 import {toTokenCount} from "./TokenCount";
 import {
     commandExecutionUsesTerminalOutput,
+    createCollabAgentToolCallCompleteUpdate,
+    createCollabAgentToolCallUpdate,
     createCommandExecutionUpdate,
     createDynamicToolCallUpdate,
     createFileChangeUpdate,
@@ -354,6 +356,7 @@ export class CodexEventHandler {
                 this.activeImageGenerationItems.add(event.item.id);
                 return createImageGenerationStartUpdate(event.item);
             case "collabAgentToolCall":
+                return createCollabAgentToolCallUpdate(event.item);
             case "subAgentActivity":
             case "sleep":
             case "userMessage":
@@ -404,12 +407,13 @@ export class CodexEventHandler {
                 return this.createCompletedReasoningEvent(event.item);
             case "webSearch":
                 return createWebSearchCompleteUpdate(event.item);
+            case "collabAgentToolCall":
+                return createCollabAgentToolCallCompleteUpdate(event.item);
             case "exitedReviewMode":
                 return this.createExitedReviewModeEvent(event.item);
             case "contextCompaction":
                 return this.createContextCompactedEvent();
             //ignored types
-            case "collabAgentToolCall":
             case "subAgentActivity":
             case "sleep":
             case "userMessage":

src/CodexToolCallMapper.ts

@@ -9,6 +9,7 @@ import type {
     FuzzyFileSearchSessionUpdatedNotification
 } from "./app-server";
 import type {
+    CollabAgentToolCallStatus,
     CommandAction,
     CommandExecutionStatus,
     DynamicToolCallStatus,
@@ -32,12 +33,13 @@ import {
     type TerminalOutputMode,
 } from "./TerminalOutputMode";
 
-type CodexItemStatus = CommandExecutionStatus | PatchApplyStatus | McpToolCallStatus | DynamicToolCallStatus;
+type CodexItemStatus = CommandExecutionStatus | PatchApplyStatus | McpToolCallStatus | DynamicToolCallStatus | CollabAgentToolCallStatus;
 type AcpToolCallStatus = "pending" | "in_progress" | "completed" | "failed";
 type GuardianApprovalReviewNotification =
     | ItemGuardianApprovalReviewStartedNotification
     | ItemGuardianApprovalReviewCompletedNotification;
 type WebSearchItem = ThreadItem & { type: "webSearch" };
+type CollabAgentToolCallItem = ThreadItem & { type: "collabAgentToolCall" };
 type CommandExecutionItem = ThreadItem & { type: "commandExecution" };
 type AcpToolCallEvent = Extract<UpdateSessionEvent, { sessionUpdate: "tool_call" }>;
 
@@ -356,6 +358,41 @@ export function createWebSearchCompleteUpdate(
     };
 }
 
+export function createCollabAgentToolCallUpdate(
+    item: CollabAgentToolCallItem
+): UpdateSessionEvent {
+    return {
+        sessionUpdate: "tool_call",
+        toolCallId: item.id,
+        kind: "other",
+        title: item.tool,
+        status: toAcpStatus(item.status),
+        rawInput: createCollabAgentToolCallRawInput(item),
+    };
+}
+
+export function createCollabAgentToolCallCompleteUpdate(
+    item: CollabAgentToolCallItem
+): UpdateSessionEvent {
+    return {
+        sessionUpdate: "tool_call_update",
+        toolCallId: item.id,
+        title: item.tool,
+        status: toAcpStatus(item.status),
+        rawInput: createCollabAgentToolCallRawInput(item),
+    };
+}
+
+function createCollabAgentToolCallRawInput(item: CollabAgentToolCallItem) {
+    return {
+        prompt: item.prompt,
+        senderThreadId: item.senderThreadId,
+        receiverThreadIds: item.receiverThreadIds,
+        agentsStates: item.agentsStates,
+        status: item.status,
+    };
+}
+
 export function formatWebSearchTitle(item: WebSearchItem): string {
     const action = item.action;
     if (!action) {

src/__tests__/CodexACPAgent/collab-agent-events.test.ts

@@ -0,0 +1,87 @@
+import { beforeEach, describe, expect, it, vi } from "vitest";
+import type { ServerNotification } from "../../app-server";
+import type { SessionState } from "../../CodexAcpServer";
+import { AgentMode } from "../../AgentMode";
+import {
+    createCodexMockTestFixture,
+    createTestSessionState,
+    setupPromptAndSendNotifications,
+    type CodexMockTestFixture,
+} from "../acp-test-utils";
+
+describe("CodexEventHandler - collab agent tool call events", () => {
+    let mockFixture: CodexMockTestFixture;
+    const sessionId = "test-session-id";
+
+    beforeEach(() => {
+        mockFixture = createCodexMockTestFixture();
+        vi.clearAllMocks();
+    });
+
+    const sessionState: SessionState = createTestSessionState({
+        sessionId,
+        currentModelId: "model-id[effort]",
+        agentMode: AgentMode.DEFAULT_AGENT_MODE,
+    });
+
+    it("maps live collab agent tool calls to ACP tool call updates", async () => {
+        const notifications: ServerNotification[] = [
+            {
+                method: "item/started",
+                params: {
+                    threadId: sessionId,
+                    turnId: "turn-1",
+                    startedAtMs: 0,
+                    item: {
+                        type: "collabAgentToolCall",
+                        id: "call-spawn-weather",
+                        tool: "spawnAgent",
+                        status: "inProgress",
+                        senderThreadId: "thread-main",
+                        receiverThreadIds: ["thread-paris"],
+                        prompt: "Find the current weather in Paris.",
+                        model: null,
+                        reasoningEffort: null,
+                        agentsStates: {
+                            "thread-paris": {
+                                status: "running",
+                                message: "Checking weather",
+                            },
+                        },
+                    },
+                },
+            },
+            {
+                method: "item/completed",
+                params: {
+                    threadId: sessionId,
+                    turnId: "turn-1",
+                    completedAtMs: 0,
+                    item: {
+                        type: "collabAgentToolCall",
+                        id: "call-spawn-weather",
+                        tool: "spawnAgent",
+                        status: "completed",
+                        senderThreadId: "thread-main",
+                        receiverThreadIds: ["thread-paris"],
+                        prompt: "Find the current weather in Paris.",
+                        model: null,
+                        reasoningEffort: null,
+                        agentsStates: {
+                            "thread-paris": {
+                                status: "completed",
+                                message: null,
+                            },
+                        },
+                    },
+                },
+            },
+        ];
+
+        await setupPromptAndSendNotifications(mockFixture, sessionId, sessionState, notifications);
+
+        await expect(`${mockFixture.getAcpConnectionDump([])}\n`).toMatchFileSnapshot(
+            "data/collab-agent-tool-call-flow.json"
+        );
+    });
+});

src/__tests__/CodexACPAgent/data/collab-agent-tool-call-flow.json

@@ -0,0 +1,57 @@
+{
+  "method": "sessionUpdate",
+  "args": [
+    {
+      "sessionId": "test-session-id",
+      "update": {
+        "sessionUpdate": "tool_call",
+        "toolCallId": "call-spawn-weather",
+        "kind": "other",
+        "title": "spawnAgent",
+        "status": "in_progress",
+        "rawInput": {
+          "prompt": "Find the current weather in Paris.",
+          "senderThreadId": "thread-main",
+          "receiverThreadIds": [
+            "thread-paris"
+          ],
+          "agentsStates": {
+            "thread-paris": {
+              "status": "running",
+              "message": "Checking weather"
+            }
+          },
+          "status": "inProgress"
+        }
+      }
+    }
+  ]
+}
+{
+  "method": "sessionUpdate",
+  "args": [
+    {
+      "sessionId": "test-session-id",
+      "update": {
+        "sessionUpdate": "tool_call_update",
+        "toolCallId": "call-spawn-weather",
+        "title": "spawnAgent",
+        "status": "completed",
+        "rawInput": {
+          "prompt": "Find the current weather in Paris.",
+          "senderThreadId": "thread-main",
+          "receiverThreadIds": [
+            "thread-paris"
+          ],
+          "agentsStates": {
+            "thread-paris": {
+              "status": "completed",
+              "message": null
+            }
+          },
+          "status": "completed"
+        }
+      }
+    }
+  ]
+}