From 35a884c004e78a4d4cef158e89fe8bdf5e129263 Mon Sep 17 00:00:00 2001 From: tanzhenxin Date: Wed, 10 Jun 2026 13:54:51 +0800 Subject: [PATCH] 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 | 13 ++ integration-tests/cli/acp-integration.test.ts | 13 ++ .../acp-integration/session/Session.test.ts | 130 ++++++++++++++++++ .../src/acp-integration/session/Session.ts | 61 +++++++- 4 files changed, 210 insertions(+), 7 deletions(-) diff --git a/integration-tests/cli/acp-cron.test.ts b/integration-tests/cli/acp-cron.test.ts index 84eb71a01f..4b0a8c3cf3 100644 --- a/integration-tests/cli/acp-cron.test.ts +++ b/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' }, + }); } }; diff --git a/integration-tests/cli/acp-integration.test.ts b/integration-tests/cli/acp-integration.test.ts index 98a0567008..375dd57c72 100644 --- a/integration-tests/cli/acp-integration.test.ts +++ b/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' }, + }); } }; diff --git a/packages/cli/src/acp-integration/session/Session.test.ts b/packages/cli/src/acp-integration/session/Session.test.ts index c9a687123b..92b4652f75 100644 --- a/packages/cli/src/acp-integration/session/Session.test.ts +++ b/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', diff --git a/packages/cli/src/acp-integration/session/Session.ts b/packages/cli/src/acp-integration/session/Session.ts index 6d3724338b..46f7f798f2 100644 --- a/packages/cli/src/acp-integration/session/Session.ts +++ b/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 { if (this.midTurnDrainUnavailable) return []; + let drainPromise: ReturnType | 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((_, reject) => { + timeoutHandle = setTimeout( + () => reject(new MidTurnDrainTimeoutError()), + MID_TURN_QUEUE_DRAIN_TIMEOUT_MS, + ); + }); + let response: Awaited; + 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;