mirror of
https://github.com/QwenLM/qwen-code.git
synced 2026-08-25 08:33:55 +00:00
fix(core): invalidate token counts recorded for a switched model route (#9506)
* fix(core): invalidate token counts recorded for a switched model route `/model` switches rebuild the content generator but keep the same GeminiChat instance, so API-reported prompt/output token counts from the previous route survived and anchored admission, output clamping, and compression decisions for a different serialization (#9454). Attribute the counts to the route that produced them (Config.getModelRouteIdentity) and invalidate them on a route change so all safety decisions fall back to the history-walk estimate. Closes #9454 * fix(core): keep token counts scoped to producing route * fix(core): close round-2 route-scoping findings (#9454) - C1: client.test.ts session-token-limit gate test now seeds the chat mock's getLastPromptTokenCount (the gate's new source) instead of relying on the telemetry stub alone. - C2: goal-turn-integration.test.ts passes the routeKey positional added to processStreamResponse and widens the local cast. - S1: fix stale rationale comment — no session token-limit gate reads the telemetry mirror anymore; UI context counters and compression banners do. - S2: pin tryCompress's entry invalidation with a focused test; manual /compress reaches it without sendMessageStream's entry reset. - S3: direct config tests for getModelRouteIdentity — call stability, model@<sha-prefix> shape, and the guard keeping the registry baseUrl out of non-active model identities. - R2-3: sendMessageStream now resolves the request route first and invalidates counts against IT (deriving the key from the actual model param on the non-exact branch), so an active-route count can no longer anchor an exact `\0` route's output clamp; regression test added. * fix(core): close round-3 route-scoping findings (#9454) - R3-1: hard-rescue rollback now restores tokenCountsRouteKey alongside the counts; tryCompress re-stamps the key to the active route mid-rescue, which left the resurrected request-route count riding the active key past the next entry invalidation (regression test: active-route read after failed rescue must not inherit the override count). - R3-2(2): invalidation zeroes the telemetry cached-content mirror together with the prompt mirror so /context stops rendering a foreign cached count beside a zeroed prompt count; the mirror is documented as best-effort display state between a switch and the next guarded read (R3-2(1)). - R3-4: collapse the requestRouteKey ternary into one getModelRouteIdentity call (the non-exact arm passed exactly the default parameter value). - R3-5: resolve the active-route default lazily after the zero-count fast path instead of eagerly in the default parameter. - R1-3: add the missing compressFast route-invalidation test (third entrypoint; mutation-verified). * fix(core): preserve route-scoped token guards (#9454) Co-authored-by: Qwen-Coder <qwen-coder@alibabacloud.com> * fix(core): resolve full-turn route selectors at the session token limit gate (#9454) * fix(core): stamp fallback-served token counts under the request route (#9454) * fix(core): retain route-scoped token counts across route crossings (#9454) * fix(core): keep route-scoped token counts consistent through compression (#9454) * fix(core): keep rescued output counts and compression stamps route-scoped (#9454) --------- Co-authored-by: yiliang114 <yiliang114@users.noreply.github.com> Co-authored-by: Qwen-Coder <qwen-coder@alibabacloud.com>
This commit is contained in:
易良
GitHub
parent
3d96e54641
commit
3fd059368b
7 changed files with 1549 additions and 16 deletions
|
|
@ -436,3 +436,106 @@ describe('Config provider-qualified model identity', () => {
|
|||
expect(process.env['QWEN_CODE_MODEL_IDENTITY']).toBe(claimed);
|
||||
});
|
||||
});
|
||||
|
||||
describe('Config.getModelRouteIdentity (#9454 route key)', () => {
|
||||
it('returns an identical identity across repeated calls for one configuration', async () => {
|
||||
const { Config } = await import('./config.js');
|
||||
const { resolveContentGeneratorConfigWithSources, AuthType } = await import(
|
||||
'../core/contentGenerator.js'
|
||||
);
|
||||
vi.mocked(resolveContentGeneratorConfigWithSources).mockReturnValue({
|
||||
config: {
|
||||
model: 'steady-model',
|
||||
apiKey: 'k',
|
||||
baseUrl: 'https://provider-a.example/v1',
|
||||
} as ContentGeneratorConfig,
|
||||
sources: {},
|
||||
});
|
||||
const config = new Config({ ...baseParams });
|
||||
await config.refreshAuth(AuthType.USE_GEMINI);
|
||||
|
||||
// Route-scoped caches (e.g. GeminiChat token counts) compare these
|
||||
// strings for equality — the value must not drift between calls.
|
||||
const first = config.getModelRouteIdentity();
|
||||
expect(config.getModelRouteIdentity()).toBe(first);
|
||||
expect(config.getModelRouteIdentity()).toBe(first);
|
||||
expect(first).toMatch(/^steady-model@[0-9a-f]{8}$/);
|
||||
});
|
||||
|
||||
it('keeps the `model@<sha-prefix>` digest shape for explicit route queries', async () => {
|
||||
const { Config } = await import('./config.js');
|
||||
const { resolveContentGeneratorConfigWithSources, AuthType } = await import(
|
||||
'../core/contentGenerator.js'
|
||||
);
|
||||
vi.mocked(resolveContentGeneratorConfigWithSources).mockReturnValue({
|
||||
config: {
|
||||
model: 'active-model',
|
||||
apiKey: 'k',
|
||||
baseUrl: 'https://provider-a.example/v1',
|
||||
} as ContentGeneratorConfig,
|
||||
sources: {},
|
||||
});
|
||||
const config = new Config({ ...baseParams });
|
||||
await config.refreshAuth(AuthType.USE_GEMINI);
|
||||
|
||||
// The readable model id stays the prefix; the discriminator is exactly
|
||||
// eight hex characters, stable for one (auth type, endpoint) pair.
|
||||
const explicit = config.getModelRouteIdentity('route-model', {
|
||||
model: 'route-model',
|
||||
authType: 'openai',
|
||||
baseUrl: 'https://route.example/v1',
|
||||
} as ContentGeneratorConfig);
|
||||
expect(explicit).toMatch(/^route-model@[0-9a-f]{8}$/);
|
||||
expect(config.getModelRouteIdentity('route-model', {
|
||||
model: 'route-model',
|
||||
authType: 'openai',
|
||||
baseUrl: 'https://route.example/v1',
|
||||
} as ContentGeneratorConfig)).toBe(explicit);
|
||||
});
|
||||
|
||||
it('does not mix the registry base URL into a non-active model identity', async () => {
|
||||
// The registry-baseUrl fallback qualifies the ACTIVE model's route when
|
||||
// its own generator config carries no baseUrl. A foreign model queried
|
||||
// with its own configuration must not pick that fallback up — hashing
|
||||
// the registry endpoint into another model's identity would invalidate
|
||||
// its route-scoped state on unrelated registry changes.
|
||||
const { Config } = await import('./config.js');
|
||||
const { resolveContentGeneratorConfigWithSources, AuthType } = await import(
|
||||
'../core/contentGenerator.js'
|
||||
);
|
||||
vi.mocked(resolveContentGeneratorConfigWithSources).mockReturnValue({
|
||||
config: {
|
||||
model: 'active-model',
|
||||
apiKey: 'k',
|
||||
// No baseUrl of its own: the active model falls back to the
|
||||
// registry base URL below.
|
||||
} as ContentGeneratorConfig,
|
||||
sources: {},
|
||||
});
|
||||
const config = new Config({ ...baseParams });
|
||||
await config.refreshAuth(AuthType.USE_GEMINI);
|
||||
const registrySpy = vi.spyOn(config, 'getCurrentModelRegistryBaseUrl');
|
||||
|
||||
const foreignGeneratorConfig = {
|
||||
model: 'foreign-model',
|
||||
authType: 'openai',
|
||||
} as ContentGeneratorConfig;
|
||||
|
||||
registrySpy.mockReturnValue('https://registry.example/v1');
|
||||
const activeWithRegistry = config.getModelRouteIdentity();
|
||||
const foreignWithRegistry =
|
||||
config.getModelRouteIdentity('foreign-model', foreignGeneratorConfig);
|
||||
|
||||
registrySpy.mockReturnValue(null);
|
||||
const activeWithoutRegistry = config.getModelRouteIdentity();
|
||||
const foreignWithoutRegistry =
|
||||
config.getModelRouteIdentity('foreign-model', foreignGeneratorConfig);
|
||||
|
||||
// The fallback is load-bearing for the ACTIVE model…
|
||||
expect(activeWithRegistry).toMatch(/^active-model@[0-9a-f]{8}$/);
|
||||
expect(activeWithRegistry).not.toBe(activeWithoutRegistry);
|
||||
// …but the foreign model's identity ignores the registry entirely.
|
||||
expect(foreignWithRegistry).toMatch(/^foreign-model@[0-9a-f]{8}$/);
|
||||
expect(foreignWithRegistry).toBe(foreignWithoutRegistry);
|
||||
});
|
||||
});
|
||||
|
|
|
|||
|
|
@ -4392,12 +4392,16 @@ export class Config {
|
|||
* configurations without publishing where they point. The bare id stays the
|
||||
* readable half, so a mismatch still names the model a human recognises.
|
||||
*/
|
||||
private resolvedModelIdentity(): string {
|
||||
const model = this.getModel();
|
||||
const authType = this.getContentGeneratorConfig()?.authType ?? '';
|
||||
private resolvedModelIdentity(
|
||||
model = this.getModel(),
|
||||
generatorConfig = this.getContentGeneratorConfig(),
|
||||
): string {
|
||||
const authType = generatorConfig?.authType ?? '';
|
||||
const baseUrl =
|
||||
this.getContentGeneratorConfig()?.baseUrl ??
|
||||
this.getCurrentModelRegistryBaseUrl() ??
|
||||
generatorConfig?.baseUrl ??
|
||||
(model === this.getModel()
|
||||
? this.getCurrentModelRegistryBaseUrl()
|
||||
: undefined) ??
|
||||
'';
|
||||
if (authType === '' && baseUrl === '') return model;
|
||||
const digest = createHash('sha256')
|
||||
|
|
@ -4407,6 +4411,19 @@ export class Config {
|
|||
return `${model}@${digest}`;
|
||||
}
|
||||
|
||||
/**
|
||||
* Identity of the currently active model route for consumers that cache
|
||||
* route-specific state and must invalidate it when a model/auth/endpoint
|
||||
* switch swaps the content generator — e.g. GeminiChat's API-reported
|
||||
* token counts (#9454). Same identity ⇒ same serialization target.
|
||||
*/
|
||||
getModelRouteIdentity(
|
||||
model?: string,
|
||||
generatorConfig?: ContentGeneratorConfig,
|
||||
): string {
|
||||
return this.resolvedModelIdentity(model, generatorConfig);
|
||||
}
|
||||
|
||||
/**
|
||||
* Returns the configured fast model selector when it resolves to an available
|
||||
* model. Bare selectors stay bare and authType-qualified selectors keep their
|
||||
|
|
|
|||
|
|
@ -619,7 +619,7 @@ describe('Gemini Client (client.ts)', () => {
|
|||
toolResultsThresholdMinutes: 60,
|
||||
toolResultsNumToKeep: 5,
|
||||
}),
|
||||
getSessionTokenLimit: vi.fn().mockReturnValue(32000),
|
||||
getSessionTokenLimit: vi.fn().mockReturnValue(0),
|
||||
getNoBrowser: vi.fn().mockReturnValue(false),
|
||||
getUsageStatisticsEnabled: vi.fn().mockReturnValue(true),
|
||||
getTelemetryIncludeSensitiveSpanAttributes: vi
|
||||
|
|
@ -654,6 +654,7 @@ describe('Gemini Client (client.ts)', () => {
|
|||
.mockReturnValue('/test/project/root/.gemini/projects/test-project'),
|
||||
},
|
||||
getContentGenerator: vi.fn().mockReturnValue(mockContentGenerator),
|
||||
getModelRouteIdentity: vi.fn().mockReturnValue('test-route'),
|
||||
getEffectiveInputModalities: vi.fn().mockReturnValue({}),
|
||||
getBaseLlmClient: vi.fn(),
|
||||
getSkipLoopDetection: vi.fn().mockReturnValue(false),
|
||||
|
|
@ -3147,6 +3148,7 @@ describe('Gemini Client (client.ts)', () => {
|
|||
addHistory,
|
||||
getHistory: vi.fn().mockReturnValue([]),
|
||||
getHistoryLength: vi.fn().mockReturnValue(1),
|
||||
getLastPromptTokenCount: vi.fn().mockReturnValue(101),
|
||||
// Send is skipped, so the push counter never advances → restore.
|
||||
getUserContentPushCount: vi.fn().mockReturnValue(0),
|
||||
stripOrphanedUserEntriesFromHistory: vi
|
||||
|
|
@ -3172,6 +3174,170 @@ describe('Gemini Client (client.ts)', () => {
|
|||
expect(mockTurnRunFn).not.toHaveBeenCalled();
|
||||
expect(addHistory).toHaveBeenCalledWith(retryEntry);
|
||||
});
|
||||
|
||||
it('invalidates a foreign route count before the session limit gate', async () => {
|
||||
let route = 'route-a';
|
||||
let telemetryCount = 691_000;
|
||||
vi.mocked(mockConfig.getModelRouteIdentity).mockImplementation(
|
||||
() => route,
|
||||
);
|
||||
vi.mocked(mockConfig.getSessionTokenLimit).mockReturnValue(100_000);
|
||||
vi.mocked(uiTelemetryService.getLastPromptTokenCount).mockImplementation(
|
||||
() => telemetryCount,
|
||||
);
|
||||
vi.mocked(uiTelemetryService.setLastPromptTokenCount).mockImplementation(
|
||||
(count) => {
|
||||
telemetryCount = count;
|
||||
},
|
||||
);
|
||||
client.getChat().setLastPromptTokenCount(telemetryCount);
|
||||
route = 'route-b';
|
||||
mockTurnRunFn.mockReturnValue(
|
||||
(async function* () {
|
||||
yield { type: GeminiEventType.Content, value: 'response' };
|
||||
})(),
|
||||
);
|
||||
|
||||
const events = await fromAsync(
|
||||
client.sendMessageStream(
|
||||
[{ text: 'new route' }],
|
||||
new AbortController().signal,
|
||||
'prompt-route-switch',
|
||||
),
|
||||
);
|
||||
|
||||
expect(events).not.toContainEqual(
|
||||
expect.objectContaining({
|
||||
type: GeminiEventType.SessionTokenLimitExceeded,
|
||||
}),
|
||||
);
|
||||
expect(telemetryCount).toBe(0);
|
||||
});
|
||||
|
||||
it('applies the session limit to the requested override route', async () => {
|
||||
vi.mocked(mockConfig.getModelRouteIdentity).mockImplementation((model) =>
|
||||
model ? `${model}@route` : 'override-model@route',
|
||||
);
|
||||
vi.mocked(mockConfig.getSessionTokenLimit).mockReturnValue(100);
|
||||
client.getChat().setLastPromptTokenCount(101);
|
||||
vi.mocked(mockConfig.getModelRouteIdentity).mockImplementation((model) =>
|
||||
model ? `${model}@route` : 'active-model@route',
|
||||
);
|
||||
|
||||
const events = await fromAsync(
|
||||
client.sendMessageStream(
|
||||
[{ text: 'override route' }],
|
||||
new AbortController().signal,
|
||||
'prompt-override-limit',
|
||||
{
|
||||
type: SendMessageType.UserQuery,
|
||||
modelOverride: 'override-model',
|
||||
},
|
||||
),
|
||||
);
|
||||
|
||||
expect(events).toContainEqual({
|
||||
type: GeminiEventType.SessionTokenLimitExceeded,
|
||||
value: expect.objectContaining({ currentTokens: 101, limit: 100 }),
|
||||
});
|
||||
expect(mockTurnRunFn).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('applies the session limit to a resolved full-turn route selector', async () => {
|
||||
// The vision-bridge full-turn selector `${id}\0${baseUrl}\0` arrives as
|
||||
// modelOverride. GeminiChat.sendMessageStream resolves it and stamps
|
||||
// counts under the RESOLVED route's identity, so the gate must resolve
|
||||
// the selector before keying — the raw selector key (always containing
|
||||
// a NUL) can never match a stamped count (#9454).
|
||||
vi.mocked(mockConfig.getModelRouteIdentity).mockReturnValue(
|
||||
'vision-agent@route',
|
||||
);
|
||||
vi.mocked(mockConfig.getSessionTokenLimit).mockReturnValue(100);
|
||||
client.getChat().setLastPromptTokenCount(101);
|
||||
const resolveForModel = vi.fn().mockResolvedValue({
|
||||
model: 'vision-agent',
|
||||
contentGeneratorConfig: undefined,
|
||||
});
|
||||
vi.mocked(mockConfig.getBaseLlmClient).mockReturnValue({
|
||||
resolveForModel,
|
||||
} as unknown as ReturnType<Config['getBaseLlmClient']>);
|
||||
vi.mocked(mockConfig.getModelRouteIdentity).mockImplementation((model) =>
|
||||
model ? `${model}@route` : 'active-model@route',
|
||||
);
|
||||
|
||||
const events = await fromAsync(
|
||||
client.sendMessageStream(
|
||||
[{ text: 'vision route' }],
|
||||
new AbortController().signal,
|
||||
'prompt-selector-limit',
|
||||
{
|
||||
type: SendMessageType.UserQuery,
|
||||
modelOverride: 'openai:vision-agent\0https://vision.example/v1\0',
|
||||
},
|
||||
),
|
||||
);
|
||||
|
||||
expect(resolveForModel).toHaveBeenCalledWith(
|
||||
'openai:vision-agent\0https://vision.example/v1',
|
||||
{ failClosed: true },
|
||||
);
|
||||
expect(events).toContainEqual({
|
||||
type: GeminiEventType.SessionTokenLimitExceeded,
|
||||
value: expect.objectContaining({ currentTokens: 101, limit: 100 }),
|
||||
});
|
||||
expect(mockTurnRunFn).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('keeps the session limit enforced when turns alternate routes (#9506)', async () => {
|
||||
// Counts are retained per route (#9506): an intervening turn on
|
||||
// another route must not destroy the count the gate later reads for
|
||||
// the original route. Pre-fix, the foreign-route gate read zeroed
|
||||
// the only slot, so the returning turn read 0 and was admitted
|
||||
// regardless of size — steady alternation disabled the limit.
|
||||
vi.mocked(mockConfig.getModelRouteIdentity).mockImplementation((model) =>
|
||||
model === 'route-x' ? 'route-x@route' : 'route-a',
|
||||
);
|
||||
vi.mocked(mockConfig.getSessionTokenLimit).mockReturnValue(100);
|
||||
// Route A's last response stamped an over-limit count.
|
||||
client.getChat().setLastPromptTokenCount(101);
|
||||
mockTurnRunFn.mockReturnValue(
|
||||
(async function* () {
|
||||
yield { type: GeminiEventType.Content, value: 'response' };
|
||||
})(),
|
||||
);
|
||||
|
||||
// Intervening turn on route X: no counts recorded for X yet, so the
|
||||
// gate admits it.
|
||||
const foreignEvents = await fromAsync(
|
||||
client.sendMessageStream(
|
||||
[{ text: 'foreign route turn' }],
|
||||
new AbortController().signal,
|
||||
'prompt-alternate-foreign',
|
||||
{ type: SendMessageType.UserQuery, modelOverride: 'route-x' },
|
||||
),
|
||||
);
|
||||
expect(foreignEvents).not.toContainEqual(
|
||||
expect.objectContaining({
|
||||
type: GeminiEventType.SessionTokenLimitExceeded,
|
||||
}),
|
||||
);
|
||||
|
||||
// Returning to route A must still trip the gate with the retained
|
||||
// over-limit count — the alternation must not have zeroed it.
|
||||
const events = await fromAsync(
|
||||
client.sendMessageStream(
|
||||
[{ text: 'back on route a' }],
|
||||
new AbortController().signal,
|
||||
'prompt-alternate-return',
|
||||
{ type: SendMessageType.UserQuery },
|
||||
),
|
||||
);
|
||||
expect(events).toContainEqual({
|
||||
type: GeminiEventType.SessionTokenLimitExceeded,
|
||||
value: expect.objectContaining({ currentTokens: 101, limit: 100 }),
|
||||
});
|
||||
expect(mockTurnRunFn).toHaveBeenCalledTimes(1);
|
||||
});
|
||||
});
|
||||
|
||||
/**
|
||||
|
|
@ -9266,6 +9432,7 @@ Other open files:
|
|||
const mockChat: Partial<GeminiChat> = {
|
||||
addHistory: vi.fn(),
|
||||
getHistory: vi.fn().mockReturnValue([]),
|
||||
getLastPromptTokenCount: vi.fn().mockReturnValue(9999),
|
||||
};
|
||||
client['chat'] = mockChat as GeminiChat;
|
||||
|
||||
|
|
|
|||
|
|
@ -3142,10 +3142,25 @@ export class GeminiClient {
|
|||
// via the `compressed → ChatCompressed` bridge in turn.ts. Manual /compress
|
||||
// still calls tryCompressChat directly for the full reset (env refresh +
|
||||
// forceFullIdeContext flip).
|
||||
const model = options?.modelOverride ?? this.config.getModel();
|
||||
const sessionTokenLimit = this.config.getSessionTokenLimit();
|
||||
if (sessionTokenLimit > 0) {
|
||||
// An exact `\0` full-turn route selector resolves to its route before
|
||||
// GeminiChat.sendMessageStream stamps counts under it, so the gate
|
||||
// must key the resolved route too — the raw selector key can never
|
||||
// match a stamped count. Mirrors the resolution at the top of
|
||||
// GeminiChat.sendMessageStream (#9454).
|
||||
const exactRoute = model.endsWith('\0')
|
||||
? await this.config
|
||||
.getBaseLlmClient()
|
||||
.resolveForModel(model.slice(0, -1), { failClosed: true })
|
||||
: undefined;
|
||||
const requestRouteKey = this.config.getModelRouteIdentity(
|
||||
exactRoute ? exactRoute.model : model,
|
||||
exactRoute?.contentGeneratorConfig,
|
||||
);
|
||||
const lastPromptTokenCount =
|
||||
uiTelemetryService.getLastPromptTokenCount();
|
||||
this.getChat().getLastPromptTokenCount(requestRouteKey);
|
||||
if (lastPromptTokenCount > sessionTokenLimit) {
|
||||
this.cancelPendingMemoryPrefetch('no_safe_delivery_point');
|
||||
yield {
|
||||
|
|
@ -3237,9 +3252,6 @@ export class GeminiClient {
|
|||
|
||||
const turn = new Turn(this.getChat(), prompt_id, goalPermit);
|
||||
|
||||
// Determine the model to use for this turn
|
||||
const model = options?.modelOverride ?? this.config.getModel();
|
||||
|
||||
// Assemble the outgoing request. IDE context is merged into the
|
||||
// user prompt's first text part, then on UserQuery / Cron turns
|
||||
// the system reminders block is prepended in front of everything
|
||||
|
|
|
|||
|
|
@ -115,6 +115,7 @@ vi.mock('../telemetry/loggers.js', () => ({
|
|||
vi.mock('../telemetry/uiTelemetry.js', () => ({
|
||||
uiTelemetryService: {
|
||||
setLastPromptTokenCount: vi.fn(),
|
||||
setLastCachedContentTokenCount: vi.fn(),
|
||||
},
|
||||
}));
|
||||
|
||||
|
|
@ -180,6 +181,7 @@ describe('GeminiChat', async () => {
|
|||
model: 'test-model',
|
||||
}),
|
||||
getModel: vi.fn().mockReturnValue('gemini-pro'),
|
||||
getModelRouteIdentity: vi.fn().mockReturnValue('gemini-pro@test0001'),
|
||||
setModel: vi.fn(),
|
||||
getProjectRoot: vi.fn().mockReturnValue('/test/project/root'),
|
||||
getTargetDir: vi.fn().mockReturnValue('/test/project/root'),
|
||||
|
|
@ -6986,6 +6988,7 @@ describe('GeminiChat', async () => {
|
|||
finishReason: 'STOP',
|
||||
},
|
||||
],
|
||||
usageMetadata: { promptTokenCount: 99_999 },
|
||||
} as unknown as GenerateContentResponse;
|
||||
})(),
|
||||
);
|
||||
|
|
@ -7010,6 +7013,9 @@ describe('GeminiChat', async () => {
|
|||
vi.mocked(mockConfig.getBaseLlmClient).mockReturnValue({
|
||||
resolveForModel,
|
||||
} as unknown as ReturnType<typeof mockConfig.getBaseLlmClient>);
|
||||
vi.mocked(mockConfig.getModelRouteIdentity).mockImplementation((model) =>
|
||||
model ? `${model}@route` : 'gemini-pro@test0001',
|
||||
);
|
||||
vi.mocked(mockConfig.getEffectiveInputModalities).mockReturnValue({
|
||||
pdf: true,
|
||||
});
|
||||
|
|
@ -7062,6 +7068,7 @@ describe('GeminiChat', async () => {
|
|||
for await (const _ of stream) {
|
||||
/* consume */
|
||||
}
|
||||
expect(chat.getLastPromptTokenCount()).toBe(0);
|
||||
|
||||
expect(resolveForModel).toHaveBeenCalledOnce();
|
||||
expect(resolveForModel).toHaveBeenCalledWith(routeSelector, {
|
||||
|
|
@ -7223,6 +7230,9 @@ describe('GeminiChat', async () => {
|
|||
vi.mocked(mockConfig.getBaseLlmClient).mockReturnValue({
|
||||
resolveForModel,
|
||||
} as unknown as ReturnType<typeof mockConfig.getBaseLlmClient>);
|
||||
vi.mocked(mockConfig.getModelRouteIdentity).mockImplementation((model) =>
|
||||
model ? `${model}@route` : 'gemini-pro@test0001',
|
||||
);
|
||||
|
||||
const capacityError = Object.assign(
|
||||
new StreamContentError(
|
||||
|
|
@ -7261,6 +7271,7 @@ describe('GeminiChat', async () => {
|
|||
finishReason: 'STOP',
|
||||
},
|
||||
],
|
||||
usageMetadata: { promptTokenCount: 99_999 },
|
||||
} as unknown as GenerateContentResponse;
|
||||
})(),
|
||||
);
|
||||
|
|
@ -7328,6 +7339,7 @@ describe('GeminiChat', async () => {
|
|||
).toContain('fallback-image');
|
||||
expect(fallbackAGenerateContentStream).toHaveBeenCalledTimes(1);
|
||||
expect(fallbackBGenerateContentStream).toHaveBeenCalledTimes(1);
|
||||
expect(chat.getLastPromptTokenCount()).toBe(0);
|
||||
expect(
|
||||
events.some(
|
||||
(event) =>
|
||||
|
|
@ -7338,6 +7350,88 @@ describe('GeminiChat', async () => {
|
|||
).toBe(true);
|
||||
});
|
||||
|
||||
it('stamps fallback-served counts under the request route key (#9454)', async () => {
|
||||
vi.mocked(mockConfig.getContentGeneratorConfig).mockReturnValue({
|
||||
authType: AuthType.USE_GEMINI,
|
||||
model: 'test-model',
|
||||
maxRetries: 0,
|
||||
});
|
||||
vi.mocked(mockConfig.getModelFallbacks).mockReturnValue(['fallback-b']);
|
||||
|
||||
const fallbackBGenerateContentStream = vi.fn();
|
||||
const resolveForModel = vi.fn().mockResolvedValue({
|
||||
contentGenerator: {
|
||||
generateContent: vi.fn(),
|
||||
generateContentStream: fallbackBGenerateContentStream,
|
||||
countTokens: vi.fn(),
|
||||
embedContent: vi.fn(),
|
||||
batchEmbedContents: vi.fn(),
|
||||
useSummarizedThinking: vi.fn().mockReturnValue(false),
|
||||
} as unknown as ContentGenerator,
|
||||
contentGeneratorConfig: { modalities: {} },
|
||||
retryAuthType: AuthType.USE_GEMINI,
|
||||
retryErrorCodes: undefined,
|
||||
model: 'fallback-b',
|
||||
});
|
||||
vi.mocked(mockConfig.getBaseLlmClient).mockReturnValue({
|
||||
resolveForModel,
|
||||
} as unknown as ReturnType<typeof mockConfig.getBaseLlmClient>);
|
||||
vi.mocked(mockConfig.getModelRouteIdentity).mockImplementation((model) =>
|
||||
model ? `${model}@route` : 'gemini-pro@test0001',
|
||||
);
|
||||
|
||||
const capacityError = Object.assign(
|
||||
new StreamContentError(
|
||||
'{"error":{"code":"429","message":"Throttling: TPM(1/1)"}}',
|
||||
),
|
||||
{ status: 429 },
|
||||
);
|
||||
vi.mocked(
|
||||
mockContentGenerator.generateContentStream,
|
||||
).mockResolvedValueOnce(
|
||||
(async function* () {
|
||||
yield {
|
||||
usageMetadata: { promptTokenCount: 10, totalTokenCount: 10 },
|
||||
} as GenerateContentResponse;
|
||||
throw capacityError;
|
||||
})(),
|
||||
);
|
||||
fallbackBGenerateContentStream.mockResolvedValueOnce(
|
||||
(async function* () {
|
||||
yield {
|
||||
candidates: [
|
||||
{
|
||||
content: {
|
||||
role: 'model',
|
||||
parts: [{ text: 'fallback-b ok' }],
|
||||
},
|
||||
finishReason: 'STOP',
|
||||
},
|
||||
],
|
||||
usageMetadata: { promptTokenCount: 99_999 },
|
||||
} as unknown as GenerateContentResponse;
|
||||
})(),
|
||||
);
|
||||
|
||||
const stream = await chat.sendMessageStream(
|
||||
'test-model',
|
||||
{ message: [{ text: 'test' }] },
|
||||
'prompt-fallback-route-stamp',
|
||||
);
|
||||
for await (const _ of stream) {
|
||||
/* consume */
|
||||
}
|
||||
|
||||
// The session-token-limit gate in Client reads the count keyed by the
|
||||
// REQUEST route. A fallback serves on behalf of the same request (the
|
||||
// session model never changes), so its count must survive that keyed
|
||||
// read instead of being invalidated as a foreign route's (#9454).
|
||||
expect(chat.getLastPromptTokenCount('test-model@route')).toBe(99_999);
|
||||
// The count still belongs to the serving turn's request route: a read
|
||||
// for a different route invalidates it as before.
|
||||
expect(chat.getLastPromptTokenCount('other-model@route')).toBe(0);
|
||||
});
|
||||
|
||||
it('skips a fallback alias that resolves to the current model', async () => {
|
||||
vi.mocked(mockConfig.getContentGeneratorConfig).mockReturnValue({
|
||||
authType: AuthType.USE_GEMINI,
|
||||
|
|
@ -15558,6 +15652,818 @@ describe('GeminiChat', async () => {
|
|||
});
|
||||
});
|
||||
|
||||
// Route-scoped token counts (#9454): API-reported prompt/output token
|
||||
// counts describe the serialization of the route (model + auth type +
|
||||
// endpoint) that produced them. A /model switch rebuilds the content
|
||||
// generator but keeps this GeminiChat instance, so counts recorded for the
|
||||
// previous route must be invalidated — otherwise they anchor admission,
|
||||
// clamp, and compression decisions for a different serialization.
|
||||
describe('route-scoped token counts (#9454)', () => {
|
||||
const switchRoute = (routeKey: string) => {
|
||||
vi.mocked(mockConfig.getModelRouteIdentity).mockReturnValue(routeKey);
|
||||
};
|
||||
|
||||
it('invalidates API-reported counts when the model route changes', () => {
|
||||
// Count reported by the pre-switch route (authoritative, not estimated).
|
||||
chat.setLastPromptTokenCount(691_000, false);
|
||||
expect(chat.getLastPromptTokenCount()).toBe(691_000);
|
||||
|
||||
// Simulate /model switching to a different route; the same chat
|
||||
// instance survives with its history.
|
||||
switchRoute('anthropic-model@beef1234');
|
||||
|
||||
// The stale count must not size requests for the new route: safety
|
||||
// decisions fall back to the history-walk estimate (count 0).
|
||||
expect(chat.getLastPromptTokenCount()).toBe(0);
|
||||
expect(chat.getLastOutputTokenCount()).toBe(0);
|
||||
expect(chat.isLastPromptTokenCountEstimated()).toBe(false);
|
||||
// The telemetry mirror must drop the stale count too, or the session
|
||||
// token-limit gate and compression banners keep using it.
|
||||
expect(uiTelemetryService.setLastPromptTokenCount).toHaveBeenCalledWith(
|
||||
0,
|
||||
);
|
||||
});
|
||||
|
||||
it('keeps counts authoritative while the route is unchanged', () => {
|
||||
chat.setLastPromptTokenCount(50_000, false);
|
||||
|
||||
// Repeated reads on the same route keep the API-authoritative count.
|
||||
expect(chat.getLastPromptTokenCount()).toBe(50_000);
|
||||
expect(chat.getLastOutputTokenCount()).toBe(0);
|
||||
expect(chat.isLastPromptTokenCountEstimated()).toBe(false);
|
||||
expect(chat.getLastPromptTokenCount()).toBe(50_000);
|
||||
});
|
||||
|
||||
it('keeps a foreign count intact across a keyless display read (#9506)', () => {
|
||||
// Counts stamped under one route key (e.g. the vision bridge's
|
||||
// full-turn selector route) must survive a keyless read: /context
|
||||
// calls the getters with no argument, which defaults to the ACTIVE
|
||||
// route key and used to zero the only slot before the
|
||||
// session-token-limit gate's keyed read got to it.
|
||||
chat.setLastPromptTokenCount(500_000, false);
|
||||
switchRoute('other-active@route');
|
||||
|
||||
// The foreign count must not leak to the active route...
|
||||
expect(chat.getLastPromptTokenCount()).toBe(0);
|
||||
expect(uiTelemetryService.setLastPromptTokenCount).toHaveBeenCalledWith(
|
||||
0,
|
||||
);
|
||||
// ...and the crossing must not have destroyed it: the gate's keyed
|
||||
// read for the original route restores the exact API-reported value.
|
||||
expect(chat.getLastPromptTokenCount('gemini-pro@test0001')).toBe(500_000);
|
||||
expect(chat.getLastPromptTokenCount()).toBe(0);
|
||||
});
|
||||
|
||||
it('restores retained counts when the route switches back (#9506)', () => {
|
||||
chat.seedResumeTokenCounts(321, 45, true);
|
||||
switchRoute('anthropic-model@beef1234');
|
||||
expect(chat.getLastPromptTokenCount()).toBe(0);
|
||||
expect(chat.getLastOutputTokenCount()).toBe(0);
|
||||
|
||||
// A turn returning to the original route reads its exact retained
|
||||
// counts — prompt, previous-output, and provenance — instead of the
|
||||
// destructive zero a foreign touch used to leave behind.
|
||||
switchRoute('gemini-pro@test0001');
|
||||
expect(chat.getLastPromptTokenCount()).toBe(321);
|
||||
expect(chat.getLastOutputTokenCount()).toBe(45);
|
||||
expect(chat.isLastPromptTokenCountEstimated()).toBe(true);
|
||||
});
|
||||
|
||||
it('invalidates seeded resume counts after a later route change', () => {
|
||||
chat.seedResumeTokenCounts(321, 45, false);
|
||||
expect(chat.getLastPromptTokenCount()).toBe(321);
|
||||
expect(chat.getLastOutputTokenCount()).toBe(45);
|
||||
|
||||
switchRoute('other-model@1234abcd');
|
||||
|
||||
expect(chat.getLastPromptTokenCount()).toBe(0);
|
||||
expect(chat.getLastOutputTokenCount()).toBe(0);
|
||||
});
|
||||
|
||||
it('accepts counts recorded on the new route after a switch', () => {
|
||||
chat.setLastPromptTokenCount(691_000, false);
|
||||
switchRoute('anthropic-model@beef1234');
|
||||
expect(chat.getLastPromptTokenCount()).toBe(0);
|
||||
|
||||
// First response on the new route re-establishes authoritative counts.
|
||||
chat.setLastPromptTokenCount(120_000, false);
|
||||
expect(chat.getLastPromptTokenCount()).toBe(120_000);
|
||||
expect(chat.isLastPromptTokenCountEstimated()).toBe(false);
|
||||
});
|
||||
|
||||
it('invalidates a stale count before sending on the new route', async () => {
|
||||
chat.setLastPromptTokenCount(691_000, false);
|
||||
switchRoute('anthropic-model@beef1234');
|
||||
vi.mocked(mockContentGenerator.generateContentStream).mockImplementation(
|
||||
async () => {
|
||||
expect(
|
||||
uiTelemetryService.setLastPromptTokenCount,
|
||||
).toHaveBeenCalledWith(0);
|
||||
return (async function* () {
|
||||
yield {
|
||||
candidates: [
|
||||
{
|
||||
content: { parts: [{ text: 'ok' }] },
|
||||
finishReason: 'STOP',
|
||||
},
|
||||
],
|
||||
} as unknown as GenerateContentResponse;
|
||||
})();
|
||||
},
|
||||
);
|
||||
|
||||
const stream = await chat.sendMessageStream(
|
||||
'test-model',
|
||||
{ message: 'new route' },
|
||||
'prompt-route-switch',
|
||||
);
|
||||
for await (const _ of stream) {
|
||||
/* consume */
|
||||
}
|
||||
});
|
||||
|
||||
it('invalidates a stale route count before manual compression sizing', async () => {
|
||||
// Authoritative count recorded by the pre-switch route. Manual
|
||||
// /compress reaches tryCompress without sendMessageStream's entry
|
||||
// invalidation, and tryCompress reads the count field directly, so it
|
||||
// must drop the foreign count itself before admission/sizing.
|
||||
chat.setLastPromptTokenCount(691_000, false);
|
||||
switchRoute('anthropic-model@beef1234');
|
||||
|
||||
const compressSpy = vi.spyOn(
|
||||
ChatCompressionService.prototype,
|
||||
'compress',
|
||||
);
|
||||
compressSpy.mockResolvedValue({
|
||||
newHistory: null,
|
||||
info: {
|
||||
originalTokenCount: 0,
|
||||
newTokenCount: 0,
|
||||
compressionStatus: CompressionStatus.NOOP,
|
||||
},
|
||||
});
|
||||
|
||||
await chat.tryCompress('prompt-manual-compress', true);
|
||||
|
||||
// History is empty, so the estimate path sizes the attempt at 0 — the
|
||||
// stale 691_000 must not have anchored the compression decision.
|
||||
expect(compressSpy).toHaveBeenCalledTimes(1);
|
||||
expect(compressSpy.mock.calls[0]?.[1].originalTokenCount).toBe(0);
|
||||
});
|
||||
|
||||
it('invalidates a stale route count before fast-compression sizing', () => {
|
||||
// compressFast is the third entrypoint alongside sendMessageStream
|
||||
// and tryCompress: it reads the raw count field for its apiBaseline,
|
||||
// so it must drop a pre-switch count before sizing the new route.
|
||||
vi.mocked(mockConfig.getClearContextOnIdle).mockReturnValue({
|
||||
toolResultsThresholdMinutes: 30,
|
||||
toolResultsNumToKeep: 1,
|
||||
});
|
||||
const fastChat = new GeminiChat(
|
||||
mockConfig,
|
||||
config,
|
||||
[
|
||||
{ role: 'user', parts: [{ text: 'question' }] },
|
||||
{
|
||||
role: 'model',
|
||||
parts: [
|
||||
{ text: 'reasoning '.repeat(100), thought: true },
|
||||
{ text: 'answer' },
|
||||
],
|
||||
},
|
||||
],
|
||||
{
|
||||
recordChatCompression: vi.fn(),
|
||||
} as unknown as ConstructorParameters<typeof GeminiChat>[3],
|
||||
uiTelemetryService,
|
||||
);
|
||||
fastChat.setLastPromptTokenCount(691_000, false);
|
||||
switchRoute('anthropic-model@beef1234');
|
||||
|
||||
const result = fastChat.compressFast();
|
||||
|
||||
expect(result.info.compressionStatus).toBe(CompressionStatus.COMPRESSED);
|
||||
// The stale 691_000 must not anchor sizing for the new route: the
|
||||
// baseline falls back to the history-walk estimate.
|
||||
expect(result.info.originalTokenCount).toBeLessThan(691_000);
|
||||
expect(fastChat.getLastPromptTokenCount()).toBeLessThan(691_000);
|
||||
});
|
||||
|
||||
it('zeroes the telemetry cached-content mirror when invalidating a foreign count', () => {
|
||||
// The cached content count is written for the same foreign route's
|
||||
// last response; leaving it up next to the zeroed prompt count gives
|
||||
// /context an internally inconsistent capacity picture.
|
||||
chat.setLastPromptTokenCount(691_000, false);
|
||||
switchRoute('anthropic-model@beef1234');
|
||||
|
||||
expect(chat.getLastPromptTokenCount()).toBe(0);
|
||||
expect(
|
||||
uiTelemetryService.setLastCachedContentTokenCount,
|
||||
).toHaveBeenCalledWith(0);
|
||||
});
|
||||
|
||||
it('does not mirror cached content without a route-stamped prompt count', async () => {
|
||||
vi.mocked(mockContentGenerator.generateContentStream).mockResolvedValue(
|
||||
(async function* () {
|
||||
yield {
|
||||
candidates: [
|
||||
{
|
||||
content: { parts: [{ text: 'cached' }] },
|
||||
finishReason: 'STOP',
|
||||
},
|
||||
],
|
||||
usageMetadata: {
|
||||
promptTokenCount: 0,
|
||||
totalTokenCount: 0,
|
||||
cachedContentTokenCount: 42,
|
||||
},
|
||||
} as unknown as GenerateContentResponse;
|
||||
})(),
|
||||
);
|
||||
|
||||
const stream = await chat.sendMessageStream(
|
||||
'test-model',
|
||||
{ message: 'cached-only' },
|
||||
'prompt-cached-only',
|
||||
);
|
||||
for await (const _ of stream) {
|
||||
/* consume */
|
||||
}
|
||||
switchRoute('anthropic-model@beef1234');
|
||||
expect(chat.getLastPromptTokenCount()).toBe(0);
|
||||
|
||||
expect(
|
||||
uiTelemetryService.setLastCachedContentTokenCount,
|
||||
).not.toHaveBeenCalledWith(42);
|
||||
});
|
||||
|
||||
it('mirrors cached content alongside a route-stamped prompt count', async () => {
|
||||
// The cached-content mirror's only non-zero production write lives
|
||||
// inside the prompt-count guard; deleting it must not leave the suite
|
||||
// green. Consumed without a route switch, a cached-content response
|
||||
// must reach the /context cached-tokens line (#9454).
|
||||
vi.mocked(mockContentGenerator.generateContentStream).mockResolvedValue(
|
||||
(async function* () {
|
||||
yield {
|
||||
candidates: [
|
||||
{
|
||||
content: { parts: [{ text: 'cached' }] },
|
||||
finishReason: 'STOP',
|
||||
},
|
||||
],
|
||||
usageMetadata: {
|
||||
promptTokenCount: 100,
|
||||
totalTokenCount: 100,
|
||||
cachedContentTokenCount: 42,
|
||||
},
|
||||
} as unknown as GenerateContentResponse;
|
||||
})(),
|
||||
);
|
||||
|
||||
const stream = await chat.sendMessageStream(
|
||||
'test-model',
|
||||
{ message: 'cached-happy' },
|
||||
'prompt-cached-happy',
|
||||
);
|
||||
for await (const _ of stream) {
|
||||
/* consume */
|
||||
}
|
||||
|
||||
expect(chat.getLastPromptTokenCount()).toBe(100);
|
||||
expect(
|
||||
uiTelemetryService.setLastCachedContentTokenCount,
|
||||
).toHaveBeenCalledWith(42);
|
||||
});
|
||||
|
||||
it('restores the request route key when a failed hard-rescue rolls counts back', async () => {
|
||||
// Hard-rescue only fires for non-exact sends, whose request route key
|
||||
// can differ from the active route's. tryCompress re-stamps the key
|
||||
// to the ACTIVE route mid-rescue; the rollback must restore the key
|
||||
// alongside the counts, or the resurrected override-route count rides
|
||||
// the active key past the next send's entry invalidation (#9454).
|
||||
vi.mocked(mockConfig.getModelRouteIdentity).mockImplementation((model) =>
|
||||
model ? `${model}@route` : 'active@route',
|
||||
);
|
||||
const rescueChat = new GeminiChat(
|
||||
mockConfig,
|
||||
config,
|
||||
[
|
||||
{ role: 'user', parts: [{ text: 'earlier turn' }] },
|
||||
{ role: 'model', parts: [{ text: 'ack' }] },
|
||||
],
|
||||
{
|
||||
recordAssistantTurn: vi.fn(),
|
||||
} as unknown as ConstructorParameters<typeof GeminiChat>[3],
|
||||
uiTelemetryService,
|
||||
);
|
||||
// Authoritative count recorded by an earlier override-route turn.
|
||||
vi.mocked(mockConfig.getModelRouteIdentity).mockReturnValue(
|
||||
'override-model@route',
|
||||
);
|
||||
rescueChat.setLastPromptTokenCount(176_999, false);
|
||||
vi.mocked(mockConfig.getModelRouteIdentity).mockImplementation((model) =>
|
||||
model ? `${model}@route` : 'active@route',
|
||||
);
|
||||
|
||||
vi.spyOn(
|
||||
ChatCompressionService.prototype,
|
||||
'compress',
|
||||
).mockResolvedValueOnce({
|
||||
newHistory: [
|
||||
{ role: 'user', parts: [{ text: 'still large summary' }] },
|
||||
{ role: 'model', parts: [{ text: 'ack' }] },
|
||||
],
|
||||
info: {
|
||||
originalTokenCount: 180_000,
|
||||
newTokenCount: 177_000,
|
||||
compressionStatus: CompressionStatus.COMPRESSED,
|
||||
},
|
||||
});
|
||||
|
||||
await expect(
|
||||
rescueChat.sendMessageStream(
|
||||
'override-model',
|
||||
{ message: 'continue' },
|
||||
'prompt-hard-rescue-route-key-restore',
|
||||
),
|
||||
).rejects.toThrow(/compression status: COMPRESSED/i);
|
||||
|
||||
expect(mockContentGenerator.generateContentStream).not.toHaveBeenCalled();
|
||||
// The restored count belongs to the override route: an active-route
|
||||
// read must invalidate it, not inherit it.
|
||||
expect(rescueChat.getLastPromptTokenCount()).toBe(0);
|
||||
});
|
||||
|
||||
it('restores the retention map when a failed hard-rescue rolls counts back (#9506)', async () => {
|
||||
// The rescue's own compression consumes retained map entries
|
||||
// mid-flight (the service's keyless getter reads adopt the active
|
||||
// route) and a successful compression clears the map outright. The
|
||||
// rollback must restore the pre-rescue snapshot, or the resurrected
|
||||
// route's over-limit count survives nowhere and its next
|
||||
// session-token-limit gate read passes with 0.
|
||||
vi.mocked(mockConfig.getModelRouteIdentity).mockReturnValue(
|
||||
'active@route',
|
||||
);
|
||||
const rescueChat = new GeminiChat(
|
||||
mockConfig,
|
||||
config,
|
||||
[{ role: 'user', parts: [{ text: 'x'.repeat(720_000) }] }],
|
||||
{
|
||||
recordAssistantTurn: vi.fn(),
|
||||
} as unknown as ConstructorParameters<typeof GeminiChat>[3],
|
||||
uiTelemetryService,
|
||||
);
|
||||
rescueChat.setLastPromptTokenCount(190_000, false);
|
||||
vi.mocked(mockConfig.getModelRouteIdentity).mockImplementation((model) =>
|
||||
model ? `${model}@route` : 'active@route',
|
||||
);
|
||||
|
||||
vi.spyOn(
|
||||
ChatCompressionService.prototype,
|
||||
'compress',
|
||||
).mockImplementationOnce(async (chatToCompress) => {
|
||||
// Mirror the real service's unconditional keyless reads: they
|
||||
// adopt the ACTIVE route, consuming its retained entry mid-rescue.
|
||||
chatToCompress.getLastPromptTokenCount();
|
||||
chatToCompress.isLastPromptTokenCountEstimated();
|
||||
return {
|
||||
newHistory: [
|
||||
{ role: 'user', parts: [{ text: 'still large summary' }] },
|
||||
],
|
||||
info: {
|
||||
originalTokenCount: 180_000,
|
||||
newTokenCount: 178_000,
|
||||
compressionStatus: CompressionStatus.COMPRESSED,
|
||||
},
|
||||
};
|
||||
});
|
||||
|
||||
await expect(
|
||||
rescueChat.sendMessageStream(
|
||||
'override-model',
|
||||
{ message: 'continue' },
|
||||
'prompt-rescue-retention-map-restore',
|
||||
),
|
||||
).rejects.toThrow(/compression status: COMPRESSED/i);
|
||||
|
||||
expect(mockContentGenerator.generateContentStream).not.toHaveBeenCalled();
|
||||
// The active route's over-limit count survived the failed rescue:
|
||||
// its next keyed read must still see it, not 0.
|
||||
expect(rescueChat.getLastPromptTokenCount('active@route')).toBe(190_000);
|
||||
expect(rescueChat.getLastPromptTokenCount('override-model@route')).toBe(
|
||||
0,
|
||||
);
|
||||
});
|
||||
|
||||
it('restores the output token count when a failed hard-rescue rolls counts back (#9506)', async () => {
|
||||
// The rescue's COMPRESSED stamp zeroes lastOutputTokenCount via
|
||||
// setLastPromptTokenCount. The rollback restores the resurrected
|
||||
// prompt count, its provenance, the route key and the retention map
|
||||
// — it must restore the output half of the pair too, or the next
|
||||
// turn's additive prompt estimate (prompt + output + new content)
|
||||
// under-counts by the last response's size.
|
||||
vi.mocked(mockConfig.getModelRouteIdentity).mockReturnValue(
|
||||
'override-model@route',
|
||||
);
|
||||
const rescueChat = new GeminiChat(
|
||||
mockConfig,
|
||||
config,
|
||||
[
|
||||
{ role: 'user', parts: [{ text: 'earlier turn' }] },
|
||||
{ role: 'model', parts: [{ text: 'ack' }] },
|
||||
],
|
||||
{
|
||||
recordAssistantTurn: vi.fn(),
|
||||
} as unknown as ConstructorParameters<typeof GeminiChat>[3],
|
||||
uiTelemetryService,
|
||||
);
|
||||
// Authoritative count pair recorded by an earlier override-route turn.
|
||||
rescueChat.seedResumeTokenCounts(170_000, 8_000, false);
|
||||
vi.mocked(mockConfig.getModelRouteIdentity).mockImplementation((model) =>
|
||||
model ? `${model}@route` : 'active@route',
|
||||
);
|
||||
|
||||
vi.spyOn(
|
||||
ChatCompressionService.prototype,
|
||||
'compress',
|
||||
).mockResolvedValueOnce({
|
||||
newHistory: [
|
||||
{ role: 'user', parts: [{ text: 'still large summary' }] },
|
||||
{ role: 'model', parts: [{ text: 'ack' }] },
|
||||
],
|
||||
info: {
|
||||
originalTokenCount: 180_000,
|
||||
newTokenCount: 177_000,
|
||||
compressionStatus: CompressionStatus.COMPRESSED,
|
||||
},
|
||||
});
|
||||
|
||||
await expect(
|
||||
rescueChat.sendMessageStream(
|
||||
'override-model',
|
||||
{ message: 'continue' },
|
||||
'prompt-hard-rescue-output-restore',
|
||||
),
|
||||
).rejects.toThrow(/compression status: COMPRESSED/i);
|
||||
|
||||
expect(mockContentGenerator.generateContentStream).not.toHaveBeenCalled();
|
||||
// Reading through the override route (the resurrected slot's key)
|
||||
// must return the full pre-rescue pair, output half included.
|
||||
vi.mocked(mockConfig.getModelRouteIdentity).mockReturnValue(
|
||||
'override-model@route',
|
||||
);
|
||||
expect(rescueChat.getLastPromptTokenCount()).toBe(170_000);
|
||||
expect(rescueChat.getLastOutputTokenCount()).toBe(8_000);
|
||||
});
|
||||
|
||||
it('re-adopts the request route after the compression service flips the slots (#9506)', async () => {
|
||||
// ChatCompressionService.compress reads the KEYLESS count getters,
|
||||
// which adopt the ACTIVE route. On a non-exact override send whose
|
||||
// request route differs, that flips the slots back to the active
|
||||
// route's retained counts mid-rescue. When the summarization side
|
||||
// query then fails, the post-rescue stop check must size from the
|
||||
// honest history-walk estimate, not the flipped foreign count.
|
||||
vi.mocked(mockConfig.getModelRouteIdentity).mockReturnValue(
|
||||
'active@route',
|
||||
);
|
||||
const rescueChat = new GeminiChat(
|
||||
mockConfig,
|
||||
config,
|
||||
[{ role: 'user', parts: [{ text: 'x'.repeat(720_000) }] }],
|
||||
{
|
||||
recordAssistantTurn: vi.fn(),
|
||||
} as unknown as ConstructorParameters<typeof GeminiChat>[3],
|
||||
uiTelemetryService,
|
||||
);
|
||||
rescueChat.setLastPromptTokenCount(150_000, false);
|
||||
vi.mocked(mockConfig.getModelRouteIdentity).mockImplementation((model) =>
|
||||
model ? `${model}@route` : 'active@route',
|
||||
);
|
||||
|
||||
vi.spyOn(
|
||||
ChatCompressionService.prototype,
|
||||
'compress',
|
||||
).mockImplementationOnce(async (chatToCompress) => {
|
||||
chatToCompress.getLastPromptTokenCount();
|
||||
chatToCompress.isLastPromptTokenCountEstimated();
|
||||
return {
|
||||
newHistory: null,
|
||||
info: {
|
||||
originalTokenCount: 180_000,
|
||||
newTokenCount: 0,
|
||||
compressionStatus:
|
||||
CompressionStatus.COMPRESSION_FAILED_EMPTY_SUMMARY,
|
||||
},
|
||||
};
|
||||
});
|
||||
|
||||
await expect(
|
||||
rescueChat.sendMessageStream(
|
||||
'override-model',
|
||||
{ message: 'continue' },
|
||||
'prompt-rescue-flip-re-adopt',
|
||||
),
|
||||
).rejects.toThrow(/Context is too large to send safely/i);
|
||||
|
||||
expect(mockContentGenerator.generateContentStream).not.toHaveBeenCalled();
|
||||
// The active route's retained count survived the failed rescue.
|
||||
expect(rescueChat.getLastPromptTokenCount('active@route')).toBe(150_000);
|
||||
});
|
||||
|
||||
it('retains a foreign-keyed slot occupant when the usage stamp re-keys (#9506)', async () => {
|
||||
// Mid-send compression can leave the slots keyed to the ACTIVE route
|
||||
// when the response's usage report arrives for the REQUEST route.
|
||||
// The stamp must retain the displaced occupant before overwriting
|
||||
// it, or the active route's count is destroyed and its next keyed
|
||||
// read returns 0 — bypassing the session token limit.
|
||||
vi.mocked(mockConfig.getModelRouteIdentity).mockReturnValue(
|
||||
'active@route',
|
||||
);
|
||||
const stampChat = new GeminiChat(
|
||||
mockConfig,
|
||||
config,
|
||||
[{ role: 'user', parts: [{ text: 'x'.repeat(720_000) }] }],
|
||||
{
|
||||
recordAssistantTurn: vi.fn(),
|
||||
// Successful hard-rescue compression records after the
|
||||
// post-compression guard passes (deferred recording).
|
||||
recordChatCompression: vi.fn(),
|
||||
} as unknown as ConstructorParameters<typeof GeminiChat>[3],
|
||||
uiTelemetryService,
|
||||
);
|
||||
stampChat.setLastPromptTokenCount(150_000, false);
|
||||
vi.mocked(mockConfig.getModelRouteIdentity).mockImplementation((model) =>
|
||||
model ? `${model}@route` : 'active@route',
|
||||
);
|
||||
|
||||
vi.spyOn(
|
||||
ChatCompressionService.prototype,
|
||||
'compress',
|
||||
).mockImplementationOnce(async (chatToCompress) => {
|
||||
chatToCompress.getLastPromptTokenCount();
|
||||
chatToCompress.isLastPromptTokenCountEstimated();
|
||||
return {
|
||||
newHistory: [{ role: 'user', parts: [{ text: 'summary' }] }],
|
||||
info: {
|
||||
originalTokenCount: 180_000,
|
||||
newTokenCount: 60_000,
|
||||
compressionStatus: CompressionStatus.COMPRESSED,
|
||||
},
|
||||
};
|
||||
});
|
||||
vi.mocked(
|
||||
mockContentGenerator.generateContentStream,
|
||||
).mockResolvedValueOnce(
|
||||
(async function* () {
|
||||
yield {
|
||||
candidates: [
|
||||
{
|
||||
content: { parts: [{ text: 'ok' }], role: 'model' },
|
||||
finishReason: 'STOP',
|
||||
index: 0,
|
||||
},
|
||||
],
|
||||
usageMetadata: {
|
||||
promptTokenCount: 61_000,
|
||||
totalTokenCount: 62_000,
|
||||
},
|
||||
} as unknown as GenerateContentResponse;
|
||||
})(),
|
||||
);
|
||||
|
||||
const stream = await stampChat.sendMessageStream(
|
||||
'override-model',
|
||||
{ message: 'continue' },
|
||||
'prompt-stamp-retains-occupant',
|
||||
);
|
||||
for await (const _ of stream) {
|
||||
/* consume */
|
||||
}
|
||||
|
||||
// The request route's fresh API report occupies the slots...
|
||||
expect(stampChat.getLastPromptTokenCount('override-model@route')).toBe(
|
||||
61_000,
|
||||
);
|
||||
// ...and the active route's post-compression count was retained,
|
||||
// not destroyed by the re-keying stamp.
|
||||
expect(stampChat.getLastPromptTokenCount('active@route')).toBe(60_000);
|
||||
});
|
||||
|
||||
it('stamps the compressed count under the request route when the send ends without usage (#9506)', async () => {
|
||||
// In-send compression runs for the REQUEST route, but
|
||||
// setLastPromptTokenCount re-keys the fresh count to the ACTIVE
|
||||
// route. If the request then ends without a usage report (abort,
|
||||
// 400 — the reactive-overflow path exists for exactly those), the
|
||||
// request route never stamps a count of its own, and its next
|
||||
// session-token-limit gate read passes with 0 even though the
|
||||
// shared compressed history's exact measure is on record.
|
||||
vi.mocked(mockConfig.getModelRouteIdentity).mockReturnValue(
|
||||
'active@route',
|
||||
);
|
||||
const stampChat = new GeminiChat(
|
||||
mockConfig,
|
||||
config,
|
||||
[{ role: 'user', parts: [{ text: 'x'.repeat(720_000) }] }],
|
||||
{
|
||||
recordAssistantTurn: vi.fn(),
|
||||
// Successful hard-rescue compression records after the
|
||||
// post-compression guard passes (deferred recording).
|
||||
recordChatCompression: vi.fn(),
|
||||
} as unknown as ConstructorParameters<typeof GeminiChat>[3],
|
||||
uiTelemetryService,
|
||||
);
|
||||
stampChat.setLastPromptTokenCount(150_000, false);
|
||||
vi.mocked(mockConfig.getModelRouteIdentity).mockImplementation((model) =>
|
||||
model ? `${model}@route` : 'active@route',
|
||||
);
|
||||
|
||||
vi.spyOn(
|
||||
ChatCompressionService.prototype,
|
||||
'compress',
|
||||
).mockImplementationOnce(async (chatToCompress) => {
|
||||
chatToCompress.getLastPromptTokenCount();
|
||||
chatToCompress.isLastPromptTokenCountEstimated();
|
||||
return {
|
||||
newHistory: [{ role: 'user', parts: [{ text: 'summary' }] }],
|
||||
info: {
|
||||
originalTokenCount: 180_000,
|
||||
newTokenCount: 60_000,
|
||||
compressionStatus: CompressionStatus.COMPRESSED,
|
||||
},
|
||||
};
|
||||
});
|
||||
vi.mocked(
|
||||
mockContentGenerator.generateContentStream,
|
||||
).mockResolvedValueOnce(
|
||||
(async function* () {
|
||||
yield {
|
||||
candidates: [
|
||||
{
|
||||
content: { parts: [{ text: 'ok' }], role: 'model' },
|
||||
finishReason: 'STOP',
|
||||
index: 0,
|
||||
},
|
||||
],
|
||||
// No usageMetadata: the request route never stamps a count of
|
||||
// its own, so the compression stamp must be readable under it.
|
||||
} as unknown as GenerateContentResponse;
|
||||
})(),
|
||||
);
|
||||
|
||||
const stream = await stampChat.sendMessageStream(
|
||||
'override-model',
|
||||
{ message: 'continue' },
|
||||
'prompt-compression-stamps-request-route',
|
||||
);
|
||||
for await (const _ of stream) {
|
||||
/* consume */
|
||||
}
|
||||
|
||||
// The request route's keyed read sees the compressed history's
|
||||
// count instead of passing the gate with 0...
|
||||
expect(stampChat.getLastPromptTokenCount('override-model@route')).toBe(
|
||||
60_000,
|
||||
);
|
||||
// ...and the active route still sees it through the retained entry,
|
||||
// because the compressed history is shared by every route.
|
||||
expect(stampChat.getLastPromptTokenCount('active@route')).toBe(60_000);
|
||||
});
|
||||
|
||||
it('drops stale retained counts when a successful compression rewrites the history (#9506)', async () => {
|
||||
// Compression rewrites the shared history every retained entry
|
||||
// sizes. Retained pre-compression counts must not survive the
|
||||
// success path, or a later keyed read adopts one and the session-
|
||||
// token-limit gate blocks a prompt that fits the compressed history.
|
||||
chat.setLastPromptTokenCount(691_000, false);
|
||||
switchRoute('override@route');
|
||||
// The crossing retains the original route's count under its own key.
|
||||
expect(chat.getLastPromptTokenCount()).toBe(0);
|
||||
|
||||
vi.spyOn(
|
||||
ChatCompressionService.prototype,
|
||||
'compress',
|
||||
).mockResolvedValueOnce({
|
||||
newHistory: [{ role: 'user', parts: [{ text: 'summary' }] }],
|
||||
info: {
|
||||
originalTokenCount: 691_000,
|
||||
newTokenCount: 50_000,
|
||||
compressionStatus: CompressionStatus.COMPRESSED,
|
||||
},
|
||||
});
|
||||
|
||||
await chat.tryCompress('prompt-compression-drops-retained', true);
|
||||
|
||||
expect(chat.getLastPromptTokenCount()).toBe(50_000);
|
||||
expect(chat.getLastPromptTokenCount('gemini-pro@test0001')).toBe(0);
|
||||
});
|
||||
|
||||
it('drops all retained counts when fast compression rewrites the history (#9506)', () => {
|
||||
// compressFast rewrites the same shared history the other routes'
|
||||
// retained entries size; clearing only the active route's entry
|
||||
// would leave stale pre-compression counts adoptable by later keyed
|
||||
// reads.
|
||||
vi.mocked(mockConfig.getClearContextOnIdle).mockReturnValue({
|
||||
toolResultsThresholdMinutes: 30,
|
||||
toolResultsNumToKeep: 1,
|
||||
});
|
||||
const fastChat = new GeminiChat(
|
||||
mockConfig,
|
||||
config,
|
||||
[
|
||||
{ role: 'user', parts: [{ text: 'question' }] },
|
||||
{
|
||||
role: 'model',
|
||||
parts: [
|
||||
{ text: 'reasoning '.repeat(100), thought: true },
|
||||
{ text: 'answer' },
|
||||
],
|
||||
},
|
||||
],
|
||||
{
|
||||
recordChatCompression: vi.fn(),
|
||||
} as unknown as ConstructorParameters<typeof GeminiChat>[3],
|
||||
uiTelemetryService,
|
||||
);
|
||||
fastChat.setLastPromptTokenCount(691_000, false);
|
||||
// Cross routes so the count is retained under the original key.
|
||||
switchRoute('other-route@fast');
|
||||
expect(fastChat.getLastPromptTokenCount()).toBe(0);
|
||||
|
||||
const result = fastChat.compressFast();
|
||||
|
||||
expect(result.info.compressionStatus).toBe(CompressionStatus.COMPRESSED);
|
||||
// The retained pre-compression entry did not survive the rewrite.
|
||||
expect(fastChat.getLastPromptTokenCount('gemini-pro@test0001')).toBe(0);
|
||||
});
|
||||
|
||||
it('does not anchor an exact route output clamp on the active route count', async () => {
|
||||
// Authoritative count recorded by the ACTIVE route. An exact `\0`
|
||||
// route send targets a different serialization, so its output clamp
|
||||
// must not read this count — the entry invalidation has to compare
|
||||
// against the request's route, not the active one.
|
||||
chat.setLastPromptTokenCount(691_000, false);
|
||||
|
||||
const routeGenerateContentStream = vi.fn().mockResolvedValue(
|
||||
(async function* () {
|
||||
yield {
|
||||
candidates: [
|
||||
{
|
||||
content: { parts: [{ text: 'ok' }] },
|
||||
finishReason: 'STOP',
|
||||
},
|
||||
],
|
||||
} as unknown as GenerateContentResponse;
|
||||
})(),
|
||||
);
|
||||
const routeGenerator = {
|
||||
...mockContentGenerator,
|
||||
generateContentStream: routeGenerateContentStream,
|
||||
} as ContentGenerator;
|
||||
const resolveForModel = vi.fn().mockResolvedValue({
|
||||
contentGenerator: routeGenerator,
|
||||
contentGeneratorConfig: {
|
||||
model: 'vision-agent',
|
||||
authType: AuthType.USE_OPENAI,
|
||||
maxRetries: 0,
|
||||
// Large enough that the zeroed-count estimate path (history walk
|
||||
// + ESTIMATE_CLAMP_OVERHEAD_PAD + clamp margin) still leaves room
|
||||
// for the full explicit ceiling below.
|
||||
contextWindowSize: 64_000,
|
||||
modalities: {},
|
||||
},
|
||||
retryAuthType: AuthType.USE_OPENAI,
|
||||
model: 'vision-agent',
|
||||
});
|
||||
vi.mocked(mockConfig.getBaseLlmClient).mockReturnValue({
|
||||
resolveForModel,
|
||||
} as unknown as ReturnType<typeof mockConfig.getBaseLlmClient>);
|
||||
vi.mocked(mockConfig.getModelRouteIdentity).mockImplementation((model) =>
|
||||
model ? `${model}@route` : 'gemini-pro@test0001',
|
||||
);
|
||||
|
||||
const selector = 'openai:vision-agent\0https://vision.example.com/v1\0';
|
||||
const stream = await chat.sendMessageStream(
|
||||
selector,
|
||||
{
|
||||
message: 'clamp probe',
|
||||
config: { maxOutputTokens: 8_000 },
|
||||
},
|
||||
'prompt-exact-route-clamp',
|
||||
);
|
||||
for await (const _ of stream) {
|
||||
/* consume */
|
||||
}
|
||||
|
||||
// With the foreign count zeroed the estimate leaves room for the full
|
||||
// 8_000 ceiling. Had the active route's 691_000 anchored the clamp,
|
||||
// the request would have been floored at MIN_CLAMPED_OUTPUT_TOKENS.
|
||||
const routeRequest = routeGenerateContentStream.mock.calls[0]?.[0] as {
|
||||
config?: { maxOutputTokens?: number };
|
||||
};
|
||||
expect(routeRequest.config?.maxOutputTokens).toBe(8_000);
|
||||
});
|
||||
});
|
||||
|
||||
// The circuit breaker is the three-strike replacement for the old
|
||||
// single-shot hasFailedCompressionAttempt lock. After
|
||||
// MAX_CONSECUTIVE_FAILURES failures the chat stops trying to auto-compact
|
||||
|
|
|
|||
|
|
@ -469,6 +469,14 @@ interface TryCompressOptions {
|
|||
precomputedEffectiveTokens?: number;
|
||||
/** Per-request overrides needed to preserve the main request cache prefix. */
|
||||
requestGenerationConfig?: GenerateContentConfig;
|
||||
/**
|
||||
* Route the enclosing send targets. The entry adoption compares against
|
||||
* this instead of the active route, so an in-send compression never
|
||||
* re-adopts counts the active route retained while the request targets
|
||||
* another one (#9506). Omitted by between-sends callers (manual
|
||||
* `/compress`), which compress the active route's state.
|
||||
*/
|
||||
requestRouteKey?: string;
|
||||
/**
|
||||
* Delay writing the compression checkpoint until the caller has run any
|
||||
* post-compression guards that may roll the in-memory chat state back.
|
||||
|
|
@ -525,6 +533,14 @@ const TRANSPORT_STREAM_RETRY_CONFIG = {
|
|||
*/
|
||||
const ESTIMATE_CLAMP_OVERHEAD_PAD = 20_000;
|
||||
|
||||
/**
|
||||
* Cap on how many routes' token counts are retained while their route is
|
||||
* not the one owning the chat's count slots (#9506). Route identities are
|
||||
* bounded by the session's model routes, so this only guards pathological
|
||||
* selector churn; eviction is FIFO.
|
||||
*/
|
||||
const MAX_RETAINED_ROUTE_COUNTS = 8;
|
||||
|
||||
/**
|
||||
* Max recovery attempts when the escalated response is also truncated.
|
||||
* Each attempt keeps the partial response in history and injects a recovery
|
||||
|
|
@ -1856,6 +1872,35 @@ export class GeminiChat {
|
|||
*/
|
||||
private lastOutputTokenCount = 0;
|
||||
|
||||
/**
|
||||
* Route identity (model + auth type + endpoint; see
|
||||
* Config.getModelRouteIdentity) of the content generator that produced
|
||||
* the counts above. API-reported sizes are wire-specific: one route's
|
||||
* count cannot size another route's serialization (#9454). Undefined
|
||||
* until the first count is recorded.
|
||||
*/
|
||||
private tokenCountsRouteKey: string | undefined = undefined;
|
||||
|
||||
/**
|
||||
* Token counts retained for routes other than the one currently owning
|
||||
* the slots above, keyed by route identity (#9506). Crossing routes
|
||||
* retains the current slots here and adopts the target's entry back
|
||||
* instead of destroying the value: API-reported sizes are per-route
|
||||
* state that a later turn on the same route still needs — most
|
||||
* critically the session-token-limit gate, whose keyed read would
|
||||
* otherwise see 0 after any foreign-route touch between turns.
|
||||
* Invariant: never holds an entry for {@link tokenCountsRouteKey}.
|
||||
*/
|
||||
private readonly tokenCountsByRouteKey = new Map<
|
||||
string,
|
||||
{
|
||||
promptTokenCount: number;
|
||||
promptTokenCountIsEstimated: boolean;
|
||||
outputTokenCount: number;
|
||||
cachedContentTokenCount: number;
|
||||
}
|
||||
>();
|
||||
|
||||
/**
|
||||
* Number of consecutive auto-compaction failures for this chat. The
|
||||
* cheap-gate NOOPs once this reaches MAX_CONSECUTIVE_FAILURES (default 3)
|
||||
|
|
@ -1966,6 +2011,134 @@ export class GeminiChat {
|
|||
this.manualPlanExitNoticesEnabled = true;
|
||||
}
|
||||
|
||||
/**
|
||||
* Identity of the currently active model route. Optional chaining keeps
|
||||
* partial Config test mocks (`{} as Config`) from throwing on count
|
||||
* reads/writes; a missing identity degrades to one stable key, i.e. no
|
||||
* route-change invalidation.
|
||||
*/
|
||||
private currentRouteKey(): string {
|
||||
return this.config.getModelRouteIdentity?.() ?? '';
|
||||
}
|
||||
|
||||
/**
|
||||
* Make the single-slot token counters describe the route identified by
|
||||
* `targetRouteKey` (default: the active route). Counts recorded for a
|
||||
* different route must not anchor admission, output clamping, or
|
||||
* compression decisions for this one (`/model` switches rebuild the
|
||||
* content generator but keep this chat instance; #9454).
|
||||
*
|
||||
* The crossing is NON-DESTRUCTIVE (#9506): the current slots are
|
||||
* retained in {@link tokenCountsByRouteKey} under their own route key,
|
||||
* and the target's retained entry — if any — is adopted back into the
|
||||
* slots. Zeroing a foreign count outright let any foreign-route touch
|
||||
* between two turns destroy the value before the session-token-limit
|
||||
* gate (the only yield site of `SessionTokenLimitExceeded`) could read
|
||||
* it back keyed by its request route. With retention, a route with no
|
||||
* counts of its own still falls back to the history-walk estimate
|
||||
* (slots 0), with reactive overflow recovery as the safety net, while a
|
||||
* turn returning to a route that has counts reads the exact
|
||||
* API-reported values.
|
||||
*
|
||||
* Defaults to comparing against the ACTIVE route (lazy reads on the
|
||||
* getters). Send paths pass the route the upcoming request actually
|
||||
* targets so a foreign count cannot anchor that request's decisions even
|
||||
* when the active route owns it — e.g. an exact `\0` route selector, or
|
||||
* a non-exact send whose `model` param overrides the active model.
|
||||
*
|
||||
* The telemetry mirror is display-only state: it is resynchronized here
|
||||
* (adopted or zeroed alongside the slots), so between a `/model` switch
|
||||
* and the next chat touch the UI counters may briefly show the previous
|
||||
* route's counts. Decision paths never read the mirror, only the
|
||||
* route-aware chat getters above.
|
||||
*/
|
||||
private adoptTokenCountsForRoute(targetRouteKey?: string): void {
|
||||
if (
|
||||
this.lastPromptTokenCount === 0 &&
|
||||
this.lastOutputTokenCount === 0 &&
|
||||
this.tokenCountsByRouteKey.size === 0
|
||||
) {
|
||||
return;
|
||||
}
|
||||
// Resolve the active-route default only AFTER the zero-count fast path:
|
||||
// computing a route identity (SHA-256 digest + config lookups) on every
|
||||
// count read while both counts are 0 (and nothing is retained) would
|
||||
// defeat the guard above.
|
||||
targetRouteKey ??= this.currentRouteKey();
|
||||
if (this.tokenCountsRouteKey === targetRouteKey) {
|
||||
return;
|
||||
}
|
||||
const retained = this.tokenCountsByRouteKey.get(targetRouteKey);
|
||||
if (retained) {
|
||||
this.tokenCountsByRouteKey.delete(targetRouteKey);
|
||||
this.retainCurrentTokenCounts();
|
||||
debugLogger.debug(
|
||||
`[token-counts] restoring retained counts for route ${targetRouteKey}`,
|
||||
);
|
||||
this.lastPromptTokenCount = retained.promptTokenCount;
|
||||
this.lastPromptTokenCountIsEstimated =
|
||||
retained.promptTokenCountIsEstimated;
|
||||
this.lastOutputTokenCount = retained.outputTokenCount;
|
||||
this.tokenCountsRouteKey = targetRouteKey;
|
||||
this.telemetryService?.setLastPromptTokenCount(retained.promptTokenCount);
|
||||
this.telemetryService?.setLastCachedContentTokenCount(
|
||||
retained.cachedContentTokenCount,
|
||||
);
|
||||
return;
|
||||
}
|
||||
debugLogger.debug(
|
||||
`[token-counts] route changed; retaining counts recorded for ` +
|
||||
`${this.tokenCountsRouteKey ?? 'unknown'} (now ${targetRouteKey})`,
|
||||
);
|
||||
this.retainCurrentTokenCounts();
|
||||
// Raw assignment on purpose: setLastPromptTokenCount would re-attribute
|
||||
// the zero slot to the ACTIVE route. The slot is attributed to the
|
||||
// TARGET route instead so it can never collide with the just-retained
|
||||
// entry (retained under the evicted slot's key, which differs from the
|
||||
// target) — a colliding key would make the next keyed read for the
|
||||
// retained route early-return the zero slot without consulting the map.
|
||||
this.lastPromptTokenCount = 0;
|
||||
this.lastPromptTokenCountIsEstimated = false;
|
||||
this.lastOutputTokenCount = 0;
|
||||
this.tokenCountsRouteKey = targetRouteKey;
|
||||
// Keep the telemetry mirror in sync, or the UI context counters
|
||||
// and compression banners keep reading the foreign count. The cached
|
||||
// content count belongs to the same foreign route's last response.
|
||||
this.telemetryService?.setLastPromptTokenCount(0);
|
||||
this.telemetryService?.setLastCachedContentTokenCount(0);
|
||||
}
|
||||
|
||||
/**
|
||||
* Save the current slots into {@link tokenCountsByRouteKey} under their
|
||||
* owning route key so a later read keyed back to that route restores the
|
||||
* exact API-reported values. Zero slots carry nothing worth retaining;
|
||||
* the telemetry mirror still holds the owning route's cached-content
|
||||
* count at this point, so it is captured here too.
|
||||
*/
|
||||
private retainCurrentTokenCounts(): void {
|
||||
if (
|
||||
this.tokenCountsRouteKey === undefined ||
|
||||
(this.lastPromptTokenCount === 0 && this.lastOutputTokenCount === 0)
|
||||
) {
|
||||
return;
|
||||
}
|
||||
if (this.tokenCountsByRouteKey.size >= MAX_RETAINED_ROUTE_COUNTS) {
|
||||
const oldestKey = this.tokenCountsByRouteKey.keys().next().value;
|
||||
if (oldestKey !== undefined) {
|
||||
this.tokenCountsByRouteKey.delete(oldestKey);
|
||||
}
|
||||
}
|
||||
this.tokenCountsByRouteKey.set(this.tokenCountsRouteKey, {
|
||||
promptTokenCount: this.lastPromptTokenCount,
|
||||
promptTokenCountIsEstimated: this.lastPromptTokenCountIsEstimated,
|
||||
outputTokenCount: this.lastOutputTokenCount,
|
||||
// Optional chaining keeps partial telemetry test mocks from throwing
|
||||
// (same convention as currentRouteKey's Config lookups).
|
||||
cachedContentTokenCount:
|
||||
this.telemetryService?.getLastCachedContentTokenCount?.() ?? 0,
|
||||
});
|
||||
}
|
||||
|
||||
/**
|
||||
* Most recent prompt-token count reported by the model for *this* chat,
|
||||
* mirroring the value in {@link UiTelemetryService} for the main session.
|
||||
|
|
@ -1973,12 +2146,14 @@ export class GeminiChat {
|
|||
* count for compaction decisions, so this is always populated regardless
|
||||
* of whether the global telemetry is updated.
|
||||
*/
|
||||
getLastPromptTokenCount(): number {
|
||||
getLastPromptTokenCount(targetRouteKey?: string): number {
|
||||
this.adoptTokenCountsForRoute(targetRouteKey);
|
||||
return this.lastPromptTokenCount;
|
||||
}
|
||||
|
||||
/** Previous model-response tokens used by the next prompt estimate. */
|
||||
getLastOutputTokenCount(): number {
|
||||
this.adoptTokenCountsForRoute();
|
||||
return this.lastOutputTokenCount;
|
||||
}
|
||||
|
||||
|
|
@ -2053,9 +2228,16 @@ export class GeminiChat {
|
|||
this.lastPromptTokenCount = count;
|
||||
this.lastPromptTokenCountIsEstimated = isEstimated;
|
||||
this.lastOutputTokenCount = 0;
|
||||
this.tokenCountsRouteKey = this.currentRouteKey();
|
||||
// A fresh count supersedes anything this route retained while another
|
||||
// route owned the slots. Without the delete this writer alone among the
|
||||
// count writers would leave an entry for tokenCountsRouteKey behind,
|
||||
// breaking the map's documented invariant (#9506).
|
||||
this.tokenCountsByRouteKey.delete(this.tokenCountsRouteKey);
|
||||
}
|
||||
|
||||
isLastPromptTokenCountEstimated(): boolean {
|
||||
this.adoptTokenCountsForRoute();
|
||||
return this.lastPromptTokenCountIsEstimated;
|
||||
}
|
||||
|
||||
|
|
@ -2083,6 +2265,14 @@ export class GeminiChat {
|
|||
this.lastOutputTokenCount = Number.isFinite(outputTokenCount)
|
||||
? Math.max(0, outputTokenCount)
|
||||
: 0;
|
||||
// Attribute the seeded counts to the active route so a model switch
|
||||
// after resume invalidates them like any API-reported count. (Detecting
|
||||
// a route that already differed at save time requires persisting route
|
||||
// identity in the session transcript; tracked as a follow-up to #9454.)
|
||||
this.tokenCountsRouteKey = this.currentRouteKey();
|
||||
// A fresh seed supersedes any count this route retained while another
|
||||
// route owned the slots (#9506).
|
||||
this.tokenCountsByRouteKey.delete(this.tokenCountsRouteKey);
|
||||
}
|
||||
|
||||
/**
|
||||
|
|
@ -2102,6 +2292,11 @@ export class GeminiChat {
|
|||
signal?: AbortSignal,
|
||||
options?: TryCompressOptions,
|
||||
): Promise<ChatCompressionInfo> {
|
||||
// Counts from a pre-switch route must not anchor compression admission
|
||||
// or sizing for this route (#9454). In-send callers pass the request
|
||||
// route so the adoption never re-adopts the active route's retained
|
||||
// counts mid-send (#9506).
|
||||
this.adoptTokenCountsForRoute(options?.requestRouteKey);
|
||||
const originalTokenCountIsEstimated =
|
||||
options?.originalTokenCountOverride === undefined &&
|
||||
this.promptCountIsEstimateDerived();
|
||||
|
|
@ -2135,6 +2330,14 @@ export class GeminiChat {
|
|||
signal,
|
||||
});
|
||||
|
||||
// ChatCompressionService reads the keyless count getters, which adopt
|
||||
// the ACTIVE route — flipping the slots away from the request route
|
||||
// adopted above whenever the two differ (non-exact override sends).
|
||||
// Re-adopt the request route so neither the COMPRESSED stamp below nor
|
||||
// the caller's post-compression sizing anchors on the flipped
|
||||
// attribution (#9506).
|
||||
this.adoptTokenCountsForRoute(options?.requestRouteKey);
|
||||
|
||||
if (info.compressionStatus === CompressionStatus.COMPRESSED && newHistory) {
|
||||
// ChatCompressionService owns provenance. Keep a conservative fallback
|
||||
// for older/custom implementations that omit the field, but preserve an
|
||||
|
|
@ -2149,10 +2352,35 @@ export class GeminiChat {
|
|||
this.setHistory(newHistory);
|
||||
debugLogger.debug('[FILE_READ_CACHE] clear after auto tryCompress');
|
||||
this.config.getFileReadCache().clear();
|
||||
// Compression rewrote the shared history every retained entry sizes,
|
||||
// so ALL retained counts are stale — not just the current route's.
|
||||
// Drop them, or a later keyed read adopts a pre-compression count and
|
||||
// the session-token-limit gate blocks a prompt that fits the
|
||||
// compressed history (#9506).
|
||||
this.tokenCountsByRouteKey.clear();
|
||||
this.setLastPromptTokenCount(
|
||||
info.newTokenCount,
|
||||
info.newTokenCountIsEstimated,
|
||||
);
|
||||
// setLastPromptTokenCount re-keyed the fresh count to the ACTIVE
|
||||
// route, but in-send callers compress for the REQUEST route: the
|
||||
// session-token-limit gate reads by that key (client.ts's sole
|
||||
// SessionTokenLimitExceeded yield site), and a request that ends
|
||||
// without a usage report (abort, 400 — the reactive-overflow path
|
||||
// exists for exactly those) never stamps a count of its own. Re-key
|
||||
// the fresh count to the request route, retaining it under the
|
||||
// active key first: the compressed history is shared, so the count
|
||||
// must anchor BOTH routes' next gate reads (#9506).
|
||||
if (
|
||||
options?.requestRouteKey &&
|
||||
this.tokenCountsRouteKey !== options.requestRouteKey
|
||||
) {
|
||||
this.retainCurrentTokenCounts();
|
||||
this.tokenCountsRouteKey = options.requestRouteKey;
|
||||
// Same invariant as the other count writers: the fresh count
|
||||
// supersedes anything the request route retained.
|
||||
this.tokenCountsByRouteKey.delete(options.requestRouteKey);
|
||||
}
|
||||
this.telemetryService?.setLastPromptTokenCount(info.newTokenCount);
|
||||
// Reset the consecutive-failure counter on success so a forced /compress
|
||||
// (or any successful compaction) recovers a chat whose breaker had
|
||||
|
|
@ -2186,6 +2414,9 @@ export class GeminiChat {
|
|||
info: ChatCompressionInfo;
|
||||
microcompactMeta?: MicrocompactMeta;
|
||||
} {
|
||||
// A pre-switch route's count must not anchor fast-compression sizing
|
||||
// for the active route (#9454).
|
||||
this.adoptTokenCountsForRoute();
|
||||
// Use the same estimator on both sides so the NOOP gate compares
|
||||
// apples to apples. The API-authoritative lastPromptTokenCount is
|
||||
// then adjusted by the estimated delta — never replaced wholesale.
|
||||
|
|
@ -2256,6 +2487,11 @@ export class GeminiChat {
|
|||
this.setHistory(newHistory);
|
||||
this.lastPromptTokenCount = adjustedTokenCount;
|
||||
this.lastPromptTokenCountIsEstimated = true;
|
||||
this.tokenCountsRouteKey = this.currentRouteKey();
|
||||
// Fast compression rewrote the shared history every retained entry
|
||||
// sizes, so ALL retained counts are stale — the other routes' entries
|
||||
// describe the same pre-compression history (#9506).
|
||||
this.tokenCountsByRouteKey.clear();
|
||||
this.telemetryService?.setLastPromptTokenCount(adjustedTokenCount);
|
||||
this.consecutiveFailures = 0;
|
||||
|
||||
|
|
@ -2335,6 +2571,23 @@ export class GeminiChat {
|
|||
if (exactRoute) {
|
||||
model = exactRoute.model;
|
||||
}
|
||||
// Both arms are one call: for a non-exact send `exactRoute` is
|
||||
// undefined, and `resolvedModelIdentity`'s second parameter defaults to
|
||||
// `getContentGeneratorConfig()` — including when passed an explicit
|
||||
// undefined. Keeping a single call site means a future change to how
|
||||
// the request route is identified cannot drift between the arms.
|
||||
const requestRouteKey = this.config.getModelRouteIdentity(
|
||||
model,
|
||||
exactRoute?.contentGeneratorConfig,
|
||||
);
|
||||
// Counts recorded for a route other than this request's target must not
|
||||
// anchor its admission/clamp/compression decisions (#9454). Comparing
|
||||
// against the REQUEST route — resolved above — keeps an exact `\0`
|
||||
// route's decisions off the active route's counts, and a differing
|
||||
// `model` param gets its own identity instead of borrowing the active
|
||||
// route's. The crossing retains the current counts under their own
|
||||
// route key so a later turn back on that route restores them (#9506).
|
||||
this.adoptTokenCountsForRoute(requestRouteKey);
|
||||
const requestModalities =
|
||||
exactRoute?.contentGeneratorConfig.modalities ??
|
||||
this.config.getEffectiveInputModalities();
|
||||
|
|
@ -2485,6 +2738,28 @@ export class GeminiChat {
|
|||
const lastPromptTokenCountBeforeHardRescue = this.lastPromptTokenCount;
|
||||
const lastPromptTokenCountWasEstimatedBeforeHardRescue =
|
||||
this.lastPromptTokenCountIsEstimated;
|
||||
// The rescue's COMPRESSED stamp zeroes lastOutputTokenCount (via
|
||||
// setLastPromptTokenCount), so the rollback below must restore the
|
||||
// output half of the resurrected count pair alongside the prompt
|
||||
// half, or the next turn's additive prompt estimate under-counts by
|
||||
// the last response's size (#9506).
|
||||
const lastOutputTokenCountBeforeHardRescue = this.lastOutputTokenCount;
|
||||
// tryCompress re-stamps tokenCountsRouteKey to the ACTIVE route (via
|
||||
// setLastPromptTokenCount on the success path) even though this send
|
||||
// targets the REQUEST route — and hard-rescue only fires for
|
||||
// non-exact sends, whose request key can differ from the active one.
|
||||
// Capture the key so the rollback below restores the resurrected
|
||||
// count's original route attribution along with the count itself.
|
||||
const tokenCountsRouteKeyBeforeHardRescue = this.tokenCountsRouteKey;
|
||||
// Snapshot the retention map too: the rescue's compression consumes
|
||||
// retained entries mid-flight (ChatCompressionService's keyless getter
|
||||
// reads adopt the active route, deleting-and-consuming its entry) and
|
||||
// a successful compression clears the map outright. Without the
|
||||
// snapshot the rollback would restore the slots but not the map,
|
||||
// leaving the resurrected route's count nowhere (#9506).
|
||||
const retainedTokenCountsBeforeHardRescue = new Map(
|
||||
this.tokenCountsByRouteKey,
|
||||
);
|
||||
const hardRescueFailureCountBeforeHardRescue =
|
||||
this.hardRescueFailureCount;
|
||||
if (shouldForceFromHard) {
|
||||
|
|
@ -2512,6 +2787,7 @@ export class GeminiChat {
|
|||
pendingUserMessage: userContent,
|
||||
precomputedEffectiveTokens: effectiveTokens,
|
||||
requestGenerationConfig: params.config,
|
||||
requestRouteKey,
|
||||
deferChatCompressionRecord: shouldForceFromHard,
|
||||
// Hard-rescue is force=true to bypass the cheap-gate breaker
|
||||
// but it remains a semantically AUTOMATIC trigger. Tag the
|
||||
|
|
@ -2561,6 +2837,21 @@ export class GeminiChat {
|
|||
this.lastPromptTokenCount = lastPromptTokenCountBeforeHardRescue;
|
||||
this.lastPromptTokenCountIsEstimated =
|
||||
lastPromptTokenCountWasEstimatedBeforeHardRescue;
|
||||
this.lastOutputTokenCount = lastOutputTokenCountBeforeHardRescue;
|
||||
this.tokenCountsRouteKey = tokenCountsRouteKeyBeforeHardRescue;
|
||||
// Restore the retention map alongside the slots: the rescue's
|
||||
// compression consumed/cleared entries mid-flight, and without
|
||||
// the restore the resurrected route's count would survive
|
||||
// nowhere — its next gate read would pass with 0 (#9506). The
|
||||
// snapshot predates the rescue, so it already satisfies the
|
||||
// invariant (no entry for the resurrected slot key).
|
||||
this.tokenCountsByRouteKey.clear();
|
||||
for (const [
|
||||
retainedRouteKey,
|
||||
retainedCounts,
|
||||
] of retainedTokenCountsBeforeHardRescue) {
|
||||
this.tokenCountsByRouteKey.set(retainedRouteKey, retainedCounts);
|
||||
}
|
||||
this.telemetryService?.setLastPromptTokenCount(
|
||||
lastPromptTokenCountBeforeHardRescue,
|
||||
);
|
||||
|
|
@ -2914,6 +3205,7 @@ export class GeminiChat {
|
|||
params,
|
||||
prompt_id,
|
||||
requestOverrides,
|
||||
requestRouteKey,
|
||||
turnGoalContext,
|
||||
// Captured by value, so the attempt records exactly the prefix
|
||||
// `buildAttemptContents()` just asked the model to resume from,
|
||||
|
|
@ -3213,6 +3505,7 @@ export class GeminiChat {
|
|||
originalTokenCountOverride: reactiveOriginalTokenCount,
|
||||
precomputedEffectiveTokens: reactiveOriginalTokenCount,
|
||||
requestGenerationConfig: params.config,
|
||||
requestRouteKey,
|
||||
trigger: 'auto',
|
||||
},
|
||||
);
|
||||
|
|
@ -3446,6 +3739,7 @@ export class GeminiChat {
|
|||
attemptState.params,
|
||||
prompt_id,
|
||||
requestOverrides,
|
||||
requestRouteKey,
|
||||
turnGoalContext,
|
||||
);
|
||||
for await (const chunk of stream) {
|
||||
|
|
@ -3857,6 +4151,14 @@ export class GeminiChat {
|
|||
currentUserContent,
|
||||
fallbackModalities ?? {},
|
||||
);
|
||||
// Stamp the fallback-served counts under the REQUEST route
|
||||
// key: a fallback serves on behalf of the same session
|
||||
// request (the session model never changes), and the
|
||||
// session-token-limit gate in Client reads the count keyed
|
||||
// by the request route. Attributing the count to the
|
||||
// fallback's own route would make every later gate read
|
||||
// invalidate it, silently disabling the limit for any
|
||||
// session ever served through fallback (#9454).
|
||||
for await (const event of self.makeFallbackStream(
|
||||
resolvedFallbackModel,
|
||||
fallbackRequestContents,
|
||||
|
|
@ -3865,6 +4167,7 @@ export class GeminiChat {
|
|||
fallbackGenerator,
|
||||
fallbackRetryAuthType,
|
||||
fallbackRetryErrorCodes,
|
||||
requestRouteKey,
|
||||
turnGoalContext,
|
||||
)) {
|
||||
const emittedUserVisibleOutput =
|
||||
|
|
@ -4019,6 +4322,7 @@ export class GeminiChat {
|
|||
retryAuthType?: string;
|
||||
retryErrorCodes?: readonly number[];
|
||||
},
|
||||
routeKey = this.currentRouteKey(),
|
||||
goalContext?: GoalTurnPermit,
|
||||
transportContinuationPrefix?: string,
|
||||
): Promise<AsyncGenerator<GenerateContentResponse>> {
|
||||
|
|
@ -4100,6 +4404,7 @@ export class GeminiChat {
|
|||
return this.processStreamResponse(
|
||||
model,
|
||||
rejectDegradedPlaceholderResponse(streamResponse),
|
||||
routeKey,
|
||||
goalContext,
|
||||
transportContinuationPrefix,
|
||||
);
|
||||
|
|
@ -4113,6 +4418,7 @@ export class GeminiChat {
|
|||
contentGenerator: ContentGenerator,
|
||||
retryAuthType?: string,
|
||||
retryErrorCodes?: readonly number[],
|
||||
routeKey?: string,
|
||||
goalContext?: GoalTurnPermit,
|
||||
): AsyncGenerator<StreamEvent> {
|
||||
const stream = await this.makeApiCallAndProcessStream(
|
||||
|
|
@ -4121,6 +4427,7 @@ export class GeminiChat {
|
|||
params,
|
||||
prompt_id,
|
||||
{ contentGenerator, retryAuthType, retryErrorCodes },
|
||||
routeKey,
|
||||
goalContext,
|
||||
);
|
||||
|
||||
|
|
@ -4614,6 +4921,7 @@ export class GeminiChat {
|
|||
private async *processStreamResponse(
|
||||
model: string,
|
||||
streamResponse: AsyncGenerator<GenerateContentResponse>,
|
||||
routeKey: string,
|
||||
goalContext?: GoalTurnPermit,
|
||||
transportContinuationPrefix?: string,
|
||||
): AsyncGenerator<GenerateContentResponse> {
|
||||
|
|
@ -4820,6 +5128,17 @@ export class GeminiChat {
|
|||
if (lastPromptTokenCount) {
|
||||
// Always update the per-chat counter so this chat (including
|
||||
// subagents) can make its own compaction decisions.
|
||||
// Retain whatever route's counts currently occupy the slots
|
||||
// before overwriting them: a foreign-keyed slot holds another
|
||||
// route's state that its next keyed read still needs — mid-send
|
||||
// compression can leave the slots keyed to the active route
|
||||
// even though this report comes from the request route (#9506).
|
||||
if (
|
||||
this.tokenCountsRouteKey !== undefined &&
|
||||
this.tokenCountsRouteKey !== routeKey
|
||||
) {
|
||||
this.retainCurrentTokenCounts();
|
||||
}
|
||||
this.lastPromptTokenCount = lastPromptTokenCount;
|
||||
this.lastPromptTokenCountIsEstimated = false;
|
||||
this.lastOutputTokenCount = hasUsablePromptTokenCount
|
||||
|
|
@ -4830,17 +5149,23 @@ export class GeminiChat {
|
|||
thoughtsTokenCount,
|
||||
})
|
||||
: 0;
|
||||
// Attribute these counts to the route that reported them so a
|
||||
// later model switch invalidates them (#9454).
|
||||
this.tokenCountsRouteKey = routeKey;
|
||||
// A fresh API report supersedes anything retained for this
|
||||
// route while another route owned the slots (#9506).
|
||||
this.tokenCountsByRouteKey.delete(routeKey);
|
||||
// Mirror to the global telemetry only when wired — subagents
|
||||
// pass `telemetryService=undefined` to keep their context usage
|
||||
// out of the main session's UI counters.
|
||||
this.telemetryService?.setLastPromptTokenCount(
|
||||
lastPromptTokenCount,
|
||||
);
|
||||
}
|
||||
if (cachedContentTokenCount && this.telemetryService) {
|
||||
this.telemetryService.setLastCachedContentTokenCount(
|
||||
cachedContentTokenCount,
|
||||
);
|
||||
if (cachedContentTokenCount && this.telemetryService) {
|
||||
this.telemetryService.setLastCachedContentTokenCount(
|
||||
cachedContentTokenCount,
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -82,6 +82,7 @@ describe('Goal turn evidence propagation', () => {
|
|||
processStreamResponse: (
|
||||
model: string,
|
||||
stream: AsyncGenerator<GenerateContentResponse>,
|
||||
routeKey: string,
|
||||
goalContext?: GoalTurnPermit,
|
||||
) => AsyncGenerator<GenerateContentResponse>;
|
||||
pendingPartialAssistantRecord:
|
||||
|
|
@ -102,6 +103,7 @@ describe('Goal turn evidence propagation', () => {
|
|||
for await (const _ of internal.processStreamResponse(
|
||||
'test-model',
|
||||
normalStream,
|
||||
'test-route',
|
||||
permit,
|
||||
)) {
|
||||
// Consume the persisted normal assistant attempt.
|
||||
|
|
@ -136,6 +138,7 @@ describe('Goal turn evidence propagation', () => {
|
|||
for await (const _ of internal.processStreamResponse(
|
||||
'test-model',
|
||||
partialStream,
|
||||
'test-route',
|
||||
permit,
|
||||
)) {
|
||||
// Consume until the deferred partial attempt is staged.
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue