diff --git a/docs/design/2026-08-21-model-stream-attempt-state.md b/docs/design/2026-08-21-model-stream-attempt-state.md new file mode 100644 index 00000000000..f1a1393f206 --- /dev/null +++ b/docs/design/2026-08-21-model-stream-attempt-state.md @@ -0,0 +1,59 @@ +# Model stream attempt state + +## Problem + +`GeminiChat.sendMessageStream()` exposes chunks plus retry, compression, and +model-fallback control events. Its consumers currently keep their own copies of +the same per-attempt state: text, thoughts, tool calls, usage, response ids, +finish reasons, and output-truncation markers. + +The implementations have drifted. The sub-agent loop ignores model fallback +and treats continuation retries as fresh restarts. ACP background notifications +keep failed-attempt text and usage across fresh retries and fallbacks. Other ACP +loops reset tool calls but can retain stale usage. MessageDisplay also has no +way to discard text from a restarted attempt. + +## Design + +Add a small core `ModelStreamAttemptState` that consumes the existing +`StreamEvent` union and owns only protocol-derived state. It returns a compact +transition for each event so callers can keep their surface-specific side +effects. + +The state applies these rules: + +- chunks append visible text, thought text, and function calls, and replace the + latest usage, response id, and finish reason; +- every retry clears tool calls and per-attempt metadata; +- continuation retries preserve accumulated text and thought text; +- fresh retries and model fallback discard accumulated text and thought text; +- model fallback always starts a fresh attempt; +- compressed events do not change attempt state. + +`Turn`, the sub-agent reasoning loop, forked queries, speculation, and all ACP +raw-stream loops consume the same transitions. `MessageDisplayDispatcher` gains the same +`restartAttempt(preserveText)` operation already used by telemetry output +capture, so hook output follows the stream contract. + +The helper deliberately stays on the current Google response type. This change +centralizes stream lifecycle semantics without attempting the larger protocol +migration. A future protocol boundary can translate provider-neutral chunks +into the same transition model. + +## Non-goals + +- Change provider request or response protocols. +- Unify UI, ACP, and sub-agent rendering side effects. +- Change retry or fallback policy inside `GeminiChat`. +- Rework tool execution scheduling. + +## Verification + +- Pure state tests cover fresh retry, continuation retry, fallback, metadata, + and truncation reset. +- Sub-agent tests prove continuation text is preserved and fallback state is + discarded. +- Forked-query and speculation tests keep only the active attempt. +- ACP background-notification tests prove output, usage, and MessageDisplay do + not retain a failed attempt. +- Existing Turn and ACP stream tests remain green. diff --git a/packages/cli/src/acp-integration/session/Session.test.ts b/packages/cli/src/acp-integration/session/Session.test.ts index afd28a48bcd..e670edfbf9a 100644 --- a/packages/cli/src/acp-integration/session/Session.test.ts +++ b/packages/cli/src/acp-integration/session/Session.test.ts @@ -370,7 +370,12 @@ function setFakeHome(home: string): () => void { // Helper to create async generator with chunks (avoids memory leak) function createStreamWithChunks( - chunks: Array<{ type: unknown; value: unknown }>, + chunks: Array<{ + type: unknown; + value?: unknown; + info?: unknown; + isContinuation?: boolean; + }>, ) { return (async function* () { for (const chunk of chunks) { @@ -7277,6 +7282,118 @@ describe('Session', () => { ); }); + it.each([ + { + name: 'fresh retry', + event: { type: core.StreamEventType.RETRY }, + }, + { + name: 'model fallback', + event: { + type: core.StreamEventType.MODEL_FALLBACK, + info: { + fromModel: 'primary', + toModel: 'fallback', + fallbackIndex: 1, + }, + }, + }, + ])( + 'discards stale background notification state on $name', + async ({ event }) => { + const messageBus = { request: vi.fn().mockResolvedValue({}) }; + mockConfig.getMessageBus = vi.fn().mockReturnValue(messageBus); + mockConfig.getDisableAllHooks = vi.fn().mockReturnValue(false); + mockConfig.hasHooksForEvent = vi + .fn() + .mockImplementation( + (eventName: string) => eventName === 'MessageDisplay', + ); + mockChat.sendMessageStream = vi + .fn() + .mockResolvedValueOnce(createEmptyStream()) + .mockResolvedValueOnce( + createStreamWithChunks([ + { + type: core.StreamEventType.CHUNK, + value: { + candidates: [ + { content: { parts: [{ text: 'stale answer' }] } }, + ], + usageMetadata: { + promptTokenCount: 111, + candidatesTokenCount: 222, + }, + }, + }, + event, + { + type: core.StreamEventType.CHUNK, + value: { + candidates: [ + { content: { parts: [{ text: 'current answer' }] } }, + ], + }, + }, + ]), + ); + + 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 }, + ) => void; + + callback('done', '', { + agentId: 'agent-1', + status: 'completed', + }); + + await vi.waitFor(() => { + expect(mockClient.sessionUpdate).toHaveBeenCalledWith({ + sessionId: 'test-session-id', + update: expect.objectContaining({ + content: { type: 'text', text: 'current answer' }, + _meta: expect.objectContaining({ + source: 'background_notification_response', + }), + }), + }); + }); + + const responseUpdates = vi + .mocked(mockClient.sessionUpdate) + .mock.calls.map(([request]) => request.update) + .filter( + (update) => + update._meta?.['source'] === 'background_notification_response', + ); + expect(responseUpdates).toHaveLength(1); + expect(responseUpdates[0]).toEqual( + expect.objectContaining({ + content: { type: 'text', text: 'current answer' }, + }), + ); + expect(session.cumulativeUsage).toMatchObject({ + promptTokens: 0, + candidateTokens: 0, + }); + + const finalDisplay = messageBus.request.mock.calls + .map(([request]) => request) + .find( + (request) => + request.eventName === 'MessageDisplay' && request.input.is_final, + ); + expect(finalDisplay?.input.displayed_text).toBe('current answer'); + }, + ); + it('attaches structured agent metadata built from the canonical entry label', async () => { mockChat.sendMessageStream = vi .fn() diff --git a/packages/cli/src/acp-integration/session/Session.ts b/packages/cli/src/acp-integration/session/Session.ts index 2e19ef3935a..45b2bd829c0 100644 --- a/packages/cli/src/acp-integration/session/Session.ts +++ b/packages/cli/src/acp-integration/session/Session.ts @@ -73,7 +73,6 @@ import { PLAN_MODE_ENTRY_SIBLING_SKIP_MESSAGE, createDebugLogger, DiscoveredMCPTool, - StreamEventType, ToolConfirmationOutcome, generatePromptSuggestion, logPromptSuggestion, @@ -104,6 +103,7 @@ import { generateToolUseId, MessageBusType, MessageDisplayDispatcher, + ModelStreamAttemptState, getPlanModeSystemReminder, getArenaSystemReminder, getStartupContextLength, @@ -5491,12 +5491,10 @@ export class Session implements SessionContext { return { stopReason: 'cancelled' }; } - const functionCalls: FunctionCall[] = []; + const attemptState = new ModelStreamAttemptState(); const preparationTracker = new ToolCallPreparationTracker( this.toolCallEmitter, ); - let usageMetadata: GenerateContentResponseUsageMetadata | null = - null; const streamStartTime = Date.now(); const messageDisplay = this.#createMessageDisplayDispatcher( pendingSend.signal, @@ -5570,17 +5568,9 @@ export class Session implements SessionContext { return { stopReason: 'cancelled' }; } - if ( - resp.type === StreamEventType.CHUNK && - resp.value.candidates && - resp.value.candidates.length > 0 - ) { - const candidate = resp.value.candidates[0]; - for (const part of candidate.content?.parts ?? []) { - if (!part.text) { - continue; - } - + const transition = attemptState.accept(resp); + if (transition.type === 'chunk') { + for (const part of transition.textParts) { this.messageEmitter.emitMessage( part.text, 'assistant', @@ -5596,36 +5586,18 @@ export class Session implements SessionContext { } } responseCapture.agentOutput.observeFinishReason( - candidate.finishReason, + transition.finishReason, ); - } - - if ( - resp.type === StreamEventType.CHUNK && - resp.value.usageMetadata - ) { - usageMetadata = resp.value.usageMetadata; - } - - if (resp.type === StreamEventType.CHUNK) { - await preparationTracker.observe(resp.value); - if (resp.value.functionCalls) { - preparationTracker.resolve(resp.value.functionCalls); - functionCalls.push(...resp.value.functionCalls); + await preparationTracker.observe(transition.response); + if (transition.functionCalls.length > 0) { + preparationTracker.resolve(transition.functionCalls); } - } - if ( - resp.type === StreamEventType.RETRY || - resp.type === StreamEventType.MODEL_FALLBACK - ) { + } else if (transition.type === 'attempt_reset') { responseCapture.agentOutput.restartAttempt( - resp.type === StreamEventType.RETRY && - resp.isContinuation === true, + transition.preserveText, ); - if ( - resp.type === StreamEventType.MODEL_FALLBACK || - !resp.isContinuation - ) { + messageDisplay?.restartAttempt(transition.preserveText); + if (!transition.preserveText) { rewindChannelDeliveryResponseBlock( channelDeliveryResponseBlock, channelDeliveryCheckpoint, @@ -5634,9 +5606,8 @@ export class Session implements SessionContext { await finalizeToolCallPreparations( preparationTracker, true, - `main prompt ${resp.type}`, + `main prompt ${transition.reason}`, ); - functionCalls.length = 0; } } } catch (error) { @@ -5723,14 +5694,16 @@ export class Session implements SessionContext { await messageDisplay?.finish(); } + const attempt = attemptState.snapshot(); + commitChannelDeliveryResponseBlock( responseCapture, channelDeliveryResponseBlock, - functionCalls.length > 0, + attempt.functionCalls.length > 0, ); - if (usageMetadata) { - this.#recordPromptTokenCount(usageMetadata); + if (attempt.usageMetadata) { + this.#recordPromptTokenCount(attempt.usageMetadata); // Kick off rewrite in background (non-blocking, runs parallel to tools) if (this.messageRewriter) { this.messageRewriter.flushTurn(pendingSend.signal); @@ -5738,20 +5711,20 @@ export class Session implements SessionContext { const durationMs = Date.now() - streamStartTime; await this.messageEmitter.emitUsageMetadata( - usageMetadata, + attempt.usageMetadata, '', durationMs, ); } - if (functionCalls.length > 0) { + if (attempt.functionCalls.length > 0) { const toolRun = await this.#runWithFullTurnModel( fullTurnModelOverride, () => this.runToolCalls( pendingSend.signal, promptId, - functionCalls, + attempt.functionCalls, toolLoopState, onFullTurnModel, ), @@ -6283,11 +6256,10 @@ export class Session implements SessionContext { }; } - const functionCalls: FunctionCall[] = []; + const attemptState = new ModelStreamAttemptState(); const preparationTracker = new ToolCallPreparationTracker( this.toolCallEmitter, ); - let usageMetadata: GenerateContentResponseUsageMetadata | null = null; const streamStartTime = Date.now(); let streamFailed = false; let guardForThisSend = nextGuardContinuation; @@ -6639,14 +6611,9 @@ export class Session implements SessionContext { }; } - if ( - response.type === StreamEventType.CHUNK && - response.value.candidates && - response.value.candidates.length > 0 - ) { - const candidate = response.value.candidates[0]; - for (const part of candidate.content?.parts ?? []) { - if (!part.text) continue; + const transition = attemptState.accept(response); + if (transition.type === 'chunk') { + for (const part of transition.textParts) { this.messageEmitter.emitMessage( part.text, 'assistant', @@ -6662,35 +6629,18 @@ export class Session implements SessionContext { } } options.responseCapture?.agentOutput.observeFinishReason( - candidate.finishReason, + transition.finishReason, ); - } - - if ( - response.type === StreamEventType.CHUNK && - response.value.usageMetadata - ) { - usageMetadata = response.value.usageMetadata; - } - if (response.type === StreamEventType.CHUNK) { - await preparationTracker.observe(response.value); - if (response.value.functionCalls) { - preparationTracker.resolve(response.value.functionCalls); - functionCalls.push(...response.value.functionCalls); + await preparationTracker.observe(transition.response); + if (transition.functionCalls.length > 0) { + preparationTracker.resolve(transition.functionCalls); } - } - if ( - response.type === StreamEventType.RETRY || - response.type === StreamEventType.MODEL_FALLBACK - ) { + } else if (transition.type === 'attempt_reset') { options.responseCapture?.agentOutput.restartAttempt( - response.type === StreamEventType.RETRY && - response.isContinuation === true, + transition.preserveText, ); - if ( - response.type === StreamEventType.MODEL_FALLBACK || - !response.isContinuation - ) { + messageDisplay?.restartAttempt(transition.preserveText); + if (!transition.preserveText) { rewindChannelDeliveryResponseBlock( channelDeliveryResponseBlock, channelDeliveryCheckpoint, @@ -6699,9 +6649,8 @@ export class Session implements SessionContext { await finalizeToolCallPreparations( preparationTracker, true, - `daemon continuation ${response.type}`, + `daemon continuation ${transition.reason}`, ); - functionCalls.length = 0; } } } catch (error) { @@ -6770,30 +6719,32 @@ export class Session implements SessionContext { } } + const attempt = attemptState.snapshot(); + commitChannelDeliveryResponseBlock( options.responseCapture, channelDeliveryResponseBlock, - functionCalls.length > 0, + attempt.functionCalls.length > 0, ); - if (usageMetadata) { - this.#recordPromptTokenCount(usageMetadata); + if (attempt.usageMetadata) { + this.#recordPromptTokenCount(attempt.usageMetadata); const durationMs = Date.now() - streamStartTime; await this.messageEmitter.emitUsageMetadata( - usageMetadata, + attempt.usageMetadata, '', durationMs, ); } - if (functionCalls.length > 0) { + if (attempt.functionCalls.length > 0) { const toolRun = await this.#runWithFullTurnModel( options.getModelOverride?.(), () => this.runToolCalls( pendingSend.signal, toolPromptId, - functionCalls, + attempt.functionCalls, toolLoopState, options.onFullTurnModel, ), @@ -8320,12 +8271,10 @@ export class Session implements SessionContext { return; } - const functionCalls: FunctionCall[] = []; + const attemptState = new ModelStreamAttemptState(); const preparationTracker = new ToolCallPreparationTracker( this.toolCallEmitter, ); - let usageMetadata: GenerateContentResponseUsageMetadata | null = - null; const streamStartTime = Date.now(); const sendResult = await this.#sendMessageStreamWithAutoCompression( @@ -8371,14 +8320,9 @@ export class Session implements SessionContext { return; } - if ( - resp.type === StreamEventType.CHUNK && - resp.value.candidates && - resp.value.candidates.length > 0 - ) { - const candidate = resp.value.candidates[0]; - for (const part of candidate.content?.parts ?? []) { - if (!part.text) continue; + const transition = attemptState.accept(resp); + if (transition.type === 'chunk') { + for (const part of transition.textParts) { this.messageEmitter.emitMessage( part.text, 'assistant', @@ -8394,36 +8338,18 @@ export class Session implements SessionContext { } } responseCapture.agentOutput.observeFinishReason( - candidate.finishReason, + transition.finishReason, ); - } - - if ( - resp.type === StreamEventType.CHUNK && - resp.value.usageMetadata - ) { - usageMetadata = resp.value.usageMetadata; - } - - if (resp.type === StreamEventType.CHUNK) { - await preparationTracker.observe(resp.value); - if (resp.value.functionCalls) { - preparationTracker.resolve(resp.value.functionCalls); - functionCalls.push(...resp.value.functionCalls); + await preparationTracker.observe(transition.response); + if (transition.functionCalls.length > 0) { + preparationTracker.resolve(transition.functionCalls); } - } - if ( - resp.type === StreamEventType.RETRY || - resp.type === StreamEventType.MODEL_FALLBACK - ) { + } else if (transition.type === 'attempt_reset') { responseCapture.agentOutput.restartAttempt( - resp.type === StreamEventType.RETRY && - resp.isContinuation === true, + transition.preserveText, ); - if ( - resp.type === StreamEventType.MODEL_FALLBACK || - !resp.isContinuation - ) { + messageDisplay?.restartAttempt(transition.preserveText); + if (!transition.preserveText) { rewindChannelDeliveryResponseBlock( channelDeliveryResponseBlock, channelDeliveryCheckpoint, @@ -8432,9 +8358,8 @@ export class Session implements SessionContext { await finalizeToolCallPreparations( preparationTracker, true, - `cron/loop tick ${resp.type}`, + `cron/loop tick ${transition.reason}`, ); - functionCalls.length = 0; } } } catch (error) { @@ -8454,30 +8379,32 @@ export class Session implements SessionContext { } } + const attempt = attemptState.snapshot(); + commitChannelDeliveryResponseBlock( responseCapture, channelDeliveryResponseBlock, - functionCalls.length > 0, + attempt.functionCalls.length > 0, ); - if (usageMetadata) { - this.#recordPromptTokenCount(usageMetadata); + if (attempt.usageMetadata) { + this.#recordPromptTokenCount(attempt.usageMetadata); if (this.messageRewriter) { this.messageRewriter.flushTurn(ac.signal); } const durationMs = Date.now() - streamStartTime; await this.messageEmitter.emitUsageMetadata( - usageMetadata, + attempt.usageMetadata, '', durationMs, ); } - if (functionCalls.length > 0) { + if (attempt.functionCalls.length > 0) { const toolRun = await this.runToolCalls( ac.signal, promptId, - functionCalls, + attempt.functionCalls, toolLoopState, ); if (toolRun.stopAfterPermissionCancel || ac.signal.aborted) { @@ -9012,13 +8939,10 @@ export class Session implements SessionContext { return; } - const functionCalls: FunctionCall[] = []; + const attemptState = new ModelStreamAttemptState(); const preparationTracker = new ToolCallPreparationTracker( this.toolCallEmitter, ); - let usageMetadata: GenerateContentResponseUsageMetadata | null = - null; - let responseText = ''; const streamStartTime = Date.now(); const sendResult = await this.#sendMessageStreamWithAutoCompression( @@ -9053,14 +8977,9 @@ export class Session implements SessionContext { return; } - if ( - resp.type === StreamEventType.CHUNK && - resp.value.candidates && - resp.value.candidates.length > 0 - ) { - const candidate = resp.value.candidates[0]; - for (const part of candidate.content?.parts ?? []) { - if (!part.text) continue; + const transition = attemptState.accept(resp); + if (transition.type === 'chunk') { + for (const part of transition.textParts) { if (part.thought) { await this.messageEmitter.emitMessage( part.text, @@ -9068,36 +8987,20 @@ export class Session implements SessionContext { true, ); } else { - responseText += part.text; messageDisplay?.addChunk(part.text); } } - } - - if ( - resp.type === StreamEventType.CHUNK && - resp.value.usageMetadata - ) { - usageMetadata = resp.value.usageMetadata; - } - - if (resp.type === StreamEventType.CHUNK) { - await preparationTracker.observe(resp.value); - if (resp.value.functionCalls) { - preparationTracker.resolve(resp.value.functionCalls); - functionCalls.push(...resp.value.functionCalls); + await preparationTracker.observe(transition.response); + if (transition.functionCalls.length > 0) { + preparationTracker.resolve(transition.functionCalls); } - } - if ( - resp.type === StreamEventType.RETRY || - resp.type === StreamEventType.MODEL_FALLBACK - ) { + } else if (transition.type === 'attempt_reset') { + messageDisplay?.restartAttempt(transition.preserveText); await finalizeToolCallPreparations( preparationTracker, true, - `background notification ${resp.type}`, + `background notification ${transition.reason}`, ); - functionCalls.length = 0; } } } catch (error) { @@ -9117,10 +9020,12 @@ export class Session implements SessionContext { } } - if (responseText.length > 0) { + const attempt = attemptState.snapshot(); + + if (attempt.text.length > 0) { await this.#emitBackgroundNotificationResponse( item, - responseText, + attempt.text, ac.signal, ); } @@ -9129,21 +9034,21 @@ export class Session implements SessionContext { await this.messageRewriter.flushTurn(ac.signal); } - if (usageMetadata) { - this.#recordPromptTokenCount(usageMetadata); + if (attempt.usageMetadata) { + this.#recordPromptTokenCount(attempt.usageMetadata); const durationMs = Date.now() - streamStartTime; await this.messageEmitter.emitUsageMetadata( - usageMetadata, + attempt.usageMetadata, '', durationMs, ); } - if (functionCalls.length > 0) { + if (attempt.functionCalls.length > 0) { const toolRun = await this.runToolCalls( ac.signal, promptId, - functionCalls, + attempt.functionCalls, toolLoopState, ); if (toolRun.stopAfterPermissionCancel || ac.signal.aborted) { diff --git a/packages/core/src/agents/forkedAgent.cache.test.ts b/packages/core/src/agents/forkedAgent.cache.test.ts index 5da024b4e76..08853235675 100644 --- a/packages/core/src/agents/forkedAgent.cache.test.ts +++ b/packages/core/src/agents/forkedAgent.cache.test.ts @@ -403,6 +403,59 @@ describe('runForkedAgent (cache path)', () => { expect(result.model).toBe('test-model'); }); + it('discards failed-attempt text and usage after model fallback', async () => { + saveCacheSafeParams({}, [], 'test-model'); + const mockSendMessageStream = vi.fn(() => + Promise.resolve( + (async function* () { + yield { + type: StreamEventType.CHUNK, + value: { + candidates: [{ content: { parts: [{ text: 'stale answer' }] } }], + usageMetadata: { + promptTokenCount: 10, + candidatesTokenCount: 5, + }, + }, + }; + yield { + type: StreamEventType.MODEL_FALLBACK, + info: { + fromModel: 'primary', + toModel: 'fallback', + fallbackIndex: 1, + }, + }; + yield { + type: StreamEventType.CHUNK, + value: { + candidates: [ + { content: { parts: [{ text: 'current answer' }] } }, + ], + }, + }; + })(), + ), + ); + vi.mocked(LlmChat).mockImplementation( + () => + ({ sendMessageStream: mockSendMessageStream }) as unknown as LlmChat, + ); + + const result = await runForkedAgent({ + config: {} as Config, + userMessage: 'suggest something', + cacheSafeParams: getCacheSafeParams()!, + }); + + expect(result.text).toBe('current answer'); + expect(result.usage).toEqual({ + inputTokens: 0, + outputTokens: 0, + cacheHitTokens: 0, + }); + }); + it('preserves tools: [] even when jsonSchema is provided', async () => { saveCacheSafeParams( { diff --git a/packages/core/src/agents/forkedAgent.ts b/packages/core/src/agents/forkedAgent.ts index 3572c9afa0c..4620acd5b24 100644 --- a/packages/core/src/agents/forkedAgent.ts +++ b/packages/core/src/agents/forkedAgent.ts @@ -40,7 +40,8 @@ import { type RuntimeContentGeneratorView, } from './runtime/agent-context.js'; import { ApprovalMode, type Config } from '../config/config.js'; -import { LlmChat, StreamEventType } from '../core/llm-chat.js'; +import { LlmChat } from '../core/llm-chat.js'; +import { ModelStreamAttemptState } from '../core/model-stream-attempt-state.js'; import { createRuntimeContentGeneratorView } from '../models/content-generator-config.js'; import { createApprovalModeOverride } from '../tools/agent/agent.js'; import { createDebugLogger } from '../utils/debugLogger.js'; @@ -550,16 +551,12 @@ export async function runForkedAgent( ) : await chat.sendMessageStream(model, sendParams, 'forked_query'); - let fullText = ''; - let usage: ForkedQueryResult['usage'] = { - inputTokens: 0, - outputTokens: 0, - cacheHitTokens: 0, - }; + const attemptState = new ModelStreamAttemptState(); for await (const event of stream) { - if (event.type !== StreamEventType.CHUNK) continue; - const response = event.value; + const transition = attemptState.accept(event); + if (transition.type !== 'chunk') continue; + const response = transition.response; const parts = response.candidates?.[0]?.content?.parts ?? []; // Defensive: when preserveTools is true the model could produce @@ -572,18 +569,17 @@ export async function runForkedAgent( 'Cache-path forked query received functionCall with preserveTools; discarding.', ); } - - const text = parts - .filter((p) => !(p as Record)['thought']) - .filter((p) => !(p as Record)['functionCall']) - .map((p) => p.text ?? '') - .join(''); - if (text) fullText += text; - if (response.usageMetadata) - usage = extractQueryUsage(response.usageMetadata); } - const trimmed = fullText.trim() || null; + const attempt = attemptState.snapshot(); + const trimmed = attempt.text.trim() || null; + const usage: ForkedQueryResult['usage'] = attempt.usageMetadata + ? extractQueryUsage(attempt.usageMetadata) + : { + inputTokens: 0, + outputTokens: 0, + cacheHitTokens: 0, + }; let jsonResult: Record | undefined; if (jsonSchema && trimmed) { try { diff --git a/packages/core/src/agents/runtime/agent-core.test.ts b/packages/core/src/agents/runtime/agent-core.test.ts index be031c2e9c3..c7997278c8c 100644 --- a/packages/core/src/agents/runtime/agent-core.test.ts +++ b/packages/core/src/agents/runtime/agent-core.test.ts @@ -8,7 +8,11 @@ import { describe, it, expect, vi } from 'vitest'; import * as fs from 'node:fs'; import * as os from 'node:os'; import * as path from 'node:path'; -import type { FunctionDeclaration, GenerateContentConfig } from '@google/genai'; +import type { + FunctionDeclaration, + GenerateContentConfig, + GenerateContentResponse, +} from '@google/genai'; import { AgentCore, extractParentToolNames, @@ -52,7 +56,7 @@ import { runWithInvocationContext, type InvocationContextV1, } from '../../utils/invocation-context.js'; -import { LlmChat } from '../../core/llm-chat.js'; +import { LlmChat, StreamEventType } from '../../core/llm-chat.js'; import { ContextState } from './agent-headless.js'; import type { ToolResultBoundaryObservation } from '../../tools/tool-result-boundary-diagnostics.js'; import { @@ -468,6 +472,111 @@ describe('AgentCore.runInAgentFrames', () => { }); }); +describe('AgentCore model stream attempts', () => { + function createCore(): AgentCore { + const runtimeContext = { + getSessionId: () => 'session', + getSkipLoopDetection: () => true, + getDebugLogger: () => ({ debug: vi.fn() }), + } as unknown as Config; + return new AgentCore( + 'attempt-agent', + runtimeContext, + { renderedSystemPrompt: 'system', initialMessages: [] }, + { model: 'test-model' }, + { max_turns: 1 }, + ); + } + + function createChat(events: unknown[]): LlmChat { + return { + getHistoryFunctionResponseIds: () => new Set(), + getHistoryToolCallFingerprints: () => new Map(), + sendMessageStream: vi.fn().mockResolvedValue( + (async function* () { + for (const event of events) yield event; + })(), + ), + } as unknown as LlmChat; + } + + it('preserves accumulated text for continuation retries', async () => { + const core = createCore(); + const chat = createChat([ + { + type: StreamEventType.CHUNK, + value: { + candidates: [{ content: { parts: [{ text: 'first half ' }] } }], + } as GenerateContentResponse, + }, + { type: StreamEventType.RETRY, isContinuation: true }, + { + type: StreamEventType.CHUNK, + value: { + candidates: [{ content: { parts: [{ text: 'second half' }] } }], + } as GenerateContentResponse, + }, + ]); + + const result = await core.runReasoningLoop( + chat, + [{ role: 'user', parts: [{ text: 'continue' }] }], + [], + new AbortController(), + ); + + expect(result.text).toBe('first half second half'); + }); + + it('discards the previous attempt when the model falls back', async () => { + const core = createCore(); + const chat = createChat([ + { + type: StreamEventType.CHUNK, + value: { + responseId: 'primary-response', + candidates: [ + { + content: { + parts: [ + { text: 'stale answer' }, + { text: 'stale thought', thought: true }, + ], + }, + }, + ], + } as GenerateContentResponse, + }, + { + type: StreamEventType.MODEL_FALLBACK, + info: { + fromModel: 'primary', + toModel: 'fallback', + fallbackIndex: 1, + }, + }, + { + type: StreamEventType.CHUNK, + value: { + candidates: [{ content: { parts: [{ text: 'fallback answer' }] } }], + } as GenerateContentResponse, + }, + ]); + + const result = await core.runReasoningLoop( + chat, + [{ role: 'user', parts: [{ text: 'fallback' }] }], + [], + new AbortController(), + ); + + expect(result.text).toBe('fallback answer'); + expect(core.getMessages()).toEqual([ + expect.objectContaining({ content: 'fallback answer' }), + ]); + }); +}); + describe('AgentCore approval response deduplication', () => { function buildApprovalCore(): { core: AgentCore; diff --git a/packages/core/src/agents/runtime/agent-core.ts b/packages/core/src/agents/runtime/agent-core.ts index a9772be6364..a20ae5e9265 100644 --- a/packages/core/src/agents/runtime/agent-core.ts +++ b/packages/core/src/agents/runtime/agent-core.ts @@ -68,7 +68,7 @@ import { toolResultBoundaryArtifact, toolResultPartDiagnosticValues, } from '../../tools/tool-result-boundary-diagnostics.js'; -import { FinishReason } from '../../core/genai-compat.js'; +import { ModelStreamAttemptState } from '../../core/model-stream-attempt-state.js'; import type { Content, Part, @@ -993,13 +993,7 @@ export class AgentCore { timestamp: Date.now(), } as AgentRoundEvent); - const functionCalls: FunctionCall[] = []; - let roundText = ''; - let roundThoughtText = ''; - let lastUsage: GenerateContentResponseUsageMetadata | undefined = - undefined; - let currentResponseId: string | undefined = undefined; - let wasOutputTruncated = false; + const attemptState = new ModelStreamAttemptState(); let loopDetectedInStream = false; for await (const streamEvent of responseStream) { @@ -1011,139 +1005,107 @@ export class AgentCore { }; } - // Handle retry events — reset all per-attempt state so a successful - // retry does not inherit stale data (e.g. wasOutputTruncated) from a - // previous attempt that may have hit MAX_TOKENS. - if (streamEvent.type === 'retry') { - if ( - checkSubagentLoop({ - type: LlmEventType.Retry, - ...('isContinuation' in streamEvent - ? { isContinuation: streamEvent.isContinuation } - : {}), - }) - ) { + const transition = attemptState.accept(streamEvent); + if (transition.type === 'attempt_reset') { + const resetEvent: ServerLlmStreamEvent = + transition.reason === 'retry' + ? { + type: LlmEventType.Retry, + // Preserve the upstream event shape: only continuation + // retries carry `isContinuation`; fresh retries must not + // grow the property (consumers rely on its absence). + ...('isContinuation' in transition.retryInfo + ? { isContinuation: transition.retryInfo.isContinuation } + : {}), + } + : { + type: LlmEventType.ModelFallback, + fromModel: transition.info.fromModel, + toModel: transition.info.toModel, + statusCode: transition.info.statusCode, + fallbackIndex: transition.info.fallbackIndex, + }; + if (checkSubagentLoop(resetEvent)) { terminateMode = AgentTerminateMode.LOOP_DETECTED; loopDetectedInStream = true; break; } - if (streamEvent.maxOutputTokensEscalated !== undefined) { - stickyMaxOutputTokens = streamEvent.maxOutputTokensEscalated; + if ( + transition.reason === 'retry' && + transition.retryInfo.maxOutputTokensEscalated !== undefined + ) { + stickyMaxOutputTokens = + transition.retryInfo.maxOutputTokensEscalated; } - functionCalls.length = 0; - roundText = ''; - roundThoughtText = ''; - lastUsage = undefined; - currentResponseId = undefined; - wasOutputTruncated = false; continue; } // LlmChat already mutated its own history; surface to the debug // log so subagent compactions show up alongside the main session's. - if (streamEvent.type === 'compressed') { + if (transition.type === 'compressed') { this.runtimeContext .getDebugLogger() .debug( `[AGENT-COMPACT] subagent=${this.subagentId} round=${turnCounter} ` + - `tokens ${streamEvent.info.originalTokenCount} -> ${streamEvent.info.newTokenCount}`, + `tokens ${transition.info.originalTokenCount} -> ${transition.info.newTokenCount}`, ); continue; } - // Handle chunk events - if (streamEvent.type === 'chunk') { - const resp = streamEvent.value; - // Track the response ID for tool call correlation - if (resp.responseId) { - currentResponseId = resp.responseId; - } - const chunkFunctionCalls = resp.functionCalls ?? []; - functionCalls.push(...chunkFunctionCalls); - if ( - resp.candidates?.[0]?.finishReason === FinishReason.MAX_TOKENS - ) { - wasOutputTruncated = true; - } - const content = resp.candidates?.[0]?.content; - const parts = content?.parts || []; - for (const p of parts) { - const txt = p.text; - const isThought = p.thought ?? false; - if (txt && isThought) roundThoughtText += txt; - if (txt && !isThought) roundText += txt; - if (txt) - this.eventEmitter?.emit(AgentEventType.STREAM_TEXT, { - subagentId: this.subagentId, - runId, - round: turnCounter, - text: txt, - thought: isThought, - timestamp: Date.now(), - }); - } - if (resp.usageMetadata) lastUsage = resp.usageMetadata; - - const thoughtSummary = getThoughtSummary(resp); - if ( - thoughtSummary && - checkSubagentLoop({ - type: LlmEventType.Thought, - value: thoughtSummary, - }) - ) { - terminateMode = AgentTerminateMode.LOOP_DETECTED; - loopDetectedInStream = true; - break; - } + const resp = transition.response; + const chunkFunctionCalls = transition.functionCalls; + for (const part of transition.textParts) { + this.eventEmitter?.emit(AgentEventType.STREAM_TEXT, { + subagentId: this.subagentId, + runId, + round: turnCounter, + text: part.text, + thought: part.thought, + timestamp: Date.now(), + }); + } - const responseText = getResponseText(resp); - if ( - responseText && - checkSubagentLoop({ - type: LlmEventType.Content, - value: responseText, - }) - ) { - terminateMode = AgentTerminateMode.LOOP_DETECTED; - loopDetectedInStream = true; - break; - } + const thoughtSummary = getThoughtSummary(resp); + if ( + thoughtSummary && + checkSubagentLoop({ + type: LlmEventType.Thought, + value: thoughtSummary, + }) + ) { + terminateMode = AgentTerminateMode.LOOP_DETECTED; + loopDetectedInStream = true; + break; + } - for (const fc of chunkFunctionCalls) { - const toolName = String(fc.name); - if ( - checkSubagentLoop({ - type: LlmEventType.ToolCallRequest, - value: { - callId: fc.id ?? `${toolName}-${Date.now()}`, - providerCallId: getProviderToolCallId(fc), - name: toolName, - args: (fc.args ?? {}) as Record, - isClientInitiated: false, - prompt_id: promptId, - response_id: currentResponseId, - wasOutputTruncated, - }, - }) - ) { - terminateMode = AgentTerminateMode.LOOP_DETECTED; - loopDetectedInStream = true; - break; - } - } - if (loopDetectedInStream) { - break; - } + const responseText = getResponseText(resp); + if ( + responseText && + checkSubagentLoop({ + type: LlmEventType.Content, + value: responseText, + }) + ) { + terminateMode = AgentTerminateMode.LOOP_DETECTED; + loopDetectedInStream = true; + break; + } - const finishReason = resp.candidates?.[0]?.finishReason; + const currentAttempt = attemptState.snapshot(); + for (const fc of chunkFunctionCalls) { + const toolName = String(fc.name); if ( - finishReason && checkSubagentLoop({ - type: LlmEventType.Finished, + type: LlmEventType.ToolCallRequest, value: { - reason: finishReason, - usageMetadata: resp.usageMetadata, + callId: fc.id ?? `${toolName}-${Date.now()}`, + providerCallId: getProviderToolCallId(fc), + name: toolName, + args: (fc.args ?? {}) as Record, + isClientInitiated: false, + prompt_id: promptId, + response_id: currentAttempt.responseId, + wasOutputTruncated: currentAttempt.wasOutputTruncated, }, }) ) { @@ -1152,20 +1114,40 @@ export class AgentCore { break; } } + if (loopDetectedInStream) { + break; + } + + const finishReason = transition.finishReason; + if ( + finishReason && + checkSubagentLoop({ + type: LlmEventType.Finished, + value: { + reason: finishReason, + usageMetadata: resp.usageMetadata, + }, + }) + ) { + terminateMode = AgentTerminateMode.LOOP_DETECTED; + loopDetectedInStream = true; + break; + } } if (loopDetectedInStream) { break; } - if (roundText || roundThoughtText || lastUsage) { + const attempt = attemptState.snapshot(); + if (attempt.text || attempt.thoughtText || attempt.usageMetadata) { this.eventEmitter?.emit(AgentEventType.ROUND_TEXT, { subagentId: this.subagentId, runId, round: turnCounter, - text: roundText, - thoughtText: roundThoughtText, - usageMetadata: lastUsage, + text: attempt.text, + thoughtText: attempt.thoughtText, + usageMetadata: attempt.usageMetadata, timestamp: Date.now(), } as AgentRoundTextEvent); } @@ -1181,19 +1163,23 @@ export class AgentCore { } // Update token usage if available - if (lastUsage) { - this.recordTokenUsage(lastUsage, turnCounter, roundStreamStart); + if (attempt.usageMetadata) { + this.recordTokenUsage( + attempt.usageMetadata, + turnCounter, + roundStreamStart, + ); } - if (functionCalls.length > 0) { + if (attempt.functionCalls.length > 0) { const toolCallResult = await this.processFunctionCalls( - functionCalls, + attempt.functionCalls, roundAbortController, promptId, turnCounter, toolsList, - currentResponseId, - wasOutputTruncated, + attempt.responseId, + attempt.wasOutputTruncated, handledToolCallFingerprints, duplicateProviderToolCallResponseIds, ); @@ -1244,7 +1230,7 @@ export class AgentCore { turnCounter, ); if (waitResult.terminateMode) { - finalText = roundText.trim(); + finalText = attempt.text.trim(); terminateMode = waitResult.terminateMode; break; } @@ -1254,8 +1240,8 @@ export class AgentCore { continue; } - if (roundText && roundText.trim().length > 0) { - finalText = roundText.trim(); + if (attempt.text && attempt.text.trim().length > 0) { + finalText = attempt.text.trim(); break; } currentMessages = [ @@ -1271,8 +1257,8 @@ export class AgentCore { continue; } else { // No tool calls — treat this as the model's final answer. - if (roundText && roundText.trim().length > 0) { - finalText = roundText.trim(); + if (attempt.text && attempt.text.trim().length > 0) { + finalText = attempt.text.trim(); // Emit ROUND_END for the final round so all consumers see it. // Previously this was skipped, requiring AgentInteractive to // compensate with an explicit flushStreamBuffers() call. diff --git a/packages/core/src/core/client.test.ts b/packages/core/src/core/client.test.ts index 9823cae252c..cb34a40b836 100644 --- a/packages/core/src/core/client.test.ts +++ b/packages/core/src/core/client.test.ts @@ -11330,6 +11330,45 @@ Other open files: expect(request.input.message_id.length).toBeGreaterThan(0); }); + it('discards MessageDisplay text from a failed fresh attempt', async () => { + const mockMessageBus = { + request: vi.fn().mockResolvedValue({}), + response: vi.fn(), + }; + vi.mocked(mockConfig.getDisableAllHooks).mockReturnValue(false); + vi.mocked(mockConfig.getMessageBus).mockReturnValue( + mockMessageBus as unknown as ReturnType, + ); + vi.mocked(mockConfig.hasHooksForEvent).mockImplementation( + (event: string) => event === 'MessageDisplay', + ); + mockTurnRunFn.mockReturnValue( + (async function* () { + yield { type: LlmEventType.Content, value: 'stale answer' }; + yield { type: LlmEventType.Retry, isContinuation: false }; + yield { type: LlmEventType.Content, value: 'current answer' }; + })(), + ); + + const stream = client.sendMessageStream( + [{ text: 'Hi' }], + new AbortController().signal, + 'prompt-message-display-retry', + ); + for await (const _ of stream) { + // consume stream + } + + expect(mockMessageBus.request).toHaveBeenCalledTimes(1); + expect(mockMessageBus.request.mock.calls[0][0]).toMatchObject({ + eventName: 'MessageDisplay', + input: { + displayed_text: 'current answer', + is_final: true, + }, + }); + }); + it('fires a debounced mid-stream flush once the debounce window elapses, then a separate final flush', async () => { vi.useFakeTimers(); const mockMessageBus = { diff --git a/packages/core/src/core/client.ts b/packages/core/src/core/client.ts index f41e1694079..830bc9508f3 100644 --- a/packages/core/src/core/client.ts +++ b/packages/core/src/core/client.ts @@ -3681,10 +3681,11 @@ export class LlmClient { event.type === LlmEventType.ModelFallback ) { hasToolCalls = false; - agentOutput.restartAttempt( + const preserveText = event.type === LlmEventType.Retry && - event.isContinuation === true, - ); + event.isContinuation === true; + agentOutput.restartAttempt(preserveText); + messageDisplay?.restartAttempt(preserveText); } if (event.type === LlmEventType.Content) { agentOutput.appendText(event.value); diff --git a/packages/core/src/core/message-display-dispatcher.test.ts b/packages/core/src/core/message-display-dispatcher.test.ts index a3eb4336d34..8e67272190f 100644 --- a/packages/core/src/core/message-display-dispatcher.test.ts +++ b/packages/core/src/core/message-display-dispatcher.test.ts @@ -146,6 +146,26 @@ describe('MessageDisplayDispatcher', () => { }); }); + it('discards restarted-attempt text but preserves continuation text', async () => { + const { bus, sent, release } = createControlledBus(); + const dispatcher = createDispatcher(bus); + + dispatcher.addChunk('discarded', 0); + dispatcher.restartAttempt(false, 10); + dispatcher.addChunk('kept ', 20); + dispatcher.restartAttempt(true, 30); + dispatcher.addChunk('continuation', 40); + const finished = dispatcher.finish(); + await release(); + await finished; + + expect(sent).toHaveLength(1); + expect(sent[0]).toMatchObject({ + displayed_text: 'kept continuation', + is_final: true, + }); + }); + it('lets the final flush supersede a pending mid-stream payload, keeping is_final', async () => { const { bus, sent, release } = createControlledBus(); const dispatcher = createDispatcher(bus); diff --git a/packages/core/src/core/message-display-dispatcher.ts b/packages/core/src/core/message-display-dispatcher.ts index 4daabde6320..fde88f29569 100644 --- a/packages/core/src/core/message-display-dispatcher.ts +++ b/packages/core/src/core/message-display-dispatcher.ts @@ -119,6 +119,14 @@ export class MessageDisplayDispatcher { } } + restartAttempt(preserveText: boolean, nowMs: number = Date.now()): void { + if (this.finished || preserveText) { + return; + } + this.state = createInitialMessageDisplayState(nowMs); + this.pending = null; + } + /** * Close out this message: dispatch the `is_final: true` payload (skipped * when no text ever streamed — a tool-call-only message — or when the turn diff --git a/packages/core/src/core/model-stream-attempt-state.test.ts b/packages/core/src/core/model-stream-attempt-state.test.ts new file mode 100644 index 00000000000..defc03607b8 --- /dev/null +++ b/packages/core/src/core/model-stream-attempt-state.test.ts @@ -0,0 +1,107 @@ +/** + * @license + * Copyright 2026 Qwen Team + * SPDX-License-Identifier: Apache-2.0 + */ + +import { describe, expect, it } from 'vitest'; +import type { GenerateContentResponse } from '@google/genai'; +import { FinishReason } from './genai-compat.js'; +import { StreamEventType } from './llm-chat.js'; +import { ModelStreamAttemptState } from './model-stream-attempt-state.js'; + +describe('ModelStreamAttemptState', () => { + it('clears all attempt state on a fresh retry', () => { + const state = new ModelStreamAttemptState(); + state.accept({ + type: StreamEventType.CHUNK, + value: { + responseId: 'old-response', + candidates: [ + { + content: { + parts: [ + { text: 'old text' }, + { text: 'old thought', thought: true }, + ], + }, + finishReason: FinishReason.MAX_TOKENS, + }, + ], + functionCalls: [{ id: 'old-call', name: 'old_tool', args: {} }], + usageMetadata: { promptTokenCount: 10 }, + } as GenerateContentResponse, + }); + + const transition = state.accept({ type: StreamEventType.RETRY }); + + expect(transition).toMatchObject({ + type: 'attempt_reset', + reason: 'retry', + preserveText: false, + }); + expect(state.snapshot()).toEqual({ + text: '', + thoughtText: '', + functionCalls: [], + wasOutputTruncated: false, + }); + }); + + it('preserves text only for continuation retries', () => { + const state = new ModelStreamAttemptState(); + state.accept({ + type: StreamEventType.CHUNK, + value: { + responseId: 'old-response', + candidates: [ + { + content: { + parts: [ + { text: 'first half ' }, + { text: 'thinking ', thought: true }, + ], + }, + }, + ], + functionCalls: [{ id: 'old-call', name: 'old_tool', args: {} }], + usageMetadata: { promptTokenCount: 10 }, + } as GenerateContentResponse, + }); + + state.accept({ type: StreamEventType.RETRY, isContinuation: true }); + + expect(state.snapshot()).toEqual({ + text: 'first half ', + thoughtText: 'thinking ', + functionCalls: [], + wasOutputTruncated: false, + }); + }); + + it('treats model fallback as a fresh attempt', () => { + const state = new ModelStreamAttemptState(); + state.accept({ + type: StreamEventType.CHUNK, + value: { + candidates: [{ content: { parts: [{ text: 'old' }] } }], + } as GenerateContentResponse, + }); + + const transition = state.accept({ + type: StreamEventType.MODEL_FALLBACK, + info: { + fromModel: 'primary', + toModel: 'fallback', + fallbackIndex: 1, + }, + }); + + expect(transition).toMatchObject({ + type: 'attempt_reset', + reason: 'model_fallback', + preserveText: false, + }); + expect(state.snapshot().text).toBe(''); + }); +}); diff --git a/packages/core/src/core/model-stream-attempt-state.ts b/packages/core/src/core/model-stream-attempt-state.ts new file mode 100644 index 00000000000..9d31424938c --- /dev/null +++ b/packages/core/src/core/model-stream-attempt-state.ts @@ -0,0 +1,165 @@ +/** + * @license + * Copyright 2026 Qwen Team + * SPDX-License-Identifier: Apache-2.0 + */ + +import type { + FunctionCall, + GenerateContentResponse, + GenerateContentResponseUsageMetadata, +} from '@google/genai'; +import { FinishReason } from './genai-compat.js'; +import { + StreamEventType, + type ModelFallbackInfo, + type StreamEvent, +} from './llm-chat.js'; + +export interface ModelStreamTextPart { + text: string; + thought: boolean; +} + +export interface ModelStreamAttemptSnapshot { + text: string; + thoughtText: string; + functionCalls: FunctionCall[]; + usageMetadata?: GenerateContentResponseUsageMetadata; + responseId?: string; + finishReason?: FinishReason; + wasOutputTruncated: boolean; +} + +export type ModelStreamAttemptTransition = + | { + type: 'chunk'; + response: GenerateContentResponse; + textParts: ModelStreamTextPart[]; + functionCalls: FunctionCall[]; + finishReason?: FinishReason; + } + | { + type: 'attempt_reset'; + reason: 'retry'; + preserveText: boolean; + retryInfo: Extract; + } + | { + type: 'attempt_reset'; + reason: 'model_fallback'; + preserveText: false; + info: ModelFallbackInfo; + } + | { + type: 'compressed'; + info: Extract['info']; + }; + +export class ModelStreamAttemptState { + private text = ''; + private thoughtText = ''; + private functionCalls: FunctionCall[] = []; + private usageMetadata?: GenerateContentResponseUsageMetadata; + private responseId?: string; + private finishReason?: FinishReason; + private wasOutputTruncated = false; + + accept(event: StreamEvent): ModelStreamAttemptTransition { + switch (event.type) { + case StreamEventType.CHUNK: + return this.acceptChunk(event.value); + case StreamEventType.RETRY: { + const preserveText = event.isContinuation === true; + this.reset(preserveText); + return { + type: 'attempt_reset', + reason: 'retry', + preserveText, + retryInfo: event, + }; + } + case StreamEventType.MODEL_FALLBACK: + this.reset(false); + return { + type: 'attempt_reset', + reason: 'model_fallback', + preserveText: false, + info: event.info, + }; + case StreamEventType.COMPRESSED: + return { type: 'compressed', info: event.info }; + default: + throw new Error('Unsupported model stream event'); + } + } + + snapshot(): ModelStreamAttemptSnapshot { + return { + text: this.text, + thoughtText: this.thoughtText, + functionCalls: [...this.functionCalls], + ...(this.usageMetadata ? { usageMetadata: this.usageMetadata } : {}), + ...(this.responseId ? { responseId: this.responseId } : {}), + ...(this.finishReason ? { finishReason: this.finishReason } : {}), + wasOutputTruncated: this.wasOutputTruncated, + }; + } + + private acceptChunk( + response: GenerateContentResponse, + ): ModelStreamAttemptTransition { + const textParts = (response.candidates?.[0]?.content?.parts ?? []) + .filter( + (part): part is typeof part & { text: string } => + typeof part.text === 'string' && part.text.length > 0, + ) + .map((part) => ({ + text: part.text, + thought: part.thought === true, + })); + for (const part of textParts) { + if (part.thought) { + this.thoughtText += part.text; + } else { + this.text += part.text; + } + } + + const functionCalls = response.functionCalls ?? []; + this.functionCalls.push(...functionCalls); + if (response.usageMetadata) { + this.usageMetadata = response.usageMetadata; + } + if (response.responseId) { + this.responseId = response.responseId; + } + const finishReason = response.candidates?.[0]?.finishReason; + if (finishReason) { + this.finishReason = finishReason; + if (finishReason === FinishReason.MAX_TOKENS) { + this.wasOutputTruncated = true; + } + } + + return { + type: 'chunk', + response, + textParts, + functionCalls, + ...(finishReason ? { finishReason } : {}), + }; + } + + private reset(preserveText: boolean): void { + if (!preserveText) { + this.text = ''; + this.thoughtText = ''; + } + this.functionCalls = []; + this.usageMetadata = undefined; + this.responseId = undefined; + this.finishReason = undefined; + this.wasOutputTruncated = false; + } +} diff --git a/packages/core/src/core/turn.test.ts b/packages/core/src/core/turn.test.ts index 847bde64bad..db9a90097b5 100644 --- a/packages/core/src/core/turn.test.ts +++ b/packages/core/src/core/turn.test.ts @@ -1275,8 +1275,8 @@ describe('Turn', () => { }, }; yield { + type: StreamEventType.CHUNK, value: { - type: StreamEventType.CHUNK, candidates: [ { content: { parts: [{ text: 'Second part' }] }, diff --git a/packages/core/src/core/turn.ts b/packages/core/src/core/turn.ts index c3a68e4d441..87172f6adbf 100644 --- a/packages/core/src/core/turn.ts +++ b/packages/core/src/core/turn.ts @@ -44,6 +44,7 @@ import type { GoalTurnPermit, } from '../goals/goal-protocol.js'; import { getProviderToolCallId } from './toolCallIdUtils.js'; +import { ModelStreamAttemptState } from './model-stream-attempt-state.js'; const ERROR_REPORT_HISTORY_TAIL_COUNT = 8; const ERROR_REPORT_TEXT_PREVIEW_CHARS = 200; @@ -627,6 +628,7 @@ export class Turn { this.prompt_id, this.goalContext, ); + const attemptState = new ModelStreamAttemptState(); for await (const streamEvent of responseStream) { if (signal?.aborted) { @@ -634,35 +636,37 @@ export class Turn { return; } - // Handle the new RETRY event: clear accumulated state from the - // previous attempt to avoid duplicate tool calls and stale metadata. - if (streamEvent.type === 'retry') { + const transition = attemptState.accept(streamEvent); + if ( + transition.type === 'attempt_reset' && + transition.reason === 'retry' + ) { this.pendingToolCalls.length = 0; this.pendingCitations.clear(); this.finishReason = undefined; + this.currentResponseId = undefined; yield { type: LlmEventType.Retry, - retryInfo: streamEvent.retryInfo, - isContinuation: streamEvent.isContinuation, + retryInfo: transition.retryInfo.retryInfo, + isContinuation: transition.retryInfo.isContinuation, }; - continue; // Skip to the next event in the stream + continue; } - // Surface model fallback transitions from the chat stream as the - // top-level ModelFallback event. The UI uses this to notify the user - // that the system switched to a different model due to capacity issues. - if (streamEvent.type === 'model_fallback') { - // Clear accumulated state from the failed model's partial response + if ( + transition.type === 'attempt_reset' && + transition.reason === 'model_fallback' + ) { this.pendingToolCalls.length = 0; this.pendingCitations.clear(); this.finishReason = undefined; this.currentResponseId = undefined; yield { type: LlmEventType.ModelFallback, - fromModel: streamEvent.info.fromModel, - toModel: streamEvent.info.toModel, - statusCode: streamEvent.info.statusCode, - fallbackIndex: streamEvent.info.fallbackIndex, + fromModel: transition.info.fromModel, + toModel: transition.info.toModel, + statusCode: transition.info.statusCode, + fallbackIndex: transition.info.fallbackIndex, }; continue; } @@ -672,19 +676,16 @@ export class Turn { // connected. This bridge is the primary path for auto-compaction // events; manual /compress emits its own ChatCompressed in // LlmClient.tryCompressChat. - if (streamEvent.type === 'compressed') { + if (transition.type === 'compressed') { yield { type: LlmEventType.ChatCompressed, - value: streamEvent.info, + value: transition.info, }; continue; } - // Assuming other events are chunks with a `value` property - const resp = streamEvent.value as GenerateContentResponse; - if (!resp) continue; // Skip if there's no response body + const resp = transition.response; - // Track the current response ID for tool call correlation if (resp.responseId) { this.currentResponseId = resp.responseId; } @@ -708,9 +709,7 @@ export class Turn { }; } - // Handle function calls (requesting tool execution) - const functionCalls = resp.functionCalls ?? []; - for (const fnCall of functionCalls) { + for (const fnCall of transition.functionCalls) { const event = this.handlePendingFunctionCall(fnCall); if (event) { yield event; @@ -721,8 +720,7 @@ export class Turn { this.pendingCitations.add(citation); } - // Check if response was truncated or stopped for various reasons - const finishReason = resp.candidates?.[0]?.finishReason; + const finishReason = transition.finishReason; // This is the key change: Only yield 'Finished' if there is a finishReason. if (finishReason) { diff --git a/packages/core/src/followup/speculation.test.ts b/packages/core/src/followup/speculation.test.ts index eebe020d45c..057d1ce835f 100644 --- a/packages/core/src/followup/speculation.test.ts +++ b/packages/core/src/followup/speculation.test.ts @@ -64,6 +64,68 @@ afterEach(() => { }); describe('startSpeculation', () => { + it('discards failed-attempt tool calls after model fallback', async () => { + const getToolRegistry = vi.fn(); + const config = { + getApprovalMode: vi.fn().mockReturnValue(ApprovalMode.DEFAULT), + getCwd: vi.fn().mockReturnValue(process.cwd()), + getFastModel: vi.fn().mockReturnValue(undefined), + getSessionId: vi.fn().mockReturnValue('spec-session'), + getTargetDir: vi.fn().mockReturnValue('/spec/cwd'), + getToolRegistry, + } as unknown as Config; + forkedAgentMocks.runForkedAgent.mockResolvedValue({ + jsonResult: { suggestion: '' }, + }); + forkedAgentMocks.sendMessageStream.mockImplementation(async function* () { + yield { + type: 'chunk', + value: { + candidates: [ + { + content: { + parts: [ + { + functionCall: { + id: 'stale-call', + name: 'read_file', + args: { path: 'stale.ts' }, + }, + }, + ], + }, + }, + ], + }, + }; + yield { + type: 'model_fallback', + info: { + fromModel: 'primary', + toModel: 'fallback', + fallbackIndex: 1, + }, + }; + yield { + type: 'chunk', + value: { + candidates: [{ content: { parts: [{ text: 'current answer' }] } }], + }, + }; + }); + + const state = await startSpeculation(config, 'inspect the repository'); + await vi.waitFor(() => expect(state.status).toBe('completed')); + + expect(getToolRegistry).not.toHaveBeenCalled(); + expect(state.messages).toEqual([ + expect.objectContaining({ role: 'user' }), + { role: 'model', parts: [{ text: 'current answer' }] }, + ]); + + await abortSpeculation(state); + }); + it('does not start when the session-scoped lookup returns null', async () => { const config = { getSessionId: vi.fn().mockReturnValue('spec-session'), diff --git a/packages/core/src/followup/speculation.ts b/packages/core/src/followup/speculation.ts index 7220e207206..b6b6165bcec 100644 --- a/packages/core/src/followup/speculation.ts +++ b/packages/core/src/followup/speculation.ts @@ -19,7 +19,7 @@ import type { Content, Part } from '@google/genai'; import type { Config } from '../config/config.js'; import type { LlmClient } from '../core/client.js'; import type { ToolArtifact } from '../tools/tools.js'; -import { StreamEventType } from '../core/llm-chat.js'; +import { ModelStreamAttemptState } from '../core/model-stream-attempt-state.js'; import { convertToFunctionErrorResponse, convertToFunctionResponse, @@ -277,10 +277,19 @@ async function runSpeculativeLoop( ); const modelParts: Part[] = []; + const attemptState = new ModelStreamAttemptState(); for await (const event of stream) { if (state.abortController?.signal.aborted) break; - if (event.type !== StreamEventType.CHUNK) continue; - const response = event.value; + const transition = attemptState.accept(event); + if (transition.type === 'attempt_reset') { + const preservedTextParts = transition.preserveText + ? modelParts.filter((part) => part.text !== undefined) + : []; + modelParts.splice(0, modelParts.length, ...preservedTextParts); + continue; + } + if (transition.type !== 'chunk') continue; + const response = transition.response; const parts = response.candidates?.[0]?.content?.parts ?? []; for (const part of parts) { // Skip thought/reasoning parts — only capture visible text + function calls diff --git a/packages/core/src/index.ts b/packages/core/src/index.ts index d3d5b975b6d..25af8635804 100644 --- a/packages/core/src/index.ts +++ b/packages/core/src/index.ts @@ -89,6 +89,7 @@ export * from './core/inlineMediaLimit.js'; export * from './core/insightProtocol.js'; export * from './core/logger.js'; export * from './core/message-display-dispatcher.js'; +export * from './core/model-stream-attempt-state.js'; export * from './core/nonInteractiveToolExecutor.js'; export * from './core/prompts.js'; export * from './core/output-styles.js'; diff --git a/packages/core/src/services/loopDetectionService.test.ts b/packages/core/src/services/loopDetectionService.test.ts index 59ba5929537..2ffc02abd67 100644 --- a/packages/core/src/services/loopDetectionService.test.ts +++ b/packages/core/src/services/loopDetectionService.test.ts @@ -191,6 +191,24 @@ describe('LoopDetectionService', () => { expect(loggers.logLoopDetected).not.toHaveBeenCalled(); }); + it('resets the consecutive tool-call counter on model fallback', () => { + const event = createToolCallRequestEvent('testTool', { param: 'value' }); + for (let i = 0; i < TOOL_CALL_LOOP_THRESHOLD - 1; i++) { + expect(service.checkAlwaysOnSafeties(event)).toBe(false); + } + + expect( + service.checkAlwaysOnSafeties({ + type: LlmEventType.ModelFallback, + } as ServerLlmStreamEvent), + ).toBe(false); + + for (let i = 0; i < TOOL_CALL_LOOP_THRESHOLD - 1; i++) { + expect(service.checkAlwaysOnSafeties(event)).toBe(false); + } + expect(loggers.logLoopDetected).not.toHaveBeenCalled(); + }); + it('should expose the current consecutive tool-call count', () => { const event = createToolCallRequestEvent('testTool', { param: 'value' }); for (let i = 0; i < TOOL_CALL_LOOP_THRESHOLD - 1; i++) { diff --git a/packages/core/src/services/loopDetectionService.ts b/packages/core/src/services/loopDetectionService.ts index 3f1173387db..96bd8ade9f3 100644 --- a/packages/core/src/services/loopDetectionService.ts +++ b/packages/core/src/services/loopDetectionService.ts @@ -433,15 +433,18 @@ export class LoopDetectionService { return false; } - // A retry re-streams the failed attempt's tool calls, which would + // A retry or fallback re-streams the failed attempt's tool calls, which would // double-count against both always-on guards. Roll the per-turn cap back // to the last committed round-trip (never below it — prior round-trips // stay) and drop the consecutive-identical streak so the replayed attempt // cannot push it over the threshold. The adaptive cap's repeat tracker is // cleared (consistent with how the heuristic path clears - // globalToolCallCounts on retry): the replayed calls re-populate it, and a + // globalToolCallCounts on reset): the replayed calls re-populate it, and a // stuck pattern simply re-accumulates toward the threshold. - if (event.type === LlmEventType.Retry) { + if ( + event.type === LlmEventType.Retry || + event.type === LlmEventType.ModelFallback + ) { this.turnToolCallTotal = this.turnToolCallTotalCommitted; this.resetToolCallCount(); this.capKeyCounts.clear();