From 6a6c3eb0fd437206cdb2127c7f7651f87c39a7a9 Mon Sep 17 00:00:00 2001 From: David East Date: Tue, 7 Jul 2026 22:41:41 -0400 Subject: [PATCH] fix gemini thinking-only stop handling --- packages/agent/src/index.ts | 1 + packages/agent/src/retrieval.ts | 2 +- packages/agent/src/session.ts | 9 +- packages/agent/src/strategy.ts | 34 +++- packages/agent/src/types/llm.ts | 8 + packages/agent/src/types/session.ts | 4 +- packages/agent/src/types/strategy.ts | 4 +- packages/agent/test/session.test.ts | 36 ++++ packages/agent/test/strategy.test.ts | 181 ++++++++++++++++++ packages/model/src/contract.ts | 10 +- packages/model/src/index.ts | 1 + packages/model/src/providers/gemini.ts | 137 ++++++++++--- packages/model/test/contract.test.ts | 10 +- .../model/test/providers/gemini-retry.test.ts | 93 ++++++++- .../test/providers/gemini-toolcall.test.ts | 83 ++++++++ packages/relay/src/index.ts | 1 + packages/relay/src/types.ts | 8 + 17 files changed, 590 insertions(+), 32 deletions(-) diff --git a/packages/agent/src/index.ts b/packages/agent/src/index.ts index f25b4f4..4e79b5a 100644 --- a/packages/agent/src/index.ts +++ b/packages/agent/src/index.ts @@ -40,6 +40,7 @@ export type { export type { ModelClient, + ModelErrorEvent, ModelRequest, ModelEvent, ModelMessage, diff --git a/packages/agent/src/retrieval.ts b/packages/agent/src/retrieval.ts index 38138d5..a1c3091 100644 --- a/packages/agent/src/retrieval.ts +++ b/packages/agent/src/retrieval.ts @@ -291,7 +291,7 @@ export function createRetrievalStrategy(opts: RetrievalStrategyOpts = {}): Agent } else if (ev.kind === 'usage') { turnUsage = ev.usage; } else if (ev.kind === 'error') { - yield { kind: 'error', message: ev.message }; + yield ev; return; } // `tool_call` events are impossible here (toolUseEnabled:false, diff --git a/packages/agent/src/session.ts b/packages/agent/src/session.ts index 90ff9e8..a203a72 100644 --- a/packages/agent/src/session.ts +++ b/packages/agent/src/session.ts @@ -139,7 +139,14 @@ export function createAgentSession(config: AgentSessionConfig): AgentSession { history = [...history, assistantMsg]; yield { kind: 'turn_completed', turnId, metrics, details: ev.details }; } else if (ev.kind === 'error') { - yield { kind: 'error', turnId, message: ev.message }; + yield { + kind: 'error', + turnId, + message: ev.message, + ...(ev.code ? { code: ev.code } : {}), + ...(ev.retryable !== undefined ? { retryable: ev.retryable } : {}), + ...(ev.details ? { details: ev.details } : {}), + }; return; } else if (ev.kind === 'custom') { yield { kind: 'strategy_event', name: ev.name, data: ev.data }; diff --git a/packages/agent/src/strategy.ts b/packages/agent/src/strategy.ts index af05c29..7205ea6 100644 --- a/packages/agent/src/strategy.ts +++ b/packages/agent/src/strategy.ts @@ -32,6 +32,14 @@ import type { ToolHandler, ToolResult } from './types/tools.js'; interface ReactLoopOptions { /** Cap on loop iterations to avoid runaway tool-call ping-pong. Default 24. */ maxTurns?: number; + /** + * Default strict behavior preserves existing ReAct semantics: any + * model error fails the run. `complete-with-warning` is for hosts + * where prior successful tool effects can be authoritative even when + * a later final model call fails with a known retryable no-output + * provider error. + */ + toolProgressErrorPolicy?: 'strict' | 'complete-with-warning'; /** * Opt-in: when `true`, tool calls produced in a single turn are partitioned * by the handler's `parallelSafe` tag. Parallel-safe calls run concurrently @@ -59,6 +67,7 @@ const DEFAULT_CRITIQUE_SYSTEM_PROMPT = export function createReactLoopStrategy(options: ReactLoopOptions = {}): AgentStrategy { const maxTurns = options.maxTurns ?? 24; const parallelDispatch = options.parallelDispatch === true; + const toolProgressErrorPolicy = options.toolProgressErrorPolicy ?? 'strict'; const reflexionEnabled = options.reflexion?.enabled === true; const reflexionMaxRetries = Math.max(0, options.reflexion?.maxRetries ?? 1); const critiqueSystemPrompt = @@ -72,6 +81,7 @@ export function createReactLoopStrategy(options: ReactLoopOptions = {}): AgentSt // critique branch is skipped entirely and the strategy returns // on the first no-tool-calls turn exactly as before. let reflexionRetriesRemaining = reflexionMaxRetries; + let successfulToolResultsThisRun = 0; for (let turn = 0; turn < maxTurns; turn++) { if (signal.aborted) { @@ -160,7 +170,23 @@ export function createReactLoopStrategy(options: ReactLoopOptions = {}): AgentSt } else if (ev.kind === 'usage') { turnUsage = ev.usage; } else if (ev.kind === 'error') { - yield { kind: 'error', message: ev.message }; + const modelError = ev as Extract; + if ( + toolProgressErrorPolicy === 'complete-with-warning' && + successfulToolResultsThisRun > 0 && + isCompletableToolProgressError(modelError) + ) { + yield { + kind: 'custom', + name: 'react_loop_model_warning', + data: { + error: modelError, + successfulToolResultCount: successfulToolResultsThisRun, + }, + }; + return; + } + yield modelError; return; } } @@ -451,6 +477,7 @@ export function createReactLoopStrategy(options: ReactLoopOptions = {}): AgentSt resultJson: JSON.stringify(safeSerializable(result)), text: '', }); + successfulToolResultsThisRun += 1; } } else { // Default behavior: byte-for-byte identical to the pre-change @@ -469,6 +496,7 @@ export function createReactLoopStrategy(options: ReactLoopOptions = {}): AgentSt resultJson: JSON.stringify(safeSerializable(result)), text: '', }); + successfulToolResultsThisRun += 1; } } // Close the tool-dispatch segment for this iteration. Pair @@ -500,6 +528,10 @@ export function createReactLoopStrategy(options: ReactLoopOptions = {}): AgentSt }; } +function isCompletableToolProgressError(ev: Extract): boolean { + return ev.code === 'gemini.thinking_only_stop' && ev.retryable === true; +} + function buildMessages(input: StrategyRunInput): ModelMessage[] { const out: ModelMessage[] = []; out.push({ role: 'system', text: input.systemPrompt }); diff --git a/packages/agent/src/types/llm.ts b/packages/agent/src/types/llm.ts index 9d8f5b0..d665a42 100644 --- a/packages/agent/src/types/llm.ts +++ b/packages/agent/src/types/llm.ts @@ -20,6 +20,14 @@ export type { ReasoningEffort, } from '@inbrowser/model'; +export interface ModelErrorEvent { + kind: 'error'; + message: string; + code?: string; + retryable?: boolean; + details?: Record; +} + export interface LlmConfig { apiKey?: string; model: string; diff --git a/packages/agent/src/types/session.ts b/packages/agent/src/types/session.ts index 8d54270..7034462 100644 --- a/packages/agent/src/types/session.ts +++ b/packages/agent/src/types/session.ts @@ -9,7 +9,7 @@ */ import type { ChatMessage, TurnDetails, TurnMetrics } from './chat.js'; -import type { ModelClient } from './llm.js'; +import type { ModelClient, ModelErrorEvent } from './llm.js'; import type { MetricsCollector } from './metrics.js'; import type { RuntimeState } from './runtime.js'; import type { AgentStrategy } from './strategy.js'; @@ -71,7 +71,7 @@ export type SessionEvent = | { kind: 'workspace_changed'; workspace: Workspace } | { kind: 'runtime_changed'; runtime: RuntimeState } | { kind: 'turn_completed'; turnId: string; metrics: TurnMetrics; details: TurnDetails } - | { kind: 'error'; turnId?: string; message: string } + | (ModelErrorEvent & { turnId?: string }) | { kind: 'completed' } /** Strategy-emitted milestones (planner phases, branch expansions, …) * — generic envelope so new strategies can surface custom events diff --git a/packages/agent/src/types/strategy.ts b/packages/agent/src/types/strategy.ts index 064c6e8..c011177 100644 --- a/packages/agent/src/types/strategy.ts +++ b/packages/agent/src/types/strategy.ts @@ -8,7 +8,7 @@ */ import type { ChatMessage, TurnDetails } from './chat.js'; -import type { ModelClient, ModelUsage } from './llm.js'; +import type { ModelClient, ModelErrorEvent, ModelUsage } from './llm.js'; import type { RuntimeState } from './runtime.js'; import type { ToolContext, ToolDispatch, ToolHandler, ToolResult } from './tools.js'; import type { Tracer } from './trace.js'; @@ -61,7 +61,7 @@ export type StrategyEvent = | { kind: 'tool_call'; id: string; name: string; args: unknown; signature?: string } | { kind: 'tool_result'; id: string; result: ToolResult } | { kind: 'turn_complete'; usage: ModelUsage; details: TurnDetails } - | { kind: 'error'; message: string } + | ModelErrorEvent /** Custom milestone — name + arbitrary payload, surfaced as * `SessionEvent.kind === 'strategy_event'` to the host. */ | { kind: 'custom'; name: string; data?: unknown }; diff --git a/packages/agent/test/session.test.ts b/packages/agent/test/session.test.ts index 5627a3c..c33e572 100644 --- a/packages/agent/test/session.test.ts +++ b/packages/agent/test/session.test.ts @@ -86,6 +86,42 @@ describe('createAgentSession', () => { expect(kinds[kinds.length - 1]).toBe('completed'); }); + test('preserves structured strategy error metadata on session errors', async () => { + const strategy: AgentStrategy = { + id: 'structured-error', + async *run() { + yield { + kind: 'error', + message: 'Gemini produced no output', + code: 'gemini.thinking_only_stop', + retryable: true, + details: { finishReason: 'STOP' }, + }; + }, + }; + const session = createAgentSession({ + strategy, + llm: fakeLlm([]), + tools: createDispatch(createToolRegistry()), + toolList: [], + toolContext: fakeCtx, + systemPromptBuilder: () => 'system', + metrics: createMetricsCollector(), + history: [], + }); + + const events = await collect(session.submit('hi', new AbortController().signal)); + const error = events.find((e) => e.kind === 'error'); + + expect(error).toMatchObject({ + kind: 'error', + message: 'Gemini produced no output', + code: 'gemini.thinking_only_stop', + retryable: true, + details: { finishReason: 'STOP' }, + }); + }); + test('applies tool result patches to workspace + runtime + emits change events', async () => { const writeRulesTool: ToolHandler<{ source: string }> = { name: 'writeRules', diff --git a/packages/agent/test/strategy.test.ts b/packages/agent/test/strategy.test.ts index 277a208..faeffb3 100644 --- a/packages/agent/test/strategy.test.ts +++ b/packages/agent/test/strategy.test.ts @@ -192,6 +192,187 @@ describe('createReactLoopStrategy', () => { } }); + test('default strict mode preserves structured model errors after tool progress', async () => { + const echoTool: ToolHandler<{ msg: string }, { msg: string }> = { + name: 'echo', + description: 'echo', + parameters: { type: 'object' }, + async execute({ msg }) { + return { ok: true, summary: msg, data: { msg } }; + }, + }; + const registry = createToolRegistry(); + registry.register(echoTool); + const error = { + kind: 'error' as const, + message: 'Gemini produced no output', + code: 'gemini.thinking_only_stop', + retryable: true, + details: { finishReason: 'STOP' }, + }; + const llm = fakeLlm([ + [ + { kind: 'tool_call', id: 'c1', name: 'echo', args: { msg: 'done' } }, + { kind: 'usage', usage: { promptTokens: 5, outputTokens: 1 } }, + ], + [error], + ]); + + const events = await collect( + createReactLoopStrategy().run( + { + prompt: 'use the tool', + history: [], + workspace: EMPTY_WORKSPACE, + runtime: EMPTY_RUNTIME, + llm, + tools: createDispatch(registry), + toolList: [echoTool], + toolContext: fakeCtx, + systemPrompt: 'You may call tools.', + }, + new AbortController().signal, + ), + ); + + const finalError = events.find((e) => e.kind === 'error'); + expect(finalError).toMatchObject(error); + }); + + test('progress-aware mode completes with warning after tool progress and Gemini thinking-only STOP', async () => { + const echoTool: ToolHandler<{ msg: string }, { msg: string }> = { + name: 'echo', + description: 'echo', + parameters: { type: 'object' }, + async execute({ msg }) { + return { ok: true, summary: msg, data: { msg } }; + }, + }; + const registry = createToolRegistry(); + registry.register(echoTool); + const providerError = { + kind: 'error' as const, + message: 'Gemini produced no output', + code: 'gemini.thinking_only_stop', + retryable: true, + details: { finishReason: 'STOP' }, + }; + const llm = fakeLlm([ + [ + { kind: 'tool_call', id: 'c1', name: 'echo', args: { msg: 'done' } }, + { kind: 'usage', usage: { promptTokens: 5, outputTokens: 1 } }, + ], + [providerError], + ]); + + const events = await collect( + createReactLoopStrategy({ toolProgressErrorPolicy: 'complete-with-warning' }).run( + { + prompt: 'use the tool', + history: [], + workspace: EMPTY_WORKSPACE, + runtime: EMPTY_RUNTIME, + llm, + tools: createDispatch(registry), + toolList: [echoTool], + toolContext: fakeCtx, + systemPrompt: 'You may call tools.', + }, + new AbortController().signal, + ), + ); + + expect(events.some((e) => e.kind === 'error')).toBe(false); + expect(events.filter((e) => e.kind === 'turn_complete')).toHaveLength(1); + const warning = events.find( + (e) => e.kind === 'custom' && e.name === 'react_loop_model_warning', + ); + expect(warning?.kind).toBe('custom'); + if (warning?.kind === 'custom') { + expect(warning.data).toMatchObject({ + error: providerError, + successfulToolResultCount: 1, + }); + } + }); + + test('progress-aware mode does not swallow unknown provider errors', async () => { + const echoTool: ToolHandler<{ msg: string }, { msg: string }> = { + name: 'echo', + description: 'echo', + parameters: { type: 'object' }, + async execute({ msg }) { + return { ok: true, summary: msg, data: { msg } }; + }, + }; + const registry = createToolRegistry(); + registry.register(echoTool); + const unknownError = { + kind: 'error' as const, + message: 'provider exploded', + code: 'provider.unknown', + retryable: true, + }; + const llm = fakeLlm([ + [ + { kind: 'tool_call', id: 'c1', name: 'echo', args: { msg: 'done' } }, + { kind: 'usage', usage: { promptTokens: 5, outputTokens: 1 } }, + ], + [unknownError], + ]); + + const events = await collect( + createReactLoopStrategy({ toolProgressErrorPolicy: 'complete-with-warning' }).run( + { + prompt: 'use the tool', + history: [], + workspace: EMPTY_WORKSPACE, + runtime: EMPTY_RUNTIME, + llm, + tools: createDispatch(registry), + toolList: [echoTool], + toolContext: fakeCtx, + systemPrompt: 'You may call tools.', + }, + new AbortController().signal, + ), + ); + + expect(events.find((e) => e.kind === 'error')).toMatchObject(unknownError); + }); + + test('progress-aware mode does not complete when no tool result was dispatched', async () => { + const providerError = { + kind: 'error' as const, + message: 'Gemini produced no output', + code: 'gemini.thinking_only_stop', + retryable: true, + }; + const llm = fakeLlm([[providerError]]); + + const events = await collect( + createReactLoopStrategy({ toolProgressErrorPolicy: 'complete-with-warning' }).run( + { + prompt: 'answer directly', + history: [], + workspace: EMPTY_WORKSPACE, + runtime: EMPTY_RUNTIME, + llm, + tools: createDispatch(createToolRegistry()), + toolList: [], + toolContext: fakeCtx, + systemPrompt: 'You are helpful.', + }, + new AbortController().signal, + ), + ); + + expect(events.find((e) => e.kind === 'error')).toMatchObject(providerError); + expect(events.some((e) => e.kind === 'custom' && e.name === 'react_loop_model_warning')).toBe( + false, + ); + }); + test('aborts when the signal fires before the turn starts', async () => { const controller = new AbortController(); controller.abort(); diff --git a/packages/model/src/contract.ts b/packages/model/src/contract.ts index 335a20b..4d982c2 100644 --- a/packages/model/src/contract.ts +++ b/packages/model/src/contract.ts @@ -73,12 +73,20 @@ export interface ModelUsage { * can therefore rely on exactly one of {a `usage` event, an `error` event} per * turn. */ +export interface ModelErrorEvent { + kind: 'error'; + message: string; + code?: string; + retryable?: boolean; + details?: Record; +} + export type ModelEvent = | { kind: 'text'; text: string } | { kind: 'thinking'; text: string } | { kind: 'tool_call'; id: string; name: string; args: unknown; signature?: string } | { kind: 'usage'; usage: ModelUsage } - | { kind: 'error'; message: string }; + | ModelErrorEvent; /** * The one model-call contract. Implemented by the cloud providers and the diff --git a/packages/model/src/index.ts b/packages/model/src/index.ts index 5780e11..873ed15 100644 --- a/packages/model/src/index.ts +++ b/packages/model/src/index.ts @@ -143,6 +143,7 @@ export type { // it; relay + agent consume it). Type-only; importing it pulls no runtime. export type { ModelClient, + ModelErrorEvent, ModelEvent, ModelMessage, ModelRequest, diff --git a/packages/model/src/providers/gemini.ts b/packages/model/src/providers/gemini.ts index b05b63e..3ea79b5 100644 --- a/packages/model/src/providers/gemini.ts +++ b/packages/model/src/providers/gemini.ts @@ -1,4 +1,11 @@ -import type { ModelClient, ModelEvent, ModelRequest, ToolSpec } from '../contract.js'; +import type { + ModelClient, + ModelErrorEvent, + ModelEvent, + ModelRequest, + ReasoningEffort, + ToolSpec, +} from '../contract.js'; import { readSseDataLines } from '../sse.js'; import type { CloudProviderConfig } from './types.js'; /** @@ -43,6 +50,30 @@ interface GeminiBody { generationConfig?: Record; } +const GEMINI_25_THINKING_BUDGET: Record, number> = { + low: 1024, + medium: 4096, + high: 8192, +}; + +function buildGeminiThinkingConfig( + model: string, + effort: ReasoningEffort | undefined, +): Record | undefined { + if (!effort || effort === 'off') return undefined; + + const thinkingConfig: Record = { includeThoughts: true }; + const normalized = model.toLowerCase(); + + if (normalized.includes('gemini-3.5-') || normalized.includes('gemini-3-flash')) { + thinkingConfig.thinkingLevel = effort; + } else if (normalized.includes('gemini-2.5-')) { + thinkingConfig.thinkingBudget = GEMINI_25_THINKING_BUDGET[effort]; + } + + return thinkingConfig; +} + function toGeminiBody(config: CloudProviderConfig, req: ModelRequest): GeminiBody { const contents: GeminiContent[] = []; let systemText = ''; @@ -109,7 +140,6 @@ function toGeminiBody(config: CloudProviderConfig, req: ModelRequest): GeminiBod } const gen: Record = { - thinkingConfig: { includeThoughts: true }, // Generous output budget. Left unset, the model can truncate a // large tool-call argument — writeApp/writeCode emit whole source // files as a string arg — and a truncated call is exactly what @@ -119,6 +149,8 @@ function toGeminiBody(config: CloudProviderConfig, req: ModelRequest): GeminiBod // 400, not silently. maxOutputTokens: 65536, }; + const thinkingConfig = buildGeminiThinkingConfig(config.model, req.reasoningEffort); + if (thinkingConfig) gen.thinkingConfig = thinkingConfig; // Per-request temperature wins; otherwise fall back to the // construction-time default (the docs agent pins 0.2; the relay sets // neither, preserving "send only what the client did"). @@ -140,6 +172,7 @@ interface GeminiStreamChunk { promptTokenCount?: number; candidatesTokenCount?: number; cachedContentTokenCount?: number; + thoughtsTokenCount?: number; }; } @@ -229,13 +262,16 @@ export async function* geminiEventsFromResponse( let promptTokens = 0; let completionTokens = 0; let cachedTokens = 0; + let reasoningTokens = 0; // Diagnostics for the "thinking-only, no output" case: Gemini can // end a response after the thinking phase having produced nothing // visible. `finishReason` on the last chunk names why (MAX_TOKENS / // SAFETY / RECITATION); a missing one means the stream was simply // truncated. `sawOutput` tracks whether any *visible* output (text // or a tool call — not thinking) actually came through. - let sawOutput = false; + let sawThinking = false; + let sawVisibleText = false; + let sawFunctionCall = false; let lastFinishReason: string | undefined; // Function calls accumulate here and are flushed once, after the // stream closes (see the per-part merge below and the flush loop). @@ -270,14 +306,15 @@ export async function* geminiEventsFromResponse( for (const p of parts) { if (typeof p.text === 'string' && p.text.length > 0) { if (p.thought === true) { + sawThinking = true; yield { kind: 'thinking', text: p.text }; } else { - sawOutput = true; + sawVisibleText = true; yield { kind: 'text', text: p.text }; } } if (p.functionCall) { - sawOutput = true; + sawFunctionCall = true; const slot = fnOrdinal++; let call = pending.get(slot); if (!call) { @@ -322,6 +359,9 @@ export async function* geminiEventsFromResponse( if (typeof usage.cachedContentTokenCount === 'number') { cachedTokens = usage.cachedContentTokenCount; } + if (typeof usage.thoughtsTokenCount === 'number') { + reasoningTokens = usage.thoughtsTokenCount; + } } } } catch (e) { @@ -333,13 +373,13 @@ export async function* geminiEventsFromResponse( // Stream ended cleanly but the model never produced visible output — // only thinking. Surface why: a non-STOP `finishReason` names it, // `none` means the stream was truncated before one arrived. - if (!sawOutput) { - yield { - kind: 'error', - message: `Gemini produced no output — finishReason=${ - lastFinishReason ?? 'none' - } (response ended after thinking only)`, - }; + if (!sawVisibleText && !sawFunctionCall) { + yield geminiNoOutputError({ + finishReason: lastFinishReason, + sawThinking, + sawVisibleText, + sawFunctionCall, + }); return; } @@ -369,6 +409,47 @@ export async function* geminiEventsFromResponse( promptTokens, outputTokens: completionTokens, ...(cachedTokens > 0 ? { cachedTokens } : {}), + ...(reasoningTokens > 0 ? { reasoningTokens } : {}), + }, + }; +} + +function geminiNoOutputError(opts: { + finishReason: string | undefined; + sawThinking: boolean; + sawVisibleText: boolean; + sawFunctionCall: boolean; +}): ModelErrorEvent { + const finishReason = opts.finishReason ?? 'none'; + const message = `Gemini produced no output — finishReason=${finishReason} (${ + opts.sawThinking + ? 'response ended after thinking only' + : 'response ended with no visible output' + })`; + + let code = 'gemini.no_output'; + let retryable = false; + if (opts.finishReason === undefined) { + code = 'gemini.truncated_no_output'; + retryable = true; + } else if (opts.finishReason === 'MALFORMED_FUNCTION_CALL') { + code = 'gemini.malformed_function_call'; + retryable = true; + } else if (opts.finishReason === 'STOP' && opts.sawThinking) { + code = 'gemini.thinking_only_stop'; + retryable = true; + } + + return { + kind: 'error', + message, + code, + retryable, + details: { + finishReason, + sawThinking: opts.sawThinking, + sawVisibleText: opts.sawVisibleText, + sawFunctionCall: opts.sawFunctionCall, }, }; } @@ -397,16 +478,32 @@ export async function* geminiEventsFromResponse( const MAX_GEMINI_ATTEMPTS = 3; const RETRY_DELAY_MS = 500; -/** Substrings that identify a retryable provider error. Matched - * against `ModelEvent.message` from `geminiEventsFromResponse`. */ +/** Substrings kept as a fallback for older message-only provider errors. */ const RETRYABLE_ERROR_MARKERS = [ 'MALFORMED_FUNCTION_CALL', 'finishReason=STOP', 'finishReason=none', ]; -function isRetryableError(message: string): boolean { - return RETRYABLE_ERROR_MARKERS.some((m) => message.includes(m)); +function isRetryableError(event: ModelErrorEvent): boolean { + if (event.retryable === true) return true; + if (event.retryable === false) return false; + return RETRYABLE_ERROR_MARKERS.some((m) => event.message.includes(m)); +} + +function withAttemptMetadata( + event: ModelErrorEvent, + attempt: number, + maxAttempts: number, +): ModelErrorEvent { + return { + ...event, + details: { + ...(event.details ?? {}), + attempt, + maxAttempts, + }, + }; } /** @@ -438,15 +535,11 @@ export function geminiModelClient(config: GeminiConfig): ModelClient { // The non-retryable kinds (SAFETY, RECITATION, MAX_TOKENS, // network/parse failures) fall straight through and surface. // Final attempt always yields whatever it produces. - if ( - evt.kind === 'error' && - isRetryableError(evt.message) && - attempt < MAX_GEMINI_ATTEMPTS - ) { + if (evt.kind === 'error' && isRetryableError(evt) && attempt < MAX_GEMINI_ATTEMPTS) { retry = true; break; } - yield evt; + yield evt.kind === 'error' ? withAttemptMetadata(evt, attempt, MAX_GEMINI_ATTEMPTS) : evt; } if (!retry) return; diff --git a/packages/model/test/contract.test.ts b/packages/model/test/contract.test.ts index 3ceb300..308d9a0 100644 --- a/packages/model/test/contract.test.ts +++ b/packages/model/test/contract.test.ts @@ -14,7 +14,13 @@ describe('model contract', () => { { kind: 'thinking', text: 'hmm' }, { kind: 'tool_call', id: 'c1', name: 'search', args: { q: 'x' } }, { kind: 'usage', usage: { promptTokens: 10, outputTokens: 5 } }, - { kind: 'error', message: 'nope' }, + { + kind: 'error', + message: 'nope', + code: 'provider.nope', + retryable: false, + details: { reason: 'test' }, + }, ]; expect(events.map((e) => e.kind)).toEqual(['text', 'thinking', 'tool_call', 'usage', 'error']); @@ -24,6 +30,8 @@ describe('model contract', () => { if (tc.kind === 'tool_call') expect(tc.id).toBe('c1'); // `id`, not `callId` const u = events[3]; if (u.kind === 'usage') expect(u.usage.outputTokens).toBe(5); // nested usage; `outputTokens` + const err = events[4]; + if (err.kind === 'error') expect(err.code).toBe('provider.nope'); }); test('a minimal ModelClient is structurally valid; the turn ends by returning', async () => { diff --git a/packages/model/test/providers/gemini-retry.test.ts b/packages/model/test/providers/gemini-retry.test.ts index ad76563..3fce6c3 100644 --- a/packages/model/test/providers/gemini-retry.test.ts +++ b/packages/model/test/providers/gemini-retry.test.ts @@ -15,7 +15,7 @@ */ import { afterEach, beforeEach, describe, expect, it } from 'bun:test'; import type { ModelEvent, ModelRequest } from '../../src/contract'; -import { geminiModelClient } from '../../src/providers/gemini'; +import { geminiEventsFromResponse, geminiModelClient } from '../../src/providers/gemini'; function makeSseResponse(chunks: unknown[]): Response { const body = chunks.map((c) => `data: ${JSON.stringify(c)}\n\n`).join(''); @@ -165,6 +165,8 @@ describe('gemini ModelClient retry', () => { expect(error?.kind).toBe('error'); if (error?.kind === 'error') { expect(error.message).toContain('finishReason=SAFETY'); + expect(error.code).toBe('gemini.no_output'); + expect(error.retryable).toBe(false); } }); @@ -189,6 +191,95 @@ describe('gemini ModelClient retry', () => { expect(error?.kind).toBe('error'); if (error?.kind === 'error') { expect(error.message).toContain('finishReason=STOP'); + expect(error.code).toBe('gemini.thinking_only_stop'); + expect(error.retryable).toBe(true); + expect(error.details).toMatchObject({ + finishReason: 'STOP', + sawThinking: true, + sawVisibleText: false, + sawFunctionCall: false, + attempt: 3, + maxAttempts: 3, + }); + } + }); + + it('marks truncated no-output streams as retryable structured errors', async () => { + const events = await collect( + geminiEventsFromResponse( + makeSseResponse([ + { + candidates: [{ content: { parts: [{ text: 'starting', thought: true }] } }], + }, + ]), + ), + ); + + const error = events.find((e) => e.kind === 'error'); + expect(error?.kind).toBe('error'); + if (error?.kind === 'error') { + expect(error.code).toBe('gemini.truncated_no_output'); + expect(error.retryable).toBe(true); + expect(error.details).toMatchObject({ + finishReason: 'none', + sawThinking: true, + sawVisibleText: false, + sawFunctionCall: false, + }); + } + }); + + it('marks MALFORMED_FUNCTION_CALL no-output streams as retryable structured errors', async () => { + const events = await collect( + geminiEventsFromResponse( + makeSseResponse([ + { + candidates: [ + { + content: { parts: [{ text: 'thinking', thought: true }] }, + finishReason: 'MALFORMED_FUNCTION_CALL', + }, + ], + }, + ]), + ), + ); + + const error = events.find((e) => e.kind === 'error'); + expect(error?.kind).toBe('error'); + if (error?.kind === 'error') { + expect(error.code).toBe('gemini.malformed_function_call'); + expect(error.retryable).toBe(true); + expect(error.details).toMatchObject({ + finishReason: 'MALFORMED_FUNCTION_CALL', + }); + } + }); + + it('marks MAX_TOKENS no-output streams as non-retryable', async () => { + const events = await collect( + geminiEventsFromResponse( + makeSseResponse([ + { + candidates: [ + { + content: { parts: [{ text: 'thinking', thought: true }] }, + finishReason: 'MAX_TOKENS', + }, + ], + }, + ]), + ), + ); + + const error = events.find((e) => e.kind === 'error'); + expect(error?.kind).toBe('error'); + if (error?.kind === 'error') { + expect(error.code).toBe('gemini.no_output'); + expect(error.retryable).toBe(false); + expect(error.details).toMatchObject({ + finishReason: 'MAX_TOKENS', + }); } }); }); diff --git a/packages/model/test/providers/gemini-toolcall.test.ts b/packages/model/test/providers/gemini-toolcall.test.ts index 9d8bc1c..79a9c3e 100644 --- a/packages/model/test/providers/gemini-toolcall.test.ts +++ b/packages/model/test/providers/gemini-toolcall.test.ts @@ -56,6 +56,67 @@ function toolCalls(events: ModelEvent[]) { >[]; } +function baseReq(overrides: Partial = {}): ModelRequest { + return { + messages: [{ role: 'user', text: 'hi' }], + tools: [], + toolUseEnabled: false, + ...overrides, + }; +} + +async function geminiBody(model: string, req: ModelRequest): Promise> { + return JSON.parse(await buildGeminiRequest({ apiKey: 'sk-test', model }, req).text()) as Record< + string, + unknown + >; +} + +function generationConfig(body: Record): Record { + return body.generationConfig as Record; +} + +describe('buildGeminiRequest — thinking config', () => { + it('omits thought summaries when reasoningEffort is omitted', async () => { + const body = await geminiBody('gemini-3.5-flash', baseReq()); + expect(generationConfig(body).thinkingConfig).toBeUndefined(); + }); + + it("omits thought summaries when reasoningEffort is 'off'", async () => { + const body = await geminiBody('gemini-3.5-flash', baseReq({ reasoningEffort: 'off' })); + expect(generationConfig(body).thinkingConfig).toBeUndefined(); + }); + + it('sets includeThoughts and thinkingLevel for Gemini 3.5 models when requested', async () => { + const body = await geminiBody('gemini-3.5-flash', baseReq({ reasoningEffort: 'medium' })); + expect(generationConfig(body).thinkingConfig).toEqual({ + includeThoughts: true, + thinkingLevel: 'medium', + }); + }); + + it('sets includeThoughts and thinkingLevel for Gemini 3 Flash models when requested', async () => { + const body = await geminiBody('gemini-3-flash-preview', baseReq({ reasoningEffort: 'high' })); + expect(generationConfig(body).thinkingConfig).toEqual({ + includeThoughts: true, + thinkingLevel: 'high', + }); + }); + + it('sets includeThoughts and thinkingBudget for Gemini 2.5 models when requested', async () => { + const body = await geminiBody('gemini-2.5-flash', baseReq({ reasoningEffort: 'low' })); + expect(generationConfig(body).thinkingConfig).toEqual({ + includeThoughts: true, + thinkingBudget: 1024, + }); + }); + + it('requests summaries without model-specific controls for unknown Gemini families', async () => { + const body = await geminiBody('gemini-experimental', baseReq({ reasoningEffort: 'high' })); + expect(generationConfig(body).thinkingConfig).toEqual({ includeThoughts: true }); + }); +}); + describe('geminiEventsFromResponse — function-call accumulation', () => { it('collapses a call re-sent across chunks into one event with complete args', async () => { // One logical `bash` call: empty-arg partial, then the complete args @@ -266,6 +327,28 @@ describe('geminiEventsFromResponse — function-call accumulation', () => { expect(firstCall).toBeGreaterThan(lastText); }); + it('maps thoughtsTokenCount to reasoningTokens in usage', async () => { + const events = await collect( + geminiEventsFromResponse( + makeSseResponse([ + chunk([{ text: 'answer' }], 'STOP'), + { + usageMetadata: { + promptTokenCount: 10, + candidatesTokenCount: 5, + thoughtsTokenCount: 3, + }, + }, + ]), + ), + ); + + expect(events.find((e) => e.kind === 'usage')).toEqual({ + kind: 'usage', + usage: { promptTokens: 10, outputTokens: 5, reasoningTokens: 3 }, + }); + }); + it('yields exactly N events for an N-call turn flowing through the ModelClient', async () => { // End-to-end through the retry-bearing client to confirm the flush // survives the provider layer that piebox actually consumes. diff --git a/packages/relay/src/index.ts b/packages/relay/src/index.ts index f8e453f..c73096f 100644 --- a/packages/relay/src/index.ts +++ b/packages/relay/src/index.ts @@ -60,6 +60,7 @@ export type { // The shared model-call contract, re-exported from the relay's import // site (sourced from `@inbrowser/model/contract`). There is no // relay-local `InferenceEvent` / `ChatMessage` / `ToolDecl` anymore. + ModelErrorEvent, ModelEvent, ModelMessage, ModelRequest, diff --git a/packages/relay/src/types.ts b/packages/relay/src/types.ts index d049f98..38a14a6 100644 --- a/packages/relay/src/types.ts +++ b/packages/relay/src/types.ts @@ -33,6 +33,14 @@ export type { ToolSpec, } from '@inbrowser/model'; +export interface ModelErrorEvent { + kind: 'error'; + message: string; + code?: string; + retryable?: boolean; + details?: Record; +} + /** * The wire shape the relay accepts at `handleStart`. It is the shared * `ModelRequest` plus the relay-only transport concerns: `provider` (the