fix(acp): prevent session/prompt hang when client ignores mid-turn drain requests (#4925)

The mid-turn queue drain request sent after every tool batch awaited the
client's response with no deadline. A client that silently drops unknown
methods instead of rejecting with JSON-RPC -32601 never answers, wedging
every tool-calling prompt turn. Race the drain against a 2s timeout and
latch the feature off after three consecutive timeouts.

Also make the integration-test ACP client spec-conforming: reply -32601
to unknown agent->client requests instead of ignoring them.

integration-tests/cli/acp-cron.test.ts

@@ -199,6 +199,19 @@ function setupAcpCronTest(rig: TestRig) {
       } catch (e) {
         sendResponse(msg.id, { message: (e as Error).message });
       }
+      return;
+    }
+
+    // JSON-RPC requires every request to get a response. Reject unknown
+    // agent->client requests (e.g. optional extension methods like
+    // craft/drainMidTurnQueue) with -32601 so the agent fails fast instead
+    // of awaiting a reply that never comes.
+    if (typeof msg.id === 'number' && typeof msg.method === 'string') {
+      send({
+        jsonrpc: '2.0',
+        id: msg.id,
+        error: { code: -32601, message: 'Method not found' },
+      });
     }
   };
 

integration-tests/cli/acp-integration.test.ts

@@ -233,6 +233,19 @@ function setupAcpTest(
       } catch (e) {
         sendResponse(msg.id, { message: (e as Error).message });
       }
+      return;
+    }
+
+    // JSON-RPC requires every request to get a response. Reject unknown
+    // agent->client requests (e.g. optional extension methods like
+    // craft/drainMidTurnQueue) with -32601 so the agent fails fast instead
+    // of awaiting a reply that never comes.
+    if (typeof msg.id === 'number' && typeof msg.method === 'string') {
+      send({
+        jsonrpc: '2.0',
+        id: msg.id,
+        error: { code: -32601, message: 'Method not found' },
+      });
     }
   };
 

packages/cli/src/acp-integration/session/Session.test.ts

@@ -2320,6 +2320,136 @@ describe('Session', () => {
         expect(drainCalls).toHaveLength(1);
       });
 
