From b5abf5beda99eb0b5dec28402a5b8488d9ccc86d Mon Sep 17 00:00:00 2001 From: Garry Tan Date: Wed, 3 Jun 2026 01:29:34 -0700 Subject: [PATCH] fix(ai): bound every outbound AI call so a stalled provider can't hang (#1762/#1775) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit withDefaultTimeout composes a per-touchpoint default deadline (chat 300s, embed/multimodal 60s) with any caller signal via AbortSignal.any. Applied at the SDK call layer (chat generateText, expand generateObject, OCR, per-sub-batch embed) — covers native-anthropic + retries — plus per-request multimodal fetch. embedQuery forwards abortSignal. Env: GBRAIN_AI_{CHAT,EMBED,MULTIMODAL}_TIMEOUT_MS. --- src/core/ai/gateway.ts | 65 ++++++++++++++++++++++++++++++++++++++---- src/core/embedding.ts | 2 +- 2 files changed, 60 insertions(+), 7 deletions(-) diff --git a/src/core/ai/gateway.ts b/src/core/ai/gateway.ts index f871bc5d8..38a4dfede 100644 --- a/src/core/ai/gateway.ts +++ b/src/core/ai/gateway.ts @@ -53,6 +53,42 @@ import { hasAnthropicKey } from './anthropic-key.ts'; import { AIConfigError, AITransientError, normalizeAIError } from './errors.ts'; import { runGuardrails, hasGuardrails, type GuardrailHook } from '../guardrails.ts'; +// ---- Gateway-wide AI-HTTP timeout (v0.42.11.0, #1762/#1775) ---- +// +// Plain `fetch` (Bun/Node) has NO default request timeout, so a stalled provider +// socket makes an `await` never settle — which hangs `gbrain capture`/`search` +// and, on PGLite, pins the single-writer lock. The AI SDK's `maxRetries` only +// fires on a SETTLED error; a half-open socket never settles. So we bound at the +// SDK CALL layer: default an `abortSignal` into every generateText / generateObject +// / embed call. This (1) covers EVERY provider — including `native-anthropic` +// (the default chat model + the facts:absorb Haiku), which the AI SDK forwards +// the signal to as `fetch(url, {signal})`; and (2) bounds the WHOLE call incl. +// internal retries, not one attempt. Direct-`fetch` paths (multimodal) get the +// signal explicitly. Rerank is already bounded by its recipe `default_timeout_ms`. +function resolveAiTimeoutMs(envVar: string, fallback: number): number { + const raw = process.env[envVar]; + if (raw === undefined) return fallback; + const n = Number(raw); + return Number.isFinite(n) && n > 0 ? n : fallback; +} +/** chat / expansion / OCR — generous; only catches true hangs (non-streaming generateText). */ +const AI_CHAT_TIMEOUT_MS = resolveAiTimeoutMs('GBRAIN_AI_CHAT_TIMEOUT_MS', 300_000); +/** embed sub-batch (per SDK call, NOT per whole import). */ +const AI_EMBED_TIMEOUT_MS = resolveAiTimeoutMs('GBRAIN_AI_EMBED_TIMEOUT_MS', 60_000); +/** multimodal per request. */ +const AI_MULTIMODAL_TIMEOUT_MS = resolveAiTimeoutMs('GBRAIN_AI_MULTIMODAL_TIMEOUT_MS', 60_000); + +/** + * Compose a caller signal with a default wall-clock timeout. When the caller + * supplies its own (Fix 3's 6s query deadline, the facts queue's shutdown abort, + * a budget signal), `AbortSignal.any` makes whichever fires FIRST win — so a + * shorter caller deadline always takes precedence over the default backstop. + */ +function withDefaultTimeout(caller: AbortSignal | undefined, timeoutMs: number): AbortSignal { + const timeout = AbortSignal.timeout(timeoutMs); + return caller ? AbortSignal.any([caller, timeout]) : timeout; +} + const MAX_CHARS = 8000; // v0.36.0.0 (D3 + D4): ZeroEntropy zembed-1 at 1280d via Matryoshka is the // new default for embedding. Real-corpus benchmark across 20 queries: @@ -1440,10 +1476,11 @@ async function embedSubBatch( model, values: texts, providerOptions: providerOpts, - // v0.33.4: caller-supplied abortSignal + maxRetries passthrough. - // Undefined fields are ignored by the AI SDK so the call shape stays - // identical for production callers that don't opt in. - ...(opts?.abortSignal !== undefined && { abortSignal: opts.abortSignal }), + // v0.42.11.0 — default a per-SUB-BATCH embed timeout (codex #3: bounding + // once at embed() top would cap a whole multi-batch import; this is the + // per-SDK-call scope). Composes with a caller signal (Fix 3's 6s query + // deadline) — shorter wins. + abortSignal: withDefaultTimeout(opts?.abortSignal, AI_EMBED_TIMEOUT_MS), ...(opts?.maxRetries !== undefined && { maxRetries: opts.maxRetries }), }); @@ -1503,12 +1540,16 @@ export async function embedOne(text: string): Promise { */ export async function embedQuery( text: string, - opts?: { embeddingModel?: string; dimensions?: number }, + opts?: { embeddingModel?: string; dimensions?: number; abortSignal?: AbortSignal }, ): Promise { const [v] = await embed([text], { inputType: 'query', embeddingModel: opts?.embeddingModel, dimensions: opts?.dimensions, + // v0.42.11.0 (Fix 3) — forward a caller deadline so the query-time embed + // can be bounded BELOW the CLI force-exit; composes with the gateway embed + // default via withDefaultTimeout (shorter wins). + abortSignal: opts?.abortSignal, }); return v; } @@ -1642,6 +1683,9 @@ export async function embedMultimodal( Authorization: `Bearer ${apiKey}`, }, body: JSON.stringify(body), + // v0.42.11.0 (codex #4) — per-request multimodal timeout (direct fetch + // bypasses the SDK abortSignal). + signal: AbortSignal.timeout(AI_MULTIMODAL_TIMEOUT_MS), }); } catch (err) { throw normalizeAIError(err, `embedMultimodal(${recipe.id}:${parsed.modelId})`); @@ -1784,6 +1828,8 @@ async function embedMultimodalOpenAICompat( [authResult.headerName]: authResult.token, }, body: JSON.stringify(body), + // v0.42.11.0 (codex #4) — per-request multimodal timeout (direct fetch). + signal: AbortSignal.timeout(AI_MULTIMODAL_TIMEOUT_MS), }); } catch (err) { throw normalizeAIError(err, `embedMultimodal(${recipe.id}:${modelId})`); @@ -2034,6 +2080,9 @@ export async function expand(query: string): Promise { const result = await generateObject({ model, schema: ExpansionSchema, + // v0.42.11.0 (codex P0) — expansion had NO abortSignal; same stalled-socket + // class as chat. Default the chat timeout. + abortSignal: withDefaultTimeout(undefined, AI_CHAT_TIMEOUT_MS), prompt: [ 'Rewrite the search query below into 3-4 different, related queries that would help find relevant documents.', 'Return ONLY the JSON object. Do NOT include the original query in the result.', @@ -2084,6 +2133,8 @@ export async function generateOcrText(imageBytes: Buffer, mime: string): Promise const base64 = imageBytes.toString('base64'); const result = await generateText({ model, + // v0.42.11.0 (codex) — OCR is a 5th unbounded generateText entry point. + abortSignal: withDefaultTimeout(undefined, AI_CHAT_TIMEOUT_MS), messages: [ { role: 'system', @@ -2561,7 +2612,9 @@ export async function chat(opts: ChatOpts): Promise { messages: opts.messages as any, tools: opts.tools && opts.tools.length > 0 ? tools : undefined, maxOutputTokens: opts.maxTokens ?? 4096, - abortSignal: opts.abortSignal, + // v0.42.11.0 — default a chat timeout (composes with the caller's signal, + // shorter wins). Covers native-anthropic (the default provider + facts Haiku). + abortSignal: withDefaultTimeout(opts.abortSignal, AI_CHAT_TIMEOUT_MS), providerOptions: Object.keys(providerOptions).length > 0 ? providerOptions : undefined, }); diff --git a/src/core/embedding.ts b/src/core/embedding.ts index 7fe3c9148..20515d313 100644 --- a/src/core/embedding.ts +++ b/src/core/embedding.ts @@ -53,7 +53,7 @@ export async function embed(text: string): Promise { */ export async function embedQuery( text: string, - opts?: { embeddingModel?: string; dimensions?: number }, + opts?: { embeddingModel?: string; dimensions?: number; abortSignal?: AbortSignal }, ): Promise { return gatewayEmbedQuery(text, opts); }