diff --git a/packages/cli/src/acp-integration/session/Session.test.ts b/packages/cli/src/acp-integration/session/Session.test.ts index afd28a48bcd..993de4f5d8a 100644 --- a/packages/cli/src/acp-integration/session/Session.test.ts +++ b/packages/cli/src/acp-integration/session/Session.test.ts @@ -368,9 +368,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) { @@ -379,6 +380,66 @@ 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()); +} + +/** + * 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, @@ -12725,118 +12786,63 @@ describe('Session', () => { ); }); - it('returns cancelled when automatic compression is aborted', async () => { + 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); - mockLlmClient.tryCompressChat.mockImplementation( - async (_promptId: string, _force: boolean, signal: AbortSignal) => - new Promise((_, reject) => { - signal.addEventListener('abort', () => { - const abortError = new Error('aborted'); - abortError.name = 'AbortError'; - reject(abortError); - }); - }), - ); - mockChat.sendMessageStream = vi - .fn() - .mockResolvedValue(createEmptyStream()); - - const promptPromise = session.prompt({ - sessionId: 'test-session-id', - prompt: [{ type: 'text', text: 'hello' }], - }); - await vi.waitFor(() => { - expect(mockLlmClient.tryCompressChat).toHaveBeenCalled(); + // ACP model switches keep the same LlmChat 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. + mockLlmClient.tryCompressChat.mockResolvedValueOnce({ + originalTokenCount: 50, + newTokenCount: 50, + compressionStatus: core.CompressionStatus.NOOP, }); + mockChat.sendMessageStream = createOverLimitUsageSendStream(); - await session.cancelPendingPrompt(); + await expect( + session.prompt({ + sessionId: 'test-session-id', + prompt: [{ type: 'text', text: 'first' }], + }), + ).resolves.toEqual({ stopReason: 'end_turn' }); - await expect(promptPromise).resolves.toEqual({ - stopReason: 'cancelled', - }); - expect(mockChat.sendMessageStream).not.toHaveBeenCalled(); - 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.', - }, - }, - }); - }); + // Switch to route B on the same chat instance. + routeIdentity = 'route-b'; - it('surfaces an automatic compression AbortError without an aborted signal', async () => { - const error = new Error('compression transport aborted unexpectedly'); - error.name = 'AbortError'; - mockLlmClient.tryCompressChat.mockRejectedValueOnce(error); - mockChat.sendMessageStream = vi - .fn() - .mockResolvedValue(createEmptyStream()); + // 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. + mockLlmClient.tryCompressChat.mockRejectedValueOnce( + new Error('compression rate limited'), + ); await expect( session.prompt({ sessionId: 'test-session-id', - prompt: [{ type: 'text', text: 'hello' }], + prompt: [{ type: 'text', text: 'second' }], }), - ).rejects.toThrow('compression transport aborted unexpectedly'); + ).resolves.toEqual({ stopReason: 'end_turn' }); - expect(mockChat.sendMessageStream).not.toHaveBeenCalled(); + expect(mockChat.sendMessageStream).toHaveBeenCalledTimes(2); }); - it('uses compression token info instead of global UI telemetry for the session limit', async () => { + it('still intercepts a same-route send whose recorded count exceeds the session token limit (#9529)', async () => { mockConfig.getSessionTokenLimit = vi.fn().mockReturnValue(100); - vi.spyOn( - core.uiTelemetryService, - 'getLastPromptTokenCount', - ).mockReturnValue(999); + // 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'); + mockLlmClient.tryCompressChat.mockResolvedValueOnce({ originalTokenCount: 50, newTokenCount: 50, compressionStatus: core.CompressionStatus.NOOP, }); - mockChat.sendMessageStream = vi - .fn() - .mockResolvedValue(createEmptyStream()); - - await session.prompt({ - sessionId: 'test-session-id', - prompt: [{ type: 'text', text: 'hello' }], - }); - - expect(mockChat.sendMessageStream).toHaveBeenCalledTimes(1); - }); - - it('falls back to the previous prompt token count when compression returns zero token info', async () => { - mockConfig.getSessionTokenLimit = vi.fn().mockReturnValue(100); - mockLlmClient.tryCompressChat.mockResolvedValue({ - originalTokenCount: 0, - 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({ @@ -12844,6 +12850,13 @@ describe('Session', () => { 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. + mockLlmClient.tryCompressChat.mockRejectedValueOnce( + new Error('compression rate limited'), + ); + await expect( session.prompt({ sessionId: 'test-session-id', @@ -12851,22 +12864,26 @@ describe('Session', () => { }), ).resolves.toEqual({ stopReason: 'max_tokens' }); + expect(debugLoggerWarnSpy).toHaveBeenCalledWith( + expect.stringContaining( + 'requestRoute=route-a, activeModel=qwen3-code-plus', + ), + ); expect(mockChat.sendMessageStream).toHaveBeenCalledTimes(1); }); - it('falls back to the previous prompt token count when compressed token info is zero', async () => { + 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); - mockLlmClient.tryCompressChat - .mockResolvedValueOnce({ - originalTokenCount: 50, - newTokenCount: 50, - compressionStatus: core.CompressionStatus.NOOP, - }) - .mockResolvedValueOnce({ - originalTokenCount: 1200, - newTokenCount: 0, - compressionStatus: core.CompressionStatus.COMPRESSED, - }); + 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. + mockLlmClient.tryCompressChat.mockResolvedValueOnce({ + originalTokenCount: 50, + newTokenCount: 50, + compressionStatus: core.CompressionStatus.NOOP, + }); mockChat.sendMessageStream = vi .fn() .mockResolvedValueOnce( @@ -12882,6 +12899,25 @@ describe('Session', () => { }, ]), ) + // Prompt 2 on route B trips an in-send compression inside + // LlmChat.sendMessageStream (hard-tier rescue / reactive + // overflow). LlmChat 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( @@ -12890,105 +12926,865 @@ describe('Session', () => { 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'; + mockLlmClient.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: 'max_tokens' }); + ).resolves.toEqual({ stopReason: 'end_turn' }); - expect(mockChat.sendMessageStream).toHaveBeenCalledTimes(1); + // 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'; + mockLlmClient.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('records prompt token count instead of total token count for later session-limit checks', async () => { + it('does not drop a send using another route token count (#9529)', async () => { mockConfig.getSessionTokenLimit = vi.fn().mockReturnValue(100); - mockLlmClient.tryCompressChat - .mockResolvedValueOnce({ - originalTokenCount: 0, - newTokenCount: 0, - compressionStatus: core.CompressionStatus.NOOP, - }) - .mockRejectedValueOnce(new Error('compression unavailable')); - mockChat.sendMessageStream = vi - .fn() - .mockResolvedValueOnce( - createStreamWithChunks([ - { - type: core.StreamEventType.CHUNK, - value: { - usageMetadata: { - totalTokenCount: 500, - promptTokenCount: 50, - }, - }, - }, - ]), - ) - .mockResolvedValueOnce(createEmptyStream()); + let routeIdentity = 'route-a'; + mockConfig.getModelRouteIdentity = vi.fn(() => routeIdentity); + + mockLlmClient.tryCompressChat.mockResolvedValueOnce({ + originalTokenCount: 50, + newTokenCount: 50, + compressionStatus: core.CompressionStatus.NOOP, + }); + mockChat.sendMessageStream = createOverLimitUsageSendStream(); await expect( session.prompt({ sessionId: 'test-session-id', - prompt: [{ type: 'text', text: 'long response' }], + prompt: [{ type: 'text', text: 'first' }], }), ).resolves.toEqual({ stopReason: 'end_turn' }); + + routeIdentity = 'route-b'; await expect( session.prompt({ sessionId: 'test-session-id', - prompt: [{ type: 'text', text: 'next prompt' }], + prompt: [{ type: 'text', text: 'second' }], }), ).resolves.toEqual({ stopReason: 'end_turn' }); expect(mockChat.sendMessageStream).toHaveBeenCalledTimes(2); }); - it('resets the session-local token count when the active chat instance changes', async () => { - const clearedChat = { - sendMessageStream: vi.fn().mockResolvedValue(createEmptyStream()), - addHistory: vi.fn(), - getHistory: vi.fn().mockReturnValue([]), - getHistoryShallow: vi.fn().mockReturnValue([]), - getLastModelMessageText: vi.fn().mockReturnValue(''), - } as unknown as LlmChat; - mockConfig.getSessionTokenLimit = vi.fn().mockReturnValue(100); + it('does not drop a runtime-scoped override using the active route count (#9529)', async () => { + const resolveForModel = setupVisionRouteOverrideMocks(mockConfig); + mockLlmClient.tryCompressChat .mockResolvedValueOnce({ originalTokenCount: 50, newTokenCount: 50, compressionStatus: core.CompressionStatus.NOOP, }) - .mockRejectedValueOnce(new Error('compression unavailable')); - mockChat.sendMessageStream = vi.fn().mockResolvedValueOnce( - createStreamWithChunks([ - { - type: core.StreamEventType.CHUNK, - value: { - usageMetadata: { - totalTokenCount: 500, - promptTokenCount: 101, - }, - }, - }, - ]), - ); + .mockRejectedValueOnce(new Error('compression rate limited')); + mockChat.sendMessageStream = createOverLimitUsageSendStream(); await expect( session.prompt({ sessionId: 'test-session-id', - prompt: [{ type: 'text', text: 'before clear' }], + prompt: [{ type: 'text', text: 'primary route' }], }), ).resolves.toEqual({ stopReason: 'end_turn' }); - mockLlmClient.getChat.mockReturnValue(clearedChat); - await expect( session.prompt({ sessionId: 'test-session-id', - prompt: [{ type: 'text', text: 'after clear' }], + 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 () => { + setupVisionRouteOverrideMocks( + mockConfig, + 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 () => { + // 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. + setupVisionRouteOverrideMocks(mockConfig); + + 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. + mockLlmClient.tryCompressChat.mockResolvedValueOnce({ + originalTokenCount: 50, + newTokenCount: 50, + compressionStatus: core.CompressionStatus.NOOP, + }); + mockChat.sendMessageStream = createOverLimitUsageSendStream(); + + 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. + mockLlmClient.tryCompressChat.mockRejectedValueOnce( + new Error('compression rate limited'), + ); + + await expect(session.prompt(visionPrompt)).resolves.toEqual({ + stopReason: 'max_tokens', + }); + + 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 () => { + 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. + 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. + mockLlmClient.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. + mockLlmClient.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'; + mockConfig.getModelRouteIdentity = vi.fn(() => routeIdentity); + + mockLlmClient.tryCompressChat.mockResolvedValueOnce({ + originalTokenCount: 50, + newTokenCount: 50, + compressionStatus: core.CompressionStatus.NOOP, + }); + mockChat.sendMessageStream = createOverLimitUsageSendStream(); + + await expect( + session.prompt({ + sessionId: 'test-session-id', + prompt: [{ type: 'text', text: 'first' }], + }), + ).resolves.toEqual({ stopReason: 'end_turn' }); + + routeIdentity = 'route-b'; + mockLlmClient.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'; + mockLlmClient.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('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'; + mockConfig.getModelRouteIdentity = vi.fn(() => routeIdentity); + + // 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 = 1; i <= 9; i += 1) { + routeIdentity = `route-${i}`; + mockLlmClient.tryCompressChat.mockResolvedValueOnce({ + originalTokenCount: 50, + newTokenCount: 50, + compressionStatus: core.CompressionStatus.NOOP, + }); + await expect( + session.prompt({ + sessionId: 'test-session-id', + prompt: [{ type: 'text', text: `prompt ${i}` }], + }), + ).resolves.toEqual({ stopReason: 'end_turn' }); + } + + // 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'; + mockLlmClient.tryCompressChat.mockRejectedValueOnce( + new Error('compression rate limited'), + ); + await expect( + session.prompt({ + sessionId: 'test-session-id', + prompt: [{ type: 'text', text: 'evicted route' }], + }), + ).resolves.toEqual({ stopReason: 'end_turn' }); + + // A retained route (route-2) still trips the gate from the cache. + routeIdentity = 'route-2'; + mockLlmClient.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 () => { + mockConfig.getSessionTokenLimit = vi.fn().mockReturnValue(100); + mockLlmClient.tryCompressChat.mockImplementation( + async (_promptId: string, _force: boolean, signal: AbortSignal) => + new Promise((_, reject) => { + signal.addEventListener('abort', () => { + const abortError = new Error('aborted'); + abortError.name = 'AbortError'; + reject(abortError); + }); + }), + ); + mockChat.sendMessageStream = vi + .fn() + .mockResolvedValue(createEmptyStream()); + + const promptPromise = session.prompt({ + sessionId: 'test-session-id', + prompt: [{ type: 'text', text: 'hello' }], + }); + await vi.waitFor(() => { + expect(mockLlmClient.tryCompressChat).toHaveBeenCalled(); + }); + + await session.cancelPendingPrompt(); + + await expect(promptPromise).resolves.toEqual({ + stopReason: 'cancelled', + }); + expect(mockChat.sendMessageStream).not.toHaveBeenCalled(); + 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('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=' }, + ], + }; + + mockLlmClient.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. + mockLlmClient.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'; + mockLlmClient.tryCompressChat.mockRejectedValueOnce(error); + mockChat.sendMessageStream = vi + .fn() + .mockResolvedValue(createEmptyStream()); + + await expect( + session.prompt({ + sessionId: 'test-session-id', + prompt: [{ type: 'text', text: 'hello' }], + }), + ).rejects.toThrow('compression transport aborted unexpectedly'); + + expect(mockChat.sendMessageStream).not.toHaveBeenCalled(); + }); + + it('uses compression token info instead of global UI telemetry for the session limit', async () => { + mockConfig.getSessionTokenLimit = vi.fn().mockReturnValue(100); + vi.spyOn( + core.uiTelemetryService, + 'getLastPromptTokenCount', + ).mockReturnValue(999); + mockLlmClient.tryCompressChat.mockResolvedValueOnce({ + originalTokenCount: 50, + newTokenCount: 50, + compressionStatus: core.CompressionStatus.NOOP, + }); + mockChat.sendMessageStream = vi + .fn() + .mockResolvedValue(createEmptyStream()); + + await session.prompt({ + sessionId: 'test-session-id', + prompt: [{ type: 'text', text: 'hello' }], + }); + + expect(mockChat.sendMessageStream).toHaveBeenCalledTimes(1); + }); + + it('falls back to the previous prompt token count when compression returns zero token info', async () => { + mockConfig.getSessionTokenLimit = vi.fn().mockReturnValue(100); + mockLlmClient.tryCompressChat.mockResolvedValue({ + originalTokenCount: 0, + newTokenCount: 0, + compressionStatus: core.CompressionStatus.NOOP, + }); + mockChat.sendMessageStream = createOverLimitUsageSendStream(); + + 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' }], + }), + ).resolves.toEqual({ stopReason: 'max_tokens' }); + + expect(mockChat.sendMessageStream).toHaveBeenCalledTimes(1); + }); + + it('does not gate on the pre-compression count when compressed token info is zero', async () => { + mockConfig.getSessionTokenLimit = vi.fn().mockReturnValue(100); + mockLlmClient.tryCompressChat + .mockResolvedValueOnce({ + originalTokenCount: 50, + newTokenCount: 50, + compressionStatus: core.CompressionStatus.NOOP, + }) + .mockResolvedValueOnce({ + originalTokenCount: 1200, + newTokenCount: 0, + compressionStatus: core.CompressionStatus.COMPRESSED, + }); + mockChat.sendMessageStream = createOverLimitUsageSendStream(); + + await expect( + session.prompt({ + sessionId: 'test-session-id', + 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 + // LlmChat, which clears its keyed counts on COMPRESSED (#9529). + await expect( + session.prompt({ + sessionId: 'test-session-id', + prompt: [{ type: 'text', text: 'second' }], + }), + ).resolves.toEqual({ stopReason: 'end_turn' }); + + expect(mockChat.sendMessageStream).toHaveBeenCalledTimes(2); + }); + + it('records prompt token count instead of total token count for later session-limit checks', async () => { + mockConfig.getSessionTokenLimit = vi.fn().mockReturnValue(100); + mockLlmClient.tryCompressChat + .mockResolvedValueOnce({ + originalTokenCount: 0, + newTokenCount: 0, + compressionStatus: core.CompressionStatus.NOOP, + }) + .mockRejectedValueOnce(new Error('compression unavailable')); + mockChat.sendMessageStream = vi + .fn() + .mockResolvedValueOnce( + createStreamWithChunks([ + { + type: core.StreamEventType.CHUNK, + value: { + usageMetadata: { + totalTokenCount: 500, + promptTokenCount: 50, + }, + }, + }, + ]), + ) + .mockResolvedValueOnce(createEmptyStream()); + + await expect( + session.prompt({ + sessionId: 'test-session-id', + prompt: [{ type: 'text', text: 'long response' }], + }), + ).resolves.toEqual({ stopReason: 'end_turn' }); + await expect( + session.prompt({ + sessionId: 'test-session-id', + prompt: [{ type: 'text', text: 'next prompt' }], + }), + ).resolves.toEqual({ stopReason: 'end_turn' }); + + expect(mockChat.sendMessageStream).toHaveBeenCalledTimes(2); + }); + + it('resets session-local route counts when the active chat instance changes', async () => { + const clearedChat = { + sendMessageStream: vi.fn().mockResolvedValue(createEmptyStream()), + addHistory: vi.fn(), + getHistory: vi.fn().mockReturnValue([]), + getHistoryShallow: vi.fn().mockReturnValue([]), + getLastModelMessageText: vi.fn().mockReturnValue(''), + } as unknown as LlmChat; + mockConfig.getSessionTokenLimit = vi.fn().mockReturnValue(100); + let routeIdentity = 'route-a'; + mockConfig.getModelRouteIdentity = vi.fn(() => routeIdentity); + mockLlmClient.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([ + { + type: core.StreamEventType.CHUNK, + value: { + usageMetadata: { + totalTokenCount: 500, + promptTokenCount: 101, + }, + }, + }, + ]), + ); + + await expect( + session.prompt({ + sessionId: 'test-session-id', + prompt: [{ type: 'text', text: 'before clear' }], + }), + ).resolves.toEqual({ stopReason: 'end_turn' }); + + mockLlmClient.getChat.mockReturnValue(clearedChat); + routeIdentity = 'route-b'; + + await expect( + session.prompt({ + sessionId: 'test-session-id', + prompt: [{ type: 'text', text: 'after clear' }], + }), + ).resolves.toEqual({ stopReason: 'end_turn' }); + + 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(1); + 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 2e19ef3935a..d5a27d79677 100644 --- a/packages/cli/src/acp-integration/session/Session.ts +++ b/packages/cli/src/acp-integration/session/Session.ts @@ -374,6 +374,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'; @@ -484,7 +485,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( @@ -1938,6 +1943,14 @@ export class Session implements SessionContext { private cronDisabledByTokenLimit = false; private lastPromptTokenCount = 0; private lastPromptTokenCountChat: LlmChat | null = null; + // Private ACP fallback cache, bounded like LlmChat 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 LlmChat, 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 @@ -5505,6 +5518,10 @@ 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 { // Set where the model request is actually issued, not at @@ -5555,6 +5572,7 @@ export class Session implements SessionContext { if (restorePostAnswerNoticesAttached) { this.#clearPendingRestoreNotices(); } + requestRouteKey = sendResult.requestRouteKey; const responseStream = sendResult.responseStream; nextMessage = null; channelDeliveryResponseBlock = @@ -5638,6 +5656,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; @@ -5730,7 +5757,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); @@ -6314,6 +6341,10 @@ export class Session implements SessionContext { let channelDeliveryCheckpoint = 0; let providerSendChat: LlmChat | 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 { const sendResult = await this.#sendMessageStreamWithAutoCompression( @@ -6607,6 +6638,7 @@ export class Session implements SessionContext { }; } + requestRouteKey = sendResult.requestRouteKey; const responseStream = sendResult.responseStream; nextMessage = null; channelDeliveryResponseBlock = beginChannelDeliveryResponseBlock( @@ -6703,6 +6735,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; @@ -6777,7 +6815,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, @@ -7184,7 +7222,6 @@ export class Session implements SessionContext { abortSignal, ); compressionInfo = compressed; - this.#recordCompressionTokenCount(compressed); compressionFailed = isCompressionFailureStatus( compressed.compressionStatus, ); @@ -7237,18 +7274,38 @@ 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); + 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). + if (compressionInfo) { + this.#recordCompressionTokenCount(compressionInfo, requestRouteKey); + } else { + 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}: ` + - `${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. ` + @@ -7299,10 +7356,6 @@ export class Session implements SessionContext { } const chat = this.#getCurrentChat(); - const model = - options.getModelOverride?.() ?? - options.modelOverride ?? - this.config.getModel(); const request = { message, config: { @@ -7313,7 +7366,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 { @@ -7513,31 +7566,75 @@ export class Session implements SessionContext { }; } - #recordCompressionTokenCount(info: ChatCompressionInfo): void { - this.#syncPromptTokenCountWithCurrentChat(); + #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.lastPromptTokenCount = 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 LlmChat clearing its keyed counts in the + * COMPRESSED branch of tryCompress; without this, in-send compressions + * (LlmChat.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(), ): 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; } @@ -7557,15 +7654,71 @@ export class Session implements SessionContext { return tokenCount; } - #syncPromptTokenCountWithCurrentChat(): void { - const chat = this.#getCurrentChat(); + #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 LlmChat.currentRouteKey, #9454). + return this.config.getModelRouteIdentity?.() ?? ''; + } + + 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; if ( - this.lastPromptTokenCountChat && - this.lastPromptTokenCountChat !== chat + !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); + } + + #syncPromptTokenCountWithCurrentChat( + routeKey = this.#currentRouteKey(), + ): void { + const chat = this.#getCurrentChat(); + const chatChanged = + this.lastPromptTokenCountChat && this.lastPromptTokenCountChat !== chat; + 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; } #isAbortError(error: unknown): boolean { @@ -8345,6 +8498,7 @@ export class Session implements SessionContext { return; } const responseStream = sendResult.responseStream; + const requestRouteKey = sendResult.requestRouteKey; const channelDeliveryResponseBlock: | ChannelDeliveryResponseBlock | undefined = @@ -8436,6 +8590,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; @@ -8461,7 +8624,7 @@ export class Session implements SessionContext { ); if (usageMetadata) { - this.#recordPromptTokenCount(usageMetadata); + this.#recordPromptTokenCount(usageMetadata, requestRouteKey); if (this.messageRewriter) { this.messageRewriter.flushTurn(ac.signal); } @@ -9039,6 +9202,7 @@ export class Session implements SessionContext { } const responseStream = sendResult.responseStream; + const requestRouteKey = sendResult.requestRouteKey; nextMessage = null; const messageDisplay = this.#createMessageDisplayDispatcher( ac.signal, @@ -9099,6 +9263,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; @@ -9130,7 +9300,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,