From c1bfb37b4f5e8124e99a4c476623a73c91270560 Mon Sep 17 00:00:00 2001 From: "jinjing.zzj" Date: Mon, 24 Aug 2026 13:59:45 +0800 Subject: [PATCH 01/14] fix(acp): route-scope the session token-limit cache in Session.ts The ACP Session keeps a private `lastPromptTokenCount` fed from streamed `usageMetadata`, reset only when the chat instance changes (#syncPromptTokenCountWithCurrentChat). ACP model switches (unstable_setSessionModel -> setModel -> config.switchModel) rebuild the content generator but keep the same GeminiChat, so a count recorded on the previous route survived and anchored the session-token-limit gate for the new route: any modelOverride send (compression skipped) or any send whose compression attempt throws reached #getPostCompressionTokenCount(null) with the stale pre-switch count and was wrongly dropped with SessionTokenLimitExceeded / stopReason 'max_tokens'. Attribute the cached count to the route that produced it (Config.getModelRouteIdentity) and invalidate it on a route change, mirroring the route-scoping #9506 applied to the GeminiChat counts. Same-route counting and the chat-instance reset are unchanged. Fixes #9529 --- .../acp-integration/session/Session.test.ts | 111 ++++++++++++++++++ .../src/acp-integration/session/Session.ts | 29 ++++- 2 files changed, 136 insertions(+), 4 deletions(-) diff --git a/packages/cli/src/acp-integration/session/Session.test.ts b/packages/cli/src/acp-integration/session/Session.test.ts index 090ae123347..613b9b6652d 100644 --- a/packages/cli/src/acp-integration/session/Session.test.ts +++ b/packages/cli/src/acp-integration/session/Session.test.ts @@ -11989,6 +11989,117 @@ describe('Session', () => { ); }); + it('does not drop a route-B send using a stale route-A token count after a model switch (#9529)', async () => { + mockConfig.getSessionTokenLimit = vi.fn().mockReturnValue(100); + // ACP model switches keep the same GeminiChat instance, so only the + // route identity changes — the chat-instance reset in the session's + // token cache never fires (#9529, follow-up to #9454/#9506). + let routeIdentity = 'route-a'; + mockConfig.getModelRouteIdentity = vi.fn(() => routeIdentity); + + // Prompt 1 on route A: an API-reported count lands in the session's + // private token cache. + mockGeminiClient.tryCompressChat.mockResolvedValueOnce({ + originalTokenCount: 50, + newTokenCount: 50, + compressionStatus: core.CompressionStatus.NOOP, + }); + mockChat.sendMessageStream = vi + .fn() + .mockResolvedValueOnce( + createStreamWithChunks([ + { + type: core.StreamEventType.CHUNK, + value: { + usageMetadata: { + totalTokenCount: 101, + promptTokenCount: 101, + }, + }, + }, + ]), + ) + .mockResolvedValueOnce(createEmptyStream()); + + await expect( + session.prompt({ + sessionId: 'test-session-id', + prompt: [{ type: 'text', text: 'first' }], + }), + ).resolves.toEqual({ stopReason: 'end_turn' }); + + // Switch to route B on the same chat instance. + routeIdentity = 'route-b'; + + // Prompt 2 on route B reaches the session-token-limit gate without + // compression info (compression throws), so the gate falls back to the + // cached count. The stale route-A count (101 > 100) must not drop a + // route-B send. + mockGeminiClient.tryCompressChat.mockRejectedValueOnce( + new Error('compression rate limited'), + ); + + await expect( + session.prompt({ + sessionId: 'test-session-id', + prompt: [{ type: 'text', text: 'second' }], + }), + ).resolves.toEqual({ stopReason: 'end_turn' }); + + expect(mockChat.sendMessageStream).toHaveBeenCalledTimes(2); + }); + + it('still intercepts a same-route send whose recorded count exceeds the session token limit (#9529)', async () => { + mockConfig.getSessionTokenLimit = vi.fn().mockReturnValue(100); + // The route never changes, so the route-scoped cache must keep the + // recorded count and let the gate trip exactly as it did before #9529. + mockConfig.getModelRouteIdentity = vi.fn(() => 'route-a'); + + mockGeminiClient.tryCompressChat.mockResolvedValueOnce({ + originalTokenCount: 50, + newTokenCount: 50, + compressionStatus: core.CompressionStatus.NOOP, + }); + mockChat.sendMessageStream = vi + .fn() + .mockResolvedValueOnce( + createStreamWithChunks([ + { + type: core.StreamEventType.CHUNK, + value: { + usageMetadata: { + totalTokenCount: 101, + promptTokenCount: 101, + }, + }, + }, + ]), + ) + .mockResolvedValueOnce(createEmptyStream()); + + await expect( + session.prompt({ + sessionId: 'test-session-id', + prompt: [{ type: 'text', text: 'first' }], + }), + ).resolves.toEqual({ stopReason: 'end_turn' }); + + // Same route, compression throws → the gate reads the recorded + // same-route count (101 > 100) and must still drop the send. + mockGeminiClient.tryCompressChat.mockRejectedValueOnce( + new Error('compression rate limited'), + ); + + await expect( + session.prompt({ + sessionId: 'test-session-id', + prompt: [{ type: 'text', text: 'second' }], + }), + ).resolves.toEqual({ stopReason: 'max_tokens' }); + + expect(mockChat.sendMessageStream).toHaveBeenCalledTimes(1); + }); + it('returns cancelled when automatic compression is aborted', async () => { mockConfig.getSessionTokenLimit = vi.fn().mockReturnValue(100); mockGeminiClient.tryCompressChat.mockImplementation( diff --git a/packages/cli/src/acp-integration/session/Session.ts b/packages/cli/src/acp-integration/session/Session.ts index ba1c6ee4ec4..ce1f365f53e 100644 --- a/packages/cli/src/acp-integration/session/Session.ts +++ b/packages/cli/src/acp-integration/session/Session.ts @@ -1849,6 +1849,11 @@ export class Session implements SessionContext { private cronDisabledByTokenLimit = false; private lastPromptTokenCount = 0; private lastPromptTokenCountChat: GeminiChat | null = null; + // The model route that produced `lastPromptTokenCount` (Config + // .getModelRouteIdentity). ACP model switches keep the same GeminiChat, so + // the chat-instance check alone never invalidates the count on a route + // change (#9529, follow-up to #9454/#9506). + private lastPromptTokenCountRouteKey: string | undefined = undefined; private midTurnDrainUnavailable = false; private midTurnDrainTimeoutStrikes = 0; // ACP can continue one logical conversation through prompt, cron, and @@ -7100,15 +7105,31 @@ export class Session implements SessionContext { return tokenCount; } + #currentRouteKey(): string { + // Optional chaining keeps partial Config test mocks from throwing; a + // missing identity degrades to one stable key, i.e. no route-change + // invalidation (mirrors GeminiChat.currentRouteKey, #9454). + return this.config.getModelRouteIdentity?.() ?? ''; + } + #syncPromptTokenCountWithCurrentChat(): void { const chat = this.#getCurrentChat(); - if ( - this.lastPromptTokenCountChat && - this.lastPromptTokenCountChat !== chat - ) { + const routeKey = this.#currentRouteKey(); + const chatChanged = + this.lastPromptTokenCountChat && this.lastPromptTokenCountChat !== chat; + // A model switch rebuilds the content generator but keeps the same + // GeminiChat, so the chat-instance check above never fires on a route + // change. A count recorded for the previous route must not anchor the + // session-token-limit gate for the new one — reset it so the gate falls + // back to the fresh send's own count (#9529, follow-up to #9454/#9506). + const routeChanged = + this.lastPromptTokenCountRouteKey !== undefined && + this.lastPromptTokenCountRouteKey !== routeKey; + if (chatChanged || routeChanged) { this.lastPromptTokenCount = 0; } this.lastPromptTokenCountChat = chat; + this.lastPromptTokenCountRouteKey = routeKey; } #isAbortError(error: unknown): boolean { From c121781382e4ee5040daa54b961b1b871ccc5684 Mon Sep 17 00:00:00 2001 From: yiliang114 Date: Mon, 24 Aug 2026 16:19:26 +0800 Subject: [PATCH 02/14] fix(cli): retain acp token counts per route --- .../acp-integration/session/Session.test.ts | 110 ++++++++++++++++++ .../src/acp-integration/session/Session.ts | 96 ++++++++++----- 2 files changed, 178 insertions(+), 28 deletions(-) diff --git a/packages/cli/src/acp-integration/session/Session.test.ts b/packages/cli/src/acp-integration/session/Session.test.ts index 613b9b6652d..044e79b1524 100644 --- a/packages/cli/src/acp-integration/session/Session.test.ts +++ b/packages/cli/src/acp-integration/session/Session.test.ts @@ -12100,6 +12100,116 @@ describe('Session', () => { expect(mockChat.sendMessageStream).toHaveBeenCalledTimes(1); }); + it('does not drop a modelOverride send using the active route token count (#9529)', async () => { + mockConfig.getSessionTokenLimit = vi.fn().mockReturnValue(100); + mockConfig.getModelRouteIdentity = vi.fn((model?: string) => + model === 'route-b-model' ? 'route-b' : 'route-a', + ); + + mockGeminiClient.tryCompressChat.mockResolvedValueOnce({ + originalTokenCount: 50, + newTokenCount: 50, + compressionStatus: core.CompressionStatus.NOOP, + }); + mockChat.sendMessageStream = vi + .fn() + .mockResolvedValueOnce( + createStreamWithChunks([ + { + type: core.StreamEventType.CHUNK, + value: { + usageMetadata: { + totalTokenCount: 101, + promptTokenCount: 101, + }, + }, + }, + ]), + ) + .mockResolvedValueOnce(createEmptyStream()); + + await expect( + session.prompt({ + sessionId: 'test-session-id', + prompt: [{ type: 'text', text: 'first' }], + }), + ).resolves.toEqual({ stopReason: 'end_turn' }); + + await expect( + session.prompt({ + sessionId: 'test-session-id', + prompt: [{ type: 'text', text: 'second' }], + model: 'route-b-model', + }), + ).resolves.toEqual({ stopReason: 'end_turn' }); + + expect(mockChat.sendMessageStream).toHaveBeenCalledTimes(2); + expect(mockChat.sendMessageStream).toHaveBeenLastCalledWith( + 'route-b-model', + expect.any(Object), + expect.any(String), + ); + }); + + it('retains a returning route token count after an A-B-A switch (#9529)', async () => { + mockConfig.getSessionTokenLimit = vi.fn().mockReturnValue(100); + let routeIdentity = 'route-a'; + mockConfig.getModelRouteIdentity = vi.fn(() => routeIdentity); + + mockGeminiClient.tryCompressChat.mockResolvedValueOnce({ + originalTokenCount: 50, + newTokenCount: 50, + compressionStatus: core.CompressionStatus.NOOP, + }); + mockChat.sendMessageStream = vi + .fn() + .mockResolvedValueOnce( + createStreamWithChunks([ + { + type: core.StreamEventType.CHUNK, + value: { + usageMetadata: { + totalTokenCount: 101, + promptTokenCount: 101, + }, + }, + }, + ]), + ) + .mockResolvedValueOnce(createEmptyStream()); + + await expect( + session.prompt({ + sessionId: 'test-session-id', + prompt: [{ type: 'text', text: 'first' }], + }), + ).resolves.toEqual({ stopReason: 'end_turn' }); + + routeIdentity = 'route-b'; + mockGeminiClient.tryCompressChat.mockRejectedValueOnce( + new Error('compression rate limited'), + ); + await expect( + session.prompt({ + sessionId: 'test-session-id', + prompt: [{ type: 'text', text: 'second' }], + }), + ).resolves.toEqual({ stopReason: 'end_turn' }); + + routeIdentity = 'route-a'; + mockGeminiClient.tryCompressChat.mockRejectedValueOnce( + new Error('compression rate limited'), + ); + await expect( + session.prompt({ + sessionId: 'test-session-id', + prompt: [{ type: 'text', text: 'third' }], + }), + ).resolves.toEqual({ stopReason: 'max_tokens' }); + + expect(mockChat.sendMessageStream).toHaveBeenCalledTimes(2); + }); + it('returns cancelled when automatic compression is aborted', async () => { mockConfig.getSessionTokenLimit = vi.fn().mockReturnValue(100); mockGeminiClient.tryCompressChat.mockImplementation( diff --git a/packages/cli/src/acp-integration/session/Session.ts b/packages/cli/src/acp-integration/session/Session.ts index ce1f365f53e..0d35292a214 100644 --- a/packages/cli/src/acp-integration/session/Session.ts +++ b/packages/cli/src/acp-integration/session/Session.ts @@ -476,7 +476,11 @@ function maskApiKeyForDisplay(apiKey: string | undefined): string { } type AutoCompressionSendResult = - | { responseStream: AsyncGenerator; stopReason?: never } + | { + responseStream: AsyncGenerator; + requestRouteKey: string; + stopReason?: never; + } | { responseStream: null; stopReason: PromptResponse['stopReason'] }; function getAbortAwareEndTurnStopReason( @@ -1849,6 +1853,7 @@ export class Session implements SessionContext { private cronDisabledByTokenLimit = false; private lastPromptTokenCount = 0; private lastPromptTokenCountChat: GeminiChat | null = null; + private readonly lastPromptTokenCountsByRouteKey = new Map(); // The model route that produced `lastPromptTokenCount` (Config // .getModelRouteIdentity). ACP model switches keep the same GeminiChat, so // the chat-instance check alone never invalidates the count on a route @@ -6357,7 +6362,7 @@ export class Session implements SessionContext { ); if (usageMetadata) { - this.#recordPromptTokenCount(usageMetadata); + this.#recordPromptTokenCount(usageMetadata, sendResult.requestRouteKey); const durationMs = Date.now() - streamStartTime; await this.messageEmitter.emitUsageMetadata( usageMetadata, @@ -6795,14 +6800,19 @@ export class Session implements SessionContext { return { responseStream: null, stopReason: 'cancelled' }; } - if (!compressionInfo) { - this.#syncPromptTokenCountWithCurrentChat(); - } + const model = + options.getModelOverride?.() ?? + options.modelOverride ?? + this.config.getModel(); + const requestRouteKey = await this.#requestRouteKeyForModel(model); + this.#syncPromptTokenCountWithCurrentChat(requestRouteKey); const sessionTokenLimit = this.config.getSessionTokenLimit(); if (sessionTokenLimit > 0) { - const lastPromptTokenCount = - this.#getPostCompressionTokenCount(compressionInfo); + const lastPromptTokenCount = this.#getPostCompressionTokenCount( + compressionInfo, + requestRouteKey, + ); if (lastPromptTokenCount > sessionTokenLimit) { debugLogger.warn( `Session token limit exceeded for prompt ${promptId}: ` + @@ -6847,10 +6857,6 @@ export class Session implements SessionContext { } const chat = this.#getCurrentChat(); - const model = - options.getModelOverride?.() ?? - options.modelOverride ?? - this.config.getModel(); const request = { message, config: { @@ -6861,7 +6867,7 @@ export class Session implements SessionContext { const responseStream = goalPermit ? await chat.sendMessageStream(model, request, promptId, goalPermit) : await chat.sendMessageStream(model, request, promptId); - return { responseStream }; + return { responseStream, requestRouteKey }; } #clearPendingRestoreNotices(): void { @@ -7062,30 +7068,36 @@ export class Session implements SessionContext { } #recordCompressionTokenCount(info: ChatCompressionInfo): void { - this.#syncPromptTokenCountWithCurrentChat(); + const routeKey = this.#currentRouteKey(); + this.#syncPromptTokenCountWithCurrentChat(routeKey); const tokenCount = this.#extractCompressionTokenCount(info); if (tokenCount !== null && tokenCount > 0) { - this.lastPromptTokenCount = tokenCount; + this.#setLastPromptTokenCount(routeKey, tokenCount); } } #recordPromptTokenCount( usageMetadata: GenerateContentResponseUsageMetadata, + routeKey = this.#currentRouteKey(), ): void { - this.#syncPromptTokenCountWithCurrentChat(); + this.#syncPromptTokenCountWithCurrentChat(routeKey); const tokenCount = usageMetadata.promptTokenCount ?? usageMetadata.totalTokenCount; if (tokenCount !== undefined && tokenCount > 0) { - this.lastPromptTokenCount = tokenCount; + this.#setLastPromptTokenCount(routeKey, tokenCount); } } - #getPostCompressionTokenCount(info: ChatCompressionInfo | null): number { + #getPostCompressionTokenCount( + info: ChatCompressionInfo | null, + routeKey = this.#currentRouteKey(), + ): number { const tokenCount = this.#extractCompressionTokenCount(info); if (tokenCount !== null) { return tokenCount; } + this.#syncPromptTokenCountWithCurrentChat(routeKey); return this.lastPromptTokenCount; } @@ -7112,21 +7124,49 @@ export class Session implements SessionContext { return this.config.getModelRouteIdentity?.() ?? ''; } - #syncPromptTokenCountWithCurrentChat(): void { + async #requestRouteKeyForModel(model: string): Promise { + if (!this.config.getModelRouteIdentity) { + return ''; + } + if (!model.endsWith('\0')) { + return this.config.getModelRouteIdentity(model); + } + const runtimeView = await this.config + .getBaseLlmClient() + .resolveForModel(model.slice(0, -1), { failClosed: true }); + return this.config.getModelRouteIdentity( + runtimeView.model, + runtimeView.contentGeneratorConfig, + ); + } + + #setLastPromptTokenCount(routeKey: string, tokenCount: number): void { + this.lastPromptTokenCount = tokenCount; + this.lastPromptTokenCountRouteKey = routeKey; + this.lastPromptTokenCountsByRouteKey.set(routeKey, tokenCount); + } + + #syncPromptTokenCountWithCurrentChat( + routeKey = this.#currentRouteKey(), + ): void { const chat = this.#getCurrentChat(); - const routeKey = this.#currentRouteKey(); const chatChanged = this.lastPromptTokenCountChat && this.lastPromptTokenCountChat !== chat; - // A model switch rebuilds the content generator but keeps the same - // GeminiChat, so the chat-instance check above never fires on a route - // change. A count recorded for the previous route must not anchor the - // session-token-limit gate for the new one — reset it so the gate falls - // back to the fresh send's own count (#9529, follow-up to #9454/#9506). - const routeChanged = - this.lastPromptTokenCountRouteKey !== undefined && - this.lastPromptTokenCountRouteKey !== routeKey; - if (chatChanged || routeChanged) { + if (chatChanged) { + this.lastPromptTokenCountsByRouteKey.clear(); this.lastPromptTokenCount = 0; + } else if (this.lastPromptTokenCountRouteKey !== routeKey) { + if ( + this.lastPromptTokenCountRouteKey !== undefined && + this.lastPromptTokenCount > 0 + ) { + this.lastPromptTokenCountsByRouteKey.set( + this.lastPromptTokenCountRouteKey, + this.lastPromptTokenCount, + ); + } + this.lastPromptTokenCount = + this.lastPromptTokenCountsByRouteKey.get(routeKey) ?? 0; } this.lastPromptTokenCountChat = chat; this.lastPromptTokenCountRouteKey = routeKey; From 8ff5a94b14794142a4ca4a76d8a8e117cc9051d3 Mon Sep 17 00:00:00 2001 From: yiliang114 Date: Mon, 24 Aug 2026 18:26:48 +0800 Subject: [PATCH 03/14] fix(cli): retain token counts for request route --- .../src/acp-integration/session/Session.test.ts | 14 ++++---------- .../cli/src/acp-integration/session/Session.ts | 12 ++++++++++-- 2 files changed, 14 insertions(+), 12 deletions(-) diff --git a/packages/cli/src/acp-integration/session/Session.test.ts b/packages/cli/src/acp-integration/session/Session.test.ts index 044e79b1524..37e7061589a 100644 --- a/packages/cli/src/acp-integration/session/Session.test.ts +++ b/packages/cli/src/acp-integration/session/Session.test.ts @@ -12100,11 +12100,10 @@ describe('Session', () => { expect(mockChat.sendMessageStream).toHaveBeenCalledTimes(1); }); - it('does not drop a modelOverride send using the active route token count (#9529)', async () => { + it('does not drop a send using another route token count (#9529)', async () => { mockConfig.getSessionTokenLimit = vi.fn().mockReturnValue(100); - mockConfig.getModelRouteIdentity = vi.fn((model?: string) => - model === 'route-b-model' ? 'route-b' : 'route-a', - ); + let routeIdentity = 'route-a'; + mockConfig.getModelRouteIdentity = vi.fn(() => routeIdentity); mockGeminiClient.tryCompressChat.mockResolvedValueOnce({ originalTokenCount: 50, @@ -12135,20 +12134,15 @@ describe('Session', () => { }), ).resolves.toEqual({ stopReason: 'end_turn' }); + routeIdentity = 'route-b'; await expect( session.prompt({ sessionId: 'test-session-id', prompt: [{ type: 'text', text: 'second' }], - model: 'route-b-model', }), ).resolves.toEqual({ stopReason: 'end_turn' }); expect(mockChat.sendMessageStream).toHaveBeenCalledTimes(2); - expect(mockChat.sendMessageStream).toHaveBeenLastCalledWith( - 'route-b-model', - expect.any(Object), - expect.any(String), - ); }); it('retains a returning route token count after an A-B-A switch (#9529)', async () => { diff --git a/packages/cli/src/acp-integration/session/Session.ts b/packages/cli/src/acp-integration/session/Session.ts index 0d35292a214..341e7e9034e 100644 --- a/packages/cli/src/acp-integration/session/Session.ts +++ b/packages/cli/src/acp-integration/session/Session.ts @@ -5136,6 +5136,9 @@ export class Session implements SessionContext { | ChannelDeliveryResponseBlock | undefined; let channelDeliveryCheckpoint = 0; + let requestRouteKey = this.config.getModelRouteIdentity( + this.config.getModel(), + ); try { // Set where the model request is actually issued, not at @@ -5181,6 +5184,7 @@ export class Session implements SessionContext { if (restorePostAnswerNoticesAttached) { this.#clearPendingRestoreNotices(); } + requestRouteKey = sendResult.requestRouteKey; const responseStream = sendResult.responseStream; nextMessage = null; channelDeliveryResponseBlock = @@ -5356,7 +5360,7 @@ export class Session implements SessionContext { ); if (usageMetadata) { - this.#recordPromptTokenCount(usageMetadata); + this.#recordPromptTokenCount(usageMetadata, requestRouteKey); // Kick off rewrite in background (non-blocking, runs parallel to tools) if (this.messageRewriter) { this.messageRewriter.flushTurn(pendingSend.signal); @@ -5899,6 +5903,9 @@ export class Session implements SessionContext { let channelDeliveryCheckpoint = 0; let providerSendChat: GeminiChat | undefined; let userContentPushCountBeforeSend = 0; + let requestRouteKey = this.config.getModelRouteIdentity( + this.config.getModel(), + ); try { const sendResult = await this.#sendMessageStreamWithAutoCompression( @@ -6192,6 +6199,7 @@ export class Session implements SessionContext { }; } + requestRouteKey = sendResult.requestRouteKey; const responseStream = sendResult.responseStream; nextMessage = null; channelDeliveryResponseBlock = beginChannelDeliveryResponseBlock( @@ -6362,7 +6370,7 @@ export class Session implements SessionContext { ); if (usageMetadata) { - this.#recordPromptTokenCount(usageMetadata, sendResult.requestRouteKey); + this.#recordPromptTokenCount(usageMetadata, requestRouteKey); const durationMs = Date.now() - streamStartTime; await this.messageEmitter.emitUsageMetadata( usageMetadata, From 97996823fe7c99d73859c1a3a4ec853ae904e67a Mon Sep 17 00:00:00 2001 From: "jinjing.zzj" Date: Mon, 24 Aug 2026 19:51:42 +0800 Subject: [PATCH 04/14] test(acp): pin override-route token recording for the #9529 gate MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add a #9529 regression test that drives a route override through the full-turn vision selector (fullTurnModelOverride): the first override send streams usage metadata over the session token limit, and a second same-override send whose compression falls back to the cache must then resolve max_tokens — proving the first count was recorded under the override route key, not the active route's. Reverting the record site to the default route key makes the test fail. Also make the hoisted requestRouteKey initialization in #executePromptInner and #runStopContinuation use optional chaining (this.config.getModelRouteIdentity?.(...) ?? ''), matching the #currentRouteKey convention for partial Config mocks; the unguarded call threw for every prompt in the ~340 Session tests whose mock config does not define getModelRouteIdentity. --- .../acp-integration/session/Session.test.ts | 81 +++++++++++++++++++ .../src/acp-integration/session/Session.ts | 11 ++- 2 files changed, 86 insertions(+), 6 deletions(-) diff --git a/packages/cli/src/acp-integration/session/Session.test.ts b/packages/cli/src/acp-integration/session/Session.test.ts index 37e7061589a..e8edbb1ce12 100644 --- a/packages/cli/src/acp-integration/session/Session.test.ts +++ b/packages/cli/src/acp-integration/session/Session.test.ts @@ -12145,6 +12145,87 @@ describe('Session', () => { expect(mockChat.sendMessageStream).toHaveBeenCalledTimes(2); }); + it('records an override send usage under the override route so the next same-override send trips the gate (#9529)', async () => { + mockConfig.getSessionTokenLimit = vi.fn().mockReturnValue(100); + mockConfig.getModelRouteIdentity = vi.fn((model?: string) => + model === 'vision-agent' ? 'route-vision' : 'route-primary', + ); + // Full-turn vision selector: an image turn is sent under the \0 exact + // route override — a different route than the active one, driven + // through fullTurnModelOverride, which Session.prompt consumes. + mockConfig.getEffectiveInputModalities = vi.fn().mockReturnValue({}); + mockConfig.getDefaultVisionBridgeModel = vi.fn().mockReturnValue({ + id: 'vision-agent', + baseUrl: 'https://vision.example.com/v1', + agentCapable: true, + }); + const resolveForModel = vi.fn().mockResolvedValue({ + contentGenerator: {}, + contentGeneratorConfig: { model: 'vision-agent' }, + model: 'vision-agent', + }); + mockConfig.getBaseLlmClient = vi.fn().mockReturnValue({ + resolveForModel, + }); + + const visionPrompt: PromptRequest = { + sessionId: 'test-session-id', + prompt: [ + { type: 'text', text: 'look at this' }, + { type: 'image', mimeType: 'image/png', data: 'iVBORw0KGgo=' }, + ], + }; + + // First override send goes out (compression info under the limit) and + // streams usage metadata over the limit; the count must be recorded + // under the override route key, not the active route's. + mockGeminiClient.tryCompressChat.mockResolvedValueOnce({ + originalTokenCount: 50, + newTokenCount: 50, + compressionStatus: core.CompressionStatus.NOOP, + }); + mockChat.sendMessageStream = vi + .fn() + .mockResolvedValueOnce( + createStreamWithChunks([ + { + type: core.StreamEventType.CHUNK, + value: { + usageMetadata: { + totalTokenCount: 101, + promptTokenCount: 101, + }, + }, + }, + ]), + ) + .mockResolvedValueOnce(createEmptyStream()); + + await expect(session.prompt(visionPrompt)).resolves.toEqual({ + stopReason: 'end_turn', + }); + expect(mockChat.sendMessageStream).toHaveBeenCalledWith( + 'vision-agent\0https://vision.example.com/v1\0', + expect.any(Object), + expect.any(String), + ); + + // Second same-override send: compression throws, so the gate falls + // back to the cached count — which must be the first send's 101 + // recorded under the override route (101 > 100 → drop). If the record + // had gone under the active route instead, the override-route cache + // would be empty and this send would wrongly go out. + mockGeminiClient.tryCompressChat.mockRejectedValueOnce( + new Error('compression rate limited'), + ); + + await expect(session.prompt(visionPrompt)).resolves.toEqual({ + stopReason: 'max_tokens', + }); + + expect(mockChat.sendMessageStream).toHaveBeenCalledTimes(1); + }); + it('retains a returning route token count after an A-B-A switch (#9529)', async () => { mockConfig.getSessionTokenLimit = vi.fn().mockReturnValue(100); let routeIdentity = 'route-a'; diff --git a/packages/cli/src/acp-integration/session/Session.ts b/packages/cli/src/acp-integration/session/Session.ts index 341e7e9034e..fe562b843fe 100644 --- a/packages/cli/src/acp-integration/session/Session.ts +++ b/packages/cli/src/acp-integration/session/Session.ts @@ -5136,9 +5136,9 @@ export class Session implements SessionContext { | ChannelDeliveryResponseBlock | undefined; let channelDeliveryCheckpoint = 0; - let requestRouteKey = this.config.getModelRouteIdentity( - this.config.getModel(), - ); + let requestRouteKey = + this.config.getModelRouteIdentity?.(this.config.getModel()) ?? + ''; try { // Set where the model request is actually issued, not at @@ -5903,9 +5903,8 @@ export class Session implements SessionContext { let channelDeliveryCheckpoint = 0; let providerSendChat: GeminiChat | undefined; let userContentPushCountBeforeSend = 0; - let requestRouteKey = this.config.getModelRouteIdentity( - this.config.getModel(), - ); + let requestRouteKey = + this.config.getModelRouteIdentity?.(this.config.getModel()) ?? ''; try { const sendResult = await this.#sendMessageStreamWithAutoCompression( From 8e0db9bdf212a73316e2ffc91fb2c63726cfa5a0 Mon Sep 17 00:00:00 2001 From: yiliang114 Date: Mon, 24 Aug 2026 20:19:34 +0800 Subject: [PATCH 05/14] fix(cli): bound acp route token cache --- .../acp-integration/session/Session.test.ts | 124 +++++++++++++++++- .../src/acp-integration/session/Session.ts | 19 ++- 2 files changed, 140 insertions(+), 3 deletions(-) diff --git a/packages/cli/src/acp-integration/session/Session.test.ts b/packages/cli/src/acp-integration/session/Session.test.ts index e8edbb1ce12..183b9c6d3ce 100644 --- a/packages/cli/src/acp-integration/session/Session.test.ts +++ b/packages/cli/src/acp-integration/session/Session.test.ts @@ -12097,6 +12097,11 @@ describe('Session', () => { }), ).resolves.toEqual({ stopReason: 'max_tokens' }); + expect(debugLoggerWarnSpy).toHaveBeenCalledWith( + expect.stringContaining( + 'requestRoute=route-a, activeModel=qwen3-code-plus', + ), + ); expect(mockChat.sendMessageStream).toHaveBeenCalledTimes(1); }); @@ -12145,6 +12150,109 @@ describe('Session', () => { expect(mockChat.sendMessageStream).toHaveBeenCalledTimes(2); }); + it('does not drop a runtime-scoped override using the active route count (#9529)', async () => { + mockConfig.getSessionTokenLimit = vi.fn().mockReturnValue(100); + mockConfig.getModelRouteIdentity = vi.fn((model?: string) => + model === 'vision-agent' ? 'route-vision' : 'route-primary', + ); + mockConfig.getEffectiveInputModalities = vi.fn().mockReturnValue({}); + mockConfig.getDefaultVisionBridgeModel = vi.fn().mockReturnValue({ + id: 'vision-agent', + baseUrl: 'https://vision.example.com/v1', + agentCapable: true, + }); + const resolveForModel = vi.fn().mockResolvedValue({ + contentGenerator: {}, + contentGeneratorConfig: { model: 'vision-agent' }, + model: 'vision-agent', + }); + mockConfig.getBaseLlmClient = vi.fn().mockReturnValue({ + resolveForModel, + }); + + mockGeminiClient.tryCompressChat + .mockResolvedValueOnce({ + originalTokenCount: 50, + newTokenCount: 50, + compressionStatus: core.CompressionStatus.NOOP, + }) + .mockRejectedValueOnce(new Error('compression rate limited')); + mockChat.sendMessageStream = vi + .fn() + .mockResolvedValueOnce( + createStreamWithChunks([ + { + type: core.StreamEventType.CHUNK, + value: { + usageMetadata: { + totalTokenCount: 101, + promptTokenCount: 101, + }, + }, + }, + ]), + ) + .mockResolvedValueOnce(createEmptyStream()); + + await expect( + session.prompt({ + sessionId: 'test-session-id', + prompt: [{ type: 'text', text: 'primary route' }], + }), + ).resolves.toEqual({ stopReason: 'end_turn' }); + + await expect( + session.prompt({ + sessionId: 'test-session-id', + prompt: [ + { type: 'text', text: 'look at this' }, + { type: 'image', mimeType: 'image/png', data: 'iVBORw0KGgo=' }, + ], + }), + ).resolves.toEqual({ stopReason: 'end_turn' }); + + expect(resolveForModel).toHaveBeenCalledWith( + 'vision-agent\0https://vision.example.com/v1', + { + failClosed: true, + }, + ); + expect(mockChat.sendMessageStream).toHaveBeenCalledTimes(2); + }); + + it('fails closed when runtime-scoped route resolution rejects (#9529)', async () => { + mockConfig.getSessionTokenLimit = vi.fn().mockReturnValue(100); + mockConfig.getModelRouteIdentity = vi.fn((model?: string) => + model === 'vision-agent' ? 'route-vision' : 'route-primary', + ); + mockConfig.getEffectiveInputModalities = vi.fn().mockReturnValue({}); + mockConfig.getDefaultVisionBridgeModel = vi.fn().mockReturnValue({ + id: 'vision-agent', + baseUrl: 'https://vision.example.com/v1', + agentCapable: true, + }); + mockConfig.getBaseLlmClient = vi.fn().mockReturnValue({ + resolveForModel: vi + .fn() + .mockRejectedValue(new Error('runtime unavailable')), + }); + mockChat.sendMessageStream = vi + .fn() + .mockResolvedValue(createEmptyStream()); + + await expect( + session.prompt({ + sessionId: 'test-session-id', + prompt: [ + { type: 'text', text: 'look at this' }, + { type: 'image', mimeType: 'image/png', data: 'iVBORw0KGgo=' }, + ], + }), + ).rejects.toThrow('runtime unavailable'); + + expect(mockChat.sendMessageStream).not.toHaveBeenCalled(); + }); + it('records an override send usage under the override route so the next same-override send trips the gate (#9529)', async () => { mockConfig.getSessionTokenLimit = vi.fn().mockReturnValue(100); mockConfig.getModelRouteIdentity = vi.fn((model?: string) => @@ -12502,7 +12610,7 @@ describe('Session', () => { expect(mockChat.sendMessageStream).toHaveBeenCalledTimes(2); }); - it('resets the session-local token count when the active chat instance changes', async () => { + it('resets session-local route counts when the active chat instance changes', async () => { const clearedChat = { sendMessageStream: vi.fn().mockResolvedValue(createEmptyStream()), addHistory: vi.fn(), @@ -12511,12 +12619,15 @@ describe('Session', () => { getLastModelMessageText: vi.fn().mockReturnValue(''), } as unknown as GeminiChat; mockConfig.getSessionTokenLimit = vi.fn().mockReturnValue(100); + let routeIdentity = 'route-a'; + mockConfig.getModelRouteIdentity = vi.fn(() => routeIdentity); mockGeminiClient.tryCompressChat .mockResolvedValueOnce({ originalTokenCount: 50, newTokenCount: 50, compressionStatus: core.CompressionStatus.NOOP, }) + .mockRejectedValueOnce(new Error('compression unavailable')) .mockRejectedValueOnce(new Error('compression unavailable')); mockChat.sendMessageStream = vi.fn().mockResolvedValueOnce( createStreamWithChunks([ @@ -12540,6 +12651,7 @@ describe('Session', () => { ).resolves.toEqual({ stopReason: 'end_turn' }); mockGeminiClient.getChat.mockReturnValue(clearedChat); + routeIdentity = 'route-b'; await expect( session.prompt({ @@ -12548,7 +12660,15 @@ describe('Session', () => { }), ).resolves.toEqual({ stopReason: 'end_turn' }); - expect(clearedChat.sendMessageStream).toHaveBeenCalledTimes(1); + routeIdentity = 'route-a'; + await expect( + session.prompt({ + sessionId: 'test-session-id', + prompt: [{ type: 'text', text: 'back to route A' }], + }), + ).resolves.toEqual({ stopReason: 'end_turn' }); + + expect(clearedChat.sendMessageStream).toHaveBeenCalledTimes(2); }); it('continues sending when the compression notification fails', async () => { diff --git a/packages/cli/src/acp-integration/session/Session.ts b/packages/cli/src/acp-integration/session/Session.ts index fe562b843fe..68222322ae2 100644 --- a/packages/cli/src/acp-integration/session/Session.ts +++ b/packages/cli/src/acp-integration/session/Session.ts @@ -366,6 +366,7 @@ const permissionRequestTails = new WeakMap< AgentSideConnection, Promise >(); +const MAX_RETAINED_SESSION_ROUTE_COUNTS = 8; const USER_CANCEL_ABORT_REASON = 'qwen:user-cancel'; const NEW_PROMPT_ABORT_REASON = 'qwen:new-prompt'; const SESSION_DISPOSE_ABORT_REASON = 'qwen:session-dispose'; @@ -1853,6 +1854,8 @@ export class Session implements SessionContext { private cronDisabledByTokenLimit = false; private lastPromptTokenCount = 0; private lastPromptTokenCountChat: GeminiChat | null = null; + // Private ACP fallback cache, bounded like GeminiChat without exposing a + // cross-package route resolver just for this closeout. private readonly lastPromptTokenCountsByRouteKey = new Map(); // The model route that produced `lastPromptTokenCount` (Config // .getModelRouteIdentity). ACP model switches keep the same GeminiChat, so @@ -6823,7 +6826,9 @@ export class Session implements SessionContext { if (lastPromptTokenCount > sessionTokenLimit) { debugLogger.warn( `Session token limit exceeded for prompt ${promptId}: ` + - `${lastPromptTokenCount} > ${sessionTokenLimit}. Send dropped.`, + `${lastPromptTokenCount} > ${sessionTokenLimit}. ` + + `requestRoute=${requestRouteKey}, activeModel=${this.config.getModel()}. ` + + 'Send dropped.', ); await this.#emitAgentDiagnosticMessageSafely( `Session token limit exceeded: ${lastPromptTokenCount} tokens > ${sessionTokenLimit} limit. ` + @@ -7150,6 +7155,18 @@ export class Session implements SessionContext { #setLastPromptTokenCount(routeKey: string, tokenCount: number): void { this.lastPromptTokenCount = tokenCount; this.lastPromptTokenCountRouteKey = routeKey; + if ( + !this.lastPromptTokenCountsByRouteKey.has(routeKey) && + this.lastPromptTokenCountsByRouteKey.size >= + MAX_RETAINED_SESSION_ROUTE_COUNTS + ) { + const oldestKey = this.lastPromptTokenCountsByRouteKey + .keys() + .next().value; + if (oldestKey !== undefined) { + this.lastPromptTokenCountsByRouteKey.delete(oldestKey); + } + } this.lastPromptTokenCountsByRouteKey.set(routeKey, tokenCount); } From c30055ab136b97a090edeb92c348576bca4b006b Mon Sep 17 00:00:00 2001 From: yiliang114 Date: Tue, 25 Aug 2026 00:22:15 +0800 Subject: [PATCH 06/14] fix(cli): cover acp route token eviction --- .../acp-integration/session/Session.test.ts | 47 +++++++++++++++++++ .../src/acp-integration/session/Session.ts | 7 +-- 2 files changed, 49 insertions(+), 5 deletions(-) diff --git a/packages/cli/src/acp-integration/session/Session.test.ts b/packages/cli/src/acp-integration/session/Session.test.ts index 183b9c6d3ce..64f47d0eae8 100644 --- a/packages/cli/src/acp-integration/session/Session.test.ts +++ b/packages/cli/src/acp-integration/session/Session.test.ts @@ -12393,6 +12393,53 @@ describe('Session', () => { expect(mockChat.sendMessageStream).toHaveBeenCalledTimes(2); }); + it('evicts the oldest retained route token count after eight routes (#9529)', async () => { + mockConfig.getSessionTokenLimit = vi.fn().mockReturnValue(100); + let routeIdentity = 'route-0'; + mockConfig.getModelRouteIdentity = vi.fn(() => routeIdentity); + mockChat.sendMessageStream = vi.fn().mockImplementation(() => + createStreamWithChunks([ + { + type: core.StreamEventType.CHUNK, + value: { + usageMetadata: { + totalTokenCount: 101, + promptTokenCount: 101, + }, + }, + }, + ]), + ); + + for (let i = 0; i < 9; i++) { + routeIdentity = `route-${i}`; + mockGeminiClient.tryCompressChat.mockResolvedValueOnce({ + originalTokenCount: 50, + newTokenCount: 50, + compressionStatus: core.CompressionStatus.NOOP, + }); + await expect( + session.prompt({ + sessionId: 'test-session-id', + prompt: [{ type: 'text', text: `route ${i}` }], + }), + ).resolves.toEqual({ stopReason: 'end_turn' }); + } + + routeIdentity = 'route-0'; + mockGeminiClient.tryCompressChat.mockRejectedValueOnce( + new Error('compression rate limited'), + ); + await expect( + session.prompt({ + sessionId: 'test-session-id', + prompt: [{ type: 'text', text: 'return to evicted route' }], + }), + ).resolves.toEqual({ stopReason: 'end_turn' }); + + expect(mockChat.sendMessageStream).toHaveBeenCalledTimes(10); + }); + it('returns cancelled when automatic compression is aborted', async () => { mockConfig.getSessionTokenLimit = vi.fn().mockReturnValue(100); mockGeminiClient.tryCompressChat.mockImplementation( diff --git a/packages/cli/src/acp-integration/session/Session.ts b/packages/cli/src/acp-integration/session/Session.ts index 68222322ae2..a525d1e8f0d 100644 --- a/packages/cli/src/acp-integration/session/Session.ts +++ b/packages/cli/src/acp-integration/session/Session.ts @@ -5139,9 +5139,7 @@ export class Session implements SessionContext { | ChannelDeliveryResponseBlock | undefined; let channelDeliveryCheckpoint = 0; - let requestRouteKey = - this.config.getModelRouteIdentity?.(this.config.getModel()) ?? - ''; + let requestRouteKey = ''; try { // Set where the model request is actually issued, not at @@ -5906,8 +5904,7 @@ export class Session implements SessionContext { let channelDeliveryCheckpoint = 0; let providerSendChat: GeminiChat | undefined; let userContentPushCountBeforeSend = 0; - let requestRouteKey = - this.config.getModelRouteIdentity?.(this.config.getModel()) ?? ''; + let requestRouteKey = ''; try { const sendResult = await this.#sendMessageStreamWithAutoCompression( From d32bcd31e6475b376b3af7d215be70d6b03bc28d Mon Sep 17 00:00:00 2001 From: "jinjing.zzj" Date: Tue, 25 Aug 2026 00:24:35 +0800 Subject: [PATCH 07/14] fix(acp): drop dead request route key initializers in the send loops The hoisted requestRouteKey initializer in #executePromptInner and #runStopContinuation computed a route identity that was discarded on every turn: the null-stream paths return before any record site, and every path that reaches a record site first assigns requestRouteKey from the send result. Replace both with a plain empty initializer. --- packages/cli/src/acp-integration/session/Session.ts | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/packages/cli/src/acp-integration/session/Session.ts b/packages/cli/src/acp-integration/session/Session.ts index a525d1e8f0d..e01fe175a9c 100644 --- a/packages/cli/src/acp-integration/session/Session.ts +++ b/packages/cli/src/acp-integration/session/Session.ts @@ -5139,6 +5139,9 @@ export class Session implements SessionContext { | ChannelDeliveryResponseBlock | undefined; let channelDeliveryCheckpoint = 0; + // The send result assigns this before any read; null-stream + // paths return before the record site, so a pre-send route + // computation here would only be discarded. let requestRouteKey = ''; try { @@ -5904,6 +5907,9 @@ export class Session implements SessionContext { let channelDeliveryCheckpoint = 0; let providerSendChat: GeminiChat | undefined; let userContentPushCountBeforeSend = 0; + // The send result assigns this before any read; null-stream paths + // return before the record site, so a pre-send route computation here + // would only be discarded. let requestRouteKey = ''; try { From 800912734b96fa94af244355e859e0e0f4db2c19 Mon Sep 17 00:00:00 2001 From: "jinjing.zzj" Date: Tue, 25 Aug 2026 00:26:08 +0800 Subject: [PATCH 08/14] test(acp): factor the #9529 over-limit usage stream setup into a helper The ~25-line mock setup that streams a 101-token usage metadata chunk on the first send and an empty stream on the second was pasted verbatim in eight #9529 session-token-limit tests. Extract it into createOverLimitUsageSendStream next to the existing stream helpers and migrate all eight copies. --- .../acp-integration/session/Session.test.ts | 161 ++++-------------- 1 file changed, 33 insertions(+), 128 deletions(-) diff --git a/packages/cli/src/acp-integration/session/Session.test.ts b/packages/cli/src/acp-integration/session/Session.test.ts index 64f47d0eae8..3026f4fa28f 100644 --- a/packages/cli/src/acp-integration/session/Session.test.ts +++ b/packages/cli/src/acp-integration/session/Session.test.ts @@ -378,6 +378,31 @@ function createStreamWithChunks( })(); } +/** + * Builds a sendMessageStream mock for the #9529 session-token-limit tests: + * the first send streams usage metadata over the 100-token limit those tests + * configure (so the count lands in the session's route-scoped cache), and + * the second send returns an empty stream. + */ +function createOverLimitUsageSendStream() { + return vi + .fn() + .mockResolvedValueOnce( + createStreamWithChunks([ + { + type: core.StreamEventType.CHUNK, + value: { + usageMetadata: { + totalTokenCount: 101, + promptTokenCount: 101, + }, + }, + }, + ]), + ) + .mockResolvedValueOnce(createEmptyStream()); +} + /** Builds provider preparation metadata that arrives before complete arguments. */ function createPreparationResponse( callId: string, @@ -12004,22 +12029,7 @@ describe('Session', () => { newTokenCount: 50, compressionStatus: core.CompressionStatus.NOOP, }); - mockChat.sendMessageStream = vi - .fn() - .mockResolvedValueOnce( - createStreamWithChunks([ - { - type: core.StreamEventType.CHUNK, - value: { - usageMetadata: { - totalTokenCount: 101, - promptTokenCount: 101, - }, - }, - }, - ]), - ) - .mockResolvedValueOnce(createEmptyStream()); + mockChat.sendMessageStream = createOverLimitUsageSendStream(); await expect( session.prompt({ @@ -12060,22 +12070,7 @@ describe('Session', () => { newTokenCount: 50, compressionStatus: core.CompressionStatus.NOOP, }); - mockChat.sendMessageStream = vi - .fn() - .mockResolvedValueOnce( - createStreamWithChunks([ - { - type: core.StreamEventType.CHUNK, - value: { - usageMetadata: { - totalTokenCount: 101, - promptTokenCount: 101, - }, - }, - }, - ]), - ) - .mockResolvedValueOnce(createEmptyStream()); + mockChat.sendMessageStream = createOverLimitUsageSendStream(); await expect( session.prompt({ @@ -12115,22 +12110,7 @@ describe('Session', () => { newTokenCount: 50, compressionStatus: core.CompressionStatus.NOOP, }); - mockChat.sendMessageStream = vi - .fn() - .mockResolvedValueOnce( - createStreamWithChunks([ - { - type: core.StreamEventType.CHUNK, - value: { - usageMetadata: { - totalTokenCount: 101, - promptTokenCount: 101, - }, - }, - }, - ]), - ) - .mockResolvedValueOnce(createEmptyStream()); + mockChat.sendMessageStream = createOverLimitUsageSendStream(); await expect( session.prompt({ @@ -12177,22 +12157,7 @@ describe('Session', () => { compressionStatus: core.CompressionStatus.NOOP, }) .mockRejectedValueOnce(new Error('compression rate limited')); - mockChat.sendMessageStream = vi - .fn() - .mockResolvedValueOnce( - createStreamWithChunks([ - { - type: core.StreamEventType.CHUNK, - value: { - usageMetadata: { - totalTokenCount: 101, - promptTokenCount: 101, - }, - }, - }, - ]), - ) - .mockResolvedValueOnce(createEmptyStream()); + mockChat.sendMessageStream = createOverLimitUsageSendStream(); await expect( session.prompt({ @@ -12292,22 +12257,7 @@ describe('Session', () => { newTokenCount: 50, compressionStatus: core.CompressionStatus.NOOP, }); - mockChat.sendMessageStream = vi - .fn() - .mockResolvedValueOnce( - createStreamWithChunks([ - { - type: core.StreamEventType.CHUNK, - value: { - usageMetadata: { - totalTokenCount: 101, - promptTokenCount: 101, - }, - }, - }, - ]), - ) - .mockResolvedValueOnce(createEmptyStream()); + mockChat.sendMessageStream = createOverLimitUsageSendStream(); await expect(session.prompt(visionPrompt)).resolves.toEqual({ stopReason: 'end_turn', @@ -12344,22 +12294,7 @@ describe('Session', () => { newTokenCount: 50, compressionStatus: core.CompressionStatus.NOOP, }); - mockChat.sendMessageStream = vi - .fn() - .mockResolvedValueOnce( - createStreamWithChunks([ - { - type: core.StreamEventType.CHUNK, - value: { - usageMetadata: { - totalTokenCount: 101, - promptTokenCount: 101, - }, - }, - }, - ]), - ) - .mockResolvedValueOnce(createEmptyStream()); + mockChat.sendMessageStream = createOverLimitUsageSendStream(); await expect( session.prompt({ @@ -12536,22 +12471,7 @@ describe('Session', () => { newTokenCount: 0, compressionStatus: core.CompressionStatus.NOOP, }); - mockChat.sendMessageStream = vi - .fn() - .mockResolvedValueOnce( - createStreamWithChunks([ - { - type: core.StreamEventType.CHUNK, - value: { - usageMetadata: { - totalTokenCount: 101, - promptTokenCount: 101, - }, - }, - }, - ]), - ) - .mockResolvedValueOnce(createEmptyStream()); + mockChat.sendMessageStream = createOverLimitUsageSendStream(); await expect( session.prompt({ @@ -12582,22 +12502,7 @@ describe('Session', () => { newTokenCount: 0, compressionStatus: core.CompressionStatus.COMPRESSED, }); - mockChat.sendMessageStream = vi - .fn() - .mockResolvedValueOnce( - createStreamWithChunks([ - { - type: core.StreamEventType.CHUNK, - value: { - usageMetadata: { - totalTokenCount: 101, - promptTokenCount: 101, - }, - }, - }, - ]), - ) - .mockResolvedValueOnce(createEmptyStream()); + mockChat.sendMessageStream = createOverLimitUsageSendStream(); await expect( session.prompt({ From 3579e3ae5a4c2fda409547ceadd3ce7bdbe8854e Mon Sep 17 00:00:00 2001 From: "jinjing.zzj" Date: Tue, 25 Aug 2026 00:26:50 +0800 Subject: [PATCH 09/14] test(acp): cover route count eviction in the session token cache The evict-oldest branch in #setLastPromptTokenCount had no coverage: existing tests exercise at most three route identities, so deleting the eviction block, flipping the size comparison, or evicting the newest entry all survived silently. Drive nine distinct route identities (one past MAX_RETAINED_SESSION_ROUTE_COUNTS) through session.prompt, then assert the evicted oldest route reads back no cached count (its send goes out) while a retained route still trips the gate. --- .../acp-integration/session/Session.test.ts | 59 +++++++++++++------ 1 file changed, 42 insertions(+), 17 deletions(-) diff --git a/packages/cli/src/acp-integration/session/Session.test.ts b/packages/cli/src/acp-integration/session/Session.test.ts index 3026f4fa28f..5ee09b3c8bc 100644 --- a/packages/cli/src/acp-integration/session/Session.test.ts +++ b/packages/cli/src/acp-integration/session/Session.test.ts @@ -12328,25 +12328,34 @@ describe('Session', () => { expect(mockChat.sendMessageStream).toHaveBeenCalledTimes(2); }); - it('evicts the oldest retained route token count after eight routes (#9529)', async () => { + it('evicts the oldest route count once the retained-route budget is exhausted (#9529)', async () => { mockConfig.getSessionTokenLimit = vi.fn().mockReturnValue(100); let routeIdentity = 'route-0'; mockConfig.getModelRouteIdentity = vi.fn(() => routeIdentity); - mockChat.sendMessageStream = vi.fn().mockImplementation(() => - createStreamWithChunks([ - { - type: core.StreamEventType.CHUNK, - value: { - usageMetadata: { - totalTokenCount: 101, - promptTokenCount: 101, + + // Record an over-limit count under nine distinct route keys on the + // same chat instance — one more than the retained-route budget, so + // the oldest entry must be evicted instead of growing unbounded. + const sendStreamMock = vi.fn(); + for (let i = 1; i <= 9; i += 1) { + sendStreamMock.mockResolvedValueOnce( + createStreamWithChunks([ + { + type: core.StreamEventType.CHUNK, + value: { + usageMetadata: { + totalTokenCount: 101, + promptTokenCount: 101, + }, }, }, - }, - ]), - ); + ]), + ); + } + sendStreamMock.mockResolvedValueOnce(createEmptyStream()); + mockChat.sendMessageStream = sendStreamMock; - for (let i = 0; i < 9; i++) { + for (let i = 1; i <= 9; i += 1) { routeIdentity = `route-${i}`; mockGeminiClient.tryCompressChat.mockResolvedValueOnce({ originalTokenCount: 50, @@ -12356,23 +12365,39 @@ describe('Session', () => { await expect( session.prompt({ sessionId: 'test-session-id', - prompt: [{ type: 'text', text: `route ${i}` }], + prompt: [{ type: 'text', text: `prompt ${i}` }], }), ).resolves.toEqual({ stopReason: 'end_turn' }); } - routeIdentity = 'route-0'; + // The evicted oldest route (route-1) reads back no cached count, so + // a send on it must go out instead of tripping the gate. + routeIdentity = 'route-1'; mockGeminiClient.tryCompressChat.mockRejectedValueOnce( new Error('compression rate limited'), ); await expect( session.prompt({ sessionId: 'test-session-id', - prompt: [{ type: 'text', text: 'return to evicted route' }], + prompt: [{ type: 'text', text: 'evicted route' }], }), ).resolves.toEqual({ stopReason: 'end_turn' }); - expect(mockChat.sendMessageStream).toHaveBeenCalledTimes(10); + // A retained route (route-2) still trips the gate from the cache. + routeIdentity = 'route-2'; + mockGeminiClient.tryCompressChat.mockRejectedValueOnce( + new Error('compression rate limited'), + ); + await expect( + session.prompt({ + sessionId: 'test-session-id', + prompt: [{ type: 'text', text: 'retained route' }], + }), + ).resolves.toEqual({ stopReason: 'max_tokens' }); + + // Nine recording sends plus the evicted-route send; the + // retained-route send was dropped by the gate. + expect(sendStreamMock).toHaveBeenCalledTimes(10); }); it('returns cancelled when automatic compression is aborted', async () => { From 79b30a75ccf820663de2f0cefc97b903fad3b3b7 Mon Sep 17 00:00:00 2001 From: "jinjing.zzj" Date: Tue, 25 Aug 2026 00:27:51 +0800 Subject: [PATCH 10/14] test(acp): pin the stop-continuation token record route scope MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The existing override-route recording test only exercises the primary prompt record site in #executePromptInner; the only Stop-hook gate test drops the continuation send before streaming and never mocks getModelRouteIdentity, so the #runStopContinuation record site was unpinned. Drive a Stop-hook continuation whose send streams over-limit usage under a \0 exact-route override, then assert a second same-override send trips the gate from the cached count — reverting the continuation record site to the default route key makes the test fail. --- .../acp-integration/session/Session.test.ts | 133 ++++++++++++++++++ 1 file changed, 133 insertions(+) diff --git a/packages/cli/src/acp-integration/session/Session.test.ts b/packages/cli/src/acp-integration/session/Session.test.ts index 5ee09b3c8bc..2934f76bf0f 100644 --- a/packages/cli/src/acp-integration/session/Session.test.ts +++ b/packages/cli/src/acp-integration/session/Session.test.ts @@ -12284,6 +12284,139 @@ describe('Session', () => { expect(mockChat.sendMessageStream).toHaveBeenCalledTimes(1); }); + it('records a Stop-hook continuation usage under the continuation request route so the next same-route send trips the gate (#9529)', async () => { + mockConfig.getSessionTokenLimit = vi.fn().mockReturnValue(100); + mockConfig.getModelRouteIdentity = vi.fn((model?: string) => + model === 'vision-agent' ? 'route-vision' : 'route-primary', + ); + mockConfig.getEffectiveInputModalities = vi.fn().mockReturnValue({}); + mockConfig.getDefaultVisionBridgeModel = vi.fn().mockReturnValue({ + id: 'vision-agent', + baseUrl: 'https://vision.example.com/v1', + agentCapable: true, + }); + const resolveForModel = vi.fn().mockResolvedValue({ + contentGenerator: {}, + contentGeneratorConfig: { model: 'vision-agent' }, + model: 'vision-agent', + }); + mockConfig.getBaseLlmClient = vi.fn().mockReturnValue({ + resolveForModel, + }); + // The Stop hook blocks the first end-turn so the over-limit usage + // arrives on the continuation send inside #runStopContinuation; the + // second hook call lets the turn finish. + const messageBus = { + request: vi + .fn() + .mockResolvedValueOnce({ + success: true, + output: { + decision: 'block', + reason: 'Continue after Stop hook', + }, + }) + .mockResolvedValueOnce({ + success: true, + output: {}, + }), + }; + mockConfig.getMessageBus = vi.fn().mockReturnValue(messageBus); + mockConfig.getDisableAllHooks = vi.fn().mockReturnValue(false); + mockConfig.hasHooksForEvent = vi + .fn() + .mockImplementation((eventName: string) => eventName === 'Stop'); + mockChat.getHistory = vi + .fn() + .mockReturnValue([ + { role: 'model', parts: [{ text: 'response text' }] }, + ]); + mockChat.getLastModelMessageText = vi + .fn() + .mockReturnValue('response text'); + + const visionPrompt: PromptRequest = { + sessionId: 'test-session-id', + prompt: [ + { type: 'text', text: 'look at this' }, + { type: 'image', mimeType: 'image/png', data: 'iVBORw0KGgo=' }, + ], + }; + + // The primary override send completes without usage metadata; the + // continuation send then streams the over-limit usage, which must be + // recorded under the continuation's request route (the override + // route), not the active route's. + mockGeminiClient.tryCompressChat + .mockResolvedValueOnce({ + originalTokenCount: 50, + newTokenCount: 50, + compressionStatus: core.CompressionStatus.NOOP, + }) + .mockResolvedValueOnce({ + originalTokenCount: 50, + newTokenCount: 50, + compressionStatus: core.CompressionStatus.NOOP, + }); + mockChat.sendMessageStream = vi + .fn() + .mockResolvedValueOnce( + createStreamWithChunks([ + { + type: core.StreamEventType.CHUNK, + value: { + candidates: [ + { + content: { parts: [{ text: 'response text' }] }, + finishReason: 'STOP', + }, + ], + }, + }, + ]), + ) + .mockResolvedValueOnce( + createStreamWithChunks([ + { + type: core.StreamEventType.CHUNK, + value: { + usageMetadata: { + totalTokenCount: 101, + promptTokenCount: 101, + }, + }, + }, + ]), + ); + + await expect(session.prompt(visionPrompt)).resolves.toEqual({ + stopReason: 'end_turn', + }); + // Both sends went out under the \0 exact-route override. + expect(mockChat.sendMessageStream).toHaveBeenCalledWith( + 'vision-agent\0https://vision.example.com/v1\0', + expect.any(Object), + expect.any(String), + ); + expect(mockChat.sendMessageStream).toHaveBeenCalledTimes(2); + + // Next same-override send: compression throws, so the gate falls + // back to the cached count — which must be the continuation send's + // 101 recorded under the override route (101 > 100 → drop). If the + // continuation record had gone under the active route instead, the + // override-route cache would be empty and this send would wrongly + // go out. + mockGeminiClient.tryCompressChat.mockRejectedValueOnce( + new Error('compression rate limited'), + ); + + await expect(session.prompt(visionPrompt)).resolves.toEqual({ + stopReason: 'max_tokens', + }); + + expect(mockChat.sendMessageStream).toHaveBeenCalledTimes(2); + }); + it('retains a returning route token count after an A-B-A switch (#9529)', async () => { mockConfig.getSessionTokenLimit = vi.fn().mockReturnValue(100); let routeIdentity = 'route-a'; From f619920b53f2b372f6e31aa2508288d335fc0837 Mon Sep 17 00:00:00 2001 From: "jinjing.zzj" Date: Tue, 25 Aug 2026 04:57:00 +0800 Subject: [PATCH 11/14] test(acp): factor the #9529 vision-override mock setup into a helper --- .../acp-integration/session/Session.test.ts | 109 +++++++----------- 1 file changed, 41 insertions(+), 68 deletions(-) diff --git a/packages/cli/src/acp-integration/session/Session.test.ts b/packages/cli/src/acp-integration/session/Session.test.ts index 2934f76bf0f..e08e1204e45 100644 --- a/packages/cli/src/acp-integration/session/Session.test.ts +++ b/packages/cli/src/acp-integration/session/Session.test.ts @@ -403,6 +403,41 @@ function createOverLimitUsageSendStream() { .mockResolvedValueOnce(createEmptyStream()); } +/** + * Installs the shared vision-override mock surface for the #9529 + * override-route tests: the 100-token session limit, the vision/primary + * route identity discriminator, the modality and vision-bridge selectors, + * and a base LLM client whose `resolveForModel` resolves to the vision + * agent. Returns the `resolveForModel` mock so a test can swap in a + * rejecting one (the fail-closed path) or assert on its calls. + */ +function setupVisionRouteOverrideMocks( + mockConfig: Config, + resolveForModel?: ReturnType, +): ReturnType { + const resolver = + resolveForModel ?? + vi.fn().mockResolvedValue({ + contentGenerator: {}, + contentGeneratorConfig: { model: 'vision-agent' }, + model: 'vision-agent', + }); + mockConfig.getSessionTokenLimit = vi.fn().mockReturnValue(100); + mockConfig.getModelRouteIdentity = vi.fn((model?: string) => + model === 'vision-agent' ? 'route-vision' : 'route-primary', + ); + mockConfig.getEffectiveInputModalities = vi.fn().mockReturnValue({}); + mockConfig.getDefaultVisionBridgeModel = vi.fn().mockReturnValue({ + id: 'vision-agent', + baseUrl: 'https://vision.example.com/v1', + agentCapable: true, + }); + mockConfig.getBaseLlmClient = vi.fn().mockReturnValue({ + resolveForModel: resolver, + }); + return resolver; +} + /** Builds provider preparation metadata that arrives before complete arguments. */ function createPreparationResponse( callId: string, @@ -12131,24 +12166,7 @@ describe('Session', () => { }); it('does not drop a runtime-scoped override using the active route count (#9529)', async () => { - mockConfig.getSessionTokenLimit = vi.fn().mockReturnValue(100); - mockConfig.getModelRouteIdentity = vi.fn((model?: string) => - model === 'vision-agent' ? 'route-vision' : 'route-primary', - ); - mockConfig.getEffectiveInputModalities = vi.fn().mockReturnValue({}); - mockConfig.getDefaultVisionBridgeModel = vi.fn().mockReturnValue({ - id: 'vision-agent', - baseUrl: 'https://vision.example.com/v1', - agentCapable: true, - }); - const resolveForModel = vi.fn().mockResolvedValue({ - contentGenerator: {}, - contentGeneratorConfig: { model: 'vision-agent' }, - model: 'vision-agent', - }); - mockConfig.getBaseLlmClient = vi.fn().mockReturnValue({ - resolveForModel, - }); + const resolveForModel = setupVisionRouteOverrideMocks(mockConfig); mockGeminiClient.tryCompressChat .mockResolvedValueOnce({ @@ -12186,21 +12204,10 @@ describe('Session', () => { }); it('fails closed when runtime-scoped route resolution rejects (#9529)', async () => { - mockConfig.getSessionTokenLimit = vi.fn().mockReturnValue(100); - mockConfig.getModelRouteIdentity = vi.fn((model?: string) => - model === 'vision-agent' ? 'route-vision' : 'route-primary', + setupVisionRouteOverrideMocks( + mockConfig, + vi.fn().mockRejectedValue(new Error('runtime unavailable')), ); - mockConfig.getEffectiveInputModalities = vi.fn().mockReturnValue({}); - mockConfig.getDefaultVisionBridgeModel = vi.fn().mockReturnValue({ - id: 'vision-agent', - baseUrl: 'https://vision.example.com/v1', - agentCapable: true, - }); - mockConfig.getBaseLlmClient = vi.fn().mockReturnValue({ - resolveForModel: vi - .fn() - .mockRejectedValue(new Error('runtime unavailable')), - }); mockChat.sendMessageStream = vi .fn() .mockResolvedValue(createEmptyStream()); @@ -12219,27 +12226,10 @@ describe('Session', () => { }); it('records an override send usage under the override route so the next same-override send trips the gate (#9529)', async () => { - mockConfig.getSessionTokenLimit = vi.fn().mockReturnValue(100); - mockConfig.getModelRouteIdentity = vi.fn((model?: string) => - model === 'vision-agent' ? 'route-vision' : 'route-primary', - ); // Full-turn vision selector: an image turn is sent under the \0 exact // route override — a different route than the active one, driven // through fullTurnModelOverride, which Session.prompt consumes. - mockConfig.getEffectiveInputModalities = vi.fn().mockReturnValue({}); - mockConfig.getDefaultVisionBridgeModel = vi.fn().mockReturnValue({ - id: 'vision-agent', - baseUrl: 'https://vision.example.com/v1', - agentCapable: true, - }); - const resolveForModel = vi.fn().mockResolvedValue({ - contentGenerator: {}, - contentGeneratorConfig: { model: 'vision-agent' }, - model: 'vision-agent', - }); - mockConfig.getBaseLlmClient = vi.fn().mockReturnValue({ - resolveForModel, - }); + setupVisionRouteOverrideMocks(mockConfig); const visionPrompt: PromptRequest = { sessionId: 'test-session-id', @@ -12285,24 +12275,7 @@ describe('Session', () => { }); it('records a Stop-hook continuation usage under the continuation request route so the next same-route send trips the gate (#9529)', async () => { - mockConfig.getSessionTokenLimit = vi.fn().mockReturnValue(100); - mockConfig.getModelRouteIdentity = vi.fn((model?: string) => - model === 'vision-agent' ? 'route-vision' : 'route-primary', - ); - mockConfig.getEffectiveInputModalities = vi.fn().mockReturnValue({}); - mockConfig.getDefaultVisionBridgeModel = vi.fn().mockReturnValue({ - id: 'vision-agent', - baseUrl: 'https://vision.example.com/v1', - agentCapable: true, - }); - const resolveForModel = vi.fn().mockResolvedValue({ - contentGenerator: {}, - contentGeneratorConfig: { model: 'vision-agent' }, - model: 'vision-agent', - }); - mockConfig.getBaseLlmClient = vi.fn().mockReturnValue({ - resolveForModel, - }); + setupVisionRouteOverrideMocks(mockConfig); // The Stop hook blocks the first end-turn so the over-limit usage // arrives on the continuation send inside #runStopContinuation; the // second hook call lets the turn finish. From 8836460e9a6b9447d36b2eef9765f8df10f92710 Mon Sep 17 00:00:00 2001 From: "jinjing.zzj" Date: Tue, 25 Aug 2026 08:49:45 +0800 Subject: [PATCH 12/14] fix(acp): invalidate the session token cache on every compression rewrite MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The route-keyed fallback cache was only cleared on a chat-instance change and re-stamped by the pre-send compression hook, so compressions inside GeminiChat.sendMessageStream (hard-tier rescue, reactive overflow — surfaced as StreamEventType.COMPRESSED, which the session loops ignored) left it holding pre-compression counts sized against destroyed history. A returning route's send could then be false-dropped with 'Session token limit exceeded' when tryCompressChat failed. Handle StreamEventType.COMPRESSED in all four session send loops and clear every retained route count on any COMPRESSED result (pre-send or in-send), re-stamping the fresh count under the request route and the active route when they differ — mirroring GeminiChat clearing its keyed counts in the COMPRESSED branch of tryCompress. Move the pre-send record after request-route resolution so the invalidation keys correctly. Update the zero-newTokenCount COMPRESSED test to pin the corrected semantics: after a successful rewrite the pre-compression count must not gate the send (owner-side parity). --- .../acp-integration/session/Session.test.ts | 107 +++++++++++++++++- .../src/acp-integration/session/Session.ts | 88 +++++++++++++- 2 files changed, 184 insertions(+), 11 deletions(-) diff --git a/packages/cli/src/acp-integration/session/Session.test.ts b/packages/cli/src/acp-integration/session/Session.test.ts index e08e1204e45..17471684e1a 100644 --- a/packages/cli/src/acp-integration/session/Session.test.ts +++ b/packages/cli/src/acp-integration/session/Session.test.ts @@ -367,9 +367,10 @@ function setFakeHome(home: string): () => void { }; } -// Helper to create async generator with chunks (avoids memory leak) +// Helper to create async generator with chunks (avoids memory leak). +// COMPRESSED events carry `info` instead of `value`. function createStreamWithChunks( - chunks: Array<{ type: unknown; value: unknown }>, + chunks: Array<{ type: unknown; value?: unknown; info?: unknown }>, ) { return (async function* () { for (const chunk of chunks) { @@ -12135,6 +12136,97 @@ describe('Session', () => { expect(mockChat.sendMessageStream).toHaveBeenCalledTimes(1); }); + it('does not drop a returning route send on a stale count after in-send compression rewrote the history (#9529)', async () => { + mockConfig.getSessionTokenLimit = vi.fn().mockReturnValue(100); + let routeIdentity = 'route-a'; + mockConfig.getModelRouteIdentity = vi.fn(() => routeIdentity); + + // Prompt 1 on route A: an API-reported over-limit count (101) lands + // in the session's route-scoped cache. + mockGeminiClient.tryCompressChat.mockResolvedValueOnce({ + originalTokenCount: 50, + newTokenCount: 50, + compressionStatus: core.CompressionStatus.NOOP, + }); + mockChat.sendMessageStream = vi + .fn() + .mockResolvedValueOnce( + createStreamWithChunks([ + { + type: core.StreamEventType.CHUNK, + value: { + usageMetadata: { + totalTokenCount: 101, + promptTokenCount: 101, + }, + }, + }, + ]), + ) + // Prompt 2 on route B trips an in-send compression inside + // GeminiChat.sendMessageStream (hard-tier rescue / reactive + // overflow). GeminiChat clears its own keyed counts when it + // surfaces the COMPRESSED event; the session's fallback cache + // must be invalidated at the same point, or route A's stale + // pre-compression count survives the rewrite. + .mockResolvedValueOnce( + createStreamWithChunks([ + { + type: core.StreamEventType.COMPRESSED, + info: { + originalTokenCount: 101, + newTokenCount: 40, + compressionStatus: core.CompressionStatus.COMPRESSED, + }, + }, + ]), + ) + // Prompt 3 back on route A: an empty stream. + .mockResolvedValueOnce(createEmptyStream()); + + await expect( + session.prompt({ + sessionId: 'test-session-id', + prompt: [{ type: 'text', text: 'first' }], + }), + ).resolves.toEqual({ stopReason: 'end_turn' }); + + // Switch to route B on the same chat instance; its send compresses + // in-send. + routeIdentity = 'route-b'; + mockGeminiClient.tryCompressChat.mockResolvedValueOnce({ + originalTokenCount: 50, + newTokenCount: 50, + compressionStatus: core.CompressionStatus.NOOP, + }); + + await expect( + session.prompt({ + sessionId: 'test-session-id', + prompt: [{ type: 'text', text: 'second' }], + }), + ).resolves.toEqual({ stopReason: 'end_turn' }); + + // Back to route A while compression fails (rate limited): the gate + // falls back to the cache, which must no longer hold route A's + // stale pre-compression count (101 > 100) — the shared history was + // rewritten to 40 tokens by route B's in-send compression, so the + // send must go out. + routeIdentity = 'route-a'; + mockGeminiClient.tryCompressChat.mockRejectedValueOnce( + new Error('compression rate limited'), + ); + + await expect( + session.prompt({ + sessionId: 'test-session-id', + prompt: [{ type: 'text', text: 'third' }], + }), + ).resolves.toEqual({ stopReason: 'end_turn' }); + + expect(mockChat.sendMessageStream).toHaveBeenCalledTimes(3); + }); + it('does not drop a send using another route token count (#9529)', async () => { mockConfig.getSessionTokenLimit = vi.fn().mockReturnValue(100); let routeIdentity = 'route-a'; @@ -12620,7 +12712,7 @@ describe('Session', () => { expect(mockChat.sendMessageStream).toHaveBeenCalledTimes(1); }); - it('falls back to the previous prompt token count when compressed token info is zero', async () => { + it('does not gate on the pre-compression count when compressed token info is zero', async () => { mockConfig.getSessionTokenLimit = vi.fn().mockReturnValue(100); mockGeminiClient.tryCompressChat .mockResolvedValueOnce({ @@ -12641,14 +12733,19 @@ describe('Session', () => { prompt: [{ type: 'text', text: 'first' }], }), ).resolves.toEqual({ stopReason: 'end_turn' }); + // The second send arrives with a successful COMPRESSED whose fresh + // count is unknown (0). The compression just rewrote the shared + // history, so the previously recorded count (101) measures + // destroyed history and must not drop the send — mirroring + // GeminiChat, which clears its keyed counts on COMPRESSED (#9529). await expect( session.prompt({ sessionId: 'test-session-id', prompt: [{ type: 'text', text: 'second' }], }), - ).resolves.toEqual({ stopReason: 'max_tokens' }); + ).resolves.toEqual({ stopReason: 'end_turn' }); - expect(mockChat.sendMessageStream).toHaveBeenCalledTimes(1); + expect(mockChat.sendMessageStream).toHaveBeenCalledTimes(2); }); it('records prompt token count instead of total token count for later session-limit checks', async () => { diff --git a/packages/cli/src/acp-integration/session/Session.ts b/packages/cli/src/acp-integration/session/Session.ts index e01fe175a9c..bd7be15ae2f 100644 --- a/packages/cli/src/acp-integration/session/Session.ts +++ b/packages/cli/src/acp-integration/session/Session.ts @@ -5272,6 +5272,15 @@ export class Session implements SessionContext { ); functionCalls.length = 0; } + if (resp.type === StreamEventType.COMPRESSED) { + // In-send compression rewrote the shared history; + // invalidate every retained route count (the + // pre-send hook never sees this path). + this.#recordCompressionTokenCount( + resp.info, + requestRouteKey, + ); + } } } catch (error) { streamFailed = true; @@ -6301,6 +6310,12 @@ export class Session implements SessionContext { ); functionCalls.length = 0; } + if (response.type === StreamEventType.COMPRESSED) { + // In-send compression rewrote the shared history; invalidate + // every retained route count (the pre-send hook never sees + // this path). + this.#recordCompressionTokenCount(response.info, requestRouteKey); + } } } catch (error) { streamFailed = true; @@ -6769,7 +6784,6 @@ export class Session implements SessionContext { abortSignal, ); compressionInfo = compressed; - this.#recordCompressionTokenCount(compressed); compressionFailed = isCompressionFailureStatus( compressed.compressionStatus, ); @@ -6818,7 +6832,14 @@ export class Session implements SessionContext { options.modelOverride ?? this.config.getModel(); const requestRouteKey = await this.#requestRouteKeyForModel(model); - this.#syncPromptTokenCountWithCurrentChat(requestRouteKey); + // Recorded with the resolved request route key: a COMPRESSED result + // must invalidate every retained route count, not just the active + // route's (see #invalidateRouteTokenCountsForCompression). + if (compressionInfo) { + this.#recordCompressionTokenCount(compressionInfo, requestRouteKey); + } else { + this.#syncPromptTokenCountWithCurrentChat(requestRouteKey); + } const sessionTokenLimit = this.config.getSessionTokenLimit(); if (sessionTokenLimit > 0) { @@ -7082,15 +7103,53 @@ export class Session implements SessionContext { }; } - #recordCompressionTokenCount(info: ChatCompressionInfo): void { - const routeKey = this.#currentRouteKey(); - this.#syncPromptTokenCountWithCurrentChat(routeKey); + #recordCompressionTokenCount( + info: ChatCompressionInfo, + requestRouteKey: string, + ): void { + if (info.compressionStatus === CompressionStatus.COMPRESSED) { + this.#invalidateRouteTokenCountsForCompression(info, requestRouteKey); + return; + } + this.#syncPromptTokenCountWithCurrentChat(requestRouteKey); const tokenCount = this.#extractCompressionTokenCount(info); if (tokenCount !== null && tokenCount > 0) { - this.#setLastPromptTokenCount(routeKey, tokenCount); + this.#setLastPromptTokenCount(requestRouteKey, tokenCount); } } + /** + * Compression rewrote the shared history, so EVERY retained route-keyed + * count is stale — not just the request route's. Drop them all and + * re-record the fresh post-compression count under the request route, + * retaining it under the active route too when the two differ (the + * compressed history is shared, so the count anchors both routes' next + * gate reads). Mirrors GeminiChat clearing its keyed counts in the + * COMPRESSED branch of tryCompress; without this, in-send compressions + * (GeminiChat.sendMessageStream's hard-tier rescue and reactive-overflow + * paths, surfaced as StreamEventType.COMPRESSED) would leave this cache + * holding pre-compression sizes and the gate would drop a returning + * route's send that fits the compressed history (#9529). + */ + #invalidateRouteTokenCountsForCompression( + info: ChatCompressionInfo, + requestRouteKey: string, + ): void { + this.lastPromptTokenCountsByRouteKey.clear(); + const tokenCount = this.#extractCompressionTokenCount(info); + if (tokenCount !== null && tokenCount > 0) { + this.#setLastPromptTokenCount(requestRouteKey, tokenCount); + const activeRouteKey = this.#currentRouteKey(); + if (activeRouteKey !== requestRouteKey) { + this.lastPromptTokenCountsByRouteKey.set(activeRouteKey, tokenCount); + } + } else { + this.lastPromptTokenCount = 0; + this.lastPromptTokenCountRouteKey = requestRouteKey; + } + this.lastPromptTokenCountChat = this.#getCurrentChat(); + } + #recordPromptTokenCount( usageMetadata: GenerateContentResponseUsageMetadata, routeKey = this.#currentRouteKey(), @@ -7957,6 +8016,7 @@ export class Session implements SessionContext { return; } const responseStream = sendResult.responseStream; + const requestRouteKey = sendResult.requestRouteKey; const channelDeliveryResponseBlock: | ChannelDeliveryResponseBlock | undefined = @@ -8048,6 +8108,15 @@ export class Session implements SessionContext { ); functionCalls.length = 0; } + if (resp.type === StreamEventType.COMPRESSED) { + // In-send compression rewrote the shared history; + // invalidate every retained route count (the + // pre-send hook never sees this path). + this.#recordCompressionTokenCount( + resp.info, + requestRouteKey, + ); + } } } catch (error) { streamFailed = true; @@ -8641,6 +8710,7 @@ export class Session implements SessionContext { } const responseStream = sendResult.responseStream; + const requestRouteKey = sendResult.requestRouteKey; nextMessage = null; const messageDisplay = this.#createMessageDisplayDispatcher( ac.signal, @@ -8701,6 +8771,12 @@ export class Session implements SessionContext { ); functionCalls.length = 0; } + if (resp.type === StreamEventType.COMPRESSED) { + // In-send compression rewrote the shared history; + // invalidate every retained route count (the pre-send + // hook never sees this path). + this.#recordCompressionTokenCount(resp.info, requestRouteKey); + } } } catch (error) { streamFailed = true; From 9f9496fae8c969fecf0bed76ced5f692de49e737 Mon Sep 17 00:00:00 2001 From: "jinjing.zzj" Date: Tue, 25 Aug 2026 12:24:19 +0800 Subject: [PATCH 13/14] fix(acp): re-check abort after the route-key await (#9529) --- .../acp-integration/session/Session.test.ts | 87 +++++++++++++++++++ .../src/acp-integration/session/Session.ts | 6 ++ 2 files changed, 93 insertions(+) diff --git a/packages/cli/src/acp-integration/session/Session.test.ts b/packages/cli/src/acp-integration/session/Session.test.ts index 17471684e1a..557c4c06c5a 100644 --- a/packages/cli/src/acp-integration/session/Session.test.ts +++ b/packages/cli/src/acp-integration/session/Session.test.ts @@ -12646,6 +12646,93 @@ describe('Session', () => { }); }); + it('returns cancelled when a cancel lands in the override route-key resolution window (#9529)', async () => { + // First override send records an over-limit count under the + // override route; on the second same-override send the route-key + // resolution outlives a cancel. The abort re-check after the + // resolution must win over the session-token-limit gate (101 > 100 + // cached), which would otherwise mislabel the cancel as + // 'max_tokens' and drop the user turn from history. + setupVisionRouteOverrideMocks(mockConfig); + + const visionPrompt: PromptRequest = { + sessionId: 'test-session-id', + prompt: [ + { type: 'text', text: 'look at this' }, + { type: 'image', mimeType: 'image/png', data: 'iVBORw0KGgo=' }, + ], + }; + + mockGeminiClient.tryCompressChat.mockResolvedValueOnce({ + originalTokenCount: 50, + newTokenCount: 50, + compressionStatus: core.CompressionStatus.NOOP, + }); + mockChat.sendMessageStream = createOverLimitUsageSendStream(); + + await expect(session.prompt(visionPrompt)).resolves.toEqual({ + stopReason: 'end_turn', + }); + + // Second same-override send: compression fails so the gate falls + // back to the cached over-limit count, and the route-key + // resolution hangs until released — the cancel lands inside that + // window and the resolver settles only afterwards. + mockGeminiClient.tryCompressChat.mockRejectedValueOnce( + new Error('compression rate limited'), + ); + let releaseRouteResolution!: () => void; + const resolveForModel = vi.fn().mockImplementation( + () => + new Promise<{ + contentGenerator: Record; + contentGeneratorConfig: { model: string }; + model: string; + }>((resolve) => { + releaseRouteResolution = () => + resolve({ + contentGenerator: {}, + contentGeneratorConfig: { model: 'vision-agent' }, + model: 'vision-agent', + }); + }), + ); + mockConfig.getBaseLlmClient = vi + .fn() + .mockReturnValue({ resolveForModel }); + + const promptPromise = session.prompt(visionPrompt); + await vi.waitFor(() => { + expect(resolveForModel).toHaveBeenCalled(); + }); + + await session.cancelPendingPrompt(); + releaseRouteResolution(); + + await expect(promptPromise).resolves.toEqual({ + stopReason: 'cancelled', + }); + // The send never went out; the cancelled user turn was restored to + // history and no token-limit diagnostic was emitted. + expect(mockChat.sendMessageStream).toHaveBeenCalledTimes(1); + expect(mockChat.addHistory).toHaveBeenCalledWith({ + role: 'user', + parts: expect.any(Array), + }); + expect(mockClient.sessionUpdate).not.toHaveBeenCalledWith({ + sessionId: 'test-session-id', + update: { + sessionUpdate: 'agent_message_chunk', + content: { + type: 'text', + text: + 'Session token limit exceeded: 101 tokens > 100 limit. ' + + 'Please start a new session or increase the sessionTokenLimit in your settings.json.', + }, + }, + }); + }); + it('surfaces an automatic compression AbortError without an aborted signal', async () => { const error = new Error('compression transport aborted unexpectedly'); error.name = 'AbortError'; diff --git a/packages/cli/src/acp-integration/session/Session.ts b/packages/cli/src/acp-integration/session/Session.ts index bd7be15ae2f..00aeac7e022 100644 --- a/packages/cli/src/acp-integration/session/Session.ts +++ b/packages/cli/src/acp-integration/session/Session.ts @@ -6832,6 +6832,12 @@ export class Session implements SessionContext { options.modelOverride ?? this.config.getModel(); const requestRouteKey = await this.#requestRouteKeyForModel(model); + if (abortSignal.aborted) { + debugLogger.debug( + `Send aborted after request route key resolution for prompt ${promptId}`, + ); + return { responseStream: null, stopReason: 'cancelled' }; + } // Recorded with the resolved request route key: a COMPRESSED result // must invalidate every retained route count, not just the active // route's (see #invalidateRouteTokenCountsForCompression). From 4951028190a072a28cb39dcd4828ac85d633bd66 Mon Sep 17 00:00:00 2001 From: "jinjing.zzj" Date: Tue, 25 Aug 2026 15:40:28 +0800 Subject: [PATCH 14/14] fix(acp): key cron and background-notification usage records by the request route (#9529) The cron/loop-tick and background-notification send loops captured the request route key and threaded it into the COMPRESSED handler, but their post-stream usage record still called #recordPromptTokenCount(usageMetadata), whose default route key is the record-time active route. A model switch landing between request and record stored the outgoing route's API-reported count under the incoming route's key, so the next new-route send whose pre-send compression failed was false-dropped with 'Session token limit exceeded' (and, on the cron path, could permanently disable cron via #stopCronAfterTokenLimit). Pass the captured requestRouteKey into the usage record at both call sites, matching the interactive prompt loops. Add collocated tests pinning that each loop records usage under the request route even when the route switches mid-stream. Co-authored-by: Qwen-Coder --- .../acp-integration/session/Session.test.ts | 113 ++++++++++++++++++ .../src/acp-integration/session/Session.ts | 4 +- 2 files changed, 115 insertions(+), 2 deletions(-) diff --git a/packages/cli/src/acp-integration/session/Session.test.ts b/packages/cli/src/acp-integration/session/Session.test.ts index 557c4c06c5a..deb86c35b2b 100644 --- a/packages/cli/src/acp-integration/session/Session.test.ts +++ b/packages/cli/src/acp-integration/session/Session.test.ts @@ -12526,6 +12526,119 @@ describe('Session', () => { expect(mockChat.sendMessageStream).toHaveBeenCalledTimes(2); }); + it('records a cron tick usage under the request route even if the route switches mid-stream (#9529)', async () => { + mockConfig.getSessionTokenLimit = vi.fn().mockReturnValue(100); + let routeIdentity = 'route-a'; + mockConfig.getModelRouteIdentity = vi.fn(() => routeIdentity); + + const scheduler = { + size: 1, + hasPendingWork: true, + start: vi.fn((callback: (job: { prompt: string }) => void) => { + callback({ prompt: 'scheduled prompt' }); + }), + stop: vi.fn(), + getExitSummary: vi.fn().mockReturnValue(undefined), + }; + mockConfig.isCronEnabled = vi.fn().mockReturnValue(true); + mockConfig.getCronScheduler = vi.fn().mockReturnValue(scheduler); + + mockChat.sendMessageStream = vi + .fn() + .mockResolvedValueOnce(createEmptyStream()) + .mockResolvedValueOnce( + (async function* () { + yield { + type: core.StreamEventType.CHUNK, + value: { + usageMetadata: { + totalTokenCount: 101, + promptTokenCount: 101, + }, + }, + }; + // A model switch landing between request and record: the + // request route key was captured before this stream started, + // so the over-limit usage above must still be keyed under it. + routeIdentity = 'route-b'; + })(), + ); + + await session.prompt({ + sessionId: 'test-session-id', + prompt: [{ type: 'text', text: 'hello' }], + }); + + await vi.waitFor(() => { + expect(mockChat.sendMessageStream).toHaveBeenCalledTimes(2); + const internals = session as unknown as { + lastPromptTokenCount: number; + lastPromptTokenCountRouteKey: string | undefined; + }; + expect(internals.lastPromptTokenCountRouteKey).toBe('route-a'); + expect(internals.lastPromptTokenCount).toBe(101); + }); + }); + + it('records a background-notification usage under the request route even if the route switches mid-stream (#9529)', async () => { + mockConfig.getSessionTokenLimit = vi.fn().mockReturnValue(100); + let routeIdentity = 'route-a'; + mockConfig.getModelRouteIdentity = vi.fn(() => routeIdentity); + + mockChat.sendMessageStream = vi + .fn() + .mockResolvedValueOnce(createEmptyStream()) + .mockResolvedValueOnce( + (async function* () { + yield { + type: core.StreamEventType.CHUNK, + value: { + usageMetadata: { + totalTokenCount: 101, + promptTokenCount: 101, + }, + }, + }; + // A model switch landing between request and record: the + // request route key was captured before this stream started, + // so the over-limit usage above must still be keyed under it. + routeIdentity = 'route-b'; + })(), + ); + + await session.prompt({ + sessionId: 'test-session-id', + prompt: [{ type: 'text', text: 'start background work' }], + }); + + const callback = mockBackgroundTaskRegistry.setNotificationCallback.mock + .calls[0][0] as ( + displayText: string, + modelText: string, + meta: { agentId: string; status: string; toolUseId?: string }, + ) => void; + + callback( + 'Background agent "worker" completed.', + 'completed', + { + agentId: 'agent-1', + status: 'completed', + toolUseId: 'tool-1', + }, + ); + + await vi.waitFor(() => { + expect(mockChat.sendMessageStream).toHaveBeenCalledTimes(2); + const internals = session as unknown as { + lastPromptTokenCount: number; + lastPromptTokenCountRouteKey: string | undefined; + }; + expect(internals.lastPromptTokenCountRouteKey).toBe('route-a'); + expect(internals.lastPromptTokenCount).toBe(101); + }); + }); + it('evicts the oldest route count once the retained-route budget is exhausted (#9529)', async () => { mockConfig.getSessionTokenLimit = vi.fn().mockReturnValue(100); let routeIdentity = 'route-0'; diff --git a/packages/cli/src/acp-integration/session/Session.ts b/packages/cli/src/acp-integration/session/Session.ts index 00aeac7e022..8416f9a0ff4 100644 --- a/packages/cli/src/acp-integration/session/Session.ts +++ b/packages/cli/src/acp-integration/session/Session.ts @@ -8148,7 +8148,7 @@ export class Session implements SessionContext { ); if (usageMetadata) { - this.#recordPromptTokenCount(usageMetadata); + this.#recordPromptTokenCount(usageMetadata, requestRouteKey); if (this.messageRewriter) { this.messageRewriter.flushTurn(ac.signal); } @@ -8814,7 +8814,7 @@ export class Session implements SessionContext { } if (usageMetadata) { - this.#recordPromptTokenCount(usageMetadata); + this.#recordPromptTokenCount(usageMetadata, requestRouteKey); const durationMs = Date.now() - streamStartTime; await this.messageEmitter.emitUsageMetadata( usageMetadata,