diff --git a/src/acp-agent.ts b/src/acp-agent.ts index 3a381864..8825432a 100644 --- a/src/acp-agent.ts +++ b/src/acp-agent.ts @@ -497,17 +497,23 @@ export class ClaudeAcpAgent implements Agent { const userMessage = promptToClaude(params); + const promptUuid = randomUUID(); + userMessage.uuid = promptUuid; + + let promptReplayed = false; + if (session.promptRunning) { - const uuid = randomUUID(); - userMessage.uuid = uuid; session.input.push(userMessage); const order = session.nextPendingOrder++; const cancelled = await new Promise((resolve) => { - session.pendingMessages.set(uuid, { resolve, order }); + session.pendingMessages.set(promptUuid, { resolve, order }); }); if (cancelled) { return { stopReason: "cancelled" }; } + // The replay resolved the promise, mark in this loop too, + // so we don't treat the next result as a background task's result. + promptReplayed = true; } else { session.input.push(userMessage); } @@ -582,10 +588,6 @@ export class ClaudeAcpAgent implements Agent { } break; case "result": { - if (session.cancelled) { - return { stopReason: "cancelled" }; - } - // Accumulate usage from this result session.accumulatedUsage.inputTokens += message.usage.input_tokens; session.accumulatedUsage.outputTokens += message.usage.output_tokens; @@ -613,6 +615,18 @@ export class ClaudeAcpAgent implements Agent { }); } + if (!promptReplayed) { + // This result is from a background task that finished after + // the previous prompt loop ended. Consume it and continue + // waiting for our own prompt's result. + this.logger.log(`Session ${params.sessionId}: consuming background task result`); + break; + } + + if (session.cancelled) { + return { stopReason: "cancelled" }; + } + // Build the usage response const usage: PromptResponse["usage"] = { inputTokens: session.accumulatedUsage.inputTokens, @@ -688,8 +702,14 @@ export class ClaudeAcpAgent implements Agent { break; } - // Check for queued prompt replay + // Check for prompt replay if (message.type === "user" && "uuid" in message && message.uuid) { + if (message.uuid === promptUuid) { + // Our own prompt was replayed back — we're now processing + // our prompt's response (not a background task's). + promptReplayed = true; + break; + } const pending = session.pendingMessages.get(message.uuid as string); if (pending) { pending.resolve(false); diff --git a/src/tests/acp-agent.test.ts b/src/tests/acp-agent.test.ts index 554c6071..07c6a4ad 100644 --- a/src/tests/acp-agent.test.ts +++ b/src/tests/acp-agent.test.ts @@ -1300,15 +1300,27 @@ describe("stop reason propagation", () => { }; } - function* messageGenerator(messages: any[]) { - yield* messages; - } - function injectSession(agent: ClaudeAcpAgent, messages: any[]) { - const gen = messageGenerator(messages); + const input = new Pushable(); + async function* messageGenerator() { + // Wait for the prompt to push its user message so we can replay it + const iter = input[Symbol.asyncIterator](); + const { value: userMessage, done } = await iter.next(); + if (!done && userMessage) { + yield { + type: "user", + message: userMessage.message, + parent_tool_use_id: null, + uuid: userMessage.uuid, + session_id: "test-session", + isReplay: true, + }; + } + yield* messages; + } agent.sessions["test-session"] = { - query: gen as any, - input: new Pushable(), + query: messageGenerator() as any, + input, cancelled: false, cwd: "/test", permissionMode: "default", @@ -1392,6 +1404,79 @@ describe("stop reason propagation", () => { expect(response.stopReason).toBe("end_turn"); }); + it("should consume background task results and return the prompt's own result", async () => { + const agent = createMockAgent(); + const input = new Pushable(); + + const backgroundTaskResult = createResultMessage({ + subtype: "success", + stop_reason: null, + is_error: false, + }); + // Background task used some tokens + backgroundTaskResult.usage.input_tokens = 100; + backgroundTaskResult.usage.output_tokens = 50; + + const promptResult = createResultMessage({ + subtype: "success", + stop_reason: null, + is_error: false, + }); + + async function* messageGenerator() { + // Background task init + result arrive before our prompt's replay + yield { type: "system", subtype: "init", session_id: "test-session" }; + yield backgroundTaskResult; + + // Now the prompt's user message replay arrives + const iter = input[Symbol.asyncIterator](); + const { value: userMessage } = await iter.next(); + yield { + type: "user", + message: userMessage.message, + parent_tool_use_id: null, + uuid: userMessage.uuid, + session_id: "test-session", + isReplay: true, + }; + + // Then the prompt's own result + yield promptResult; + } + + agent.sessions["test-session"] = { + query: messageGenerator() as any, + input, + cancelled: false, + permissionMode: "default", + settingsManager: {} as any, + accumulatedUsage: { + inputTokens: 0, + outputTokens: 0, + cachedReadTokens: 0, + cachedWriteTokens: 0, + }, + configOptions: [], + promptRunning: false, + pendingMessages: new Map(), + nextPendingOrder: 0, + }; + + const response = await agent.prompt({ + sessionId: "test-session", + prompt: [{ type: "text", text: "test" }], + }); + + expect(response.stopReason).toBe("end_turn"); + // Usage should include both background task and prompt result tokens + expect(response.usage?.inputTokens).toBe( + backgroundTaskResult.usage.input_tokens + promptResult.usage.input_tokens, + ); + expect(response.usage?.outputTokens).toBe( + backgroundTaskResult.usage.output_tokens + promptResult.usage.output_tokens, + ); + }); + it("should throw internal error for success with is_error true and no max_tokens", async () => { const agent = createMockAgent(); injectSession(agent, [