admin/gbrain
1
0
Fork 0
mirror of https://github.com/garrytan/gbrain.git synced 2026-07-28 14:59:47 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix(ai): bound every outbound AI call so a stalled provider can't hang (#1762/#1775)

Browse Source 
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.
This commit is contained in:
Garry Tan
2026-06-03 01:29:34 -07:00
parent a41dfa5ff6
commit b5abf5beda
2 changed files with 60 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/core/ai/gateway.ts
+59 -6
View File
@@ -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<Float32Array> {
*/
export async function embedQuery(
text: string,
opts?: { embeddingModel?: string; dimensions?: number },
opts?: { embeddingModel?: string; dimensions?: number; abortSignal?: AbortSignal },
): Promise<Float32Array> {
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<string[]> {
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<ChatResult> {
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,
});
src/core/embedding.ts
+1 -1
View File
@@ -53,7 +53,7 @@ export async function embed(text: string): Promise<Float32Array> {
*/
export async function embedQuery(
text: string,
opts?: { embeddingModel?: string; dimensions?: number },
opts?: { embeddingModel?: string; dimensions?: number; abortSignal?: AbortSignal },
): Promise<Float32Array> {
return gatewayEmbedQuery(text, opts);
}