fix: Prune tool cache per session after results (#748)

src/acp-agent.ts

@@ -36,8 +36,6 @@ import {
   CloseSessionResponse,
   DeleteSessionRequest,
   DeleteSessionResponse,
-  TerminalHandle,
-  TerminalOutputResponse,
   WriteTextFileRequest,
   WriteTextFileResponse,
   StopReason,
@@ -195,6 +193,12 @@ type Session = {
   /** Accumulated task list for the session, keyed by task ID. Task IDs are
    *  per-session, so this state must not be shared across sessions. */
   taskState: TaskState;
+  /** Caches `tool_use` blocks by id so the matching `tool_result` can recover
+   *  the tool name/input when mapping it to a `tool_call_update`. Per-session
+   *  (tool_use ids are only unique within a session) and pruned at
+   *  `tool_result` time so a long-running session doesn't accumulate every
+   *  tool call for its whole lifetime. */
+  toolUseCache: ToolUseCache;
 };
 
 /** Compute a stable fingerprint of the session-defining params so we can
@@ -209,17 +213,6 @@ function computeSessionFingerprint(params: {
   return JSON.stringify({ cwd: params.cwd, mcpServers: servers });
 }
 
-type BackgroundTerminal =
-  | {
-      handle: TerminalHandle;
-      status: "started";
-      lastOutput: TerminalOutputResponse | null;
-    }
-  | {
-      status: "aborted" | "exited" | "killed" | "timedOut";
-      pendingOutput: TerminalOutputResponse;
-    };
-
 export type SDKMessageFilter = {
   type: string;
   subtype?: string;
@@ -548,8 +541,6 @@ export class ClaudeAcpAgent implements Agent {
     [key: string]: Session;
   };
   client: AgentSideConnection;
-  toolUseCache: ToolUseCache;
-  backgroundTerminals: { [key: string]: BackgroundTerminal } = {};
   clientCapabilities?: ClientCapabilities;
   logger: Logger;
   gatewayAuthRequest?: GatewayAuthRequest;
@@ -561,7 +552,6 @@ export class ClaudeAcpAgent implements Agent {
   constructor(client: AgentSideConnection, logger?: Logger) {
     this.sessions = {};
     this.client = client;
-    this.toolUseCache = {};
     this.logger = logger ?? console;
   }
 
@@ -1210,7 +1200,7 @@ export class ClaudeAcpAgent implements Agent {
                     message.result,
                     "assistant",
                     params.sessionId,
-                    this.toolUseCache,
+                    session.toolUseCache,
                     this.client,
                     this.logger,
                   )) {
@@ -1311,7 +1301,7 @@ export class ClaudeAcpAgent implements Agent {
             for (const notification of streamEventToAcpNotifications(
               message,
               params.sessionId,
-              this.toolUseCache,
+              session.toolUseCache,
               this.client,
               this.logger,
               {
@@ -1385,7 +1375,7 @@ export class ClaudeAcpAgent implements Agent {
                   stripped,
                   message.message.role,
                   params.sessionId,
-                  this.toolUseCache,
+                  session.toolUseCache,
                   this.client,
                   this.logger,
                   {
@@ -1447,7 +1437,7 @@ export class ClaudeAcpAgent implements Agent {
               content,
               message.message.role,
               params.sessionId,
-              this.toolUseCache,
+              session.toolUseCache,
               this.client,
               this.logger,
               {
@@ -2523,6 +2513,7 @@ export class ClaudeAcpAgent implements Agent {
       contextWindowSize:
         inferContextWindowFromModel(models.currentModelId) ?? DEFAULT_CONTEXT_WINDOW,
       taskState,
+      toolUseCache: {},
     };
 
     return {
@@ -3222,42 +3213,41 @@ export function toAcpNotifications(
         } else {
           // Only register hooks on first encounter to avoid double-firing
           if (registerHooks && !alreadyCached) {
+            // Capture the tool name in the closure rather than re-reading the
+            // cache when the hook fires. The cache entry is pruned at
+            // tool_result time, and a PostToolUse hook can fire after that, so
+            // closing over the name keeps the diff working without depending on
+            // (or pinning) the cache entry's lifetime.
+            const toolName = chunk.name;
             registerHookCallback(chunk.id, {
               onPostToolUseHook: async (toolUseId, toolInput, toolResponse) => {
-                const toolUse = toolUseCache[toolUseId];
-                if (toolUse) {
-                  // Both `Edit` and `Write` produce a structuredPatch in their
-                  // PostToolUse tool_response. For Edit the diff replaces the
-                  // optimistic content built at tool_use time. For Write the
-                  // optimistic content (built from `input.content` alone with
-                  // `oldText: null`) shows "creation" semantics regardless of
-                  // whether the file existed; the structuredPatch from the
-                  // hook lets us emit the real diff for `type: "update"`. The
-                  // helper returns `{}` if the response shape isn't usable.
-                  const editDiff =
-                    toolUse.name === "Edit" || toolUse.name === "Write"
-                      ? toolUpdateFromDiffToolResponse(toolResponse)
-                      : {};
-                  const update: SessionNotification["update"] = {
-                    _meta: {
-                      claudeCode: {
-                        toolResponse,
-                        toolName: toolUse.name,
-                      },
-                    } satisfies ToolUpdateMeta,
-                    toolCallId: toolUseId,
-                    sessionUpdate: "tool_call_update",
-                    ...editDiff,
-                  };
-                  await client.sessionUpdate({
-                    sessionId,
-                    update,
-                  });
-                } else {
-                  logger.error(
-                    `[claude-agent-acp] Got a tool response for tool use that wasn't tracked: ${toolUseId}`,
-                  );
-                }
+                // Both `Edit` and `Write` produce a structuredPatch in their
+                // PostToolUse tool_response. For Edit the diff replaces the
+                // optimistic content built at tool_use time. For Write the
+                // optimistic content (built from `input.content` alone with
+                // `oldText: null`) shows "creation" semantics regardless of
+                // whether the file existed; the structuredPatch from the
+                // hook lets us emit the real diff for `type: "update"`. The
+                // helper returns `{}` if the response shape isn't usable.
+                const editDiff =
+                  toolName === "Edit" || toolName === "Write"
+                    ? toolUpdateFromDiffToolResponse(toolResponse)
+                    : {};
+                const update: SessionNotification["update"] = {
+                  _meta: {
+                    claudeCode: {
+                      toolResponse,
+                      toolName,
+                    },
+                  } satisfies ToolUpdateMeta,
+                  toolCallId: toolUseId,
+                  sessionUpdate: "tool_call_update",
+                  ...editDiff,
+                };
+                await client.sessionUpdate({
+                  sessionId,
+                  update,
+                });
               },
             });
           }
@@ -3394,6 +3384,11 @@ export function toAcpNotifications(
             ...toolUpdate,
           };
         }
+        // The tool_use is fully resolved now — drop it so a long session doesn't
+        // retain every tool call. The PostToolUse hook (Edit/Write diffs) closes
+        // over the tool name and no longer reads the cache, so pruning here is
+        // safe regardless of hook/result ordering.
+        delete toolUseCache[chunk.tool_use_id];
         break;
       }
 

src/tests/acp-agent.test.ts

@@ -1548,6 +1548,7 @@ describe("stop reason propagation", () => {
       emitRawSDKMessages: false,
       contextWindowSize: 200000,
       taskState: new Map(),
+      toolUseCache: {},
     };
   }
 
@@ -1692,6 +1693,7 @@ describe("stop reason propagation", () => {
       emitRawSDKMessages: false,
       contextWindowSize: 200000,
       taskState: new Map(),
+      toolUseCache: {},
     };
 
     const response = await agent.prompt({
@@ -1851,6 +1853,7 @@ describe("session/close", () => {
       emitRawSDKMessages: false,
       contextWindowSize: 200000,
       taskState: new Map(),
+      toolUseCache: {},
     };
     return agent.sessions[sessionId]!;
   }
@@ -1935,6 +1938,7 @@ describe("session/delete", () => {
       emitRawSDKMessages: false,
       contextWindowSize: 200000,
       taskState: new Map(),
+      toolUseCache: {},
     };
     return agent.sessions[sessionId]!;
   }
@@ -2036,6 +2040,7 @@ describe("getOrCreateSession param change detection", () => {
       emitRawSDKMessages: false,
       contextWindowSize: 200000,
       taskState: new Map(),
+      toolUseCache: {},
     };
     return agent.sessions[sessionId]!;
   }
@@ -2271,6 +2276,7 @@ describe("usage_update computation", () => {
       emitRawSDKMessages: false,
       contextWindowSize: 200000,
       taskState: new Map(),
+      toolUseCache: {},
     };
   }
 
@@ -3214,6 +3220,7 @@ describe("emitRawSDKMessages", () => {
       emitRawSDKMessages,
       contextWindowSize: 200000,
       taskState: new Map(),
+      toolUseCache: {},
     };
   }
 
@@ -3442,6 +3449,7 @@ describe("result origin handling", () => {
       emitRawSDKMessages: false,
       contextWindowSize: 200000,
       taskState: new Map(),
+      toolUseCache: {},
     };
   }
 
@@ -3617,6 +3625,7 @@ describe("memory_recall handling", () => {
       emitRawSDKMessages: false,
       contextWindowSize: 200000,
       taskState: new Map(),
+      toolUseCache: {},
     };
   }
 
@@ -3847,6 +3856,7 @@ describe("post-error recovery", () => {
       emitRawSDKMessages: false,
       contextWindowSize: 200000,
       taskState: new Map(),
+      toolUseCache: {},
     };
     return { interrupt };
   }
@@ -3992,6 +4002,7 @@ describe("session/cancel wedge recovery (issue #680)", () => {
       emitRawSDKMessages: false,
       contextWindowSize: 200000,
       taskState: new Map(),
+      toolUseCache: {},
     };
     return { interrupt };
   }

src/tests/tools.test.ts

@@ -1,4 +1,4 @@
-import { describe, it, expect } from "vitest";
+import { describe, it, expect, beforeEach } from "vitest";
 import { AgentSideConnection, ClientCapabilities } from "@agentclientprotocol/sdk";
 import { ImageBlockParam, ToolResultBlockParam } from "@anthropic-ai/sdk/resources";
 import {
@@ -673,14 +673,19 @@ describe("Bash terminal output", () => {
   });
 
   describe("toAcpNotifications with clientCapabilities", () => {
-    const toolUseCache: ToolUseCache = {
-      toolu_bash: {
-        type: "tool_use",
-        id: "toolu_bash",
-        name: "Bash",
-        input: { command: "ls -la" },
-      },
-    };
+    // Reset before each test: toAcpNotifications prunes the cache entry once it
+    // maps the tool_result, so a shared object would be empty by the 2nd test.
+    let toolUseCache: ToolUseCache;
+    beforeEach(() => {
+      toolUseCache = {
+        toolu_bash: {
+          type: "tool_use",
+          id: "toolu_bash",
+          name: "Bash",
+          input: { command: "ls -la" },
+        },
+      };
+    });
 
     const bashResult: BetaBashCodeExecutionResultBlock = {
       type: "bash_code_execution_result",
@@ -792,11 +797,21 @@ describe("Bash terminal output", () => {
         { clientCapabilities: { _meta: { terminal_output: true } } },
       );
 
+      // Fresh cache: the withSupport call above already consumed the entry,
+      // since toAcpNotifications prunes the tool_use once it maps the result.
+      const toolUseCacheWithoutSupport: ToolUseCache = {
+        toolu_bash: {
+          type: "tool_use",
+          id: "toolu_bash",
+          name: "Bash",
+          input: { command: "ls -la" },
+        },
+      };
       const withoutSupport = toAcpNotifications(
         [toolResult],
         "assistant",
         "test-session",
-        toolUseCache,
+        toolUseCacheWithoutSupport,
         mockClient,
         mockLogger,
       );
@@ -848,6 +863,31 @@ describe("Bash terminal output", () => {
     });
   });
 
+  describe("toolUseCache pruning", () => {
+    it("retains the tool_use entry until its result, then prunes it", () => {
+      const toolUseCache: ToolUseCache = {};
+      const toolUse = {
+        type: "tool_use" as const,
+        id: "toolu_read",
+        name: "Read",
+        input: { file_path: "/tmp/x.txt" },
+      };
+
+      // tool_use is cached and kept (the matching result hasn't arrived yet).
+      toAcpNotifications([toolUse], "assistant", "s", toolUseCache, mockClient, mockLogger);
+      expect(toolUseCache.toolu_read).toBeDefined();
+
+      // tool_result resolves it, so the entry is pruned to bound memory.
+      const toolResult: ToolResultBlockParam = {
+        type: "tool_result",
+        tool_use_id: "toolu_read",
+        content: [{ type: "text", text: "hello" }],
+      };
+      toAcpNotifications([toolResult], "assistant", "s", toolUseCache, mockClient, mockLogger);
+      expect(toolUseCache.toolu_read).toBeUndefined();
+    });
+  });
+
   describe("post-tool-use hook sends diff content for Edit tool", () => {
     it("should include content and locations from structuredPatch in hook update", async () => {
       const toolUseCache: ToolUseCache = {};