From e68ddc2e3cf4cffcbe3855ac8785ce7a6c6807c7 Mon Sep 17 00:00:00 2001 From: qqqys Date: Mon, 17 Aug 2026 11:13:35 +0800 Subject: [PATCH 01/10] feat(goal): account the tokens a Goal spends MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A Goal reported how many turns it had run and how long it had been active, but never what it cost. That is the one number a user needs to decide whether an autonomous run is worth continuing, and the one number every future limit has to be expressed in — a budget cannot be enforced against a figure nobody keeps. `GoalRecord` now carries `tokensUsed`, summed across the Goal's turns by `reduceGoalTurnFinished`, and `get_goal` reports it in the unpermitted `lastGoal` summary alongside the turn count. The figure is the same one `/stats` shows, read from the session's own model metrics, so a Goal's spend and the session's spend are one measurement rather than two definitions. The runtime pulls the reading rather than having hosts push it. `finishTurn` is called from three separate hosts on the normal path and from core on the interrupted paths, so a pushed count would have to be threaded through four call sites and would go missing wherever it was forgotten; a meter injected once at construction cannot be. A session with no meter, or a meter that throws, bills the turn zero rather than guessing, and never fails the turn. No limit is introduced here — this only counts. Goals recovered from a transcript written before the field existed restore with zero spend. Co-Authored-By: Claude Opus 5 --- .../cli/src/acp-integration/acpAgent.test.ts | 2 + .../acp-integration/session/Session.test.ts | 24 ++++++ .../session/emitters/MessageEmitter.test.ts | 1 + .../session/history-replay-page.test.ts | 2 + .../session/recovered-goal-update.test.ts | 1 + .../io/StreamJsonOutputAdapter.test.ts | 1 + packages/cli/src/serve/routes/goals.test.ts | 1 + .../cli/src/ui/commands/goalCommand.test.ts | 1 + .../cli/src/ui/components/GoalPill.test.tsx | 1 + .../ui/components/HistoryItemDisplay.test.tsx | 1 + .../messages/GoalStatusMessage.test.tsx | 1 + .../src/ui/utils/resumeHistoryUtils.test.ts | 7 ++ packages/core/src/config/config.test.ts | 1 + packages/core/src/config/config.ts | 10 +++ packages/core/src/core/client-goal.test.ts | 1 + packages/core/src/goals/goal-evidence.test.ts | 1 + .../src/goals/goal-legacy-projection.test.ts | 1 + .../core/src/goals/goal-persistence.test.ts | 2 + packages/core/src/goals/goal-persistence.ts | 1 + packages/core/src/goals/goal-protocol.ts | 8 ++ packages/core/src/goals/goal-reducer.test.ts | 40 ++++++++++ packages/core/src/goals/goal-reducer.ts | 9 +++ packages/core/src/goals/goal-runtime.test.ts | 77 +++++++++++++++++++ packages/core/src/goals/goal-runtime.ts | 54 +++++++++++++ packages/core/src/goals/goal-tools.test.ts | 12 +++ packages/core/src/goals/goal-tools.ts | 11 ++- .../src/services/chatRecordingService.test.ts | 1 + .../session-transcript-reader.test.ts | 8 ++ 28 files changed, 278 insertions(+), 2 deletions(-) diff --git a/packages/cli/src/acp-integration/acpAgent.test.ts b/packages/cli/src/acp-integration/acpAgent.test.ts index 6ab12b650db..c6ce8d60f03 100644 --- a/packages/cli/src/acp-integration/acpAgent.test.ts +++ b/packages/cli/src/acp-integration/acpAgent.test.ts @@ -999,6 +999,7 @@ function goalSnapshot( evidenceCursor: { recordId: 'cursor-1' }, turnCount: 0, activeTimeMs: 0, + tokensUsed: 0, createdAt: 123, updatedAt: 123, ...overrides, @@ -17103,6 +17104,7 @@ describe('QwenAgent loadSession / unstable_resumeSession', () => { evidenceCursor: { recordId: 'goal-state' }, turnCount: 0, activeTimeMs: 0, + tokensUsed: 0, createdAt: 1, updatedAt: 1, }, diff --git a/packages/cli/src/acp-integration/session/Session.test.ts b/packages/cli/src/acp-integration/session/Session.test.ts index 2c5bda0ec47..05859bc2226 100644 --- a/packages/cli/src/acp-integration/session/Session.test.ts +++ b/packages/cli/src/acp-integration/session/Session.test.ts @@ -15156,6 +15156,7 @@ describe('Session', () => { evidenceCursor: { recordId: 'cursor-1' }, turnCount: 0, activeTimeMs: 0, + tokensUsed: 0, createdAt: 1234, updatedAt: 1234, }, @@ -15198,6 +15199,7 @@ describe('Session', () => { evidenceCursor: { recordId: 'cursor-1' }, turnCount: 0, activeTimeMs: 0, + tokensUsed: 0, createdAt: 1234, updatedAt: 1234, }, @@ -15261,6 +15263,7 @@ describe('Session', () => { evidenceCursor: { recordId: 'cursor-1' }, turnCount: 0, activeTimeMs: 0, + tokensUsed: 0, createdAt: 1234, updatedAt: 1234, }, @@ -15368,6 +15371,7 @@ describe('Session', () => { evidenceCursor: { recordId: 'cursor-1' }, turnCount: 0, activeTimeMs: 0, + tokensUsed: 0, createdAt: 1234, updatedAt: 1234, }, @@ -15574,6 +15578,7 @@ describe('Session', () => { evidenceCursor: { recordId: 'cursor-1' }, turnCount: 0, activeTimeMs: 0, + tokensUsed: 0, createdAt: 1234, updatedAt: 1234, }; @@ -15621,6 +15626,7 @@ describe('Session', () => { evidenceCursor: { recordId: 'cursor-1' }, turnCount: 0, activeTimeMs: 0, + tokensUsed: 0, createdAt: 1234, updatedAt: 1234, }; @@ -15655,6 +15661,7 @@ describe('Session', () => { evidenceCursor: { recordId: 'cursor-1' }, turnCount: 0, activeTimeMs: 0, + tokensUsed: 0, createdAt: 1234, updatedAt: 1234, }, @@ -15725,6 +15732,7 @@ describe('Session', () => { evidenceCursor: { recordId: 'cursor-1' }, turnCount: 0, activeTimeMs: 0, + tokensUsed: 0, createdAt: 1234, updatedAt: 1234, }, @@ -15781,6 +15789,7 @@ describe('Session', () => { evidenceCursor: { recordId: 'cursor-1' }, turnCount: 0, activeTimeMs: 0, + tokensUsed: 0, createdAt: 1234, updatedAt: 1234, }, @@ -15912,6 +15921,7 @@ describe('Session', () => { evidenceCursor: { recordId: 'cursor-1' }, turnCount: 0, activeTimeMs: 0, + tokensUsed: 0, createdAt: 1234, updatedAt: 1234, }, @@ -15986,6 +15996,7 @@ describe('Session', () => { evidenceCursor: { recordId: 'cursor-1' }, turnCount: 0, activeTimeMs: 0, + tokensUsed: 0, createdAt: 1234, updatedAt: 1234, }, @@ -16045,6 +16056,7 @@ describe('Session', () => { evidenceCursor: { recordId: 'cursor-1' }, turnCount: 0, activeTimeMs: 0, + tokensUsed: 0, createdAt: 1234, updatedAt: 1234, }, @@ -16086,6 +16098,7 @@ describe('Session', () => { evidenceCursor: { recordId: 'cursor-1' }, turnCount: 0, activeTimeMs: 0, + tokensUsed: 0, createdAt: 1234, updatedAt: 1234, }, @@ -16147,6 +16160,7 @@ describe('Session', () => { evidenceCursor: { recordId: 'cursor-1' }, turnCount: 0, activeTimeMs: 0, + tokensUsed: 0, createdAt: 1234, updatedAt: 1234, }, @@ -16194,6 +16208,7 @@ describe('Session', () => { evidenceCursor: { recordId: 'cursor-1' }, turnCount: 0, activeTimeMs: 0, + tokensUsed: 0, createdAt: 1234, updatedAt: 1234, }, @@ -16241,6 +16256,7 @@ describe('Session', () => { evidenceCursor: { recordId: 'cursor-1' }, turnCount: 0, activeTimeMs: 0, + tokensUsed: 0, createdAt: 1234, updatedAt: 1234, }, @@ -16312,6 +16328,7 @@ describe('Session', () => { evidenceCursor: { recordId: 'cursor-1' }, turnCount: 0, activeTimeMs: 0, + tokensUsed: 0, createdAt: 1234, updatedAt: 1234, }, @@ -16382,6 +16399,7 @@ describe('Session', () => { evidenceCursor: { recordId: 'cursor-1' }, turnCount: 0, activeTimeMs: 0, + tokensUsed: 0, createdAt: 1234, updatedAt: 1234, }, @@ -16439,6 +16457,7 @@ describe('Session', () => { evidenceCursor: { recordId: 'cursor-1' }, turnCount: 0, activeTimeMs: 0, + tokensUsed: 0, createdAt: 1234, updatedAt: 1234, }, @@ -16554,6 +16573,7 @@ describe('Session', () => { evidenceCursor: { recordId: 'cursor-1' }, turnCount: 0, activeTimeMs: 0, + tokensUsed: 0, createdAt: 1234, updatedAt: 1234, }, @@ -16664,6 +16684,7 @@ describe('Session', () => { evidenceCursor: { recordId: 'cursor-1' }, turnCount: 0, activeTimeMs: 0, + tokensUsed: 0, createdAt: 1234, updatedAt: 1234, }; @@ -16790,6 +16811,7 @@ describe('Session', () => { evidenceCursor: { recordId: 'cursor-1' }, turnCount: 0, activeTimeMs: 0, + tokensUsed: 0, createdAt: 1234, updatedAt: 1234, }, @@ -20254,6 +20276,7 @@ describe('Session', () => { evidenceCursor: { recordId: 'cursor-1' }, turnCount: 0, activeTimeMs: 0, + tokensUsed: 0, createdAt: 1234, updatedAt: 1234, }, @@ -26716,6 +26739,7 @@ describe('Session', () => { evidenceCursor: { recordId: 'cursor-1' }, turnCount: 0, activeTimeMs: 0, + tokensUsed: 0, createdAt: 1234, updatedAt: 1234, }, diff --git a/packages/cli/src/acp-integration/session/emitters/MessageEmitter.test.ts b/packages/cli/src/acp-integration/session/emitters/MessageEmitter.test.ts index 9b3192e10b3..a6ba6d0777b 100644 --- a/packages/cli/src/acp-integration/session/emitters/MessageEmitter.test.ts +++ b/packages/cli/src/acp-integration/session/emitters/MessageEmitter.test.ts @@ -156,6 +156,7 @@ describe('MessageEmitter', () => { evidenceCursor: { recordId: 'cursor-1' }, turnCount: 0, activeTimeMs: 0, + tokensUsed: 0, createdAt: 1234, updatedAt: 1234, }, diff --git a/packages/cli/src/acp-integration/session/history-replay-page.test.ts b/packages/cli/src/acp-integration/session/history-replay-page.test.ts index b018d7ba6c8..57c4d973033 100644 --- a/packages/cli/src/acp-integration/session/history-replay-page.test.ts +++ b/packages/cli/src/acp-integration/session/history-replay-page.test.ts @@ -38,6 +38,7 @@ const GOAL_STATE: GoalSnapshotV2 = { evidenceCursor: { recordId: 'goal-state' }, turnCount: 2, activeTimeMs: 1000, + tokensUsed: 0, createdAt: 1, updatedAt: 2, }, @@ -628,6 +629,7 @@ describe('history replay page', () => { evidenceCursor: { recordId: 'goal-state' }, turnCount: 3, activeTimeMs: 1234, + tokensUsed: 0, createdAt: 10, updatedAt: 20, }, diff --git a/packages/cli/src/acp-integration/session/recovered-goal-update.test.ts b/packages/cli/src/acp-integration/session/recovered-goal-update.test.ts index c587c83a4d7..04e03992756 100644 --- a/packages/cli/src/acp-integration/session/recovered-goal-update.test.ts +++ b/packages/cli/src/acp-integration/session/recovered-goal-update.test.ts @@ -23,6 +23,7 @@ const hiddenSnapshot: GoalSnapshotV2 = { evidenceCursor: { recordId: 'hidden-record' }, turnCount: 1, activeTimeMs: 10, + tokensUsed: 0, createdAt: 1, updatedAt: 2, }, diff --git a/packages/cli/src/nonInteractive/io/StreamJsonOutputAdapter.test.ts b/packages/cli/src/nonInteractive/io/StreamJsonOutputAdapter.test.ts index c968f36c08f..f2ebc70d41b 100644 --- a/packages/cli/src/nonInteractive/io/StreamJsonOutputAdapter.test.ts +++ b/packages/cli/src/nonInteractive/io/StreamJsonOutputAdapter.test.ts @@ -37,6 +37,7 @@ const goalSnapshot: GoalSnapshotV2 = { evidenceCursor: { recordId: 'record-1' }, turnCount: 3, activeTimeMs: 12_000, + tokensUsed: 0, createdAt: 1, updatedAt: 2, lastReason: 'keep going', diff --git a/packages/cli/src/serve/routes/goals.test.ts b/packages/cli/src/serve/routes/goals.test.ts index fcff6888d12..1ffde6ede9a 100644 --- a/packages/cli/src/serve/routes/goals.test.ts +++ b/packages/cli/src/serve/routes/goals.test.ts @@ -55,6 +55,7 @@ const activeGoal = ( evidenceCursor: { recordId: 'cursor-1' }, turnCount: active.iterations, activeTimeMs: 0, + tokensUsed: 0, createdAt: active.setAt, updatedAt: active.setAt, ...(active.lastReason ? { lastReason: active.lastReason } : {}), diff --git a/packages/cli/src/ui/commands/goalCommand.test.ts b/packages/cli/src/ui/commands/goalCommand.test.ts index 6b1de0dca56..cdb64e7f336 100644 --- a/packages/cli/src/ui/commands/goalCommand.test.ts +++ b/packages/cli/src/ui/commands/goalCommand.test.ts @@ -49,6 +49,7 @@ function goalSnapshot( evidenceCursor: { recordId: 'cursor-1' }, turnCount: 3, activeTimeMs: 1_000, + tokensUsed: 0, createdAt: 10, updatedAt: 20, ...overrides, diff --git a/packages/cli/src/ui/components/GoalPill.test.tsx b/packages/cli/src/ui/components/GoalPill.test.tsx index 7aaee245f88..ca8cdf7ab9c 100644 --- a/packages/cli/src/ui/components/GoalPill.test.tsx +++ b/packages/cli/src/ui/components/GoalPill.test.tsx @@ -39,6 +39,7 @@ function snapshot( evidenceCursor: { recordId: null }, turnCount: 3, activeTimeMs: 2_000, + tokensUsed: 0, createdAt: 1_000, updatedAt: 7_000, ...overrides, diff --git a/packages/cli/src/ui/components/HistoryItemDisplay.test.tsx b/packages/cli/src/ui/components/HistoryItemDisplay.test.tsx index fce85d88a77..64a0df90d94 100644 --- a/packages/cli/src/ui/components/HistoryItemDisplay.test.tsx +++ b/packages/cli/src/ui/components/HistoryItemDisplay.test.tsx @@ -169,6 +169,7 @@ describe('', () => { evidenceCursor: { recordId: 'record-1' }, turnCount: 2, activeTimeMs: 4_000, + tokensUsed: 0, createdAt: 1_000, updatedAt: 5_000, lastReason: 'waiting for approval', diff --git a/packages/cli/src/ui/components/messages/GoalStatusMessage.test.tsx b/packages/cli/src/ui/components/messages/GoalStatusMessage.test.tsx index d18289c1af0..cc6b05e163c 100644 --- a/packages/cli/src/ui/components/messages/GoalStatusMessage.test.tsx +++ b/packages/cli/src/ui/components/messages/GoalStatusMessage.test.tsx @@ -26,6 +26,7 @@ function snapshot( evidenceCursor: { recordId: 'record-1' }, turnCount: 4, activeTimeMs: 12_000, + tokensUsed: 0, createdAt: 1_000, updatedAt: 13_000, ...(lastReason ? { lastReason } : {}), diff --git a/packages/cli/src/ui/utils/resumeHistoryUtils.test.ts b/packages/cli/src/ui/utils/resumeHistoryUtils.test.ts index 44487572302..ca63d397021 100644 --- a/packages/cli/src/ui/utils/resumeHistoryUtils.test.ts +++ b/packages/cli/src/ui/utils/resumeHistoryUtils.test.ts @@ -55,6 +55,7 @@ describe('resumeHistoryUtils', () => { evidenceCursor: { recordId: 'goal-create' }, turnCount: 0, activeTimeMs: 0, + tokensUsed: 0, createdAt: 1, updatedAt: 1, }; @@ -119,6 +120,7 @@ describe('resumeHistoryUtils', () => { evidenceCursor: { recordId: 'goal-create' }, turnCount: 0, activeTimeMs: 0, + tokensUsed: 0, createdAt: 1, updatedAt: 1, }; @@ -145,12 +147,14 @@ describe('resumeHistoryUtils', () => { ...goal, turnCount: 1, activeTimeMs: 10, + tokensUsed: 0, updatedAt: 2, }; const rejected = { ...turned, lastReason: 'More work remains', activeTimeMs: 20, + tokensUsed: 0, updatedAt: 3, }; const checkpointed = { @@ -169,6 +173,7 @@ describe('resumeHistoryUtils', () => { ], }, activeTimeMs: 30, + tokensUsed: 0, updatedAt: 4, }; const limited = { @@ -176,6 +181,7 @@ describe('resumeHistoryUtils', () => { status: 'usage_limited' as const, lastReason: 'provider failed', activeTimeMs: 40, + tokensUsed: 0, updatedAt: 5, }; const conversation = { @@ -327,6 +333,7 @@ describe('resumeHistoryUtils', () => { evidenceCursor: { recordId: 'goal-create' }, turnCount: 0, activeTimeMs: 0, + tokensUsed: 0, createdAt: 1, updatedAt: 1, }; diff --git a/packages/core/src/config/config.test.ts b/packages/core/src/config/config.test.ts index cc95c2d53e1..f0d77cea7c0 100644 --- a/packages/core/src/config/config.test.ts +++ b/packages/core/src/config/config.test.ts @@ -2375,6 +2375,7 @@ describe('Server Config (config.ts)', () => { evidenceCursor: { recordId: 'goal-active' }, turnCount: 1, activeTimeMs: 10, + tokensUsed: 0, createdAt: 1, updatedAt: 2, }, diff --git a/packages/core/src/config/config.ts b/packages/core/src/config/config.ts index 9698a8e478a..9e103d766f6 100644 --- a/packages/core/src/config/config.ts +++ b/packages/core/src/config/config.ts @@ -139,6 +139,7 @@ import { logRipgrepFallback, RipgrepFallbackEvent, StartSessionEvent, + uiTelemetryService, type TelemetryTarget, } from '../telemetry/index.js'; import { @@ -7703,11 +7704,20 @@ export class Config { return; } const recorder = this.chatRecordingService; + const sessionId = this.getSessionId(); const runtime = createGoalRuntime({ journal: recorder, evidenceSource: recorder, verifier: createGoalVerifier(this), checkpointVerifier: createGoalCheckpointVerifier(this), + tokenMeter: { + // The same figure `/stats` reports, so a Goal's spend and the + // session's spend are the same number measured once. + readSessionTokens: () => + Object.values( + uiTelemetryService.getMetricsForSession(sessionId).models, + ).reduce((total, model) => total + model.tokens.total, 0), + }, }); this.goalRuntime = runtime; if (this.goalTurnHost) { diff --git a/packages/core/src/core/client-goal.test.ts b/packages/core/src/core/client-goal.test.ts index 296130219d1..b21783661cd 100644 --- a/packages/core/src/core/client-goal.test.ts +++ b/packages/core/src/core/client-goal.test.ts @@ -123,6 +123,7 @@ function setupGoalClient() { evidenceCursor: { recordId: 'create-record' }, turnCount: 0, activeTimeMs: 0, + tokensUsed: 0, createdAt: 1, updatedAt: 1, }, diff --git a/packages/core/src/goals/goal-evidence.test.ts b/packages/core/src/goals/goal-evidence.test.ts index fdafd4f1da3..e291a3e7adf 100644 --- a/packages/core/src/goals/goal-evidence.test.ts +++ b/packages/core/src/goals/goal-evidence.test.ts @@ -89,6 +89,7 @@ function goal(cursor: string | null = 'cursor'): GoalRecord { evidenceCursor: { recordId: cursor }, turnCount: 2, activeTimeMs: 100, + tokensUsed: 0, createdAt: 1, updatedAt: 2, }; diff --git a/packages/core/src/goals/goal-legacy-projection.test.ts b/packages/core/src/goals/goal-legacy-projection.test.ts index 910fccf778d..d07ed52a9a4 100644 --- a/packages/core/src/goals/goal-legacy-projection.test.ts +++ b/packages/core/src/goals/goal-legacy-projection.test.ts @@ -20,6 +20,7 @@ const GOAL: GoalRecord = { evidenceCursor: { recordId: 'state-1' }, turnCount: 4, activeTimeMs: 2000, + tokensUsed: 0, createdAt: 100, updatedAt: 200, lastReason: 'continuing', diff --git a/packages/core/src/goals/goal-persistence.test.ts b/packages/core/src/goals/goal-persistence.test.ts index 9280ce6fcb5..2805ca6d10e 100644 --- a/packages/core/src/goals/goal-persistence.test.ts +++ b/packages/core/src/goals/goal-persistence.test.ts @@ -26,6 +26,7 @@ const ACTIVE_PAYLOAD: GoalStateRecordPayloadV2 = { evidenceCursor: { recordId: 'state-1' }, turnCount: 3, activeTimeMs: 1500, + tokensUsed: 0, createdAt: 100, updatedAt: 200, }, @@ -229,6 +230,7 @@ describe('legacy migration', () => { evidenceCursor: { recordId: 'migration-record' }, turnCount: 0, activeTimeMs: 0, + tokensUsed: 0, createdAt: 1000, updatedAt: 1000, }, diff --git a/packages/core/src/goals/goal-persistence.ts b/packages/core/src/goals/goal-persistence.ts index 8ba9b05101e..324521592c0 100644 --- a/packages/core/src/goals/goal-persistence.ts +++ b/packages/core/src/goals/goal-persistence.ts @@ -221,6 +221,7 @@ export function createMigratedGoalState( evidenceCursor: { recordId: input.recordUuid }, turnCount: 0, activeTimeMs: 0, + tokensUsed: 0, createdAt: input.now, updatedAt: input.now, }, diff --git a/packages/core/src/goals/goal-protocol.ts b/packages/core/src/goals/goal-protocol.ts index ed5d2bcf6a1..dad740a0e51 100644 --- a/packages/core/src/goals/goal-protocol.ts +++ b/packages/core/src/goals/goal-protocol.ts @@ -104,6 +104,14 @@ export interface GoalRecord { evidenceCursor: TranscriptCursor; turnCount: number; activeTimeMs: number; + /** + * Model tokens billed to this Goal so far, summed across its turns. + * + * Measured the way the session already measures itself, so the number a Goal + * reports and the number `/stats` reports mean the same thing. Zero on Goals + * recovered from a transcript written before the field existed. + */ + tokensUsed: number; createdAt: number; updatedAt: number; evidenceCheckpoint?: GoalEvidenceCheckpoint; diff --git a/packages/core/src/goals/goal-reducer.test.ts b/packages/core/src/goals/goal-reducer.test.ts index 8e602b8e6ff..89e38f21c87 100644 --- a/packages/core/src/goals/goal-reducer.test.ts +++ b/packages/core/src/goals/goal-reducer.test.ts @@ -35,6 +35,7 @@ const goalRecord = (overrides: Partial = {}): GoalRecord => ({ evidenceCursor: { recordId: 'r-100' }, turnCount: 0, activeTimeMs: 0, + tokensUsed: 0, createdAt: 100, updatedAt: 100, ...overrides, @@ -523,6 +524,7 @@ describe('goal reducer', () => { revision: 1, turnCount: 150, activeTimeMs: 150, + tokensUsed: 0, evidenceCursor: { recordId: 'r-100' }, }); }); @@ -533,6 +535,7 @@ describe('goal reducer', () => { status: 'paused', turnCount: 2, activeTimeMs: 60, + tokensUsed: 0, updatedAt: 160, }); @@ -544,10 +547,47 @@ describe('goal reducer', () => { evidenceCursor: { recordId: 'r-100' }, turnCount: 3, activeTimeMs: 60, + tokensUsed: 0, updatedAt: 225, }); }); + it('accumulates per-turn token spend across finished turns', () => { + let goal = goalRecord({ tokensUsed: 0 }); + + goal = reduceGoalTurnFinished(goal, { now: 200, tokensUsed: 1_200 }); + goal = reduceGoalTurnFinished(goal, { now: 300, tokensUsed: 800 }); + + expect(goal).toMatchObject({ turnCount: 2, tokensUsed: 2_000 }); + }); + + it.each([ + ['an unmetered turn', undefined], + ['a meter that went backwards', -50], + ])('adds nothing for %s', (_label, tokensUsed) => { + const finished = reduceGoalTurnFinished(goalRecord({ tokensUsed: 700 }), { + now: 200, + ...(tokensUsed === undefined ? {} : { tokensUsed }), + }); + + expect(finished).toMatchObject({ turnCount: 1, tokensUsed: 700 }); + }); + + it('migrates a snapshot persisted before spend was recorded', () => { + const goal = goalRecord(); + delete (goal as Partial).tokensUsed; + + expect(parseGoalSnapshotV2(snapshot(goal))).toMatchObject({ + goal: { tokensUsed: 0 }, + }); + }); + + it('rejects a snapshot carrying negative spend', () => { + expect( + parseGoalSnapshotV2(snapshot(goalRecord({ tokensUsed: -1 }))), + ).toBeUndefined(); + }); + it.each(['blocked', 'usage_limited', 'complete'] as const)( 'rejects finishing a turn for a %s goal', (status) => { diff --git a/packages/core/src/goals/goal-reducer.ts b/packages/core/src/goals/goal-reducer.ts index 270d374a4f9..6aec5d32974 100644 --- a/packages/core/src/goals/goal-reducer.ts +++ b/packages/core/src/goals/goal-reducer.ts @@ -35,6 +35,8 @@ export interface GoalControlTransition { export interface GoalTurnFinishedTransition { now: number; lastReason?: string; + /** Tokens billed to the turn that just finished. */ + tokensUsed?: number; } export class GoalConflictError extends Error { @@ -158,6 +160,7 @@ export function reduceGoalTurnFinished( } return transitionGoal(current, transition.now, { turnCount: current.turnCount + 1, + tokensUsed: current.tokensUsed + Math.max(0, transition.tokensUsed ?? 0), ...(transition.lastReason === undefined ? {} : { lastReason: transition.lastReason }), @@ -307,6 +310,7 @@ function createGoal( evidenceCursor: copyCursor(cursor), turnCount: 0, activeTimeMs: 0, + tokensUsed: 0, createdAt: now, updatedAt: now, }; @@ -407,6 +411,7 @@ function parseGoalRecord(value: unknown): GoalRecord | undefined { 'evidenceCursor', 'turnCount', 'activeTimeMs', + 'tokensUsed', 'createdAt', 'updatedAt', 'evidenceCheckpoint', @@ -423,6 +428,8 @@ function parseGoalRecord(value: unknown): GoalRecord | undefined { !isTranscriptCursor(value['evidenceCursor']) || !isNonNegativeInteger(value['turnCount']) || !isNonNegativeNumber(value['activeTimeMs']) || + (value['tokensUsed'] !== undefined && + !isNonNegativeNumber(value['tokensUsed'])) || !isFiniteNumber(value['createdAt']) || !isFiniteNumber(value['updatedAt']) || !isGoalEvidenceCheckpoint(value['evidenceCheckpoint']) || @@ -449,6 +456,8 @@ function parseGoalRecord(value: unknown): GoalRecord | undefined { evidenceCursor: copyCursor(value['evidenceCursor']), turnCount: value['turnCount'], activeTimeMs: value['activeTimeMs'], + // Goals persisted before `tokensUsed` existed carry no spend to restore. + tokensUsed: value['tokensUsed'] ?? 0, createdAt: value['createdAt'], updatedAt: value['updatedAt'], ...(value['evidenceCheckpoint'] === undefined diff --git a/packages/core/src/goals/goal-runtime.test.ts b/packages/core/src/goals/goal-runtime.test.ts index 50db7af7904..d31ab87eb05 100644 --- a/packages/core/src/goals/goal-runtime.test.ts +++ b/packages/core/src/goals/goal-runtime.test.ts @@ -269,6 +269,72 @@ describe('goal runtime', () => { }); }); + it('bills a finished turn the tokens the session spent during it', async () => { + const journal = fakeGoalJournal(); + const host = fakeGoalTurnHost(); + let sessionTokens = 10_000; + const runtime = createGoalRuntime({ + journal, + tokenMeter: { readSessionTokens: () => sessionTokens }, + }); + runtime.bindHost(host); + await runtime.dispatch({ action: 'create', objective: 'ship' }); + + const first = host.started[0]!; + sessionTokens = 12_500; + await runtime.finishTurn(first); + + expect(runtime.getSnapshot().goal).toMatchObject({ + turnCount: 1, + tokensUsed: 2_500, + }); + + const second = host.started[1]!; + sessionTokens = 13_000; + await runtime.finishTurn(second); + + expect(runtime.getSnapshot().goal).toMatchObject({ + turnCount: 2, + tokensUsed: 3_000, + }); + }); + + it('bills nothing when the session is unmetered', async () => { + const journal = fakeGoalJournal(); + const host = fakeGoalTurnHost(); + const runtime = createGoalRuntime({ journal }); + runtime.bindHost(host); + await runtime.dispatch({ action: 'create', objective: 'ship' }); + + await runtime.finishTurn(host.started[0]!); + + expect(runtime.getSnapshot().goal).toMatchObject({ + turnCount: 1, + tokensUsed: 0, + }); + }); + + it('finishes the turn when the token meter throws', async () => { + const journal = fakeGoalJournal(); + const host = fakeGoalTurnHost(); + const runtime = createGoalRuntime({ + journal, + tokenMeter: { + readSessionTokens: () => { + throw new Error('metrics are unavailable'); + }, + }, + }); + runtime.bindHost(host); + await runtime.dispatch({ action: 'create', objective: 'ship' }); + + await expect(runtime.finishTurn(host.started[0]!)).resolves.toBeUndefined(); + expect(runtime.getSnapshot().goal).toMatchObject({ + turnCount: 1, + tokensUsed: 0, + }); + }); + it('persists verifier acceptance before completing a verified proposal', async () => { const journal = fakeGoalJournal(); let records: readonly RuntimeRecord[] = []; @@ -889,6 +955,7 @@ describe('goal runtime', () => { expect(runtime.getSnapshot()).toMatchObject({ goal: { activeTimeMs: 4_000, + tokensUsed: 0, evidenceCheckpoint: { checkpointId: expect.any(String) }, }, }); @@ -2494,6 +2561,7 @@ describe('goal runtime', () => { evidenceCursor: { recordId: 'limit-record' }, turnCount: FORMER_GOAL_CONTINUATION_LIMIT, activeTimeMs: 1_000, + tokensUsed: 0, createdAt: 1, updatedAt: 2, }, @@ -2693,6 +2761,7 @@ describe('goal runtime', () => { evidenceCursor: { recordId: 'create-record' }, turnCount: 2, activeTimeMs: 10, + tokensUsed: 0, createdAt: 1, updatedAt: 2, }, @@ -2718,6 +2787,7 @@ describe('goal runtime', () => { evidenceCursor: { recordId: 'create-record' }, turnCount: 2, activeTimeMs: 10, + tokensUsed: 0, createdAt: 1, updatedAt: 2, }, @@ -2749,6 +2819,7 @@ describe('goal runtime', () => { evidenceCursor: { recordId: 'create-record' }, turnCount: 2, activeTimeMs: 10, + tokensUsed: 0, createdAt: 1, updatedAt: 2, }, @@ -2974,6 +3045,7 @@ describe('goal runtime', () => { evidenceCursor: { recordId: 'create-record' }, turnCount: 0, activeTimeMs: 0, + tokensUsed: 0, createdAt: 1, updatedAt: 1, }, @@ -3070,6 +3142,7 @@ describe('goal runtime', () => { evidenceCursor: { recordId: 'create-record' }, turnCount: 0, activeTimeMs: 0, + tokensUsed: 0, createdAt: 1, updatedAt: 1, }, @@ -3350,6 +3423,7 @@ describe('goal runtime', () => { evidenceCursor: { recordId: 'create-record' }, turnCount: 3, activeTimeMs: 0, + tokensUsed: 0, createdAt: 1, updatedAt: 2, }, @@ -3594,6 +3668,7 @@ describe('goal runtime', () => { evidenceCursor: { recordId: 'create-record' }, turnCount: 0, activeTimeMs: 0, + tokensUsed: 0, createdAt: 1, updatedAt: 1, }, @@ -3617,6 +3692,7 @@ describe('goal runtime', () => { evidenceCursor: { recordId: 'create-record' }, turnCount: 1, activeTimeMs: 1, + tokensUsed: 0, createdAt: 1, updatedAt: 2, }, @@ -3691,6 +3767,7 @@ describe('goal runtime', () => { evidenceCursor: { recordId: 'restore-record' }, turnCount: 1, activeTimeMs: 10, + tokensUsed: 0, createdAt: 1, updatedAt: 2, }, diff --git a/packages/core/src/goals/goal-runtime.ts b/packages/core/src/goals/goal-runtime.ts index 9e460af861b..73546b23f79 100644 --- a/packages/core/src/goals/goal-runtime.ts +++ b/packages/core/src/goals/goal-runtime.ts @@ -69,6 +69,20 @@ export interface CreateGoalRuntimeOptions { evidenceSource?: GoalEvidenceSource; verifier?: GoalVerifier; checkpointVerifier?: GoalCheckpointVerifier; + tokenMeter?: GoalTokenMeter; +} + +/** + * A monotonic read of the session's billed tokens. + * + * The runtime takes the difference across a turn rather than having each host + * push a number in: the normal path calls `finishTurn` from three separate + * hosts and the interrupted paths call it from core, so a pushed count would + * have to be threaded through all four and would go missing wherever it was + * forgotten. Pulling it needs one wiring point and cannot be forgotten. + */ +export interface GoalTokenMeter { + readSessionTokens(): number; } export interface GoalEvidenceSource { @@ -188,6 +202,7 @@ export function createGoalRuntime( let currentPermit: GoalTurnPermit | undefined; let currentPermitHost: GoalTurnHost | undefined; let currentTurnKey: string | undefined; + let currentTurnTokensAtStart: number | undefined; let queuedTurnKey: string | undefined; let continuationQueued = false; let currentProposal: @@ -255,6 +270,33 @@ export function createGoalRuntime( } : undefined; + const readSessionTokens = (): number | undefined => { + if (!options.tokenMeter) return undefined; + try { + const total = options.tokenMeter.readSessionTokens(); + return Number.isFinite(total) ? total : undefined; + } catch { + // Goal accounting is bookkeeping. A meter that cannot answer costs the + // Goal its spend figure for this turn, not the turn. + return undefined; + } + }; + + /** + * The tokens the finishing turn billed, consuming the reading it opened with. + * + * Undefined at either end — no meter, an unreadable meter, or a turn whose + * opening reading was never taken — yields zero rather than a guess, so an + * unmetered session reports no spend instead of a wrong one. + */ + const takeTurnTokens = (): number => { + const opened = currentTurnTokensAtStart; + currentTurnTokensAtStart = undefined; + const closed = readSessionTokens(); + if (opened === undefined || closed === undefined) return 0; + return Math.max(0, closed - opened); + }; + const assertAvailable = () => { if (disposed) throw new Error(GOAL_RUNTIME_DISPOSED_MESSAGE); }; @@ -309,6 +351,7 @@ export function createGoalRuntime( turnId: randomUUID(), }; currentPermitHost = scheduledHost; + currentTurnTokensAtStart = readSessionTokens(); currentTurnKey = `goal-runtime:${currentPermit.turnId}`; const startedPermit = structuredClone(currentPermit); snapshot = { ...snapshot, activity: 'running' }; @@ -318,6 +361,7 @@ export function createGoalRuntime( if (isCurrentPermit(startedPermit)) { const nextTurnKey = queuedTurnKey; currentPermit = undefined; + currentTurnTokensAtStart = undefined; currentPermitHost = undefined; currentTurnKey = undefined; currentProposal = undefined; @@ -333,6 +377,7 @@ export function createGoalRuntime( turnId: randomUUID(), }; currentPermitHost = host; + currentTurnTokensAtStart = readSessionTokens(); currentTurnKey = nextTurnKey; currentTurnFeedback = nextVerifierFeedback; nextVerifierFeedback = undefined; @@ -476,6 +521,7 @@ export function createGoalRuntime( turnId: randomUUID(), }; currentPermitHost = host; + currentTurnTokensAtStart = readSessionTokens(); currentTurnKey = nextTurnKey; currentTurnFeedback = nextVerifierFeedback; nextVerifierFeedback = undefined; @@ -1124,6 +1170,7 @@ export function createGoalRuntime( turnId: randomUUID(), }; currentPermitHost = host; + currentTurnTokensAtStart = readSessionTokens(); currentTurnKey = turnKey; currentTurnFeedback = nextVerifierFeedback; nextVerifierFeedback = undefined; @@ -1144,6 +1191,7 @@ export function createGoalRuntime( nextVerifierFeedback ??= currentTurnFeedback; } currentPermit = undefined; + currentTurnTokensAtStart = undefined; currentPermitHost = undefined; currentTurnKey = undefined; currentTurnFeedback = undefined; @@ -1170,6 +1218,7 @@ export function createGoalRuntime( turnId: randomUUID(), }; currentPermitHost = host; + currentTurnTokensAtStart = readSessionTokens(); currentTurnKey = nextTurnKey; currentTurnFeedback = nextVerifierFeedback; nextVerifierFeedback = undefined; @@ -1208,6 +1257,7 @@ export function createGoalRuntime( const recordUuid = randomUUID(); const nextGoal = reduceGoalTurnFinished(snapshot.goal, { now: Date.now(), + tokensUsed: takeTurnTokens(), }); const persistedSnapshot: GoalSnapshotV2 = { v: GOAL_STATE_VERSION, @@ -1271,6 +1321,7 @@ export function createGoalRuntime( activity: verifying ? 'verifying' : 'idle', }; currentPermit = undefined; + currentTurnTokensAtStart = undefined; currentPermitHost = undefined; currentTurnKey = undefined; currentTurnFeedback = undefined; @@ -1284,6 +1335,7 @@ export function createGoalRuntime( turnId: randomUUID(), }; currentPermitHost = host; + currentTurnTokensAtStart = readSessionTokens(); currentTurnKey = nextTurnKey; currentTurnFeedback = nextVerifierFeedback; nextVerifierFeedback = undefined; @@ -1453,6 +1505,7 @@ export function createGoalRuntime( } if (invalidatesPermit) { currentPermit = undefined; + currentTurnTokensAtStart = undefined; currentPermitHost = undefined; currentTurnKey = undefined; queuedTurnKey = undefined; @@ -1491,6 +1544,7 @@ export function createGoalRuntime( disposed = true; const invalidatedHost = currentPermitHost ?? host; currentPermit = undefined; + currentTurnTokensAtStart = undefined; currentPermitHost = undefined; currentTurnKey = undefined; queuedTurnKey = undefined; diff --git a/packages/core/src/goals/goal-tools.test.ts b/packages/core/src/goals/goal-tools.test.ts index 364edc843e5..2fab5e440ec 100644 --- a/packages/core/src/goals/goal-tools.test.ts +++ b/packages/core/src/goals/goal-tools.test.ts @@ -143,6 +143,7 @@ describe('GetGoalTool', () => { evidenceCursor: { recordId: 'record-1' }, turnCount: 27, activeTimeMs: 1_763_705, + tokensUsed: 4_096, createdAt: 1, updatedAt: 2, lastReason: GOAL_EVIDENCE_CATALOG_EXHAUSTED_REASON, @@ -173,6 +174,7 @@ describe('GetGoalTool', () => { status: 'usage_limited', turnCount: 27, activeTimeMs: 1_763_705, + tokensUsed: 4_096, lastReason: GOAL_EVIDENCE_CATALOG_EXHAUSTED_REASON, }, }); @@ -196,6 +198,7 @@ describe('GetGoalTool', () => { evidenceCursor: { recordId: 'record-1' }, turnCount: 1, activeTimeMs: 750, + tokensUsed: 4_096, createdAt: 1, updatedAt: 2, }, @@ -212,6 +215,7 @@ describe('GetGoalTool', () => { status: 'paused', turnCount: 1, activeTimeMs: 750, + tokensUsed: 4_096, }, }); expect(result.returnDisplay).toBe( @@ -233,6 +237,7 @@ describe('GetGoalTool', () => { evidenceCursor: { recordId: 'record-1' }, turnCount: 2, activeTimeMs: 10, + tokensUsed: 4_096, createdAt: 1, updatedAt: 2, evidenceCheckpoint: { @@ -263,6 +268,7 @@ describe('GetGoalTool', () => { status: 'complete', turnCount: 2, activeTimeMs: 10, + tokensUsed: 4_096, }, }); }); @@ -306,6 +312,7 @@ describe('GetGoalTool', () => { evidenceCursor: { recordId: 'cursor-1' }, turnCount: 4, activeTimeMs: 120, + tokensUsed: 4_096, createdAt: 10, updatedAt: 20, }, @@ -448,6 +455,7 @@ describe('UpdateGoalTool', () => { evidenceCursor: { recordId: 'goal-created' }, turnCount: 3, activeTimeMs: 100, + tokensUsed: 4_096, createdAt: 1, updatedAt: 2, }, @@ -523,6 +531,7 @@ describe('UpdateGoalTool', () => { evidenceCursor: { recordId: 'goal-created' }, turnCount: 3, activeTimeMs: 100, + tokensUsed: 4_096, createdAt: 1, updatedAt: 2, }, @@ -591,6 +600,7 @@ describe('UpdateGoalTool', () => { evidenceCursor: { recordId: 'goal-created' }, turnCount: 1, activeTimeMs: 0, + tokensUsed: 4_096, createdAt: 1, updatedAt: 1, }, @@ -1022,6 +1032,7 @@ describe('UpdateGoalTool', () => { evidenceCursor: { recordId: 'cursor-1' }, turnCount: 1, activeTimeMs: 0, + tokensUsed: 4_096, createdAt: 1, updatedAt: 1, }, @@ -1111,6 +1122,7 @@ describe('UpdateGoalTool', () => { evidenceCursor: { recordId: 'cursor-1' }, turnCount: 1, activeTimeMs: 0, + tokensUsed: 4_096, createdAt: 1, updatedAt: 1, }, diff --git a/packages/core/src/goals/goal-tools.ts b/packages/core/src/goals/goal-tools.ts index ddcee87f074..822e76532fb 100644 --- a/packages/core/src/goals/goal-tools.ts +++ b/packages/core/src/goals/goal-tools.ts @@ -45,7 +45,13 @@ export type GoalToolResult = ToolResult; type LastGoalSummary = Pick< GoalRecord, - 'goalId' | 'revision' | 'status' | 'turnCount' | 'activeTimeMs' | 'lastReason' + | 'goalId' + | 'revision' + | 'status' + | 'turnCount' + | 'activeTimeMs' + | 'tokensUsed' + | 'lastReason' >; type GetGoalRuntime = Pick & { @@ -108,7 +114,7 @@ export class GetGoalTool extends BaseDeclarativeTool< super( GetGoalTool.Name, ToolDisplayNames.GET_GOAL, - 'Read the current Goal identity, objective, evidence cursor, and bounded evidence-reference catalog for this permitted Goal turn. Outside a permitted Goal turn it reports "active": false together with "lastGoal", a scalar summary (goalId, revision, status, turnCount, activeTimeMs, and lastReason when one was recorded) of the session\'s most recent Goal, so a Goal that has already stopped can still be inspected. It never returns uncited transcript history or changes Goal state. Use the result silently; do not narrate or acknowledge the retrieval to the user.', + 'Read the current Goal identity, objective, evidence cursor, and bounded evidence-reference catalog for this permitted Goal turn. Outside a permitted Goal turn it reports "active": false together with "lastGoal", a scalar summary (goalId, revision, status, turnCount, activeTimeMs, tokensUsed, and lastReason when one was recorded) of the session\'s most recent Goal, so a Goal that has already stopped can still be inspected. It never returns uncited transcript history or changes Goal state. Use the result silently; do not narrate or acknowledge the retrieval to the user.', Kind.Read, { type: 'object', @@ -159,6 +165,7 @@ export class GetGoalTool extends BaseDeclarativeTool< status: goal.status, turnCount: goal.turnCount, activeTimeMs: goal.activeTimeMs, + tokensUsed: goal.tokensUsed, ...(goal.lastReason === undefined ? {} : { lastReason: goal.lastReason }), }; } diff --git a/packages/core/src/services/chatRecordingService.test.ts b/packages/core/src/services/chatRecordingService.test.ts index cd3dac199a0..440d8144aec 100644 --- a/packages/core/src/services/chatRecordingService.test.ts +++ b/packages/core/src/services/chatRecordingService.test.ts @@ -832,6 +832,7 @@ describe('ChatRecordingService', () => { evidenceCursor: { recordId: 'goal-record' }, turnCount: 0, activeTimeMs: 0, + tokensUsed: 0, createdAt: 100, updatedAt: 100, }, diff --git a/packages/core/src/services/session-transcript-reader.test.ts b/packages/core/src/services/session-transcript-reader.test.ts index 77d16a4b6e2..a06dc6cb61d 100644 --- a/packages/core/src/services/session-transcript-reader.test.ts +++ b/packages/core/src/services/session-transcript-reader.test.ts @@ -433,6 +433,7 @@ describe('SessionTranscriptReader', () => { evidenceCursor: { recordId: null }, turnCount: 0, activeTimeMs: 0, + tokensUsed: 0, createdAt: 1, updatedAt: 1, }, @@ -500,6 +501,7 @@ describe('SessionTranscriptReader', () => { evidenceCursor: { recordId: null }, turnCount: 0, activeTimeMs: 0, + tokensUsed: 0, createdAt: 1, updatedAt: 1, }, @@ -528,6 +530,7 @@ describe('SessionTranscriptReader', () => { evidenceCursor: { recordId: null }, turnCount: 0, activeTimeMs: 0, + tokensUsed: 0, createdAt: 1, updatedAt: 1, }, @@ -787,6 +790,7 @@ describe('SessionTranscriptReader', () => { evidenceCursor: { recordId: 'goal' }, turnCount: 1, activeTimeMs: 0, + tokensUsed: 0, createdAt: 1, updatedAt: 1, }, @@ -1204,6 +1208,7 @@ describe('SessionTranscriptReader', () => { evidenceCursor: { recordId: 'goal' }, turnCount: 1, activeTimeMs: 0, + tokensUsed: 0, createdAt: 1, updatedAt: 1, }, @@ -1643,6 +1648,7 @@ describe('SessionTranscriptReader', () => { evidenceCursor: { recordId: 'u1' }, turnCount: 0, activeTimeMs: 0, + tokensUsed: 0, createdAt: 1, updatedAt: 1, }, @@ -1910,6 +1916,7 @@ describe('SessionTranscriptReader', () => { evidenceCursor: { recordId: 'cursor' }, turnCount: 1, activeTimeMs: 0, + tokensUsed: 0, createdAt: 1, updatedAt: 2, }, @@ -1992,6 +1999,7 @@ describe('SessionTranscriptReader', () => { evidenceCursor: { recordId: 'missing-cursor' }, turnCount: 1, activeTimeMs: 0, + tokensUsed: 0, createdAt: 1, updatedAt: 2, }, From d78703df091e0353c7403dabe6a9e860c5614799 Mon Sep 17 00:00:00 2001 From: qqqys <266654365+qqqys@users.noreply.github.com> Date: Mon, 17 Aug 2026 12:12:50 +0800 Subject: [PATCH 02/10] fix(acp-bridge): include Goal tokens in replay fixture --- packages/acp-bridge/src/transcript-replay.test.ts | 1 + 1 file changed, 1 insertion(+) diff --git a/packages/acp-bridge/src/transcript-replay.test.ts b/packages/acp-bridge/src/transcript-replay.test.ts index f9e5707157f..6a942d19d14 100644 --- a/packages/acp-bridge/src/transcript-replay.test.ts +++ b/packages/acp-bridge/src/transcript-replay.test.ts @@ -23,6 +23,7 @@ const GOAL: GoalRecord = { status: 'active', evidenceCursor: { recordId: 'record-0' }, turnCount: 4, + tokensUsed: 0, activeTimeMs: 2000, createdAt: 100, updatedAt: 200, From 2c651e135123aa9999e8697e39dffabae8389738 Mon Sep 17 00:00:00 2001 From: qqqys <266654365+qqqys@users.noreply.github.com> Date: Mon, 17 Aug 2026 15:24:16 +0800 Subject: [PATCH 03/10] fix(goal): clarify and test token accounting --- packages/core/src/config/config.test.ts | 58 ++++++++++++++++++++ packages/core/src/config/config.ts | 5 +- packages/core/src/goals/goal-protocol.ts | 8 +-- packages/core/src/goals/goal-runtime.test.ts | 27 +++++++++ 4 files changed, 92 insertions(+), 6 deletions(-) diff --git a/packages/core/src/config/config.test.ts b/packages/core/src/config/config.test.ts index f0d77cea7c0..28a94aad6ef 100644 --- a/packages/core/src/config/config.test.ts +++ b/packages/core/src/config/config.test.ts @@ -35,6 +35,7 @@ import { refreshSessionContext, logStartSession, logSessionEnd, + uiTelemetryService, } from '../telemetry/index.js'; import type { ContentGenerator, @@ -55,6 +56,11 @@ import { ShellTool } from '../tools/shell.js'; import { canUseRipgrep } from '../utils/ripgrepUtils.js'; import { getSessionProjectDir } from '../utils/sessionIdContext.js'; import { logRipgrepFallback } from '../telemetry/loggers.js'; +import { + EVENT_API_RESPONSE, + UiTelemetryService, +} from '../telemetry/uiTelemetry.js'; +import type { ApiResponseEvent } from '../telemetry/types.js'; import { RipgrepFallbackEvent } from '../telemetry/types.js'; import { ToolRegistry } from '../tools/tool-registry.js'; import { ToolNames } from '../tools/tool-names.js'; @@ -318,6 +324,7 @@ vi.mock('../telemetry/index.js', async (importOriginal) => { refreshSessionContext: vi.fn(), uiTelemetryService: { getLastPromptTokenCount: vi.fn(), + getMetricsForSession: vi.fn(), }, }; }); @@ -2614,6 +2621,57 @@ describe('Server Config (config.ts)', () => { ).rejects.toThrow('Goal runtime has been disposed'); }); + it('bills only the current session telemetry to its Goal', async () => { + const sessionId = 'metered-session'; + const telemetry = new UiTelemetryService(); + vi.mocked(uiTelemetryService.getMetricsForSession).mockImplementation( + (requestedSessionId) => + telemetry.getMetricsForSession(requestedSessionId), + ); + const config = new Config({ + ...baseParams, + chatRecording: true, + sessionId, + }); + const recorder = config.getChatRecordingService(); + if (!recorder) throw new Error('expected a chat recording service'); + vi.spyOn(recorder, 'recordGoalState').mockResolvedValue({} as ChatRecord); + const runtime = config.getGoalRuntime(); + await runtime.dispatch({ action: 'create', objective: 'ship' }); + const permit = runtime.beginTurn('turn-1'); + if (!permit) throw new Error('expected a Goal turn permit'); + const apiResponse = ( + model: string, + totalTokens: number, + ): ApiResponseEvent & { 'event.name': typeof EVENT_API_RESPONSE } => + ({ + 'event.name': EVENT_API_RESPONSE, + model, + duration_ms: 1, + input_token_count: totalTokens, + output_token_count: 0, + total_token_count: totalTokens, + cached_content_token_count: 0, + thoughts_token_count: 0, + }) as ApiResponseEvent & { + 'event.name': typeof EVENT_API_RESPONSE; + }; + telemetry.addEvent(apiResponse('model-a', 40), sessionId); + telemetry.addEvent(apiResponse('model-b', 60), sessionId); + telemetry.addEvent(apiResponse('decoy-model', 1_000), 'other-session'); + + await runtime.finishTurn(permit); + + expect(runtime.getSnapshot().goal?.tokensUsed).toBe(100); + expect( + vi + .mocked(uiTelemetryService.getMetricsForSession) + .mock.calls.every(([requestedSessionId]) => + Object.is(requestedSessionId, sessionId), + ), + ).toBe(true); + }); + it('rebinds the current Goal host to every replacement runtime', async () => { const config = new Config({ ...baseParams, diff --git a/packages/core/src/config/config.ts b/packages/core/src/config/config.ts index 9e103d766f6..829af3d27a1 100644 --- a/packages/core/src/config/config.ts +++ b/packages/core/src/config/config.ts @@ -7711,8 +7711,9 @@ export class Config { verifier: createGoalVerifier(this), checkpointVerifier: createGoalCheckpointVerifier(this), tokenMeter: { - // The same figure `/stats` reports, so a Goal's spend and the - // session's spend are the same number measured once. + // The figure `/stats` reports within Goal turn windows. Spend from + // verification and checkpoint side queries lands between turns and + // is not billed here. readSessionTokens: () => Object.values( uiTelemetryService.getMetricsForSession(sessionId).models, diff --git a/packages/core/src/goals/goal-protocol.ts b/packages/core/src/goals/goal-protocol.ts index dad740a0e51..db09bb2ca55 100644 --- a/packages/core/src/goals/goal-protocol.ts +++ b/packages/core/src/goals/goal-protocol.ts @@ -105,11 +105,11 @@ export interface GoalRecord { turnCount: number; activeTimeMs: number; /** - * Model tokens billed to this Goal so far, summed across its turns. + * Model tokens billed to this Goal so far, summed across its turn windows. * - * Measured the way the session already measures itself, so the number a Goal - * reports and the number `/stats` reports mean the same thing. Zero on Goals - * recovered from a transcript written before the field existed. + * Measured from the same session token source as `/stats`. Verification and + * checkpoint side queries run between turn windows and are not included. + * Zero on Goals recovered from a transcript written before the field existed. */ tokensUsed: number; createdAt: number; diff --git a/packages/core/src/goals/goal-runtime.test.ts b/packages/core/src/goals/goal-runtime.test.ts index d31ab87eb05..3c8ddb4c3fd 100644 --- a/packages/core/src/goals/goal-runtime.test.ts +++ b/packages/core/src/goals/goal-runtime.test.ts @@ -335,6 +335,33 @@ describe('goal runtime', () => { }); }); + it.each([Number.NaN, Number.POSITIVE_INFINITY])( + 'finishes the turn when the token meter returns %s', + async (invalidTotal) => { + const journal = fakeGoalJournal(); + const host = fakeGoalTurnHost(); + const runtime = createGoalRuntime({ + journal, + tokenMeter: { + readSessionTokens: vi + .fn() + .mockReturnValueOnce(0) + .mockReturnValue(invalidTotal), + }, + }); + runtime.bindHost(host); + await runtime.dispatch({ action: 'create', objective: 'ship' }); + + await expect( + runtime.finishTurn(host.started[0]!), + ).resolves.toBeUndefined(); + expect(runtime.getSnapshot().goal).toMatchObject({ + turnCount: 1, + tokensUsed: 0, + }); + }, + ); + it('persists verifier acceptance before completing a verified proposal', async () => { const journal = fakeGoalJournal(); let records: readonly RuntimeRecord[] = []; From d74015c81f3079771817743e41420b0904ff84e7 Mon Sep 17 00:00:00 2001 From: qqqys <266654365+qqqys@users.noreply.github.com> Date: Mon, 17 Aug 2026 22:27:29 +0800 Subject: [PATCH 04/10] fix(goal): seed resume token meter before activation --- .../cli/src/ui/hooks/useBranchCommand.test.ts | 15 +++++++ packages/cli/src/ui/hooks/useBranchCommand.ts | 5 +++ .../cli/src/ui/hooks/useResumeCommand.test.ts | 12 +++++- packages/cli/src/ui/hooks/useResumeCommand.ts | 5 +++ .../core/src/services/sessionService.test.ts | 42 +++++++++++++++++++ packages/core/src/services/sessionService.ts | 15 +++++-- 6 files changed, 89 insertions(+), 5 deletions(-) diff --git a/packages/cli/src/ui/hooks/useBranchCommand.test.ts b/packages/cli/src/ui/hooks/useBranchCommand.test.ts index e742b0da524..258d13f780c 100644 --- a/packages/cli/src/ui/hooks/useBranchCommand.test.ts +++ b/packages/cli/src/ui/hooks/useBranchCommand.test.ts @@ -9,6 +9,17 @@ import { renderHook, act } from '@testing-library/react'; import { useBranchCommand } from './useBranchCommand.js'; import type { LoadedSettings } from '../../config/settings.js'; +const replayUiTelemetryEventsMock = vi.hoisted(() => vi.fn()); + +vi.mock('@qwen-code/qwen-code-core', async (importOriginal) => { + const original = + await importOriginal(); + return { + ...original, + replayUiTelemetryEventsFromConversation: replayUiTelemetryEventsMock, + }; +}); + const mockSettings = { merged: { ui: { history: { collapseOnResume: false } } }, } as unknown as LoadedSettings; @@ -80,6 +91,7 @@ describe('useBranchCommand', () => { }); beforeEach(() => { + replayUiTelemetryEventsMock.mockClear(); forkSession = vi .fn() .mockResolvedValue({ filePath: '/tmp/new.jsonl', copiedCount: 2 }); @@ -283,6 +295,9 @@ describe('useBranchCommand', () => { await result.current.handleBranch('my-branch'); }); expect(getGoalRuntimeReady).toHaveBeenCalledTimes(1); + expect( + replayUiTelemetryEventsMock.mock.invocationCallOrder[0], + ).toBeLessThan(getGoalRuntimeReady.mock.invocationCallOrder[0]!); }); it('rolls core back when the fork contains malformed Goal state', async () => { diff --git a/packages/cli/src/ui/hooks/useBranchCommand.ts b/packages/cli/src/ui/hooks/useBranchCommand.ts index 4a56c32e457..a70291e03b9 100644 --- a/packages/cli/src/ui/hooks/useBranchCommand.ts +++ b/packages/cli/src/ui/hooks/useBranchCommand.ts @@ -12,6 +12,7 @@ import { type ResumedSessionData, SessionStartSource, computeUniqueBranchTitle, + replayUiTelemetryEventsFromConversation, } from '@qwen-code/qwen-code-core'; import { buildResumedHistoryItems, @@ -199,6 +200,10 @@ export function useBranchCommand( // the parent, silently recording user input into an orphan. config.startNewSession(newSessionId, resumed); coreSwapped = true; + replayUiTelemetryEventsFromConversation( + resumed.conversation, + newSessionId, + ); await waitForGoalRuntime(config); await config.getGeminiClient()?.initialize?.(SessionStartSource.Branch); diff --git a/packages/cli/src/ui/hooks/useResumeCommand.test.ts b/packages/cli/src/ui/hooks/useResumeCommand.test.ts index 8bfe6a67858..50ebfa99098 100644 --- a/packages/cli/src/ui/hooks/useResumeCommand.test.ts +++ b/packages/cli/src/ui/hooks/useResumeCommand.test.ts @@ -15,6 +15,8 @@ import { useHistory } from './useHistoryManager.js'; import type { Content } from '@google/genai'; import type { LoadedSettings } from '../../config/settings.js'; +const replayUiTelemetryEventsMock = vi.hoisted(() => vi.fn()); + const mockSettings = { merged: { ui: { @@ -105,6 +107,7 @@ vi.mock('@qwen-code/qwen-code-core', async (importOriginal) => { return { ...original, SessionService, + replayUiTelemetryEventsFromConversation: replayUiTelemetryEventsMock, }; }); @@ -231,6 +234,7 @@ describe('useResumeCommand', () => { }); it('handleResume closes the dialog immediately and restores session state', async () => { + replayUiTelemetryEventsMock.mockClear(); resumeMocks.reset(); resumeMocks.createPendingLoadSession(); @@ -244,6 +248,7 @@ describe('useResumeCommand', () => { const geminiClient = { initialize: vi.fn().mockResolvedValue(undefined), }; + const getGoalRuntimeReady = vi.fn().mockResolvedValue({}); const resetMonitorRegistry = vi.fn(); const config = { @@ -251,7 +256,7 @@ describe('useResumeCommand', () => { getTargetDir: () => '/tmp', getGeminiClient: () => geminiClient, startNewSession: vi.fn(), - getGoalRuntimeReady: vi.fn().mockResolvedValue({}), + getGoalRuntimeReady, getBackgroundTaskRegistry: () => ({ hasRunningTasks: vi.fn().mockReturnValue(false), reset: vi.fn(), @@ -332,7 +337,10 @@ describe('useResumeCommand', () => { historyManager.loadHistory.mock.invocationCallOrder[0]!, ); expect(resetMonitorRegistry).toHaveBeenCalledTimes(1); - expect(config.getGoalRuntimeReady).toHaveBeenCalledTimes(1); + expect(getGoalRuntimeReady).toHaveBeenCalledTimes(1); + expect( + replayUiTelemetryEventsMock.mock.invocationCallOrder[0], + ).toBeLessThan(getGoalRuntimeReady.mock.invocationCallOrder[0]!); }); it('adds a recovery notice when resuming an interrupted tool turn', async () => { diff --git a/packages/cli/src/ui/hooks/useResumeCommand.ts b/packages/cli/src/ui/hooks/useResumeCommand.ts index c85314524a4..b61511601c3 100644 --- a/packages/cli/src/ui/hooks/useResumeCommand.ts +++ b/packages/cli/src/ui/hooks/useResumeCommand.ts @@ -8,6 +8,7 @@ import { useState, useCallback } from 'react'; import { SessionService, buildSessionRecoveryPlan, + replayUiTelemetryEventsFromConversation, type Config, type SessionListItem, } from '@qwen-code/qwen-code-core'; @@ -164,6 +165,10 @@ export function useResumeCommand( resetBackgroundStateForSessionSwitch(config); config.startNewSession(sessionId, sessionData); coreSwapped = true; + replayUiTelemetryEventsFromConversation( + sessionData.conversation, + sessionId, + ); await waitForGoalRuntime(config); // Rebuild turn boundary tracking so rewind works within resumed sessions. config diff --git a/packages/core/src/services/sessionService.test.ts b/packages/core/src/services/sessionService.test.ts index 0dd6a978434..4e0112d0da5 100644 --- a/packages/core/src/services/sessionService.test.ts +++ b/packages/core/src/services/sessionService.test.ts @@ -24,6 +24,7 @@ import { buildApiHistoryFromConversation, getResumePromptTokenCount, getResumeTokenCounts, + replayUiTelemetryEventsFromConversation, type ConversationRecord, } from './sessionService.js'; import { @@ -38,6 +39,8 @@ import { SessionOrganizationService } from './session-organization-service.js'; import { CompressionStatus } from '../core/turn.js'; import type { ChatRecord } from './chatRecordingService.js'; import * as jsonl from '../utils/jsonl-utils.js'; +import { EVENT_API_RESPONSE, uiTelemetryService } from '../telemetry/index.js'; +import type { ApiResponseEvent } from '../telemetry/types.js'; vi.mock('./usageHistoryService.js', () => ({ persistUsageBeforeTranscriptDeletion: vi.fn().mockResolvedValue(true), @@ -2726,6 +2729,45 @@ describe('SessionService', () => { }); }); + it('replays historical token metrics for a resumed session', () => { + const sessionId = 'resumed-goal-meter'; + const uiEvent = { + 'event.name': EVENT_API_RESPONSE, + model: 'model-a', + duration_ms: 1, + input_token_count: 200_000, + output_token_count: 0, + total_token_count: 200_000, + cached_content_token_count: 0, + thoughts_token_count: 0, + } as ApiResponseEvent & { 'event.name': typeof EVENT_API_RESPONSE }; + const telemetryRecord: ChatRecord = { + ...recordA1, + uuid: 'telemetry-1', + sessionId, + type: 'system', + subtype: 'ui_telemetry', + systemPayload: { uiEvent }, + }; + + replayUiTelemetryEventsFromConversation( + { + sessionId, + projectHash: 'test-project-hash', + startTime: '2024-01-01T00:00:00Z', + lastUpdated: '2024-01-01T00:00:00Z', + messages: [telemetryRecord], + }, + sessionId, + ); + + expect( + Object.values(uiTelemetryService.getMetricsForSession(sessionId).models) + .map((model) => model.tokens.total) + .reduce((total, tokens) => total + tokens, 0), + ).toBe(200_000); + }); + describe('getResumePromptTokenCount', () => { const baseRecord: ChatRecord = { uuid: 'r1', diff --git a/packages/core/src/services/sessionService.ts b/packages/core/src/services/sessionService.ts index 4e778e43864..90148ef9c88 100644 --- a/packages/core/src/services/sessionService.ts +++ b/packages/core/src/services/sessionService.ts @@ -2618,12 +2618,11 @@ function collectReferencedFileHistoryBackupNames( /** * Replays stored UI telemetry events to rebuild metrics when resuming a session. - * Also restores the last prompt token count from the best available source. */ -export function replayUiTelemetryFromConversation( +export function replayUiTelemetryEventsFromConversation( conversation: ConversationRecord, sessionId?: string, -): ResumeTokenCounts | undefined { +): void { if (sessionId) { uiTelemetryService.resetSession(sessionId); } else { @@ -2642,6 +2641,16 @@ export function replayUiTelemetryFromConversation( uiTelemetryService.addEvent(uiEvent, sessionId); } } +} + +/** + * Replays stored UI telemetry and restores the last prompt token count. + */ +export function replayUiTelemetryFromConversation( + conversation: ConversationRecord, + sessionId?: string, +): ResumeTokenCounts | undefined { + replayUiTelemetryEventsFromConversation(conversation, sessionId); const resumeTokenCounts = getResumeTokenCounts(conversation); if (resumeTokenCounts !== undefined) { From 12672cdb3f139d01e60b922e57ed61ed24102d64 Mon Sep 17 00:00:00 2001 From: qqqys Date: Tue, 18 Aug 2026 18:41:59 +0800 Subject: [PATCH 05/10] fix(goal): replay resumed telemetry once per session swap MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Round-3 review found the same Critical in two places (R3-1 in useResumeCommand.ts, R3-2 in useBranchCommand.ts): both hooks replay the stored UI telemetry, and `GeminiClient.initialize()` replays it again a few lines later. `uiTelemetryService.addEvent` accumulates into both the process-wide metrics and the per-session bucket, but the `resetSession` leading each replay clears only the bucket — so the aggregate ends up carrying one extra copy of the whole session history. At process exit or on /clear, `persistSessionUsage` writes that doubled figure out under the resumed session id, permanently inflating cross-session usage reports, while in-session /stats reads the session-keyed bucket and stays correct — which is why manual testing never surfaces it. For /branch the same mechanism can reach a third copy when the fork fails after the replay and the rollback's own initialize() replays the parent. Replay exactly once per swap. The hook-level call stays where it is — its position before waitForGoalRuntime is load-bearing, it is what makes the Goal meter open on the restored totals — and it now hands off to the client via a one-shot marker on Config. initialize() consumes the marker and restores only the token counts, so its chat seeding is unchanged. startNewSession() clears the marker, so resuming the same session again later still replays exactly once, and a rollback still replays for the session it rolls back to. Also addressed from the same round: - R3-3 / R3-4: the order tests pinned only the call order of replayUiTelemetryEventsFromConversation, never its arguments, so dropping the optional sessionId survived all 38 hook tests. Dropping it sends sessionService down the global reset() branch, clearing every live session's bucket and keying no history under the resumed session. Both tests now assert the argument and the new hand-off. - R2-1: the config meter fixture set output_token_count to 0, making tokens.prompt equal tokens.total so the test could not tell a prompt-only meter from a total meter. The fixture now splits the total. - R1-2: of the six turn-start sites that open a meter reading, only two ran under a meter in the tests. Added metered tests for the finishTurn promotion, the releaseTurn promotion, and the post-rejection promoteQueuedUserTurn path; all four sites the review named now fail a suite when their reading is deleted. - R1-5: every meter failure was swallowed with zero logging, making a persistent fault indistinguishable from "no API calls happened". Added one breadcrumb per runtime, covering both the throw and the non-finite reading. Every change above is mutation-verified: reverting it turns at least one named test red. Verified with vitest on client, sessionService, config and goal-runtime (1164 passed) plus the two hook suites (38 passed). --- .../cli/src/ui/hooks/useBranchCommand.test.ts | 16 ++ packages/cli/src/ui/hooks/useBranchCommand.ts | 5 + .../cli/src/ui/hooks/useResumeCommand.test.ts | 22 +++ packages/cli/src/ui/hooks/useResumeCommand.ts | 5 + packages/core/src/config/config.test.ts | 8 +- packages/core/src/config/config.ts | 35 ++++ packages/core/src/core/client.test.ts | 116 +++++++++++++ packages/core/src/core/client.ts | 23 ++- packages/core/src/goals/goal-runtime.test.ts | 152 ++++++++++++++++++ packages/core/src/goals/goal-runtime.ts | 30 +++- packages/core/src/services/sessionService.ts | 21 ++- 11 files changed, 420 insertions(+), 13 deletions(-) diff --git a/packages/cli/src/ui/hooks/useBranchCommand.test.ts b/packages/cli/src/ui/hooks/useBranchCommand.test.ts index 258d13f780c..f8893810250 100644 --- a/packages/cli/src/ui/hooks/useBranchCommand.test.ts +++ b/packages/cli/src/ui/hooks/useBranchCommand.test.ts @@ -32,6 +32,7 @@ describe('useBranchCommand', () => { let finalize: ReturnType; let flush: ReturnType; let startNewSessionConfig: ReturnType; + let markUiTelemetryEventsReplayed: ReturnType; let getGoalRuntimeReady: ReturnType; let startNewSessionUI: ReturnType; let clearPendingState: ReturnType; @@ -108,6 +109,7 @@ describe('useBranchCommand', () => { flush = vi.fn().mockResolvedValue(undefined); findSessionTitlesByPrefix = vi.fn().mockResolvedValue([]); startNewSessionConfig = vi.fn(); + markUiTelemetryEventsReplayed = vi.fn(); getGoalRuntimeReady = vi.fn().mockResolvedValue({}); startNewSessionUI = vi.fn(); clearPendingState = vi.fn(); @@ -152,6 +154,7 @@ describe('useBranchCommand', () => { getBackgroundShellRegistry: () => backgroundShellRegistry, getWorkflowRunRegistry: () => workflowRunRegistry, startNewSession: startNewSessionConfig, + markUiTelemetryEventsReplayed, getGoalRuntimeReady, getDebugLogger: () => ({ warn: vi.fn() }), }; @@ -295,9 +298,22 @@ describe('useBranchCommand', () => { await result.current.handleBranch('my-branch'); }); expect(getGoalRuntimeReady).toHaveBeenCalledTimes(1); + expect(replayUiTelemetryEventsMock).toHaveBeenCalledTimes(1); + // The session id is what keys the inherited history onto the forked + // session. Dropping it sends sessionService down the global reset() + // branch, which clears every live session's bucket and keys nothing + // under the fork, so the new Goal meter reads zero history. + const forkedSessionId = replayUiTelemetryEventsMock.mock.calls[0][1]; + expect(forkedSessionId).toEqual(expect.any(String)); expect( replayUiTelemetryEventsMock.mock.invocationCallOrder[0], ).toBeLessThan(getGoalRuntimeReady.mock.invocationCallOrder[0]!); + // Hand the replay off to the client so initialize() does not replay the + // same history a second time into the process-wide usage aggregate. + expect(markUiTelemetryEventsReplayed).toHaveBeenCalledWith(forkedSessionId); + expect( + replayUiTelemetryEventsMock.mock.invocationCallOrder[0], + ).toBeLessThan(markUiTelemetryEventsReplayed.mock.invocationCallOrder[0]!); }); it('rolls core back when the fork contains malformed Goal state', async () => { diff --git a/packages/cli/src/ui/hooks/useBranchCommand.ts b/packages/cli/src/ui/hooks/useBranchCommand.ts index a70291e03b9..60b9c6d0025 100644 --- a/packages/cli/src/ui/hooks/useBranchCommand.ts +++ b/packages/cli/src/ui/hooks/useBranchCommand.ts @@ -200,10 +200,15 @@ export function useBranchCommand( // the parent, silently recording user input into an orphan. config.startNewSession(newSessionId, resumed); coreSwapped = true; + // Replay before the Goal runtime opens its meter so the meter reads + // the inherited totals, then tell the client it is already done — + // otherwise initialize() replays the forked history a second time and + // the process-wide usage aggregate carries two copies of it. replayUiTelemetryEventsFromConversation( resumed.conversation, newSessionId, ); + config.markUiTelemetryEventsReplayed(newSessionId); await waitForGoalRuntime(config); await config.getGeminiClient()?.initialize?.(SessionStartSource.Branch); diff --git a/packages/cli/src/ui/hooks/useResumeCommand.test.ts b/packages/cli/src/ui/hooks/useResumeCommand.test.ts index 50ebfa99098..e5dd82520f8 100644 --- a/packages/cli/src/ui/hooks/useResumeCommand.test.ts +++ b/packages/cli/src/ui/hooks/useResumeCommand.test.ts @@ -256,6 +256,7 @@ describe('useResumeCommand', () => { getTargetDir: () => '/tmp', getGeminiClient: () => geminiClient, startNewSession: vi.fn(), + markUiTelemetryEventsReplayed: vi.fn(), getGoalRuntimeReady, getBackgroundTaskRegistry: () => ({ hasRunningTasks: vi.fn().mockReturnValue(false), @@ -338,9 +339,26 @@ describe('useResumeCommand', () => { ); expect(resetMonitorRegistry).toHaveBeenCalledTimes(1); expect(getGoalRuntimeReady).toHaveBeenCalledTimes(1); + expect(replayUiTelemetryEventsMock).toHaveBeenCalledTimes(1); + // The session id is what keys the replayed history onto the resumed + // session. Dropping it sends sessionService down the global reset() + // branch, which clears every live session's bucket and keys nothing + // under the resumed id, so the new Goal meter reads zero history. + expect(replayUiTelemetryEventsMock.mock.calls[0][1]).toBe('session-2'); expect( replayUiTelemetryEventsMock.mock.invocationCallOrder[0], ).toBeLessThan(getGoalRuntimeReady.mock.invocationCallOrder[0]!); + // Hand the replay off to the client so initialize() does not replay the + // same history a second time into the process-wide usage aggregate. + expect(config.markUiTelemetryEventsReplayed).toHaveBeenCalledWith( + 'session-2', + ); + expect( + replayUiTelemetryEventsMock.mock.invocationCallOrder[0], + ).toBeLessThan( + vi.mocked(config.markUiTelemetryEventsReplayed).mock + .invocationCallOrder[0]!, + ); }); it('adds a recovery notice when resuming an interrupted tool turn', async () => { @@ -362,6 +380,7 @@ describe('useResumeCommand', () => { getTargetDir: () => '/tmp', getGeminiClient: () => geminiClient, startNewSession: vi.fn(), + markUiTelemetryEventsReplayed: vi.fn(), getGoalRuntimeReady: vi.fn().mockResolvedValue({}), getBackgroundTaskRegistry: () => ({ hasRunningTasks: vi.fn().mockReturnValue(false), @@ -446,6 +465,7 @@ describe('useResumeCommand', () => { getTargetDir: () => '/tmp', getGeminiClient: () => geminiClient, startNewSession: vi.fn(), + markUiTelemetryEventsReplayed: vi.fn(), getGoalRuntimeReady: vi.fn().mockResolvedValue({}), getBackgroundTaskRegistry: () => ({ hasRunningTasks: vi.fn().mockReturnValue(false), @@ -542,6 +562,7 @@ describe('useResumeCommand', () => { getTargetDir: () => '/tmp', getGeminiClient: () => geminiClient, startNewSession: vi.fn(), + markUiTelemetryEventsReplayed: vi.fn(), getGoalRuntimeReady: vi.fn().mockResolvedValue({}), getBackgroundTaskRegistry: () => ({ hasRunningTasks: vi.fn().mockReturnValue(false), @@ -762,6 +783,7 @@ describe('useResumeCommand', () => { getTargetDir: () => '/tmp', getGeminiClient: () => geminiClient, startNewSession: vi.fn(), + markUiTelemetryEventsReplayed: vi.fn(), getGoalRuntimeReady: vi.fn().mockRejectedValue(goalFailure), getBackgroundTaskRegistry: () => ({ hasRunningTasks: vi.fn().mockReturnValue(false), diff --git a/packages/cli/src/ui/hooks/useResumeCommand.ts b/packages/cli/src/ui/hooks/useResumeCommand.ts index b61511601c3..676984007c1 100644 --- a/packages/cli/src/ui/hooks/useResumeCommand.ts +++ b/packages/cli/src/ui/hooks/useResumeCommand.ts @@ -165,10 +165,15 @@ export function useResumeCommand( resetBackgroundStateForSessionSwitch(config); config.startNewSession(sessionId, sessionData); coreSwapped = true; + // Replay before the Goal runtime opens its meter so the meter reads + // the restored totals, then tell the client it is already done — + // otherwise initialize() replays the same history a second time and + // the process-wide usage aggregate carries two copies of it. replayUiTelemetryEventsFromConversation( sessionData.conversation, sessionId, ); + config.markUiTelemetryEventsReplayed(sessionId); await waitForGoalRuntime(config); // Rebuild turn boundary tracking so rewind works within resumed sessions. config diff --git a/packages/core/src/config/config.test.ts b/packages/core/src/config/config.test.ts index 38a6e1aa5c5..3890474a0a2 100644 --- a/packages/core/src/config/config.test.ts +++ b/packages/core/src/config/config.test.ts @@ -2669,8 +2669,12 @@ describe('Server Config (config.ts)', () => { 'event.name': EVENT_API_RESPONSE, model, duration_ms: 1, - input_token_count: totalTokens, - output_token_count: 0, + // Split the total across input and output so prompt !== total. With + // output at 0 the two are equal and the assertion below cannot tell + // a prompt-only meter from a total meter, while real API responses + // always carry output tokens. + input_token_count: totalTokens - 15, + output_token_count: 15, total_token_count: totalTokens, cached_content_token_count: 0, thoughts_token_count: 0, diff --git a/packages/core/src/config/config.ts b/packages/core/src/config/config.ts index 0c186e7c470..859f94bc6b7 100644 --- a/packages/core/src/config/config.ts +++ b/packages/core/src/config/config.ts @@ -1750,6 +1750,17 @@ export class Config { private sessionSourceType?: string; private sessionSourceId?: string; private sessionData?: ResumedSessionData; + /** + * Session id whose stored UI telemetry a session-swap caller has already + * replayed. `/resume` and `/branch` replay before awaiting the Goal runtime + * so the meter opens on the restored totals, which lands the replay ahead of + * `GeminiClient.initialize()`. The marker lets `initialize()` recognize that + * hand-off and skip its own replay: `addEvent` accumulates into both the + * process-wide metrics and the per-session bucket, while the `resetSession` + * leading a replay clears only the bucket, so a second pass would leave one + * extra copy of the whole history in the process-wide aggregate. + */ + private uiTelemetryReplayedSessionId?: string; private pendingSessionRestoreProjection?: SessionRestoreProjection; private sessionRestoreRuntime?: SessionRuntimeResumeState; private readonly sessionRestoreProjectionSource?: () => Promise< @@ -3997,6 +4008,9 @@ export class Config { unregisterSessionModel(previousSessionId); this.publishModelEnv(); this.sessionData = sessionData; + // A swap re-arms the replay: the marker only ever covers the single + // initialize() that follows the caller's own replay for this swap. + this.uiTelemetryReplayedSessionId = undefined; this.clearSessionRestoreProjection(); this.pendingRecoveredAgentsNotice = null; this.getOwnActiveTodoReminders().clear(); @@ -4234,6 +4248,27 @@ export class Config { return this.sessionData; } + /** + * Records that the caller has already replayed the stored UI telemetry for + * `sessionId`, so `GeminiClient.initialize()` must not replay it again. + * The marker is one-shot and is cleared by the next `startNewSession()`, so + * resuming the same session again later still replays exactly once. + */ + markUiTelemetryEventsReplayed(sessionId: string): void { + this.uiTelemetryReplayedSessionId = sessionId; + } + + /** + * Returns whether `sessionId`'s stored UI telemetry was already replayed by + * the session-swap caller, clearing the marker so only the first reader + * skips its replay. + */ + consumeUiTelemetryEventsReplayed(sessionId: string): boolean { + if (this.uiTelemetryReplayedSessionId !== sessionId) return false; + this.uiTelemetryReplayedSessionId = undefined; + return true; + } + shouldLoadMemoryFromIncludeDirectories(): boolean { return this.loadMemoryFromIncludeDirectories; } diff --git a/packages/core/src/core/client.test.ts b/packages/core/src/core/client.test.ts index 1652f9855d8..a47ab1cf03c 100644 --- a/packages/core/src/core/client.test.ts +++ b/packages/core/src/core/client.test.ts @@ -781,6 +781,122 @@ describe('Gemini Client (client.ts)', () => { expect(seedResumeTokenCountsSpy).toHaveBeenCalledWith(321, 45, false); }); + it('does not replay telemetry a second time when the session swap already did', async () => { + // `/resume` and `/branch` replay before awaiting the Goal runtime, so + // the replay lands ahead of initialize(). addEvent accumulates into the + // process-wide aggregate as well as the per-session bucket, and the + // resetSession leading a replay clears only the bucket — so replaying + // again here would leave a second copy of the whole history in the + // aggregate that persistSessionUsage later writes out. + const seedResumeTokenCountsSpy = vi.spyOn( + GeminiChat.prototype, + 'seedResumeTokenCounts', + ); + const uiEvent = { 'event.name': 'api_response' }; + vi.mocked(mockConfig.getResumedSessionData).mockReturnValue({ + conversation: { + sessionId: 'test-session-id', + projectHash: 'project-hash', + startTime: new Date(0).toISOString(), + lastUpdated: new Date(0).toISOString(), + messages: [ + { + uuid: 'telemetry-1', + parentUuid: null, + sessionId: 'test-session-id', + timestamp: new Date(0).toISOString(), + type: 'system', + subtype: 'ui_telemetry', + cwd: '/test/project', + version: '1.0.0', + systemPayload: { uiEvent }, + }, + { + uuid: 'assistant-1', + parentUuid: null, + sessionId: 'test-session-id', + timestamp: new Date(0).toISOString(), + type: 'assistant', + cwd: '/test/project', + version: '1.0.0', + message: { role: 'model', parts: [{ text: 'done' }] }, + usageMetadata: { + promptTokenCount: 200, + candidatesTokenCount: 60, + thoughtsTokenCount: 20, + totalTokenCount: 280, + }, + }, + ], + }, + filePath: '/test/session.jsonl', + lastCompletedUuid: null, + } as unknown as ReturnType); + const consume = vi.fn().mockReturnValue(true); + ( + mockConfig as unknown as { + consumeUiTelemetryEventsReplayed: typeof consume; + } + ).consumeUiTelemetryEventsReplayed = consume; + + const resumedClient = new GeminiClient(mockConfig); + await resumedClient.initialize(); + + expect(consume).toHaveBeenCalledWith('test-session-id'); + expect(uiTelemetryService.addEvent).not.toHaveBeenCalled(); + expect(uiTelemetryService.resetSession).not.toHaveBeenCalled(); + // Skipping the replay must not skip the token-count restore: the chat + // still has to be seeded from the conversation's last response. + expect(seedResumeTokenCountsSpy).toHaveBeenCalledWith(200, 80, false); + expect(uiTelemetryService.setLastPromptTokenCount).toHaveBeenCalledWith( + 200, + ); + }); + + it('replays telemetry itself when the session swap did not', async () => { + const uiEvent = { 'event.name': 'api_response' }; + vi.mocked(mockConfig.getResumedSessionData).mockReturnValue({ + conversation: { + sessionId: 'test-session-id', + projectHash: 'project-hash', + startTime: new Date(0).toISOString(), + lastUpdated: new Date(0).toISOString(), + messages: [ + { + uuid: 'telemetry-1', + parentUuid: null, + sessionId: 'test-session-id', + timestamp: new Date(0).toISOString(), + type: 'system', + subtype: 'ui_telemetry', + cwd: '/test/project', + version: '1.0.0', + systemPayload: { uiEvent }, + }, + ], + }, + filePath: '/test/session.jsonl', + lastCompletedUuid: null, + } as unknown as ReturnType); + const consume = vi.fn().mockReturnValue(false); + ( + mockConfig as unknown as { + consumeUiTelemetryEventsReplayed: typeof consume; + } + ).consumeUiTelemetryEventsReplayed = consume; + + const resumedClient = new GeminiClient(mockConfig); + await resumedClient.initialize(); + + expect(uiTelemetryService.resetSession).toHaveBeenCalledWith( + 'test-session-id', + ); + expect(uiTelemetryService.addEvent).toHaveBeenCalledWith( + uiEvent, + 'test-session-id', + ); + }); + it('seeds resumed chat with replayed prompt token count', async () => { vi.mocked(mockConfig.getResumedSessionData).mockReturnValue({ conversation: { diff --git a/packages/core/src/core/client.ts b/packages/core/src/core/client.ts index d7ee4a84e1b..dc090d33dcf 100644 --- a/packages/core/src/core/client.ts +++ b/packages/core/src/core/client.ts @@ -132,6 +132,7 @@ import type { DeferredToolSummary } from '../tools/tool-registry.js'; import { buildApiHistoryFromConversation, replayUiTelemetryFromConversation, + restoreResumeTokenCountsFromConversation, } from '../services/sessionService.js'; import { reportError } from '../utils/errorReporting.js'; import { @@ -489,10 +490,24 @@ export class GeminiClient { ); } } else if (resumedSessionData) { - const resumeTokenCounts = replayUiTelemetryFromConversation( - resumedSessionData.conversation, - this.config.getSessionId(), - ); + // `/resume` and `/branch` replay the stored telemetry themselves, before + // awaiting the Goal runtime, so the meter opens on the restored totals. + // Replaying it again here would double-count: `addEvent` accumulates + // into both the process-wide metrics and the per-session bucket, but the + // `resetSession` leading each replay clears only the bucket — so the + // second pass leaves an extra copy of the whole history in the aggregate + // that `persistSessionUsage` later writes out. Honour the caller's + // hand-off and restore only the token counts when it already replayed. + const alreadyReplayed = + this.config.consumeUiTelemetryEventsReplayed?.(sessionId) ?? false; + const resumeTokenCounts = alreadyReplayed + ? restoreResumeTokenCountsFromConversation( + resumedSessionData.conversation, + ) + : replayUiTelemetryFromConversation( + resumedSessionData.conversation, + sessionId, + ); // Convert resumed session to API history format // Each ChatRecord's message field is already a Content object const resumedHistory = buildApiHistoryFromConversation( diff --git a/packages/core/src/goals/goal-runtime.test.ts b/packages/core/src/goals/goal-runtime.test.ts index 3c8ddb4c3fd..09809fe622a 100644 --- a/packages/core/src/goals/goal-runtime.test.ts +++ b/packages/core/src/goals/goal-runtime.test.ts @@ -5,6 +5,19 @@ */ import { describe, expect, it, vi } from 'vitest'; + +const debugWarn = vi.hoisted(() => vi.fn()); +vi.mock('../utils/debugLogger.js', async (importOriginal) => { + const actual = + await importOriginal(); + return { + ...actual, + createDebugLogger: (tag?: string) => ({ + ...actual.createDebugLogger(tag), + warn: debugWarn, + }), + }; +}); import type { GoalEvidenceRecord } from './goal-evidence.js'; import type { GoalRecoveryRecord } from './goal-persistence.js'; import { @@ -335,6 +348,41 @@ describe('goal runtime', () => { }); }); + it.each([ + [ + 'throws', + () => { + throw new Error('metrics are unavailable'); + }, + 'metrics are unavailable', + ], + ['reads back NaN', () => Number.NaN, 'non-finite'], + ])( + 'leaves one breadcrumb per runtime when the token meter %s', + async (_label, readSessionTokens, expectedDetail) => { + // Failing closed to zero is the policy, but a silent failure is + // indistinguishable from "no API calls happened" when a user reports a + // Goal that says 0 tokens. The meter is read at every begin/finish, so + // the breadcrumb has to be capped at one per runtime. + debugWarn.mockClear(); + const journal = fakeGoalJournal(); + const host = fakeGoalTurnHost(); + const runtime = createGoalRuntime({ + journal, + tokenMeter: { readSessionTokens }, + }); + runtime.bindHost(host); + await runtime.dispatch({ action: 'create', objective: 'ship' }); + await runtime.finishTurn(host.started[0]!); + const permit = runtime.beginTurn('turn-2'); + if (permit) await runtime.finishTurn(permit); + + expect(debugWarn).toHaveBeenCalledTimes(1); + expect(debugWarn.mock.calls[0][0]).toContain(expectedDetail); + expect(debugWarn.mock.calls[0][0]).toContain('tokensUsed will report 0'); + }, + ); + it.each([Number.NaN, Number.POSITIVE_INFINITY])( 'finishes the turn when the token meter returns %s', async (invalidTotal) => { @@ -2031,6 +2079,54 @@ describe('goal runtime', () => { }, ); + it('bills the meter delta for user input promoted after a rejection', async () => { + // promoteQueuedUserTurn opens the promoted turn's meter reading on its own + // line, and it is the only turn-start site reachable through verification + // rejection. Without a metered test here, deleting that reading bills the + // promoted turn nothing while the whole suite stays green. + const result = deferred>>(); + const journal = fakeGoalJournal(); + let records: readonly RuntimeRecord[] = []; + const evidenceSource = fakeEvidenceSource(() => records); + const verifier: GoalVerifier = vi.fn(() => result.promise); + const host = fakeGoalTurnHost(); + let sessionTokens = 0; + const runtime = createGoalRuntime({ + journal, + evidenceSource, + verifier, + tokenMeter: { readSessionTokens: () => sessionTokens }, + }); + runtime.bindHost(host); + await runtime.dispatch({ action: 'create', objective: 'deliver result' }); + const permit = host.started[0]; + records = verifierEvidenceRecords( + permit, + runtime.getSnapshot().goal!.evidenceCursor.recordId!, + ); + runtime.recordTerminalProposal(permit, { + status: 'complete', + reason: 'Delivered', + evidenceRefs: ['assistant-evidence'], + }); + sessionTokens = 40; + const finishing = runtime.finishTurn(permit); + await vi.waitFor(() => expect(verifier).toHaveBeenCalledOnce()); + expect(runtime.beginTurn('real-user')).toBeUndefined(); + + result.resolve({ decision: 'reject', reason: 'Add the missing example' }); + await finishing; + expect(runtime.getSnapshot().goal?.tokensUsed).toBe(40); + + const userPermit = runtime.permitForTurn('real-user')!; + expect(userPermit).toBeDefined(); + sessionTokens = 95; + await runtime.finishTurn(userPermit); + + // 40 from the verified turn plus the promoted turn's own 55. + expect(runtime.getSnapshot().goal?.tokensUsed).toBe(95); + }); + it('promotes queued user input with exact verifier feedback after rejection', async () => { const result = deferred>>(); const journal = fakeGoalJournal(); @@ -2147,6 +2243,62 @@ describe('goal runtime', () => { }); }); + it('bills the meter delta for a reservation promoted by finishTurn', async () => { + // The promotion opens the new turn's meter reading on its own line. Only + // the continuation and beginTurn admit paths were metered under test, so + // deleting the reading here left takeTurnTokens with `opened === undefined` + // — the promoted turn bills nothing and every suite stays green. The + // realistic trigger is a user turn queued while a Goal turn is running. + const host = fakeGoalTurnHost(); + let sessionTokens = 0; + const runtime = createGoalRuntime({ + journal: fakeGoalJournal(), + tokenMeter: { readSessionTokens: () => sessionTokens }, + }); + runtime.bindHost(host); + await runtime.dispatch({ action: 'create', objective: 'ship' }); + const initialPermit = host.started[0]; + + expect(runtime.beginTurn('queued-user')).toBeUndefined(); + sessionTokens = 40; + await runtime.finishTurn(initialPermit); + expect(runtime.getSnapshot().goal?.tokensUsed).toBe(40); + + const promoted = runtime.permitForTurn('queued-user'); + expect(promoted).toBeDefined(); + sessionTokens = 65; + await runtime.finishTurn(promoted!); + + // 40 from the first turn plus the promoted turn's own 25 — not 40, which + // is what an unopened reading would leave behind. + expect(runtime.getSnapshot().goal?.tokensUsed).toBe(65); + }); + + it('bills the meter delta for a reservation promoted by releaseTurn', async () => { + const host = fakeGoalTurnHost(); + let sessionTokens = 0; + const runtime = createGoalRuntime({ + journal: fakeGoalJournal(), + tokenMeter: { readSessionTokens: () => sessionTokens }, + }); + runtime.bindHost(host); + await runtime.dispatch({ action: 'create', objective: 'ship' }); + const initialPermit = host.started[0]; + + expect(runtime.beginTurn('queued-user')).toBeUndefined(); + sessionTokens = 30; + await expect( + runtime.releaseTurn(`goal-runtime:${initialPermit.turnId}`), + ).resolves.toBe(true); + + const promoted = runtime.permitForTurn('queued-user'); + expect(promoted).toBeDefined(); + sessionTokens = 100; + await runtime.finishTurn(promoted!); + + expect(runtime.getSnapshot().goal?.tokensUsed).toBe(70); + }); + it('promotes a waiting reservation when the current turn is released', async () => { // The host drains continuations one at a time and the caller holding // `queued-user` is what blocks that drain, so minting a fresh diff --git a/packages/core/src/goals/goal-runtime.ts b/packages/core/src/goals/goal-runtime.ts index 73546b23f79..6397ec96d9d 100644 --- a/packages/core/src/goals/goal-runtime.ts +++ b/packages/core/src/goals/goal-runtime.ts @@ -52,6 +52,9 @@ import { recoverGoalFromRecords, type GoalRecoveryRecord, } from './goal-persistence.js'; +import { createDebugLogger } from '../utils/debugLogger.js'; + +const debugLogger = createDebugLogger('GOAL_RUNTIME'); export const GOAL_RUNTIME_DISPOSED_MESSAGE = 'Goal runtime has been disposed'; export const STALE_GOAL_TURN_MESSAGE = 'Goal turn permit is no longer valid'; @@ -270,14 +273,37 @@ export function createGoalRuntime( } : undefined; + // Failing closed to zero is the policy, but doing it silently makes a + // persistent meter fault ("my Goal says 0 tokens but I was billed for + // millions") indistinguishable from "no API calls happened" — there is no + // log line, telemetry event, or error to grep for. One breadcrumb per + // runtime is enough to tell the two apart without flooding: the meter is + // read at every beginTurn/finishTurn. + let meterFailureReported = false; + const reportMeterFailure = (detail: string): void => { + if (meterFailureReported) return; + meterFailureReported = true; + debugLogger.warn( + `Goal token meter unreadable (${detail}); tokensUsed will report 0 ` + + `for this runtime until the meter recovers. Reported once per runtime.`, + ); + }; + const readSessionTokens = (): number | undefined => { if (!options.tokenMeter) return undefined; try { const total = options.tokenMeter.readSessionTokens(); - return Number.isFinite(total) ? total : undefined; - } catch { + if (Number.isFinite(total)) return total; + // A non-finite reading is the shape that drift produces: rename + // `tokens.total` and the config meter's reduce yields NaN. + reportMeterFailure(`non-finite reading ${String(total)}`); + return undefined; + } catch (error) { // Goal accounting is bookkeeping. A meter that cannot answer costs the // Goal its spend figure for this turn, not the turn. + reportMeterFailure( + `threw ${error instanceof Error ? error.message : String(error)}`, + ); return undefined; } }; diff --git a/packages/core/src/services/sessionService.ts b/packages/core/src/services/sessionService.ts index 362c36359bf..82d38f40db2 100644 --- a/packages/core/src/services/sessionService.ts +++ b/packages/core/src/services/sessionService.ts @@ -2644,14 +2644,14 @@ export function replayUiTelemetryEventsFromConversation( } /** - * Replays stored UI telemetry and restores the last prompt token count. + * Restores the last prompt token count from a conversation without replaying + * its telemetry events. Use this when the events have already been replayed by + * the session-swap caller — replaying them twice adds a second copy of the + * whole history to the process-wide usage aggregate. */ -export function replayUiTelemetryFromConversation( +export function restoreResumeTokenCountsFromConversation( conversation: ConversationRecord, - sessionId?: string, ): ResumeTokenCounts | undefined { - replayUiTelemetryEventsFromConversation(conversation, sessionId); - const resumeTokenCounts = getResumeTokenCounts(conversation); if (resumeTokenCounts !== undefined) { uiTelemetryService.setLastPromptTokenCount( @@ -2661,6 +2661,17 @@ export function replayUiTelemetryFromConversation( return resumeTokenCounts; } +/** + * Replays stored UI telemetry and restores the last prompt token count. + */ +export function replayUiTelemetryFromConversation( + conversation: ConversationRecord, + sessionId?: string, +): ResumeTokenCounts | undefined { + replayUiTelemetryEventsFromConversation(conversation, sessionId); + return restoreResumeTokenCountsFromConversation(conversation); +} + const MAX_BRANCH_COLLISION_SCAN = 99; export async function computeUniqueBranchTitle( From 981e6cab1c304116f767c7c90365e90e43307e2c Mon Sep 17 00:00:00 2001 From: qqqys Date: Tue, 18 Aug 2026 23:44:16 +0800 Subject: [PATCH 06/10] fix(goal): keep the resume replay out of the usage aggregate it can corrupt MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Round-4 review found four Critical defects, all in the window the replay this PR added opens between a session swap starting and the swap either committing or failing. R4-1 — the replay runs before the fallible steps (`waitForGoalRuntime`, `initialize()`) and the rollback never undid it, so a failed `/resume` or `/branch` leaked one full copy of the abandoned session's history into the process-wide aggregate for the life of the process, and `persistSessionUsage` wrote the inflated figure out. `UiTelemetryService` had no compensation: `resetSession` clears one bucket and the global `reset()` would take the surviving session's live data with it. Adds `snapshotForReplay` / `restoreFromReplaySnapshot` — a narrow undo that touches only the snapshotted session — and hands the snapshot back on both hooks' rollback paths, before the rollback's own re-init runs. R4-2, R4-3 — resuming the session that is already current replayed it into itself. Nothing upstream excluded it, and there the hand-off marker is never consumed because `GeminiClient.initialize()` early-returns for an already-initialized session id: the aggregate carried the session twice, and the `resetSession` leading the replay first wiped the live bucket and refilled it from a snapshot loaded before the recorder's queued writes landed — dropping not-yet-flushed events and the internal-prompt side-query tokens that `recordUiTelemetryEventToChat` never persists at all. `/resume` now skips both the replay and the mark when `sessionId === oldSessionId`, which is what the base tree did (it replayed nothing on that path). R4-4 — on TUI `--resume` of a session interrupted with an active Goal, the restore is kicked off from the Config constructor (or, under a session-writer lease, from `activateChatRecording()`), both strictly before `initializeInternal` reaches `geminiClient.initialize()`. The continuation permit was minted — and its opening meter reading latched — against an empty bucket, so the first restored turn billed the whole replayed history: a Goal resumed from a 200k-token session billed ~200k extra, persisted into `GoalRecord.tokensUsed`. Both cold entrances now use the same prepareRestore/activateRestoredWork split the session-restore-projection path already used, with activation released once the client-init step has run. `restoreActivationPending` already holds `queueContinuation` until then. Behavior flip: a cold-start restore no longer mints its permit at construction time. `getGoalRuntimeReady()` deliberately still resolves on preparation, not activation — the Goal state is recovered and no awaiter is held — so a resumed Config that never calls `initialize()` cannot hang; only the autonomous continuation waits. `config.test.ts`'s "rebinds the current Goal host to every replacement runtime" pinned the old ordering and now drives `initialize()`, asserting no turn starts before it. The old ordering was the defect. Also fixes R3-3: the branch test's sessionId assertion was self-referential (`expect.any(String)` on the mock's own argument), so keying both calls onto the parent id passed all 26 tests while double-counting in production. It now pins the forked id, exactly the reviewer's witness. Every fix is mutation-verified: reverting each one turns at least one test red (3 tests for the snapshot API, 1 each for the same-session guard, the two rollback paths, the cold-start deferral, and the sessionId pin). Deferred D4-4 is not fixed, but the test added here resets the shared meter mock it introduces rather than widening the leak. --- .../cli/src/ui/hooks/useBranchCommand.test.ts | 33 ++++- packages/cli/src/ui/hooks/useBranchCommand.ts | 30 +++++ .../cli/src/ui/hooks/useResumeCommand.test.ts | 105 ++++++++++++++++ packages/cli/src/ui/hooks/useResumeCommand.ts | 58 ++++++++- packages/core/src/config/config.test.ts | 74 +++++++++++ packages/core/src/config/config.ts | 118 +++++++++++++++++- .../core/src/telemetry/uiTelemetry.test.ts | 55 ++++++++ packages/core/src/telemetry/uiTelemetry.ts | 79 ++++++++++++ 8 files changed, 543 insertions(+), 9 deletions(-) diff --git a/packages/cli/src/ui/hooks/useBranchCommand.test.ts b/packages/cli/src/ui/hooks/useBranchCommand.test.ts index f8893810250..dd0d2067854 100644 --- a/packages/cli/src/ui/hooks/useBranchCommand.test.ts +++ b/packages/cli/src/ui/hooks/useBranchCommand.test.ts @@ -8,6 +8,7 @@ import { describe, it, expect, vi, beforeEach } from 'vitest'; import { renderHook, act } from '@testing-library/react'; import { useBranchCommand } from './useBranchCommand.js'; import type { LoadedSettings } from '../../config/settings.js'; +import { uiTelemetryService } from '@qwen-code/qwen-code-core'; const replayUiTelemetryEventsMock = vi.hoisted(() => vi.fn()); @@ -304,7 +305,13 @@ describe('useBranchCommand', () => { // branch, which clears every live session's bucket and keys nothing // under the fork, so the new Goal meter reads zero history. const forkedSessionId = replayUiTelemetryEventsMock.mock.calls[0][1]; - expect(forkedSessionId).toEqual(expect.any(String)); + // Pin the exact id, not `expect.any(String)`: keying both calls onto a + // wrong-but-consistent id (the parent's, captured before the swap) would + // satisfy a self-referential assertion while double-counting in + // production — `consumeUiTelemetryEventsReplayed(newSessionId)` compares + // for exact equality, misses a parent-keyed marker, and lets initialize() + // replay the forked history a second time. + expect(forkedSessionId).toBe(startNewSessionConfig.mock.calls[0][0]); expect( replayUiTelemetryEventsMock.mock.invocationCallOrder[0], ).toBeLessThan(getGoalRuntimeReady.mock.invocationCallOrder[0]!); @@ -623,6 +630,12 @@ describe('useBranchCommand', () => { .mockRejectedValueOnce(new Error('init boom')) // fork init fails .mockResolvedValueOnce(undefined); // rollback re-init succeeds config.getGeminiClient = () => ({ initialize }); + replayUiTelemetryEventsMock.mockClear(); + const snapshotForReplay = vi.spyOn(uiTelemetryService, 'snapshotForReplay'); + const restoreFromReplaySnapshot = vi.spyOn( + uiTelemetryService, + 'restoreFromReplaySnapshot', + ); const { result } = renderHook(() => useBranchCommand(makeOptions())); await act(async () => { @@ -658,6 +671,24 @@ describe('useBranchCommand', () => { }), expect.any(Number), ); + + // The forked history was replayed into the process-wide aggregate before + // the step that failed. `resetSession` clears only the fork's bucket and + // the global `reset()` would wipe the parent's live data, so the rollback + // has to hand back the snapshot taken before the replay — otherwise the + // abandoned fork's tokens are billed to the process for good. + expect(snapshotForReplay.mock.invocationCallOrder[0]).toBeLessThan( + replayUiTelemetryEventsMock.mock.invocationCallOrder[0], + ); + expect(restoreFromReplaySnapshot).toHaveBeenCalledWith( + snapshotForReplay.mock.results[0]!.value, + ); + // Undone before the rollback re-init runs its own replay. + expect(restoreFromReplaySnapshot.mock.invocationCallOrder[0]).toBeLessThan( + startNewSessionConfig.mock.invocationCallOrder[1], + ); + snapshotForReplay.mockRestore(); + restoreFromReplaySnapshot.mockRestore(); }); it('still surfaces the error and leaves core on the parent when rollback re-init also throws', async () => { diff --git a/packages/cli/src/ui/hooks/useBranchCommand.ts b/packages/cli/src/ui/hooks/useBranchCommand.ts index 60b9c6d0025..d793fa0fe4b 100644 --- a/packages/cli/src/ui/hooks/useBranchCommand.ts +++ b/packages/cli/src/ui/hooks/useBranchCommand.ts @@ -13,6 +13,8 @@ import { SessionStartSource, computeUniqueBranchTitle, replayUiTelemetryEventsFromConversation, + uiTelemetryService, + type UiTelemetryReplaySnapshot, } from '@qwen-code/qwen-code-core'; import { buildResumedHistoryItems, @@ -131,6 +133,11 @@ export function useBranchCommand( let uiSwapped = false; let forkCreated = false; let prevSessionData: ResumedSessionData | undefined; + // Set only while a replay for this swap is outstanding, so the rollback + // below can undo it. The service has no subtraction API, so without + // this an abandoned branch leaves the forked history in the + // process-wide aggregate for good. + let telemetryReplaySnapshot: UiTelemetryReplaySnapshot | undefined; try { // 1. Flush outgoing recorder. A degraded source must not fork because @@ -204,6 +211,8 @@ export function useBranchCommand( // the inherited totals, then tell the client it is already done — // otherwise initialize() replays the forked history a second time and // the process-wide usage aggregate carries two copies of it. + telemetryReplaySnapshot = + uiTelemetryService.snapshotForReplay(newSessionId); replayUiTelemetryEventsFromConversation( resumed.conversation, newSessionId, @@ -278,6 +287,27 @@ export function useBranchCommand( // split-brain (UI on branch, recorder on parent). Post-UI-swap // failures (hook, remount, announce) are non-fatal and // surfaced as an error item without unwinding the swap. + // + // Undo the replay first, before the rollback re-init below runs + // its own. Rolling core back does not touch the usage aggregate — + // `resetSession` clears only the abandoned fork's bucket and the + // global `reset()` would wipe the parent's live data — so without + // this the process carries a full extra copy of the forked history + // until it exits, and `persistSessionUsage` writes it out. + if (telemetryReplaySnapshot) { + try { + uiTelemetryService.restoreFromReplaySnapshot( + telemetryReplaySnapshot, + ); + } catch (restoreErr) { + config + .getDebugLogger() + .warn( + `Telemetry rollback after failed /branch init failed: ${restoreErr}`, + ); + } + telemetryReplaySnapshot = undefined; + } try { config.startNewSession(oldSessionId, prevSessionData); // Re-hydrate chat history against the restored session. Best- diff --git a/packages/cli/src/ui/hooks/useResumeCommand.test.ts b/packages/cli/src/ui/hooks/useResumeCommand.test.ts index e5dd82520f8..0a4e4a9c35a 100644 --- a/packages/cli/src/ui/hooks/useResumeCommand.test.ts +++ b/packages/cli/src/ui/hooks/useResumeCommand.test.ts @@ -11,6 +11,7 @@ import { useResumeCommand, } from './useResumeCommand.js'; import { useHistory } from './useHistoryManager.js'; +import { uiTelemetryService } from '@qwen-code/qwen-code-core'; import type { Content } from '@google/genai'; import type { LoadedSettings } from '../../config/settings.js'; @@ -772,6 +773,12 @@ describe('useResumeCommand', () => { }); it('rolls core back when persisted Goal state is malformed', async () => { + replayUiTelemetryEventsMock.mockClear(); + const snapshotForReplay = vi.spyOn(uiTelemetryService, 'snapshotForReplay'); + const restoreFromReplaySnapshot = vi.spyOn( + uiTelemetryService, + 'restoreFromReplaySnapshot', + ); const startNewSession = vi.fn(); const geminiClient = { initialize: vi.fn().mockResolvedValue(undefined), @@ -860,5 +867,103 @@ describe('useResumeCommand', () => { expect.any(Number), ); expect(geminiClient.initialize).not.toHaveBeenCalled(); + + // The replay ran before the step that failed, and the service has no + // subtraction API — so the rollback has to hand back the snapshot taken + // before it. Without this the process-wide aggregate keeps a full copy of + // the abandoned session's history and `persistSessionUsage` writes it out. + expect(snapshotForReplay).toHaveBeenCalledWith('new-session-id'); + expect(snapshotForReplay.mock.invocationCallOrder[0]).toBeLessThan( + replayUiTelemetryEventsMock.mock.invocationCallOrder[0], + ); + expect(restoreFromReplaySnapshot).toHaveBeenCalledWith( + snapshotForReplay.mock.results[0]!.value, + ); + // Undone before core is put back, so the rollback's own bookkeeping is + // the last word. + expect(restoreFromReplaySnapshot.mock.invocationCallOrder[0]).toBeLessThan( + vi.mocked(config.startNewSession).mock.invocationCallOrder[1], + ); + snapshotForReplay.mockRestore(); + restoreFromReplaySnapshot.mockRestore(); + }); + + it('does not replay telemetry when resuming the session already current', async () => { + replayUiTelemetryEventsMock.mockClear(); + resumeMocks.reset(); + + const startNewSession = vi.fn(); + const geminiClient = { + initialize: vi.fn().mockResolvedValue(undefined), + }; + + const config = { + getSessionId: () => 'current-session-id', + getTargetDir: () => '/tmp', + getGeminiClient: () => geminiClient, + startNewSession: vi.fn(), + markUiTelemetryEventsReplayed: vi.fn(), + getGoalRuntimeReady: vi.fn().mockResolvedValue({}), + getBackgroundTaskRegistry: () => ({ + hasRunningTasks: vi.fn().mockReturnValue(false), + reset: vi.fn(), + }), + getBackgroundShellRegistry: () => ({ + getAll: vi.fn().mockReturnValue([]), + hasRunningEntries: vi.fn().mockReturnValue(false), + reset: vi.fn(), + }), + getMonitorRegistry: () => ({ + getRunning: vi.fn().mockReturnValue([]), + reset: vi.fn(), + }), + getWorkflowRunRegistry: () => ({ + hasRunningEntries: vi.fn().mockReturnValue(false), + reset: vi.fn(), + abortAll: vi.fn(), + }), + loadPausedBackgroundAgents: vi.fn().mockResolvedValue([]), + getChatRecordingService: () => ({ rebuildTurnBoundaries: vi.fn() }), + getDebugLogger: () => ({ + warn: vi.fn(), + debug: vi.fn(), + error: vi.fn(), + }), + } as unknown as import('@qwen-code/qwen-code-core').Config; + + const historyManager = { + addItem: vi.fn(), + clearItems: vi.fn(), + loadHistory: vi.fn(), + }; + + const { result } = renderHook(() => + useResumeCommand({ + config, + settings: mockSettings, + historyManager, + startNewSession, + }), + ); + + await act(async () => { + await result.current.handleResume('current-session-id'); + }); + + // Same-id resume is supported and the live per-session bucket is already + // the authority. Replaying would (a) add a second full copy of the + // history to the process-wide aggregate that nothing consumes the + // hand-off marker for — `initialize()` early-returns for an + // already-initialized session — and (b) wipe the live bucket and refill + // it from a snapshot that predates the recorder's queued writes and + // never contains internal-prompt side-query tokens at all. + expect(replayUiTelemetryEventsMock).not.toHaveBeenCalled(); + expect(config.markUiTelemetryEventsReplayed).not.toHaveBeenCalled(); + // The resume itself still runs. + expect(config.startNewSession).toHaveBeenCalledWith( + 'current-session-id', + expect.any(Object), + ); + expect(startNewSession).toHaveBeenCalledWith('current-session-id'); }); }); diff --git a/packages/cli/src/ui/hooks/useResumeCommand.ts b/packages/cli/src/ui/hooks/useResumeCommand.ts index 676984007c1..a070be5cd91 100644 --- a/packages/cli/src/ui/hooks/useResumeCommand.ts +++ b/packages/cli/src/ui/hooks/useResumeCommand.ts @@ -9,7 +9,9 @@ import { SessionService, buildSessionRecoveryPlan, replayUiTelemetryEventsFromConversation, + uiTelemetryService, type Config, + type UiTelemetryReplaySnapshot, type SessionListItem, } from '@qwen-code/qwen-code-core'; import { @@ -115,6 +117,11 @@ export function useResumeCommand( let coreSwapped = false; let uiSwapped = false; let recoveredBackgroundAgentsNotice: string | null = null; + // Set only while a replay for this swap is outstanding, so the rollback + // below can undo it. The service has no subtraction API, so without + // this an abandoned resume leaves the incoming session's whole history + // in the process-wide aggregate for good. + let telemetryReplaySnapshot: UiTelemetryReplaySnapshot | undefined; try { const cwd = config.getTargetDir(); @@ -169,11 +176,32 @@ export function useResumeCommand( // the restored totals, then tell the client it is already done — // otherwise initialize() replays the same history a second time and // the process-wide usage aggregate carries two copies of it. - replayUiTelemetryEventsFromConversation( - sessionData.conversation, - sessionId, - ); - config.markUiTelemetryEventsReplayed(sessionId); + // + // Resuming the session that is already current is the one case that + // must not replay. Nothing upstream excludes it (`/resume` lists the + // current session, and `startNewSession` documents same-id resume), + // and there the live per-session bucket is already the authority: + // - replaying adds a second full copy of the history to the + // process-wide aggregate, and the hand-off marker is never + // consumed because `initialize()` early-returns for an + // already-initialized session id — so `persistSessionUsage` + // writes the session out at roughly double its real usage; + // - the `resetSession` inside the replay first wipes the live + // bucket and refills it from a snapshot loaded before the + // recorder's queued writes landed, dropping both not-yet-flushed + // events and internal-prompt side-query tokens, which + // `recordUiTelemetryEventToChat` never persists at all. + // Before this feature the same-session path replayed nothing: + // `initialize()` early-returned. Keep that. + if (sessionId !== oldSessionId) { + telemetryReplaySnapshot = + uiTelemetryService.snapshotForReplay(sessionId); + replayUiTelemetryEventsFromConversation( + sessionData.conversation, + sessionId, + ); + config.markUiTelemetryEventsReplayed(sessionId); + } await waitForGoalRuntime(config); // Rebuild turn boundary tracking so rewind works within resumed sessions. config @@ -214,6 +242,26 @@ export function useResumeCommand( remount?.(); } catch (error) { if (coreSwapped && !uiSwapped) { + // Undo the replay first. Rolling core back does not touch the + // usage aggregate — `resetSession` clears only the abandoned + // session's bucket and the global `reset()` would wipe the + // surviving session's live data — so without this the process + // carries a full extra copy of the abandoned session's history + // until it exits, and `persistSessionUsage` writes it out. + if (telemetryReplaySnapshot) { + try { + uiTelemetryService.restoreFromReplaySnapshot( + telemetryReplaySnapshot, + ); + } catch (restoreErr) { + config + .getDebugLogger() + .warn( + `Telemetry rollback after failed /resume init failed: ${restoreErr}`, + ); + } + telemetryReplaySnapshot = undefined; + } // Core switched to the resumed session but UI hasn't swapped // yet — put core back on the old session, otherwise the // recorder would keep writing new user messages into the diff --git a/packages/core/src/config/config.test.ts b/packages/core/src/config/config.test.ts index 3890474a0a2..41839d8e8d0 100644 --- a/packages/core/src/config/config.test.ts +++ b/packages/core/src/config/config.test.ts @@ -105,6 +105,7 @@ import { GoalPersistenceUnavailableError, type GoalTurnHost, } from '../goals/goal-runtime.js'; +import type { GoalTurnPermit } from '../goals/goal-protocol.js'; import { getSessionWriterLockPath, SessionTranscriptChangedError, @@ -2697,6 +2698,73 @@ describe('Server Config (config.ts)', () => { ).toBe(true); }); + it('opens a cold-resumed Goal meter on the replayed totals, not an empty bucket', async () => { + const sessionId = 'cold-resume-session'; + const telemetry = new UiTelemetryService(); + vi.mocked(uiTelemetryService.getMetricsForSession).mockImplementation( + (requestedSessionId) => + telemetry.getMetricsForSession(requestedSessionId), + ); + const apiResponse = ( + totalTokens: number, + ): ApiResponseEvent & { 'event.name': typeof EVENT_API_RESPONSE } => + ({ + 'event.name': EVENT_API_RESPONSE, + model: 'model-a', + duration_ms: 1, + input_token_count: totalTokens - 15, + output_token_count: 15, + total_token_count: totalTokens, + cached_content_token_count: 0, + thoughts_token_count: 0, + }) as ApiResponseEvent & { + 'event.name': typeof EVENT_API_RESPONSE; + }; + + const config = new Config({ + ...baseParams, + chatRecording: true, + sessionId, + sessionData: resumedGoalSession('active'), + }); + const recorder = config.getChatRecordingService(); + if (!recorder) throw new Error('expected a chat recording service'); + vi.spyOn(recorder, 'recordGoalState').mockResolvedValue({} as ChatRecord); + + let permit: GoalTurnPermit | undefined; + config.bindGoalTurnHost({ + startGoalTurn: vi.fn(async (input) => { + permit = input.permit; + }), + preemptGoalTurn: vi.fn(), + }); + const runtime = await config.getGoalRuntimeReady(); + + // The constructor kicked the restore off, but the client has not + // replayed the resumed session's telemetry into the bucket yet. A + // permit minted now would open its meter on 0. + expect(permit).toBeUndefined(); + + // What `GeminiClient.initialize()` does for a resumed session, deep + // inside `initializeInternal`. + telemetry.addEvent(apiResponse(200_000), sessionId); + + await config.initialize(); + await vi.waitFor(() => expect(permit).toBeDefined()); + + // The restored turn's own spend. + telemetry.addEvent(apiResponse(500), sessionId); + await runtime.finishTurn(permit!); + + // 500, not 200_500: the replayed history belongs to the session that + // was interrupted, not to the first turn after the resume. + expect(runtime.getSnapshot().goal?.tokensUsed).toBe(500); + + // `clearAllMocks` between tests clears call history, not + // implementations — leave the shared meter mock as this file found it. + vi.mocked(uiTelemetryService.getMetricsForSession).mockReset(); + }); + it('rebinds the current Goal host to every replacement runtime', async () => { const config = new Config({ ...baseParams, @@ -2713,6 +2781,12 @@ describe('Server Config (config.ts)', () => { config.bindGoalTurnHost(host); await config.getGoalRuntimeReady(); + // A cold-start restore does not mint its permit until initialize() has + // let the client replay the resumed session's telemetry — otherwise the + // meter opens on an empty bucket and the first restored turn bills the + // whole replayed history. + expect(started).toEqual([]); + await config.initialize(); await vi.waitFor(() => expect(started).toEqual(['g-resumed'])); config.startNewSession( diff --git a/packages/core/src/config/config.ts b/packages/core/src/config/config.ts index 859f94bc6b7..e90d3a05046 100644 --- a/packages/core/src/config/config.ts +++ b/packages/core/src/config/config.ts @@ -1769,6 +1769,20 @@ export class Config { private restoredFileHistory = false; private goalRestoreActivation?: () => Promise; private rejectGoalRestoreActivation?: (reason?: unknown) => void; + /** + * A cold-start Goal restore held back until the client has replayed the + * resumed session's telemetry into the session bucket. Fired by + * {@link startDeferredGoalRestore}; see + * {@link shouldDeferGoalRestoreForTelemetryReplay} for why. + */ + private deferredGoalRestoreActivation?: () => void; + /** + * Set once `initializeInternal` has passed the client-init step, which is + * where a resumed session's stored telemetry is replayed into its bucket. + * Before that point the bucket is empty, so a Goal permit minted then would + * open its meter on 0 and bill the whole replayed history to its first turn. + */ + private clientSessionTelemetryReplayed = false; private readonly sessionRuntimeBaseDir: string; private sessionProjectDirRegistered = false; private pendingSessionWriterLease?: SessionWriterLease; @@ -3080,6 +3094,13 @@ export class Config { } else { this.debugLogger.info('Gemini client initialization skipped'); } + // The resumed session's telemetry is in its bucket now (or was never + // going to be, when client init is skipped). A cold-start Goal restore + // held back for it can activate: its first turn opens the meter on the + // restored totals instead of on an empty bucket. Unconditional so the + // skip path cannot strand `getGoalRuntimeReady()`. + this.clientSessionTelemetryReplayed = true; + this.startDeferredGoalRestore(); // Detect and capture runtime model snapshot (from CLI/ENV/credentials) this.modelsConfig.detectAndCaptureRuntimeModel(); @@ -5435,6 +5456,7 @@ export class Config { ); this.goalRestoreActivation = undefined; this.rejectGoalRestoreActivation = undefined; + this.deferredGoalRestoreActivation = undefined; this.goalTurnHostUnbind?.(); this.goalTurnHostUnbind = undefined; // Shutting down before the writer arrived: nothing will ever run @@ -7880,6 +7902,7 @@ export class Config { ); this.goalRestoreActivation = undefined; this.rejectGoalRestoreActivation = undefined; + this.deferredGoalRestoreActivation = undefined; this.goalTurnHostUnbind?.(); this.goalTurnHostUnbind = undefined; // A runtime built here supersedes any restore still waiting on the @@ -7950,6 +7973,11 @@ export class Config { this.pendingGoalRestore = { runtime, resolve, reject }; }); this.goalRuntimeReady = ready; + } else if (this.shouldDeferGoalRestoreForTelemetryReplay()) { + this.goalRuntimeReady = this.deferGoalRestoreActivation( + runtime, + records ?? [], + ); } else { this.goalRuntimeReady = runtime .restore(records ?? []) @@ -7958,6 +7986,78 @@ export class Config { void this.goalRuntimeReady.catch(() => undefined); } + /** + * Whether a Goal restore starting now must wait for the client's telemetry + * replay before it is allowed to mint a continuation permit. + * + * On a cold `--resume` start the restore is kicked off from the Config + * constructor (or, under a session-writer lease, from + * `activateChatRecording()`), both strictly before `initializeInternal` + * reaches `geminiClient.initialize()` — which is where the stored telemetry + * is replayed into the session bucket. A permit minted in that window + * samples `readSessionTokens()` on an empty bucket, so when the first + * restored turn closes, the delta it bills is the entire replayed history + * plus the turn's own spend: a Goal resumed from a 200k-token session bills + * ~200k extra on its first turn, and that figure is persisted into + * `GoalRecord.tokensUsed`. + * + * The `/resume` and `/branch` hooks, the headless path and the ACP path all + * replay before they let the runtime activate; this closes the same gap on + * the cold-start entrance. It only applies to a resumed session — a fresh + * session has nothing to replay, and a swap-time restore (`startNewSession`) + * runs after the caller already replayed. + */ + private shouldDeferGoalRestoreForTelemetryReplay(): boolean { + return ( + !this.clientSessionTelemetryReplayed && this.sessionData !== undefined + ); + } + + /** + * Prepares the restore now — recovering the Goal state so + * `getGoalRuntimeReady()` waiters see it — while holding back the + * activation that lets the runtime queue a continuation. Mirrors the split + * the session-restore-projection path already uses. + */ + private deferGoalRestoreActivation( + runtime: GoalRuntime, + records: readonly GoalRecoveryRecord[], + ): Promise { + const preparation = runtime.prepareRestore(records); + this.deferredGoalRestoreActivation = () => { + void runtime.activateRestoredWork().catch((error: unknown) => { + this.debugLogger.error( + `Deferred goal restore activation failed: ${error}`, + ); + }); + }; + // Readiness deliberately tracks preparation, not activation: the Goal + // state is recovered and every `getGoalRuntimeReady()` awaiter can + // proceed. Only the continuation permit waits — `restoreActivationPending` + // holds `queueContinuation` until activation. Gating readiness on + // activation instead would hang any resumed Config that never calls + // `initialize()`. + return preparation.then(() => runtime); + } + + /** + * Release a restore held by {@link deferGoalRestoreActivation}. Called once + * the client-init step has replayed the resumed session's telemetry, so the + * restored Goal's first turn opens its meter on the restored totals. + */ + private startDeferredGoalRestore(): void { + const activate = this.deferredGoalRestoreActivation; + this.deferredGoalRestoreActivation = undefined; + if (!activate) return; + try { + activate(); + } catch (error) { + this.debugLogger.error( + `Deferred goal restore activation failed: ${error}`, + ); + } + } + /** * Run the restore that {@link initializeGoalRuntime} deferred because the * session writer was not yet accepting writes. @@ -7979,12 +8079,20 @@ export class Config { ); return; } - void pending.runtime - .restore(this.sessionData?.conversation.messages ?? []) - .then( + const records = this.sessionData?.conversation.messages ?? []; + if (this.shouldDeferGoalRestoreForTelemetryReplay()) { + // Same cold-start window as the constructor path: this runs from + // `activateChatRecording()`, still ahead of the client's replay. + void this.deferGoalRestoreActivation(pending.runtime, records).then( () => pending.resolve(pending.runtime), (error: unknown) => pending.reject(error), ); + return; + } + void pending.runtime.restore(records).then( + () => pending.resolve(pending.runtime), + (error: unknown) => pending.reject(error), + ); } /** @@ -8017,6 +8125,10 @@ export class Config { ); this.goalRestoreActivation = undefined; this.rejectGoalRestoreActivation = undefined; + // A restore waiting on the client's telemetry replay can never activate + // once the restore is abandoned — drop it so a later swap's activation + // cannot fire against the outgoing runtime. + this.deferredGoalRestoreActivation = undefined; } private notifyChatRecordingFailure(event: ChatRecordingFailureEvent): void { diff --git a/packages/core/src/telemetry/uiTelemetry.test.ts b/packages/core/src/telemetry/uiTelemetry.test.ts index 243b0d8bb5e..6dbbced3187 100644 --- a/packages/core/src/telemetry/uiTelemetry.test.ts +++ b/packages/core/src/telemetry/uiTelemetry.test.ts @@ -1307,6 +1307,61 @@ describe('UiTelemetryService', () => { expect(global2.models['m']?.tokens.prompt).toBe(350); }); + it('snapshotForReplay/restoreFromReplaySnapshot undo an abandoned replay', () => { + // Session A is live; the user runs /resume B and the swap fails after + // the replay but before the UI commits. + service.addEvent(makeApiEvent('m', 100), SESSION_A); + const snapshot = service.snapshotForReplay(SESSION_B); + + service.resetSession(SESSION_B); + service.addEvent(makeApiEvent('m', 300), SESSION_B); + expect(service.getMetrics().models['m']?.tokens.prompt).toBe(400); + + service.restoreFromReplaySnapshot(snapshot); + + // The abandoned session's contribution is gone from the aggregate... + expect(service.getMetrics().models['m']?.tokens.prompt).toBe(100); + // ...the surviving session's live data is untouched... + expect( + service.getMetricsForSession(SESSION_A).models['m']?.tokens.prompt, + ).toBe(100); + // ...and the bucket the replay created is gone rather than left empty. + expect(service.getMetricsForSession(SESSION_B).models).toEqual({}); + }); + + it('restoreFromReplaySnapshot puts back a pre-existing bucket, not an empty one', () => { + service.addEvent(makeApiEvent('m', 40), SESSION_B); + const snapshot = service.snapshotForReplay(SESSION_B); + + service.resetSession(SESSION_B); + service.addEvent(makeApiEvent('m', 300), SESSION_B); + service.restoreFromReplaySnapshot(snapshot); + + expect( + service.getMetricsForSession(SESSION_B).models['m']?.tokens.prompt, + ).toBe(40); + expect(service.getMetrics().models['m']?.tokens.prompt).toBe(40); + }); + + it('restoreFromReplaySnapshot restores closed-session state and prompt counts', () => { + service.addEvent(makeApiEvent('m', 10), SESSION_B); + service.removeSession(SESSION_B); + service.setLastPromptTokenCount(7); + const snapshot = service.snapshotForReplay(SESSION_B); + + // A replay re-opens the closed session and moves the prompt count. + service.resetSession(SESSION_B); + service.addEvent(makeApiEvent('m', 300), SESSION_B); + service.setLastPromptTokenCount(999); + + service.restoreFromReplaySnapshot(snapshot); + + expect(service.getLastPromptTokenCount()).toBe(7); + // Closed again: a late event must not resurrect the bucket. + service.addEvent(makeApiEvent('m', 5), SESSION_B); + expect(service.getMetricsForSession(SESSION_B).models).toEqual({}); + }); + it('#closedSessions should be bounded', () => { // Add more than MAX_CLOSED_SESSIONS for (let i = 0; i < 1005; i++) { diff --git a/packages/core/src/telemetry/uiTelemetry.ts b/packages/core/src/telemetry/uiTelemetry.ts index 4c2a0bac7e7..44a81fff448 100644 --- a/packages/core/src/telemetry/uiTelemetry.ts +++ b/packages/core/src/telemetry/uiTelemetry.ts @@ -188,6 +188,33 @@ const createInitialMetrics = (): SessionMetrics => ({ skills: createInitialSkillMetrics(), }); +/** + * The slice of telemetry state a session-swap replay overwrites. + * + * `/resume` and `/branch` replay the incoming session's stored history into + * the aggregate *before* the steps that can still fail (the Goal-runtime + * wait, client init), because the Goal meter has to open on the restored + * totals. Their catch blocks roll core back to the old session, but the + * service has no subtraction API — `resetSession` clears one bucket and + * `reset()` would take the surviving session's live data with it — so an + * abandoned swap would otherwise leak a full copy of the dead session's + * history into the process-wide totals for the life of the process, and + * `persistSessionUsage` would later write that inflated figure out. + * + * Take one with {@link UiTelemetryService.snapshotForReplay} immediately + * before replaying and hand it back to + * {@link UiTelemetryService.restoreFromReplaySnapshot} on the rollback path. + */ +export interface UiTelemetryReplaySnapshot { + readonly metrics: SessionMetrics; + readonly sessionId: string; + /** Absent when the session had no bucket yet — restore removes it again. */ + readonly sessionMetrics: SessionMetrics | undefined; + readonly sessionWasClosed: boolean; + readonly lastPromptTokenCount: number; + readonly lastCachedContentTokenCount: number; +} + export class UiTelemetryService extends EventEmitter { static readonly #MAX_CLOSED_SESSIONS = 1000; #metrics: SessionMetrics = createInitialMetrics(); @@ -269,6 +296,58 @@ export class UiTelemetryService extends EventEmitter { this.#lastCachedContentTokenCount = count; } + /** + * Captures everything a session replay is about to overwrite, so a session + * swap that fails after replaying can put the aggregate back. + * + * See {@link UiTelemetryReplaySnapshot} for why a snapshot is the only + * compensation available. + */ + snapshotForReplay(sessionId: string): UiTelemetryReplaySnapshot { + const sessionMetrics = this.#sessionMetrics.get(sessionId); + return { + metrics: structuredClone(this.#metrics), + sessionId, + sessionMetrics: sessionMetrics + ? structuredClone(sessionMetrics) + : undefined, + sessionWasClosed: this.#closedSessions.has(sessionId), + lastPromptTokenCount: this.#lastPromptTokenCount, + lastCachedContentTokenCount: this.#lastCachedContentTokenCount, + }; + } + + /** + * Puts back the state {@link snapshotForReplay} captured, undoing a replay + * whose session swap was abandoned. Only the snapshotted session's bucket + * is touched — every other session's live data survives, which is what + * makes this usable on a rollback path where the old session is still live. + */ + restoreFromReplaySnapshot(snapshot: UiTelemetryReplaySnapshot): void { + this.#metrics = structuredClone(snapshot.metrics); + if (snapshot.sessionMetrics) { + this.#sessionMetrics.set( + snapshot.sessionId, + structuredClone(snapshot.sessionMetrics), + ); + } else { + // No bucket existed before the replay; the replay created one. Drop it + // rather than leave an empty bucket that reads as a live session. + this.#sessionMetrics.delete(snapshot.sessionId); + } + if (snapshot.sessionWasClosed) { + this.#closedSessions.add(snapshot.sessionId); + } else { + this.#closedSessions.delete(snapshot.sessionId); + } + this.#lastPromptTokenCount = snapshot.lastPromptTokenCount; + this.#lastCachedContentTokenCount = snapshot.lastCachedContentTokenCount; + this.emit('update', { + metrics: this.#metrics, + lastPromptTokenCount: this.#lastPromptTokenCount, + }); + } + /** * Resets metrics to the initial state (used when resuming a session). */ From 205c249c03f6d93b42807d3a24b2578469e6fc97 Mon Sep 17 00:00:00 2001 From: qqqys Date: Wed, 19 Aug 2026 06:47:20 +0800 Subject: [PATCH 07/10] fix(telemetry): keep bySource prototype-free across a replay rollback MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit R5-1 (Critical, round 5): `bySource` is built with `Object.create(null)` precisely so a subagent named after an inherited `Object` member cannot short-circuit the `!bySource[name]` check and hand back a prototype member as the bucket — the comment above `createInitialModelMetrics` names that crash class. `structuredClone` does not preserve that: it copies own properties onto a fresh object carrying `Object.prototype`. So every `snapshotForReplay` and every `restoreFromReplaySnapshot` silently re-armed the crash, permanently, for the rest of the process. Verified rather than inferred. `structuredClone({bySource: Object.create(null)})` yields a clone whose prototype IS `Object.prototype`, with `typeof clone.bySource['constructor'] === 'function'` and the truthiness check passing — so `#getOrCreateSourceMetrics` returns the `Object` function itself and `bucket.api.totalRequests++` throws `TypeError: Cannot read properties of undefined (reading 'totalRequests')`, exactly the reported witness. `constructor` is a valid subagent name per the naming regex, and uiTelemetry.test.ts already covers it on the live path — the three replay tests missed it only because their fixture sets no `subagent_name` and so never probes a colliding key. All four clone sites now go through one `cloneSessionMetrics` helper that re-homes each model's `bySource` onto a null prototype after the clone. Test: `keeps bySource prototype-free across a snapshot/restore round trip` drives a `constructor`-named event through snapshot → resetSession → restore, asserts the null prototype survives in both the aggregate and the restored session bucket, and asserts a further colliding event accumulates instead of throwing. Mutation-verified: reverting all four sites to `structuredClone` turns exactly this test red. --- .../core/src/telemetry/uiTelemetry.test.ts | 49 +++++++++++++++++++ packages/core/src/telemetry/uiTelemetry.ts | 26 ++++++++-- 2 files changed, 71 insertions(+), 4 deletions(-) diff --git a/packages/core/src/telemetry/uiTelemetry.test.ts b/packages/core/src/telemetry/uiTelemetry.test.ts index 6dbbced3187..6f399a39e44 100644 --- a/packages/core/src/telemetry/uiTelemetry.test.ts +++ b/packages/core/src/telemetry/uiTelemetry.test.ts @@ -1329,6 +1329,55 @@ describe('UiTelemetryService', () => { expect(service.getMetricsForSession(SESSION_B).models).toEqual({}); }); + it('keeps bySource prototype-free across a snapshot/restore round trip', () => { + // R5-1: `structuredClone` copies own properties onto a fresh object with + // `Object.prototype`, so a plain clone silently re-arms the crash the + // prototype-free map exists to prevent — permanently, for the rest of + // the process. `constructor` is a valid subagent name per the naming + // regex, and after a rollback the truthiness check in + // #getOrCreateSourceMetrics would hand back `Object.prototype.constructor` + // as the "bucket", so `bucket.api.totalRequests++` throws. + const constructorEvent = (inputTokens: number) => + ({ + 'event.name': EVENT_API_RESPONSE, + model: 'm', + duration_ms: 100, + input_token_count: inputTokens, + output_token_count: 10, + total_token_count: inputTokens + 10, + cached_content_token_count: 0, + thoughts_token_count: 0, + subagent_name: 'constructor', + }) as ApiResponseEvent & { 'event.name': typeof EVENT_API_RESPONSE }; + + service.addEvent(constructorEvent(10), SESSION_A); + const snapshot = service.snapshotForReplay(SESSION_B); + service.resetSession(SESSION_B); + service.addEvent(constructorEvent(300), SESSION_B); + + service.restoreFromReplaySnapshot(snapshot); + + // The guard has to survive the round trip, in both the aggregate... + expect( + Object.getPrototypeOf(service.getMetrics().models['m']!.bySource), + ).toBeNull(); + // ...and the restored session bucket. + expect( + Object.getPrototypeOf( + service.getMetricsForSession(SESSION_A).models['m']!.bySource, + ), + ).toBeNull(); + + // And the next colliding event still accumulates instead of throwing. + expect(() => + service.addEvent(constructorEvent(5), SESSION_A), + ).not.toThrow(); + const bucket = service.getMetrics().models['m']!.bySource['constructor']!; + expect(typeof bucket).toBe('object'); + expect(bucket.api.totalRequests).toBe(2); + expect(bucket.tokens.prompt).toBe(15); + }); + it('restoreFromReplaySnapshot puts back a pre-existing bucket, not an empty one', () => { service.addEvent(makeApiEvent('m', 40), SESSION_B); const snapshot = service.snapshotForReplay(SESSION_B); diff --git a/packages/core/src/telemetry/uiTelemetry.ts b/packages/core/src/telemetry/uiTelemetry.ts index 44a81fff448..51221675bab 100644 --- a/packages/core/src/telemetry/uiTelemetry.ts +++ b/packages/core/src/telemetry/uiTelemetry.ts @@ -152,6 +152,24 @@ const createInitialModelMetrics = (): ModelMetrics => ({ bySource: Object.create(null) as Record, }); +/** + * `structuredClone` copies own properties onto a FRESH object with + * `Object.prototype` — it does not preserve the null prototype above, so a + * plain clone silently re-arms the crash that comment describes, permanently, + * for every `bySource` map it touches. Every snapshot/restore clone goes + * through here so the guard survives a replay rollback. + */ +const cloneSessionMetrics = (metrics: SessionMetrics): SessionMetrics => { + const clone = structuredClone(metrics); + for (const model of Object.values(clone.models)) { + model.bySource = Object.assign( + Object.create(null) as Record, + model.bySource, + ); + } + return clone; +}; + const createInitialSkillMetrics = (): SkillMetrics => ({ totalCalls: 0, totalSuccess: 0, @@ -306,10 +324,10 @@ export class UiTelemetryService extends EventEmitter { snapshotForReplay(sessionId: string): UiTelemetryReplaySnapshot { const sessionMetrics = this.#sessionMetrics.get(sessionId); return { - metrics: structuredClone(this.#metrics), + metrics: cloneSessionMetrics(this.#metrics), sessionId, sessionMetrics: sessionMetrics - ? structuredClone(sessionMetrics) + ? cloneSessionMetrics(sessionMetrics) : undefined, sessionWasClosed: this.#closedSessions.has(sessionId), lastPromptTokenCount: this.#lastPromptTokenCount, @@ -324,11 +342,11 @@ export class UiTelemetryService extends EventEmitter { * makes this usable on a rollback path where the old session is still live. */ restoreFromReplaySnapshot(snapshot: UiTelemetryReplaySnapshot): void { - this.#metrics = structuredClone(snapshot.metrics); + this.#metrics = cloneSessionMetrics(snapshot.metrics); if (snapshot.sessionMetrics) { this.#sessionMetrics.set( snapshot.sessionId, - structuredClone(snapshot.sessionMetrics), + cloneSessionMetrics(snapshot.sessionMetrics), ); } else { // No bucket existed before the replay; the replay created one. Drop it From 76ca082c6ae674a74dd8bd8c01ab6b24907dd1bf Mon Sep 17 00:00:00 2001 From: qqqys Date: Wed, 19 Aug 2026 10:47:57 +0800 Subject: [PATCH 08/10] fix(goal): make the meter breadcrumb visible and pin the replay handoff MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Round 6 review of #9301 confirmed no new findings, but re-confirmed the two Suggestions carried over from earlier rounds. Both are cleared here. R6-12 (ledger R1-5): the one-shot meter-failure breadcrumb was delivered solely through `debugLogger.warn`, which writes nothing unless QWEN_DEBUG_LOG_FILE is set — `writeLog` returns early otherwise and the module has no console fallback. So in the default configuration a persistent meter fault still produced no log line, telemetry event or error, which is exactly the indistinguishability ("my Goal says 0 tokens but I was billed for millions" vs "no API calls happened") the breadcrumb was added to remove. It now also writes a `[warn]` line to stderr, following `quarantineCorruptFile`'s house pattern, and keeps the debugLogger line as the verbose copy. Secondary, same finding: the message promised "until the meter recovers" while `meterFailureReported` was never reset, so one transient fault permanently silenced every distinct later one. A finite reading now re-arms the flag, which makes the wording true. Crash-loop repetition of the SAME fault is still reported once, because the flag only clears on a successful read. R6-3 (ledger R4-5): the replay-handoff marker had zero real coverage — `client.test.ts` injects a fake `consumeUiTelemetryEventsReplayed` onto its mockConfig and the resume/branch hooks pass a `vi.fn()` as `markUiTelemetryEventsReplayed`, so both ends of the handshake were mocked on their own side of the territory split. Added three `config.test.ts` tests against the real methods: mark→consume is true once and false on a second read, a mismatched session id neither consumes nor clears the marker, and `startNewSession` re-arms it. Both fixes are mutation-verified. Six mutants run, six killed: dropping the stderr write, dropping the recovery re-arm, making `consume` always return false, making `mark` a no-op, making the marker non-one-shot, and dropping the `startNewSession` re-arm each turn at least one of the new tests red. Verification: `cd packages/core && npx vitest run src/goals/goal-runtime.test.ts src/config/config.test.ts src/goals/goal-reducer.test.ts` — 725 passed. eslint and prettier clean on the three touched files. --- packages/core/src/config/config.test.ts | 63 ++++++++++++++++ packages/core/src/goals/goal-runtime.test.ts | 78 ++++++++++++++++++++ packages/core/src/goals/goal-runtime.ts | 19 ++++- 3 files changed, 156 insertions(+), 4 deletions(-) diff --git a/packages/core/src/config/config.test.ts b/packages/core/src/config/config.test.ts index 41839d8e8d0..ac957f9d275 100644 --- a/packages/core/src/config/config.test.ts +++ b/packages/core/src/config/config.test.ts @@ -2320,6 +2320,69 @@ describe('Server Config (config.ts)', () => { ); }); + it('replays a marked session exactly once, then never again', async () => { + // R4-5: both ends of this handshake are mocked on their own side — + // `client.test.ts` injects a fake `consumeUiTelemetryEventsReplayed` + // onto its mockConfig, and the resume/branch hooks pass a `vi.fn()` as + // `markUiTelemetryEventsReplayed` — so a regression making the real + // `consume` always return false, or the real `mark` a no-op, would + // reintroduce the double replay this PR fixes with every suite green. + const config = new Config({ ...baseParams }); + await config.initialize({ + skipGeminiInitialization: true, + skipHooks: true, + skipMcpDiscovery: true, + skipSkillManager: true, + skipFileCheckpointing: true, + }); + + // Nothing marked: initialize() must do its own replay. + expect(config.consumeUiTelemetryEventsReplayed('session-a')).toBe(false); + + config.markUiTelemetryEventsReplayed('session-a'); + // The one initialize() that follows the caller's own replay skips it... + expect(config.consumeUiTelemetryEventsReplayed('session-a')).toBe(true); + // ...and the marker is one-shot, so a later initialize() replays again. + expect(config.consumeUiTelemetryEventsReplayed('session-a')).toBe(false); + }); + + it('does not let one session id consume another session id marker', async () => { + const config = new Config({ ...baseParams }); + await config.initialize({ + skipGeminiInitialization: true, + skipHooks: true, + skipMcpDiscovery: true, + skipSkillManager: true, + skipFileCheckpointing: true, + }); + + config.markUiTelemetryEventsReplayed('session-a'); + + expect(config.consumeUiTelemetryEventsReplayed('session-b')).toBe(false); + // And the mismatched read left session-a's marker intact. + expect(config.consumeUiTelemetryEventsReplayed('session-a')).toBe(true); + }); + + it('re-arms the replay marker on the next session swap', async () => { + // The `uiTelemetryReplayedSessionId = undefined` re-arm in + // startNewSession is what makes resuming the SAME session a second time + // replay exactly once again; dropping it leaves a marker that swallows + // the next initialize()'s legitimate replay. + const config = new Config({ ...baseParams }); + await config.initialize({ + skipGeminiInitialization: true, + skipHooks: true, + skipMcpDiscovery: true, + skipSkillManager: true, + skipFileCheckpointing: true, + }); + + config.markUiTelemetryEventsReplayed('session-a'); + config.startNewSession('session-b'); + + expect(config.consumeUiTelemetryEventsReplayed('session-a')).toBe(false); + }); + it('carries the outgoing session id when resuming a different persisted session', async () => { const config = new Config({ ...baseParams }); await config.initialize({ diff --git a/packages/core/src/goals/goal-runtime.test.ts b/packages/core/src/goals/goal-runtime.test.ts index 09809fe622a..c445ee9f41a 100644 --- a/packages/core/src/goals/goal-runtime.test.ts +++ b/packages/core/src/goals/goal-runtime.test.ts @@ -383,6 +383,84 @@ describe('goal runtime', () => { }, ); + it('surfaces the meter breadcrumb on stderr, not only the debug log', async () => { + // R1-5: `debugLogger.warn` writes nothing unless QWEN_DEBUG_LOG_FILE is + // set (`writeLog` returns early otherwise, with no console fallback), so + // on its own this breadcrumb leaves the default configuration with exactly + // the silence it exists to remove — a Goal reporting 0 tokens against a + // session billed for real usage, with nothing to grep. Same remedy as + // `quarantineCorruptFile`. + const stderr = vi + .spyOn(process.stderr, 'write') + .mockImplementation(() => true); + try { + const journal = fakeGoalJournal(); + const host = fakeGoalTurnHost(); + const runtime = createGoalRuntime({ + journal, + tokenMeter: { + readSessionTokens: () => { + throw new Error('metrics are unavailable'); + }, + }, + }); + runtime.bindHost(host); + await runtime.dispatch({ action: 'create', objective: 'ship' }); + await runtime.finishTurn(host.started[0]!); + + const written = stderr.mock.calls + .map(([chunk]) => String(chunk)) + .filter((line) => line.includes('Goal token meter unreadable')); + expect(written).toHaveLength(1); + expect(written[0]).toContain('[warn]'); + expect(written[0]).toContain('metrics are unavailable'); + } finally { + stderr.mockRestore(); + } + }); + + it('reports a later fault after the meter recovers', async () => { + // R1-5 (secondary): the message promises "until the meter recovers", but + // the one-shot flag was never reset — so one transient fault permanently + // silenced every distinct later one, and the wording was a lie. + debugWarn.mockClear(); + const stderr = vi + .spyOn(process.stderr, 'write') + .mockImplementation(() => true); + try { + let reading = 0; + const journal = fakeGoalJournal(); + const host = fakeGoalTurnHost(); + const runtime = createGoalRuntime({ + journal, + tokenMeter: { readSessionTokens: () => reading }, + }); + runtime.bindHost(host); + await runtime.dispatch({ action: 'create', objective: 'ship' }); + + // Turn 1 closes on a fault. (Turn 2 auto-starts on the same bad + // reading, and must NOT warn again — that is the crash-loop case the + // one-shot flag exists for.) + reading = Number.NaN; + await runtime.finishTurn(host.started[0]!); + expect(debugWarn).toHaveBeenCalledTimes(1); + + // Turn 2 reads fine at both ends: the meter recovered, so the flag + // must re-arm. + reading = 100; + await runtime.finishTurn(host.started[1]!); + expect(debugWarn).toHaveBeenCalledTimes(1); + + // Turn 3 faults again — a distinct fault, not a repeat of the first. + reading = Number.NaN; + await runtime.finishTurn(host.started[2]!); + + expect(debugWarn).toHaveBeenCalledTimes(2); + } finally { + stderr.mockRestore(); + } + }); + it.each([Number.NaN, Number.POSITIVE_INFINITY])( 'finishes the turn when the token meter returns %s', async (invalidTotal) => { diff --git a/packages/core/src/goals/goal-runtime.ts b/packages/core/src/goals/goal-runtime.ts index 6397ec96d9d..121121b40cb 100644 --- a/packages/core/src/goals/goal-runtime.ts +++ b/packages/core/src/goals/goal-runtime.ts @@ -283,17 +283,28 @@ export function createGoalRuntime( const reportMeterFailure = (detail: string): void => { if (meterFailureReported) return; meterFailureReported = true; - debugLogger.warn( + const message = `Goal token meter unreadable (${detail}); tokensUsed will report 0 ` + - `for this runtime until the meter recovers. Reported once per runtime.`, - ); + `until the meter recovers. Reported once per fault.`; + // `debugLogger.warn` is gated behind QWEN_DEBUG_LOG_FILE (unset for almost + // all users), so on its own it leaves the default configuration with + // nothing to grep — the exact indistinguishability this breadcrumb exists + // to remove. Surface it on stderr too, as `quarantineCorruptFile` does. + process.stderr.write(`[warn] ${message}\n`); + debugLogger.warn(message); }; const readSessionTokens = (): number | undefined => { if (!options.tokenMeter) return undefined; try { const total = options.tokenMeter.readSessionTokens(); - if (Number.isFinite(total)) return total; + if (Number.isFinite(total)) { + // Re-arm on recovery, so "until the meter recovers" is true and a + // LATER, distinct fault is reported instead of being swallowed by the + // first one's flag. + meterFailureReported = false; + return total; + } // A non-finite reading is the shape that drift produces: rename // `tokens.total` and the config meter's reduce yields NaN. reportMeterFailure(`non-finite reading ${String(total)}`); From d6348598a13e72b532d37ee3423f57f3ee6cc6a0 Mon Sep 17 00:00:00 2001 From: qqqys Date: Wed, 19 Aug 2026 15:42:40 +0800 Subject: [PATCH 09/10] test(goal): pin the four turn-start and replay readings round 7 re-confirmed MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Round 7 confirmed no new findings but re-counted four Suggestion-level findings as already-reported duplicates. All four are the same shape: a load-bearing line with no test behind it, each with a mutation witness showing the suite stays green when it is deleted. Clearing them means making the next round unable to re-derive them, so each gets the test its witness describes. R1-2 (goal-runtime.ts, `handleStartFailure` promotion reading) — the sixth and last of the `currentTurnTokensAtStart` turn-start readings with no test. All four start-failure tests build the runtime with no `tokenMeter`, and the three metered promotion tests never make `startGoalTurn` reject, so nothing in the repo combined the two. Deleting the reading left the promoted permit with `opened === undefined`, `takeTurnTokens()` returning 0, and the promoted user turn's real spend silently dropped from `tokensUsed` on every host start failure. Adds the fourth twin of the promotion meter tests. Mutation: deleting the line fails the new test with `expected +0 to be 25` -- the exact witness the finding predicted -- and no other test moves. R5-2 (config.ts, writer-lease entrance of the telemetry-replay deferral) — the deferral is load-bearing on this entrance too: `initializeOnce` awaits `activateChatRecording()` strictly before `initializeInternal` reaches `geminiClient.initialize()`, where the stored telemetry is replayed. But the only lease restore test uses a PAUSED Goal and never calls `initialize()`, so it never mints a permit and never reads the meter -- deleting the branch shipped green. Adds the lease twin of the cold-resume meter test: lease handed over via `startPendingGoalRestore()`, no permit before `initialize()`, 200k replayed into the bucket during it, then the first restored turn bills only its own 500. Mutation: deleting the branch fails the new test at `expect(permit) .toBeUndefined()` -- a permit minted on the empty bucket, which is the R4-4 regression reappearing on this entrance -- and exactly one test goes red. R5-3 (uiTelemetry.ts, `lastCachedContentTokenCount` restore) — the field rides the same replay snapshot as `lastPromptTokenCount` but had no assertion anywhere, so deleting its capture/restore pair left the suite 50/50 green while its sibling stayed guarded. `geminiChat` writes it on every live API response, so an in-flight response from the OUTGOING session landing inside a swap window is exactly what the rollback has to undo. Extends the existing sibling test rather than adding a near-duplicate. Mutation: zeroing the capture fails the test with `expected 1234 to be 42`. R5-5 (config.ts, `shouldDeferGoalRestoreForTelemetryReplay` doc) — the clause "a swap-time restore (`startNewSession`) runs after the caller already replayed" is backwards. Verified against the code: both `useResumeCommand` and `useBranchCommand` call `config.startNewSession(...)` FIRST (useResumeCommand .ts:173, useBranchCommand.ts:208), and `startNewSession` synchronously calls `initializeGoalRuntime()`. The path is safe for two different reasons -- `clientSessionTelemetryReplayed` is set once by the initial `initialize()` and never cleared, so the predicate is already false at a swap; and the hooks replay synchronously, with no `await` between `startNewSession()` and `markUiTelemetryEventsReplayed()`, so the restore's microtask activation cannot run until the replay has landed. Rewritten to state that, including the "no `await` may be inserted between those two calls" constraint a maintainer would otherwise violate on the strength of the old wording. Comment-only. Verification: - packages/core config.test.ts + goal-runtime.test.ts + uiTelemetry.test.ts: 708 tests green. - Each of the three code findings mutation-checked individually, above; every mutant turns exactly one test red, the new one. - eslint and prettier clean on all four files. The only non-test source change in this commit is the R5-5 comment. - `npm run typecheck --workspace packages/core` reports one error, `src/utils/image-view.ts` / `sharp` -- identical on the untouched branch, a dependency-typing skew in this worktree, not this change. --- packages/core/src/config/config.test.ts | 91 +++++++++++++++++++ packages/core/src/config/config.ts | 23 ++++- packages/core/src/goals/goal-runtime.test.ts | 39 ++++++++ .../core/src/telemetry/uiTelemetry.test.ts | 10 +- 4 files changed, 157 insertions(+), 6 deletions(-) diff --git a/packages/core/src/config/config.test.ts b/packages/core/src/config/config.test.ts index ac957f9d275..8c21072bce7 100644 --- a/packages/core/src/config/config.test.ts +++ b/packages/core/src/config/config.test.ts @@ -2828,6 +2828,97 @@ describe('Server Config (config.ts)', () => { vi.mocked(uiTelemetryService.getMetricsForSession).mockReset(); }); + // R5-2: the twin of the test above for the WRITER-LEASE entrance. Under a + // lease the constructor parks the restore and `activateChatRecording()` + // starts it instead — still strictly before `initializeInternal` reaches + // `geminiClient.initialize()`, where the stored telemetry is replayed. So + // the deferral is load-bearing on this entrance too, but deleting its + // branch in `startPendingGoalRestore` left every suite green: the only + // lease restore test uses a PAUSED Goal and never calls `initialize()`, so + // it never mints a permit and never reads the meter. + it('opens a lease-restored Goal meter on the replayed totals, not an empty bucket', async () => { + const sessionId = 'lease-resume-session'; + const telemetry = new UiTelemetryService(); + vi.mocked(uiTelemetryService.getMetricsForSession).mockImplementation( + (requestedSessionId) => + telemetry.getMetricsForSession(requestedSessionId), + ); + const apiResponse = ( + totalTokens: number, + ): ApiResponseEvent & { 'event.name': typeof EVENT_API_RESPONSE } => + ({ + 'event.name': EVENT_API_RESPONSE, + model: 'model-a', + duration_ms: 1, + input_token_count: totalTokens - 15, + output_token_count: 15, + total_token_count: totalTokens, + cached_content_token_count: 0, + thoughts_token_count: 0, + }) as ApiResponseEvent & { + 'event.name': typeof EVENT_API_RESPONSE; + }; + + const config = new Config({ + ...baseParams, + chatRecording: true, + experimentalZedIntegration: true, + sessionWriterLeaseEnabled: true, + sessionId, + sessionData: resumedGoalSession('active'), + }); + const recorder = config.getChatRecordingService(); + if (!recorder) throw new Error('expected a chat recording service'); + vi.spyOn(recorder, 'recordGoalState').mockResolvedValue({} as ChatRecord); + // `initialize()` re-reads the transcript under the lease before it + // activates the recorder; this file has no session store behind it. + vi.spyOn(config.getSessionService(), 'loadSession').mockResolvedValue( + resumedGoalSession('active'), + ); + // Under the lease the constructor parked the restore rather than running + // it, so there is nothing to hold back yet. + expect(recorder.hasWriteOwnership()).toBe(false); + + let permit: GoalTurnPermit | undefined; + config.bindGoalTurnHost({ + startGoalTurn: vi.fn(async (input) => { + permit = input.permit; + }), + preemptGoalTurn: vi.fn(), + }); + + // Stands in for `activateChatRecording()` handing the lease over — the + // entrance under test, and the one `initializeOnce` awaits strictly + // before the client replays. + vi.spyOn(recorder, 'hasWriteOwnership').mockReturnValue(true); + ( + config as unknown as { startPendingGoalRestore(): void } + ).startPendingGoalRestore(); + const runtime = await config.getGoalRuntimeReady(); + + // The Goal is ACTIVE and restored, but the client has not replayed the + // resumed session's telemetry into the bucket yet. A permit minted now + // would open its meter on 0. + expect(runtime.getSnapshot().goal?.status).toBe('active'); + expect(permit).toBeUndefined(); + + // What `GeminiClient.initialize()` does for a resumed session. + telemetry.addEvent(apiResponse(200_000), sessionId); + + await config.initialize(); + await vi.waitFor(() => expect(permit).toBeDefined()); + + // The restored turn's own spend. + telemetry.addEvent(apiResponse(500), sessionId); + await runtime.finishTurn(permit!); + + // 500, not 200_500 — the R4-4 regression this PR fixes for the + // constructor entrance, on the entrance a lease uses instead. + expect(runtime.getSnapshot().goal?.tokensUsed).toBe(500); + + vi.mocked(uiTelemetryService.getMetricsForSession).mockReset(); + }); + it('rebinds the current Goal host to every replacement runtime', async () => { const config = new Config({ ...baseParams, diff --git a/packages/core/src/config/config.ts b/packages/core/src/config/config.ts index e90d3a05046..a78afb657d7 100644 --- a/packages/core/src/config/config.ts +++ b/packages/core/src/config/config.ts @@ -8001,11 +8001,24 @@ export class Config { * ~200k extra on its first turn, and that figure is persisted into * `GoalRecord.tokensUsed`. * - * The `/resume` and `/branch` hooks, the headless path and the ACP path all - * replay before they let the runtime activate; this closes the same gap on - * the cold-start entrance. It only applies to a resumed session — a fresh - * session has nothing to replay, and a swap-time restore (`startNewSession`) - * runs after the caller already replayed. + * The headless path and the ACP path replay before they let the runtime + * activate; this closes the same gap on the cold-start entrance. It only + * applies to a resumed session — a fresh session has nothing to replay. + * + * A swap-time restore is not deferred either, but NOT because the caller + * already replayed: the `/resume` and `/branch` hooks call + * `startNewSession()` FIRST, and that kicks `initializeGoalRuntime()` -> + * `runtime.restore()` synchronously, before their replay block. It is safe + * for two other reasons — `clientSessionTelemetryReplayed` is already set by + * the initial `initialize()` and is never cleared, so this predicate is false + * at a swap; and the hooks replay SYNCHRONOUSLY, with no `await` between + * `startNewSession()` and `markUiTelemetryEventsReplayed()`, so the restore's + * activation — a microtask — cannot run until the replay has landed. No + * `await` may be inserted between those two calls: at that await the + * activation would run `queueContinuation` -> `flushContinuation`, minting a + * permit that samples `readSessionTokens()` on the still-empty bucket, and + * the first restored turn would bill the whole replayed history — the exact + * regression this predicate exists to prevent, on a path it does not guard. */ private shouldDeferGoalRestoreForTelemetryReplay(): boolean { return ( diff --git a/packages/core/src/goals/goal-runtime.test.ts b/packages/core/src/goals/goal-runtime.test.ts index c445ee9f41a..94d994a491d 100644 --- a/packages/core/src/goals/goal-runtime.test.ts +++ b/packages/core/src/goals/goal-runtime.test.ts @@ -3465,6 +3465,45 @@ describe('goal runtime', () => { expect(runtime.getSnapshot().activity).toBe('running'); }); + // R1-2: the fourth twin of the promotion meter tests. Of the six turn-start + // readings of `currentTurnTokensAtStart`, this was the last one no test + // reached: the promotion inside `handleStartFailure`. All four start-failure + // tests above build the runtime with NO `tokenMeter`, and the three metered + // promotion tests never make `startGoalTurn` reject — so deleting the reading + // on this path left the whole suite green while the promoted turn ran with + // `opened === undefined`, `takeTurnTokens()` returned 0, and the user turn's + // real spend vanished from `tokensUsed` on every host start failure. + it('bills the meter delta for user input promoted after a start failure', async () => { + const failedStart = deferred(); + let sessionTokens = 0; + const runtime = createGoalRuntime({ + journal: fakeGoalJournal(), + tokenMeter: { readSessionTokens: () => sessionTokens }, + }); + runtime.bindHost({ + startGoalTurn: () => failedStart.promise, + preemptGoalTurn: vi.fn(), + }); + await runtime.dispatch({ action: 'create', objective: 'ship' }); + expect(runtime.beginTurn('real-user')).toBeUndefined(); + + // Spend that lands before the host gives up. It belongs to the turn that + // never started, so the promotion must open its reading HERE, not at 0. + sessionTokens = 40; + failedStart.reject(new Error('host rejected')); + await vi.waitFor(() => + expect(runtime.permitForTurn('real-user')).toBeDefined(), + ); + + sessionTokens = 65; + await runtime.finishTurn(runtime.permitForTurn('real-user')!); + + // The promoted turn's own 25 — not 65, which is what a reading left at the + // start of the failed turn would bill, and not 0, which is what an unopened + // reading leaves behind. + expect(runtime.getSnapshot().goal?.tokensUsed).toBe(25); + }); + it('discards the rejected permit proposal before promoting queued user input', async () => { const failedStart = deferred(); const started: GoalTurnPermit[] = []; diff --git a/packages/core/src/telemetry/uiTelemetry.test.ts b/packages/core/src/telemetry/uiTelemetry.test.ts index 6f399a39e44..b845a733d52 100644 --- a/packages/core/src/telemetry/uiTelemetry.test.ts +++ b/packages/core/src/telemetry/uiTelemetry.test.ts @@ -1396,16 +1396,24 @@ describe('UiTelemetryService', () => { service.addEvent(makeApiEvent('m', 10), SESSION_B); service.removeSession(SESSION_B); service.setLastPromptTokenCount(7); + // R5-3: the cached-token field rides the same snapshot and had no + // assertion anywhere, so deleting its capture/restore pair shipped green + // while its sibling stayed guarded. `geminiChat` writes it on every live + // API response, so an in-flight response from the OUTGOING session + // landing inside a swap window is exactly what the rollback undoes. + service.setLastCachedContentTokenCount(42); const snapshot = service.snapshotForReplay(SESSION_B); - // A replay re-opens the closed session and moves the prompt count. + // A replay re-opens the closed session and moves both counts. service.resetSession(SESSION_B); service.addEvent(makeApiEvent('m', 300), SESSION_B); service.setLastPromptTokenCount(999); + service.setLastCachedContentTokenCount(1234); service.restoreFromReplaySnapshot(snapshot); expect(service.getLastPromptTokenCount()).toBe(7); + expect(service.getLastCachedContentTokenCount()).toBe(42); // Closed again: a late event must not resurrect the bucket. service.addEvent(makeApiEvent('m', 5), SESSION_B); expect(service.getMetricsForSession(SESSION_B).models).toEqual({}); From 0d5abe8ec420285ab8ecb358cc6eb38f31a9ba0d Mon Sep 17 00:00:00 2001 From: qqqys Date: Wed, 19 Aug 2026 19:40:04 +0800 Subject: [PATCH 10/10] fix(goal): keep the turn's opening meter reading until the write commits MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `finishTurn` evaluated `tokensUsed: takeTurnTokens()` while building the record, which cleared `currentTurnTokensAtStart` strictly before `await journal.recordGoalState(...)`. That write can throw — a transient writer-lease unavailability or a disk error — without mutating any state, which is exactly the contract the sibling test 'keeps turn state and the dispatch mutex usable when turn persistence fails' pins: the permit stays current so the same turn retries. On that retry the opening reading was gone, so the turn persisted with `tokensUsed: 0` while `turnCount` still advanced — the turn's real spend dropped out of the Goal's usage figure with no breadcrumb, since the meter itself read fine. The interactive TUI runs this chain (useGeminiStream finishTurn catch -> failClosedGoalTurn -> a second finishTurn on the same permit). `takeTurnTokens` becomes the non-destructive `peekTurnTokens`; the reading is cleared on the post-write path, which already did so. The new test is the metered twin of the persistence-failure contract test (that one builds its runtime with no `tokenMeter`, which is why the ordering shipped green); mutation-verified that restoring the consume turns it red at `tokensUsed: 0`. Co-Authored-By: Claude Opus 5 (1M context) --- packages/core/src/goals/goal-runtime.test.ts | 35 ++++++++++++++++++++ packages/core/src/goals/goal-runtime.ts | 19 ++++++++--- 2 files changed, 50 insertions(+), 4 deletions(-) diff --git a/packages/core/src/goals/goal-runtime.test.ts b/packages/core/src/goals/goal-runtime.test.ts index 94d994a491d..fb5ca74de25 100644 --- a/packages/core/src/goals/goal-runtime.test.ts +++ b/packages/core/src/goals/goal-runtime.test.ts @@ -3949,6 +3949,41 @@ describe('goal runtime', () => { }); }); + // The metered twin of the contract above. `finishTurn` evaluates the turn's + // token delta before its journal write, so consuming the opening reading + // there left the retry — which the sibling test exists to certify — with no + // reading at all, persisting the turn at `tokensUsed: 0` while `turnCount` + // still advanced. The TUI runs exactly this chain (finishTurn catch → + // failClosedGoalTurn → a second finishTurn on the same permit), so a + // transient fault that succeeds on retry silently dropped the turn's spend. + it('bills the retry after a failed turn write the tokens the turn spent', async () => { + const journal = fakeGoalJournal({ + appendErrors: [undefined, new Error('turn write failed')], + }); + const host = fakeGoalTurnHost(); + let sessionTokens = 10_000; + const runtime = createGoalRuntime({ + journal, + tokenMeter: { readSessionTokens: () => sessionTokens }, + }); + runtime.bindHost(host); + await runtime.dispatch({ action: 'create', objective: 'ship' }); + + const permit = host.started[0]!; + sessionTokens = 12_000; + await expect(runtime.finishTurn(permit)).rejects.toThrow( + 'turn write failed', + ); + + // The failed write mutated nothing, so the same permit retries — and the + // opening reading must still be there for it. + await runtime.finishTurn(permit); + expect(runtime.getSnapshot().goal).toMatchObject({ + turnCount: 1, + tokensUsed: 2_000, + }); + }); + it('restores active state once while stopped state remains display-only', async () => { const activeHost = fakeGoalTurnHost(); const active = createGoalRuntime({ journal: fakeGoalJournal() }); diff --git a/packages/core/src/goals/goal-runtime.ts b/packages/core/src/goals/goal-runtime.ts index 121121b40cb..ec54adb306b 100644 --- a/packages/core/src/goals/goal-runtime.ts +++ b/packages/core/src/goals/goal-runtime.ts @@ -320,15 +320,23 @@ export function createGoalRuntime( }; /** - * The tokens the finishing turn billed, consuming the reading it opened with. + * The tokens the finishing turn has billed so far, WITHOUT consuming the + * reading it opened with. + * + * Non-destructive on purpose: `finishTurn` evaluates this before its journal + * write, and that write can throw (a transient writer-lease unavailability + * or a disk error) without mutating any state — the permit stays current so + * the same turn retries. Consuming here would leave the retry with no + * opening reading, persisting the turn at `tokensUsed: 0` and dropping its + * real spend from the Goal's usage figure with no breadcrumb. The reading is + * cleared instead on the post-write path, once the state is durable. * * Undefined at either end — no meter, an unreadable meter, or a turn whose * opening reading was never taken — yields zero rather than a guess, so an * unmetered session reports no spend instead of a wrong one. */ - const takeTurnTokens = (): number => { + const peekTurnTokens = (): number => { const opened = currentTurnTokensAtStart; - currentTurnTokensAtStart = undefined; const closed = readSessionTokens(); if (opened === undefined || closed === undefined) return 0; return Math.max(0, closed - opened); @@ -1294,7 +1302,7 @@ export function createGoalRuntime( const recordUuid = randomUUID(); const nextGoal = reduceGoalTurnFinished(snapshot.goal, { now: Date.now(), - tokensUsed: takeTurnTokens(), + tokensUsed: peekTurnTokens(), }); const persistedSnapshot: GoalSnapshotV2 = { v: GOAL_STATE_VERSION, @@ -1358,6 +1366,9 @@ export function createGoalRuntime( activity: verifying ? 'verifying' : 'idle', }; currentPermit = undefined; + // Only now, past the journal write: `peekTurnTokens` deliberately + // left this in place so a write that throws can be retried on the + // same permit without losing the turn's opening reading. currentTurnTokensAtStart = undefined; currentPermitHost = undefined; currentTurnKey = undefined;