diff --git a/src/acp-agent.ts b/src/acp-agent.ts index caedb478..1ae1b865 100644 --- a/src/acp-agent.ts +++ b/src/acp-agent.ts @@ -1340,7 +1340,7 @@ export class ClaudeAcpAgent { if (session.activeTurn && !session.activeTurn.settled) { session.pendingOrphanResults = (session.pendingOrphanResults ?? 0) + 1; } - settleActive({ stopReason: "cancelled" }); + settleActive({ stopReason: "cancelled", usage: sessionUsage(session) }); // If the session is being torn down, abandon the in-flight next() // (swallowing any later rejection so it can't surface as unhandled) // and stop; otherwise re-arm and keep consuming — `pendingNext` @@ -1368,11 +1368,10 @@ export class ClaudeAcpAgent { // // Settle the turn that was in flight so its prompt() doesn't hang: // cancelled if a cancel is pending, otherwise the accumulated outcome. - settleActive( - session.cancelled - ? { stopReason: "cancelled" } - : { stopReason, usage: sessionUsage(session) }, - ); + settleActive({ + stopReason: session.cancelled ? "cancelled" : stopReason, + usage: sessionUsage(session), + }); // Queued turns the SDK never started never ran, so reject them rather // than reporting a success (end_turn) — or a misleading "cancelled" — // for a prompt that produced no output. (A cancel already settled the @@ -1517,8 +1516,14 @@ export class ClaudeAcpAgent { // the turn NOW so its session/prompt gets a terminal // response, instead of leaving it hanging until the next // prompt drains the wreckage. + // A cancelled turn still consumed tokens: its dropped result + // already fed the accumulator (the usage tally at the result + // handler runs before the `session.cancelled` guard), so + // report it — clients metering spend would otherwise lose + // the interrupted turn's tokens entirely (issue #844). Zero + // when the cancel pre-empted the result (wedge/force-cancel). if (session.cancelled && session.activeTurn && !session.activeTurn.settled) { - settleActive({ stopReason: "cancelled" }); + settleActive({ stopReason: "cancelled", usage: sessionUsage(session) }); } else if (owedTrailingIdles > 0) { // Absorb a settled turn's trailing idle. Also covers a // cancel that landed between a turn's counted result and @@ -2142,7 +2147,9 @@ export class ClaudeAcpAgent { // as the freshly-activated turn ending without a result // (which would false-fail a healthy turn — issue #825). owedTrailingIdles++; - settleActive({ stopReason: "cancelled" }); + // Before activateTurn resets the accumulator, so the + // usage still belongs to the cancelled turn. + settleActive({ stopReason: "cancelled", usage: sessionUsage(session) }); } else { settleActive({ stopReason: "end_turn", usage: sessionUsage(session) }); } @@ -2454,6 +2461,8 @@ export class ClaudeAcpAgent { for (const turn of session.turnQueue) { if (turn !== session.activeTurn && !turn.settled) { turn.settled = true; + // Deliberately no `usage`: a queued turn never ran, so the session + // accumulator (the active turn's tally) is not its spend. turn.resolve({ stopReason: "cancelled" }); orphaned++; } diff --git a/src/tests/acp-agent.test.ts b/src/tests/acp-agent.test.ts index f26f68a4..04fc47cc 100644 --- a/src/tests/acp-agent.test.ts +++ b/src/tests/acp-agent.test.ts @@ -82,6 +82,17 @@ function userEcho(u: any) { }; } +/** The `usage` a cancelled active turn settles with when the cancel pre-empted + * its result: all zeros, since only a turn's terminal result feeds the + * accumulator (issue #844). */ +const cancelledTurnUsage = { + inputTokens: 0, + outputTokens: 0, + cachedReadTokens: 0, + cachedWriteTokens: 0, + totalTokens: 0, +}; + /** Wrap a mock async generator with the `Query` methods the agent calls outside * of iteration — `close()` (teardown/closeQueryStream), `interrupt()` (cancel), * and `setModel()` — so a bare generator doesn't trip "x is not a function". */ @@ -4863,7 +4874,10 @@ describe("assembled assistant text fallback", () => { await agent.cancel({ sessionId: "test-session" }); releaseCancel(); - await expect(first).resolves.toEqual({ stopReason: "cancelled" }); + await expect(first).resolves.toEqual({ + stopReason: "cancelled", + usage: cancelledTurnUsage, + }); await agent.prompt({ sessionId: "test-session", prompt: [{ type: "text", text: "second" }] }); @@ -5791,10 +5805,55 @@ describe("post-error recovery", () => { prompt: [{ type: "text", text: "second" }], }); - await expect(first).resolves.toEqual({ stopReason: "cancelled" }); + await expect(first).resolves.toEqual({ + stopReason: "cancelled", + usage: cancelledTurnUsage, + }); await expect(second).resolves.toEqual(expect.objectContaining({ stopReason: "end_turn" })); }); + it("exposes the accumulated usage on a cancelled turn's PromptResponse (issue #844)", async () => { + // The interrupted turn's result is dropped at the `session.cancelled` + // guard, but its usage was already accumulated — the cancelled settle must + // report it so clients metering token spend don't lose the round-trips + // that completed before the cancel. + const agent = createMockAgent(); + let releaseAfterCancel!: () => void; + const afterCancel = new Promise((resolve) => (releaseAfterCancel = resolve)); + injectGeneratorSession(agent, (input) => { + async function* messageGenerator() { + const iter = input[Symbol.asyncIterator](); + const u1 = await iter.next(); + yield userEcho(u1.value); // turn 1 active + await afterCancel; // wait until the test has cancelled turn 1 + // The interrupt still yields the turn's result (usage 10/5) before the + // trailing idle that settles it cancelled. + yield createResultMessage({ subtype: "success", stop_reason: "end_turn", is_error: false }); + yield { type: "system", subtype: "session_state_changed", state: "idle" }; + } + return messageGenerator(); + }); + + const first = agent.prompt({ + sessionId: "test-session", + prompt: [{ type: "text", text: "first" }], + }); + await waitFor(() => !!agent.sessions["test-session"]?.activeTurn); + await agent.cancel({ sessionId: "test-session" }); + releaseAfterCancel(); + + await expect(first).resolves.toEqual({ + stopReason: "cancelled", + usage: { + inputTokens: 10, + outputTokens: 5, + cachedReadTokens: 0, + cachedWriteTokens: 0, + totalTokens: 15, + }, + }); + }); + it("ignores cancel() after the query stream has closed (no interrupt on a dead query)", async () => { const agent = createMockAgent(); injectGeneratorSession(agent, (input) => { @@ -5893,7 +5952,10 @@ describe("post-error recovery", () => { releaseEnd(); // stream ends -> done branch settles turn 1 + rejects turn 2 - await expect(first).resolves.toEqual({ stopReason: "cancelled" }); + await expect(first).resolves.toEqual({ + stopReason: "cancelled", + usage: cancelledTurnUsage, + }); await expect(second).rejects.toThrow(/start a new session/); }); @@ -5934,7 +5996,10 @@ describe("post-error recovery", () => { await agent.cancel({ sessionId: "test-session" }); releaseAfterCancel(); - await expect(first).resolves.toEqual({ stopReason: "cancelled" }); + await expect(first).resolves.toEqual({ + stopReason: "cancelled", + usage: cancelledTurnUsage, + }); // session.cancelled is still true here (turn 1 settled, nothing re-activated). // The /compact result must still settle via head-promotion. @@ -5997,7 +6062,12 @@ describe("post-error recovery", () => { }); afterCancelAndQueue(); - await expect(first).resolves.toEqual({ stopReason: "cancelled" }); + // Turn 1 ran (and was cancelled mid-flight) so it reports its usage; + // turn 2 never ran, so its cancelled settle carries none. + await expect(first).resolves.toEqual({ + stopReason: "cancelled", + usage: cancelledTurnUsage, + }); await expect(second).resolves.toEqual({ stopReason: "cancelled" }); const thirdResult = await third; expect(thirdResult.stopReason).toBe("end_turn"); @@ -6065,7 +6135,10 @@ describe("post-error recovery", () => { }); release(); - await expect(first).resolves.toEqual({ stopReason: "cancelled" }); + await expect(first).resolves.toEqual({ + stopReason: "cancelled", + usage: cancelledTurnUsage, + }); await expect(second).resolves.toEqual({ stopReason: "cancelled" }); const compactResult = await compact; expect(compactResult.stopReason).toBe("end_turn"); @@ -6233,7 +6306,10 @@ describe("session/cancel wedge recovery (issue #680)", () => { // Clean up the wedged prompt + long timer. await agent.closeSession({ sessionId: "test-session" }); - await expect(promptPromise).resolves.toEqual({ stopReason: "cancelled" }); + await expect(promptPromise).resolves.toEqual({ + stopReason: "cancelled", + usage: cancelledTurnUsage, + }); }); it("resolves an in-flight wedged prompt immediately when the session is closed", async () => { @@ -6252,7 +6328,10 @@ describe("session/cancel wedge recovery (issue #680)", () => { await agent.closeSession({ sessionId: "test-session" }); - await expect(promptPromise).resolves.toEqual({ stopReason: "cancelled" }); + await expect(promptPromise).resolves.toEqual({ + stopReason: "cancelled", + usage: cancelledTurnUsage, + }); expect(agent.sessions["test-session"]).toBeUndefined(); }); }); @@ -6370,7 +6449,10 @@ describe("turn abandoned by the SDK (issue #825)", () => { prompt: [{ type: "text", text: "second" }], }); - await expect(first).resolves.toEqual({ stopReason: "cancelled" }); + await expect(first).resolves.toEqual({ + stopReason: "cancelled", + usage: cancelledTurnUsage, + }); const secondResult = await second; expect(secondResult.stopReason).toBe("end_turn"); expect(secondResult.usage?.inputTokens).toBe(10); @@ -6416,7 +6498,10 @@ describe("turn abandoned by the SDK (issue #825)", () => { await agent.cancel({ sessionId: "test-session" }); releaseAfterCancel(); - await expect(first).resolves.toEqual({ stopReason: "cancelled" }); + await expect(first).resolves.toEqual({ + stopReason: "cancelled", + usage: cancelledTurnUsage, + }); await expect(second).resolves.toEqual({ stopReason: "cancelled" }); const third = await agent.prompt({ sessionId: "test-session", @@ -6461,7 +6546,10 @@ describe("turn abandoned by the SDK (issue #825)", () => { }); await waitFor(() => !!agent.sessions["test-session"]?.activeTurn); await agent.cancel({ sessionId: "test-session" }); // backstop (10ms) settles turn 1 - await expect(first).resolves.toEqual({ stopReason: "cancelled" }); + await expect(first).resolves.toEqual({ + stopReason: "cancelled", + usage: cancelledTurnUsage, + }); // Queue turn 2 BEFORE the SDK recovers, so the stale result races it. const second = agent.prompt({