+      it('latches mid-turn drain off after repeated timeouts when the client never responds', async () => {
+        const tool = {
+          name: 'read_file',
+          kind: core.Kind.Read,
+          build: vi.fn().mockReturnValue({
+            params: { path: '/tmp/test.txt' },
+            getDefaultPermission: vi.fn().mockResolvedValue('allow'),
+            getDescription: vi.fn().mockReturnValue('Read file'),
+            toolLocations: vi.fn().mockReturnValue([]),
+            execute: vi
+              .fn()
+              .mockResolvedValue({ llmContent: 'ok', returnDisplay: 'ok' }),
+          }),
+        };
+        mockToolRegistry.getTool.mockReturnValue(tool);
+        mockConfig.getApprovalMode = vi.fn().mockReturnValue(ApprovalMode.YOLO);
+        // A non-conforming client that silently drops unknown methods: the
+        // drain request never settles. The turn must not hang on it.
+        mockClient.extMethod = vi.fn().mockReturnValue(new Promise(() => {}));
+
+        const toolCallStream = () =>
+          createStreamWithChunks([
+            {
+              type: core.StreamEventType.CHUNK,
+              value: {
+                functionCalls: [
+                  {
+                    id: 'c',
+                    name: 'read_file',
+                    args: { path: '/tmp/test.txt' },
+                  },
+                ],
+              },
+            },
+          ]);
+        // Four prompts, each with one tool batch. The first three time out
+        // (consecutive-strike budget), the fourth must skip the drain.
+        mockChat.sendMessageStream = vi
+          .fn()
+          .mockResolvedValueOnce(toolCallStream())
+          .mockResolvedValueOnce(createEmptyStream())
+          .mockResolvedValueOnce(toolCallStream())
+          .mockResolvedValueOnce(createEmptyStream())
+          .mockResolvedValueOnce(toolCallStream())
+          .mockResolvedValueOnce(createEmptyStream())
+          .mockResolvedValueOnce(toolCallStream())
+          .mockResolvedValueOnce(createEmptyStream());
+
+        const prompt = {
+          sessionId: 'test-session-id',
+          prompt: [{ type: 'text' as const, text: 'read file' }],
+        };
+        await session.prompt(prompt);
+        await session.prompt(prompt);
+        await session.prompt(prompt);
+        await session.prompt(prompt);
+
+        // Three consecutive timeouts trip the latch, so the never-answered
+        // extMethod is attempted on the first three tool batches only.
+        const drainCalls = vi
+          .mocked(mockClient.extMethod)
+          .mock.calls.filter((call) => call[0] === 'craft/drainMidTurnQueue');
+        expect(drainCalls).toHaveLength(3);
+      }, 20_000);
+
+      it('resets the timeout strike count when a drain succeeds', async () => {
+        const tool = {
+          name: 'read_file',
+          kind: core.Kind.Read,
+          build: vi.fn().mockReturnValue({
+            params: { path: '/tmp/test.txt' },
+            getDefaultPermission: vi.fn().mockResolvedValue('allow'),
+            getDescription: vi.fn().mockReturnValue('Read file'),
+            toolLocations: vi.fn().mockReturnValue([]),
+            execute: vi
+              .fn()
+              .mockResolvedValue({ llmContent: 'ok', returnDisplay: 'ok' }),
+          }),
+        };
+        mockToolRegistry.getTool.mockReturnValue(tool);
+        mockConfig.getApprovalMode = vi.fn().mockReturnValue(ApprovalMode.YOLO);
+        // Timeout, success, then timeouts: the success must reset the strike
+        // count, so the latch needs three NEW consecutive timeouts to trip.
+        mockClient.extMethod = vi
+          .fn()
+          .mockReturnValueOnce(new Promise(() => {}))
+          .mockResolvedValueOnce({ messages: [] })
+          .mockReturnValue(new Promise(() => {}));
+
+        const toolCallStream = () =>
+          createStreamWithChunks([
+            {
+              type: core.StreamEventType.CHUNK,
+              value: {
+                functionCalls: [
+                  {
+                    id: 'c',
+                    name: 'read_file',
+                    args: { path: '/tmp/test.txt' },
+                  },
+                ],
+              },
+            },
+          ]);
+        const streamMock = vi.fn();
+        for (let i = 0; i < 5; i++) {
+          streamMock
+            .mockResolvedValueOnce(toolCallStream())
+            .mockResolvedValueOnce(createEmptyStream());
+        }
+        mockChat.sendMessageStream = streamMock;
+
+        const prompt = {
+          sessionId: 'test-session-id',
+          prompt: [{ type: 'text' as const, text: 'read file' }],
+        };
+        for (let i = 0; i < 5; i++) {
+          await session.prompt(prompt);
+        }
+
+        // Strikes: timeout(1), success(reset to 0), timeout(1), timeout(2),
+        // timeout(3 -> latch). All five batches attempt the drain; without
+        // the reset the latch would trip on the fourth batch and the fifth
+        // attempt would be skipped.
+        const drainCalls = vi
+          .mocked(mockClient.extMethod)
+          .mock.calls.filter((call) => call[0] === 'craft/drainMidTurnQueue');
+        expect(drainCalls).toHaveLength(5);
+      }, 30_000);
+
       it('keeps mid-turn drain enabled after a transient error', async () => {
         const tool = {
           name: 'read_file',

packages/cli/src/acp-integration/session/Session.ts

@@ -139,6 +139,24 @@ type AutoCompressionSendResult =
   | { responseStream: null; stopReason: PromptResponse['stopReason'] };
 
 const MID_TURN_QUEUE_DRAIN_METHOD = 'craft/drainMidTurnQueue';
+// The drain is served from an in-memory queue, so a conforming client answers
+// near-instantly (or rejects with -32601). No response within this window
+// means the client silently drops unknown methods; without a deadline the
+// await would wedge the prompt turn forever.
+const MID_TURN_QUEUE_DRAIN_TIMEOUT_MS = 2_000;
+// Latch the drain off only after this many consecutive timeouts: one slow
+// answer must not permanently disable mid-turn messages for a
+// conforming-but-busy client, while a client that never answers stops
+// costing a stall per tool batch after a few batches.
+const MID_TURN_QUEUE_DRAIN_MAX_TIMEOUT_STRIKES = 3;
+
+class MidTurnDrainTimeoutError extends Error {
+  constructor() {
+    super(
+      `mid-turn queue drain got no response within ${MID_TURN_QUEUE_DRAIN_TIMEOUT_MS}ms`,
+    );
+  }
+}
 
 interface BackgroundNotificationQueueItem {
   displayText: string;
@@ -373,6 +391,7 @@ export class Session implements SessionContext {
   private lastPromptTokenCount = 0;
   private lastPromptTokenCountChat: GeminiChat | null = null;
   private midTurnDrainUnavailable = false;
+  private midTurnDrainTimeoutStrikes = 0;
 
   // Background notification drain state. ACP does not have the TUI's idle
   // hook, so the session serializes registry callbacks through this queue.
@@ -1512,13 +1531,25 @@ export class Session implements SessionContext {
   async #drainMidTurnUserMessages(): Promise<Part[]> {
     if (this.midTurnDrainUnavailable) return [];
 
+    let drainPromise: ReturnType<AgentSideConnection['extMethod']> | undefined;
     try {
-      const response = await this.client.extMethod(
-        MID_TURN_QUEUE_DRAIN_METHOD,
-        {
-          sessionId: this.sessionId,
-        },
-      );
+      drainPromise = this.client.extMethod(MID_TURN_QUEUE_DRAIN_METHOD, {
+        sessionId: this.sessionId,
+      });
+      let timeoutHandle: NodeJS.Timeout | undefined;
+      const timeoutPromise = new Promise<never>((_, reject) => {
+        timeoutHandle = setTimeout(
+          () => reject(new MidTurnDrainTimeoutError()),
+          MID_TURN_QUEUE_DRAIN_TIMEOUT_MS,
+        );
+      });
+      let response: Awaited<typeof drainPromise>;
+      try {
+        response = await Promise.race([drainPromise, timeoutPromise]);
+      } finally {
+        clearTimeout(timeoutHandle);
+      }
+      this.midTurnDrainTimeoutStrikes = 0;
       // A client may legally resolve with `result: null` (passed through
       // unwrapped by the ACP SDK); guard the object access so that doesn't
       // throw a TypeError and get misclassified as a transient drain error.
@@ -1557,8 +1588,24 @@ export class Session implements SessionContext {
         error && typeof error === 'object' && 'code' in error
           ? (error as { code?: unknown }).code
           : undefined;
+      const isTimeout = error instanceof MidTurnDrainTimeoutError;
+      if (isTimeout) {
+        this.midTurnDrainTimeoutStrikes += 1;
+        // The lost race leaves the request pending; if the client settles it
+        // later, a rejection must not surface as an unhandled rejection.
+        drainPromise?.catch(() => {});
+      }
+      // Repeated timeouts are also permanent: a conforming client answers
+      // (or rejects with -32601) immediately, so sustained silence means the
+      // client drops unknown methods and would stall every subsequent tool
+      // batch the same way. A single timeout is treated as transient so one
+      // slow answer doesn't disable the drain for the whole session.
       const isPermanentError =
-        errorCode === -32601 || /method not found/i.test(errorMessage);
+        errorCode === -32601 ||
+        /method not found/i.test(errorMessage) ||
+        (isTimeout &&
+          this.midTurnDrainTimeoutStrikes >=
+            MID_TURN_QUEUE_DRAIN_MAX_TIMEOUT_STRIKES);
 
       if (isPermanentError) {
         this.midTurnDrainUnavailable = true;