From dda3c52fe2e6ade71981dbc1486256d93f91a191 Mon Sep 17 00:00:00 2001 From: Daniel Salazar Date: Fri, 16 Oct 2026 20:30:49 -0400 Subject: [PATCH] Keep Gemini extra_content off other providers; report Gemini timeouts as 504 Replayed tool calls only carry extra_content for providers that opt in (Gemini), so OpenAI-compatible and Responses vendors don't see an unknown field. The native Gemini client now owns its request deadline: the SDK's own timeout aborted without a reason and read as an internal error. A deadline timeout returns 504 upstream_timeout and leaves the route in rotation. --- .../ChatCompletionDriver.routing.test.ts | 20 ++++ .../drivers/ai-chat/ChatCompletionDriver.ts | 8 +- .../providers/OpenAICompatProvider.test.ts | 27 ++++++ .../ai-chat/providers/OpenAICompatProvider.ts | 6 ++ .../gemini/GeminiChatProvider.native.test.ts | 92 ++++++++++++++++++- .../providers/gemini/GeminiChatProvider.ts | 14 +-- .../ai-chat/providers/gemini/geminiNative.ts | 52 +++++++++-- .../drivers/ai-chat/utils/OpenAIUtil.js | 22 +++-- .../drivers/ai-chat/utils/OpenAIUtil.test.ts | 26 ++++-- src/backend/drivers/util/upstreamErrors.ts | 11 +++ 10 files changed, 240 insertions(+), 38 deletions(-) diff --git a/src/backend/drivers/ai-chat/ChatCompletionDriver.routing.test.ts b/src/backend/drivers/ai-chat/ChatCompletionDriver.routing.test.ts index 95e4d43b1..aba333282 100644 --- a/src/backend/drivers/ai-chat/ChatCompletionDriver.routing.test.ts +++ b/src/backend/drivers/ai-chat/ChatCompletionDriver.routing.test.ts @@ -58,6 +58,7 @@ import { markRouteUnhealthy, } from './utils/providerHealth.js'; import { InfronProvider } from './providers/infron/InfronProvider.js'; +import { UpstreamDeadlineError } from '../util/upstreamErrors.js'; // -- OpenAI SDK mock ------------------------------------------------ // Gemini reaches Google through `new openai.OpenAI()` (default export) and @@ -1079,6 +1080,25 @@ describe('ChatCompletionDriver unhealthy-route skipping', () => { expect(isRouteUnhealthy('deepseek', 'deepseek-v4-pro')).toBe(true); }); + it('returns 504 for a request that outlived its deadline and keeps the route', async () => { + createMock.mockRejectedValue( + new UpstreamDeadlineError('Gemini request timed out'), + ); + + const err = await withTestActor(() => + driver.complete({ + model: 'deepseek-v4-pro', + messages: [{ role: 'user', content: 'hi' }], + }), + ).catch((e: unknown) => e as HttpError); + + expect(err).toMatchObject({ + statusCode: 504, + legacyCode: 'upstream_timeout', + }); + expect(isRouteUnhealthy('deepseek', 'deepseek-v4-pro')).toBe(false); + }); + describe('how long a failed route stays marked', () => { let start: number; beforeEach(() => { diff --git a/src/backend/drivers/ai-chat/ChatCompletionDriver.ts b/src/backend/drivers/ai-chat/ChatCompletionDriver.ts index 00073b9b6..7877553ee 100644 --- a/src/backend/drivers/ai-chat/ChatCompletionDriver.ts +++ b/src/backend/drivers/ai-chat/ChatCompletionDriver.ts @@ -42,6 +42,7 @@ import { isUpstreamConnectionError, isUpstreamTimeoutError, sanitizeUpstreamMessage, + UpstreamDeadlineError, } from '../util/upstreamErrors.js'; import { AlibabaProvider } from './providers/alibaba/AlibabaProvider.js'; import { AzureChatProvider } from './providers/azure/AzureChatProvider.js'; @@ -846,7 +847,12 @@ export class ChatCompletionDriver extends PuterDriver { const attempt = toAttempt(failed.id, failed.provider!, err); attempts.push(attempt); if (!isFreeModel(failed)) allModelsFree = false; - if (isRouteLevelFailure(attempt)) { + // A request that outlived its deadline says nothing about the + // route's health. + if ( + isRouteLevelFailure(attempt) && + !(err instanceof UpstreamDeadlineError) + ) { markRouteUnhealthy( failed.provider!, failed.id, diff --git a/src/backend/drivers/ai-chat/providers/OpenAICompatProvider.test.ts b/src/backend/drivers/ai-chat/providers/OpenAICompatProvider.test.ts index da77dca68..00e9ca693 100644 --- a/src/backend/drivers/ai-chat/providers/OpenAICompatProvider.test.ts +++ b/src/backend/drivers/ai-chat/providers/OpenAICompatProvider.test.ts @@ -161,6 +161,33 @@ describe('OpenAICompatProvider', () => { expect('tools' in sent[1]!).toBe(false); }); + it("drops a replayed tool call's Gemini extra_content", async () => { + const { provider, sent } = makeProvider({}); + await complete(provider, { + messages: [ + { role: 'user', content: 'Weather?' }, + { + role: 'assistant', + content: [ + { + type: 'tool_use', + id: 'fc_1', + name: 'get_weather', + input: {}, + extra_content: { + google: { thought_signature: 'c2lnLTE=' }, + }, + }, + ], + }, + ] as ICompleteArguments['messages'], + }); + const messages = sent[0]!.messages as Array<{ + tool_calls?: Array>; + }>; + expect(messages[1]!.tool_calls![0]).not.toHaveProperty('extra_content'); + }); + it('asks for stream usage only when the vendor reports it in-band', async () => { const stream = { async *[Symbol.asyncIterator]() {} }; const reporting = makeProvider({}, stream); diff --git a/src/backend/drivers/ai-chat/providers/OpenAICompatProvider.ts b/src/backend/drivers/ai-chat/providers/OpenAICompatProvider.ts index 56aecb86f..7ab78853c 100644 --- a/src/backend/drivers/ai-chat/providers/OpenAICompatProvider.ts +++ b/src/backend/drivers/ai-chat/providers/OpenAICompatProvider.ts @@ -193,6 +193,7 @@ export class OpenAICompatProvider implements IChatProvider { messages = this.prepareMessages( await OpenAIUtil.process_input_messages(messages, { preserveToolCallOrder: this.preservesToolCallOrder(model), + keepExtraContent: this.keepsToolCallExtraContent(model), }), ); @@ -318,6 +319,11 @@ export class OpenAICompatProvider implements IChatProvider { return false; } + /** Whether replayed tool calls keep their vendor `extra_content`. */ + protected keepsToolCallExtraContent(_model: IChatModel): boolean { + return false; + } + /** Last pass over the messages, after `process_input_messages`. */ protected prepareMessages(messages: PuterMessage[]): PuterMessage[] { return messages; diff --git a/src/backend/drivers/ai-chat/providers/gemini/GeminiChatProvider.native.test.ts b/src/backend/drivers/ai-chat/providers/gemini/GeminiChatProvider.native.test.ts index 815bad499..5a7b5cea9 100644 --- a/src/backend/drivers/ai-chat/providers/gemini/GeminiChatProvider.native.test.ts +++ b/src/backend/drivers/ai-chat/providers/gemini/GeminiChatProvider.native.test.ts @@ -49,7 +49,12 @@ import type { ICompleteArguments, UsageDetails } from '../../types.js'; import { normalize_messages } from '../../utils/Messages.js'; import { trackedOutputTokens } from '../../utils/pricing.js'; import { AIChatStream } from '../../utils/Streaming.js'; +import { + isUpstreamTimeoutError, + UpstreamDeadlineError, +} from '../../../util/upstreamErrors.js'; import { GeminiChatProvider } from './GeminiChatProvider.js'; +import { geminiNativeClient } from './geminiNative.js'; import { GEMINI_MODELS } from './models.js'; const { generateContent, generateContentStream, genAICtor, openAICreate } = @@ -216,13 +221,10 @@ describe('gemini-4-argon catalog and routing', () => { }); }); - it('builds the native client once with the SDK timeout', () => { + it('builds the native client once', () => { makeProvider(); expect(genAICtor).toHaveBeenCalledTimes(1); - expect(genAICtor).toHaveBeenCalledWith({ - apiKey: 'test-key', - httpOptions: { timeout: 10 * 60 * 1000 }, - }); + expect(genAICtor).toHaveBeenCalledWith({ apiKey: 'test-key' }); }); it('routes the alias to the native API under the canonical id', async () => { @@ -1303,3 +1305,83 @@ describe('gemini-4-argon tool calls', () => { expect(generateContent).not.toHaveBeenCalled(); }); }); + +describe('geminiNativeClient deadline', () => { + type Request = { config?: { abortSignal?: AbortSignal } }; + // The SDK aborts its own controller without a reason. + const abortedBy = (request: Request) => + new Promise((_resolve, reject) => { + request.config!.abortSignal!.addEventListener('abort', () => + reject( + new DOMException( + 'This operation was aborted', + 'AbortError', + ), + ), + ); + }); + const clientWith = (models: Record, timeoutMs = 10) => + geminiNativeClient({ models } as never, timeoutMs); + const params = (stream: boolean) => + ({ + model: 'gemini-4-argon', + messages: [{ role: 'user', content: 'hi' }], + stream, + }) as never; + const errorOf = async (run: () => Promise) => { + try { + await run(); + } catch (err) { + return err; + } + throw new Error('expected a rejection'); + }; + + it('reports its own deadline as a timeout', async () => { + const client = clientWith({ generateContent: abortedBy }); + const err = await errorOf(() => + client.chat.completions.create(params(false), {}), + ); + expect(err).toBeInstanceOf(UpstreamDeadlineError); + expect(isUpstreamTimeoutError(err)).toBe(true); + }); + + it('reports a stream cut by its deadline as a timeout', async () => { + const client = clientWith({ + generateContentStream: async (request: Request) => + (async function* () { + yield* []; + await abortedBy(request); + })(), + }); + const err = await errorOf(async () => { + const stream = (await client.chat.completions.create( + params(true), + {}, + )) as AsyncIterable; + for await (const _chunk of stream); + }); + expect(err).toBeInstanceOf(UpstreamDeadlineError); + }); + + it('leaves a caller abort as an abort', async () => { + const abort = new AbortController(); + const client = clientWith( + { + generateContent: (request: Request) => { + const pending = abortedBy(request); + abort.abort(); + return pending; + }, + }, + 60_000, + ); + const err = await errorOf(() => + client.chat.completions.create(params(false), { + signal: abort.signal, + }), + ); + expect((err as Error).name).toBe('AbortError'); + expect(isUpstreamTimeoutError(err)).toBe(false); + }); +}); diff --git a/src/backend/drivers/ai-chat/providers/gemini/GeminiChatProvider.ts b/src/backend/drivers/ai-chat/providers/gemini/GeminiChatProvider.ts index 04b6f406e..ceba86dd1 100644 --- a/src/backend/drivers/ai-chat/providers/gemini/GeminiChatProvider.ts +++ b/src/backend/drivers/ai-chat/providers/gemini/GeminiChatProvider.ts @@ -33,7 +33,7 @@ import type { UsageDetails, } from '../../types.js'; import { toolChoiceToWire } from '../../utils/openaiParams.js'; -import { SDK_TIMEOUT_MS, sdkClientOptions } from '../../utils/sdkClient.js'; +import { sdkClientOptions } from '../../utils/sdkClient.js'; import { type ChatCompletionsClient, type ChatProviderConfig, @@ -82,12 +82,7 @@ export class GeminiChatProvider extends OpenAICompatProvider { const native = config.apiBaseUrl === undefined || config.apiBaseUrl.replace(/\/*$/, '/') === COMPAT_BASE_URL - ? geminiNativeClient( - new GoogleGenAI({ - apiKey: config.apiKey, - httpOptions: { timeout: SDK_TIMEOUT_MS }, - }), - ) + ? geminiNativeClient(new GoogleGenAI({ apiKey: config.apiKey })) : undefined; const client: ChatCompletionsClient = { chat: { @@ -180,6 +175,11 @@ export class GeminiChatProvider extends OpenAICompatProvider { return GEMINI_NATIVE_MODEL_IDS.has(model.id); } + protected override keepsToolCallExtraContent() { + // Carries the thought signature Gemini needs replayed with each call. + return true; + } + protected override meteredUsage(source: UsageSource, model: IChatModel) { if (GEMINI_NATIVE_MODEL_IDS.has(model.id)) { return nativeMeteredUsage( diff --git a/src/backend/drivers/ai-chat/providers/gemini/geminiNative.ts b/src/backend/drivers/ai-chat/providers/gemini/geminiNative.ts index 4149c5fa6..2249acafd 100644 --- a/src/backend/drivers/ai-chat/providers/gemini/geminiNative.ts +++ b/src/backend/drivers/ai-chat/providers/gemini/geminiNative.ts @@ -36,7 +36,9 @@ import { parseDataUri, unsupportedMediaTextPart, } from '../../utils/mediaParts.js'; +import { SDK_TIMEOUT_MS } from '../../utils/sdkClient.js'; import { promoteStopForToolCalls } from '../../utils/stopReason.js'; +import { UpstreamDeadlineError } from '../../../util/upstreamErrors.js'; import type { ChatCompletionsClient } from '../OpenAICompatProvider.js'; type ChatContentPart = @@ -454,23 +456,57 @@ async function* toChatChunks(stream: AsyncIterable) { } } +async function* rethrowMapped( + stream: AsyncIterable, + map: (err: unknown) => unknown, +) { + try { + yield* stream; + } catch (err) { + throw map(err); + } +} + export const geminiNativeClient = ( client: GoogleGenAI, + timeoutMs = SDK_TIMEOUT_MS, ): ChatCompletionsClient => ({ chat: { completions: { create: async (params: never, { signal }) => { // The SDK only listens for a later abort event. signal?.throwIfAborted(); + // The SDK aborts without a reason, so its own timeout would + // look like a client disconnect. This deadline is told apart. + const deadline = AbortSignal.timeout(timeoutMs); + const timedOut = (err: unknown) => + deadline.aborted && !signal?.aborted + ? new UpstreamDeadlineError( + 'Gemini request timed out', + { cause: err }, + ) + : err; const chatParams = params as ChatParams; - const request = toGenerateContentRequest(chatParams, signal); - return chatParams.stream - ? toChatChunks( - await client.models.generateContentStream(request), - ) - : toChatCompletion( - await client.models.generateContent(request), - ); + const request = toGenerateContentRequest( + chatParams, + signal ? AbortSignal.any([signal, deadline]) : deadline, + ); + try { + return chatParams.stream + ? rethrowMapped( + toChatChunks( + await client.models.generateContentStream( + request, + ), + ), + timedOut, + ) + : toChatCompletion( + await client.models.generateContent(request), + ); + } catch (err) { + throw timedOut(err); + } }, }, }, diff --git a/src/backend/drivers/ai-chat/utils/OpenAIUtil.js b/src/backend/drivers/ai-chat/utils/OpenAIUtil.js index f9cf2d76d..4eaeb3e00 100644 --- a/src/backend/drivers/ai-chat/utils/OpenAIUtil.js +++ b/src/backend/drivers/ai-chat/utils/OpenAIUtil.js @@ -152,11 +152,12 @@ export const toOpenAIChatMessages = ( * Moves a message's `tool_use` blocks into OpenAI `tool_calls`, in place. True * when this call created `tool_calls`. `canonicalId` keeps the Responses item * id a replayed call carries. `preserveOrder` appends the hoisted calls in - * content order instead of reversed. + * content order instead of reversed. `extraContent` keeps Gemini's + * `extra_content`, which other vendors may reject as an unknown field. */ const hoistToolUseBlocks = ( msg, - { canonicalId = false, preserveOrder = false } = {}, + { canonicalId = false, preserveOrder = false, extraContent = false } = {}, ) => { const content = msg.content; const calls = []; @@ -171,7 +172,7 @@ const hoistToolUseBlocks = ( name: block.name, arguments: JSON.stringify(block.input), }, - ...(block.extra_content + ...(extraContent && block.extra_content ? { extra_content: block.extra_content } : {}), }); @@ -189,13 +190,14 @@ const hoistToolUseBlocks = ( * make changes in-place. * * @param {Message[]} messages - Array of normalized messages - * @param {{ preserveToolCallOrder?: boolean }} [opts] `preserveToolCallOrder` - * keeps a turn's tool calls in the order the model made them. + * @param {{ preserveToolCallOrder?: boolean; keepExtraContent?: boolean }} [opts] + * `preserveToolCallOrder` keeps a turn's tool calls in the order the model made + * them. `keepExtraContent` keeps each tool call's `extra_content`. * @returns {Message[]} - Array of messages in OpenAI format */ export const process_input_messages = async ( messages, - { preserveToolCallOrder = false } = {}, + { preserveToolCallOrder = false, keepExtraContent = false } = {}, ) => { for (const msg of messages) { if (!msg.content) continue; @@ -213,7 +215,10 @@ export const process_input_messages = async ( } if ( - hoistToolUseBlocks(msg, { preserveOrder: preserveToolCallOrder }) + hoistToolUseBlocks(msg, { + preserveOrder: preserveToolCallOrder, + extraContent: keepExtraContent, + }) ) msg.content = null; @@ -426,9 +431,6 @@ export const process_input_messages_responses_api = async (messages) => { id: toolCall.canonical_id, name: toolCall.function.name, arguments: toolCall.function.arguments, - ...(toolCall.extra_content - ? { extra_content: toolCall.extra_content } - : {}), }); } continue; diff --git a/src/backend/drivers/ai-chat/utils/OpenAIUtil.test.ts b/src/backend/drivers/ai-chat/utils/OpenAIUtil.test.ts index 4a19fc823..5ac4665d5 100644 --- a/src/backend/drivers/ai-chat/utils/OpenAIUtil.test.ts +++ b/src/backend/drivers/ai-chat/utils/OpenAIUtil.test.ts @@ -331,8 +331,8 @@ describe('process_input_messages', () => { ]); }); - it('preserves extra_content on hoisted tool calls', async () => { - const messages: Array> = [ + describe('tool call extra_content', () => { + const turn = (): Array> => [ { role: 'assistant', content: [ @@ -346,12 +346,24 @@ describe('process_input_messages', () => { ], }, ]; + const firstCall = (out: unknown[]) => + ( + (out[0] as Record).tool_calls as Array< + Record + > + )[0]!; - const [out] = (await process_input_messages(messages)) as Array< - Record - >; - const calls = out!.tool_calls as Array>; - expect(calls[0]?.extra_content).toEqual({ hint: 'metadata' }); + it('drops it by default', async () => { + const out = await process_input_messages(turn()); + expect(firstCall(out)).not.toHaveProperty('extra_content'); + }); + + it('keeps it when asked', async () => { + const out = await process_input_messages(turn(), { + keepExtraContent: true, + }); + expect(firstCall(out).extra_content).toEqual({ hint: 'metadata' }); + }); }); describe('parallel tool call order', () => { diff --git a/src/backend/drivers/util/upstreamErrors.ts b/src/backend/drivers/util/upstreamErrors.ts index 47b9ea309..8b1c23406 100644 --- a/src/backend/drivers/util/upstreamErrors.ts +++ b/src/backend/drivers/util/upstreamErrors.ts @@ -92,6 +92,17 @@ export const isUpstreamTimeoutError = (err: unknown): boolean => { ); }; +/** + * A request ran past its overall deadline. It reads as a timeout, but it means + * the request was long, not that the route is down. + */ +export class UpstreamDeadlineError extends Error { + constructor(message: string, options?: ErrorOptions) { + super(message, options); + this.name = 'TimeoutError'; + } +} + /** The connection to an upstream provider dropped or never opened. */ export const isUpstreamConnectionError = (err: unknown): boolean => { const e = asShape(err);