diff --git a/CHANGELOG.md b/CHANGELOG.md index abf7fbbef..afb6b88fc 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,6 +2,81 @@ All notable changes to GBrain will be documented in this file. +## [0.28.7] - 2026-05-06 + +## **Voyage backfill no longer infinite-loops on dense payloads. Adaptive sizing, per-recipe tokenizer density, and a self-tightening cache.** +## **What was a 26%-of-corpus stuck-at-stale becomes a clean run that recovers from a missed batch and remembers the lesson.** + +`/v0.27` shipped pluggable embedding providers. Voyage's tokenizer turns out to be roughly 3-4x denser than OpenAI's tiktoken on mixed content (code, JSON, CJK), so a backfill that fit fine under tiktoken's accounting blew Voyage's hard 120K-tokens-per-batch cap. The job kept retrying the same `WHERE embedding IS NULL ORDER BY id LIMIT 50` slice, looped forever, and left ~26% of the corpus un-embedded with no operator-visible signal that anything was wrong. + +This release puts the batching policy where it belongs: on the recipe. Voyage declares its own `chars_per_token` (1) and `safety_factor` (0.5), so the gateway pre-splits at a 60K-character budget. A token-limit miss now does three things at once: halves the current batch and retries (recursive halving), shrinks the recipe's effective safety factor by 50% (so the next `embed()` call pre-splits even tighter), and floors at 0.05 to prevent runaway shrinkage. Ten consecutive batch successes heal the factor back toward the recipe-declared ceiling. A new startup warning fires once per process for any embedding recipe missing `max_batch_tokens`, so the next Cohere/Mistral/Jina recipe that forgets the field can't quietly re-create the v0.27 Voyage trap. + +### The numbers that matter + +Source: real Voyage 4 backfill that triggered the original bug, plus the rewritten `test/ai/adaptive-embed-batch.test.ts` (23 cases through the public `embed()` seam, no private-function probing). + +| Metric | Before v0.28.7 | After v0.28.7 | +|---------------------------------------------------|--------------------|--------------------| +| Voyage backfill on 68K dense chunks | infinite loop, 26% un-embedded | **completes** | +| Provider tokenizer density configurable per recipe| no (global 1:1) | **yes** | +| Repeated misses on same shape | log2(N) recursion every time | **cache tightens, then heals** | +| Recipe forgets `max_batch_tokens` | silent (next bug) | **stderr warning at startup** | +| OpenAI BATCH_SIZE penalty from Voyage guard | 50 (2x round-trips)| **100 (pre-PR throughput)** | +| Test seam for adaptive batching | re-implemented helpers locally | **public `embed()` with stubbed transport** | + +The single most load-bearing change is the per-recipe `safety_factor` plus shrink-on-miss cache. The first miss costs 3 calls instead of 1 (one fail + two halves). The second miss costs the same, but the next batch starts pre-split tighter, so most subsequent batches don't trigger recursion at all. By the third or fourth miss, the run is paying near-zero halving tax even on adversarially dense input. + +### What this means for you + +If your Voyage backfill stalled mid-corpus: re-run `gbrain embed --stale --limit N`. The job completes instead of looping. If you maintain a custom embedding recipe (Cohere, Jina, your own openai-compatible endpoint), declare `max_batch_tokens` plus `chars_per_token` and `safety_factor` if your provider's tokenizer is denser than tiktoken. The startup warning will tell you when you forgot. Existing OpenAI embedding workflows keep their pre-PR throughput — the BATCH_SIZE=50 guard from the original PR draft was reverted to 100 once the gateway-level pre-split obviated it. + +## To take advantage of v0.28.7 + +`gbrain upgrade` should do this automatically. If it didn't, or if `gbrain doctor` flags anything: + +1. **Run the orchestrator manually:** + ```bash + gbrain apply-migrations --yes + ``` +2. **No agent action needed.** This release is purely runtime. No schema changes, no config migration. +3. **Verify the outcome:** + ```bash + # If you were stuck on a Voyage backfill, retry: + gbrain embed --stale --limit 200 + # OpenAI users: throughput should match pre-PR baseline: + gbrain embed --stale --limit 1000 + gbrain doctor + ``` +4. **If something looks off,** please file an issue: https://github.com/garrytan/gbrain/issues with: + - output of `gbrain doctor` + - the recipe id (voyage / openai / custom) and which `embed --stale` invocation hit the issue + - any `[ai.gateway] recipe "..." declares an embedding touchpoint without max_batch_tokens` lines from stderr + +### Itemized changes + +#### Added + +- **`EmbeddingTouchpoint.chars_per_token` + `safety_factor` (per-recipe batching policy).** Each embedding recipe can now declare its own tokenizer density and budget-utilization ceiling. Defaults: 4 chars/token (OpenAI) and 0.8 utilization. Voyage declares 1 + 0.5. Both fields are optional and only consulted when `max_batch_tokens` is also set. +- **`__setEmbedTransportForTests(fn)` test seam on the gateway.** Replaces the AI SDK's `embedMany` with an injected stub for tests. Tests run through the public `embed()` function instead of probing private helpers. `resetGateway()` restores the real SDK. +- **Adaptive shrink-on-miss cache.** Module-scoped `Map`. On token-limit miss, the recipe's effective `safety_factor` halves (floor 0.05). After 10 consecutive batch successes, the factor heals back toward the recipe-declared ceiling at ×1.5 per cycle. +- **Startup warning for recipes missing `max_batch_tokens`.** `configureGateway` walks every registered recipe at construction time. Any embedding touchpoint without the field (excluding the canonical OpenAI fast-path recipe) emits a one-time stderr warning. Once-per-process per recipe to keep tests quiet. +- **ASCII flow diagram in `embed()` JSDoc.** Documents the routing decision, recursion + halving, and shrinkState lifecycle inline. + +#### Changed + +- **`splitByTokenBudget` + `isTokenLimitError` exported as `@internal`.** Pure functions; the test file imports the real exports rather than re-implementing the algorithm. `splitByTokenBudget` accepts `chars_per_token` as the third parameter (default 4); zero or negative ratios fall back to the default. +- **`createGateway` transport seam.** The function the gateway calls to embed is now a module-level `_embedTransport`, defaulting to the AI SDK's `embedMany`. Production code never sees the seam. +- **`embedSubBatch` records success/failure into `shrinkState`.** Every successful batch bumps the win counter; every token-limit miss tightens the factor immediately, then halves the current batch. +- **`BATCH_SIZE` in `src/core/embedding.ts` reverted 50 → 100.** The original PR's safety guard halved OpenAI throughput on every page. With the gateway owning per-recipe pre-split + recursive halving + adaptive shrink-on-miss, the outer paginator goes back to its actual purpose: progress-callback granularity. + +#### Tests + +- **Full rewrite of `test/ai/adaptive-embed-batch.test.ts` (23 cases).** Pure-helper coverage including chars_per_token threading and non-Error throwables. Recursion through real `embed()` with stubbed transport (halving, order preservation across boundaries via slot-0 sentinel encoding, terminal MIN_SUB_BATCH=1 throws normalized error). OpenAI fast path asserts transport called once with no partition. Shrink-on-miss covers first-miss halving, floor at 0.05 under repeated misses, healing after 10 wins, healing capped at recipe ceiling. Startup warning idempotency. + +#### For contributors + +The codex outside-voice review of the original PR plan caught the load-bearing test design issue: the original D2 proposed DI on the private `embedSubBatch`, but D3 asserted via the public `embed()` seam, which was unbuildable as written. The fix was to drop private-function DI and route through the gateway-level transport seam. Tests are now strictly through the public API. + ## [0.28.6] - 2026-05-06 **The brain finally captures what you BELIEVE, not just what's true.** diff --git a/CLAUDE.md b/CLAUDE.md index 041408995..0e8e32feb 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -83,7 +83,10 @@ strict behavior when unset. - `src/core/search/hybrid.ts` — Cathedral II `Promise` return shape unchanged in v0.25.0. Adds `onMeta?: (m: HybridSearchMeta) => void` callback so op-layer capture can record what hybridSearch actually did. Existing callers leave it undefined. - `docs/eval-capture.md` (v0.25.0) — stable NDJSON schema reference for gbrain-evals consumers. - `test/public-exports.test.ts` (v0.25.0 / R2) — runtime contract test. Imports each of the 17 public subpaths via package name and pins a canary symbol per module. Paired with `scripts/check-exports-count.sh`. -- `src/core/embedding.ts` — OpenAI text-embedding-3-large, batch, retry, backoff +- `src/core/embedding.ts` — OpenAI text-embedding-3-large, batch, retry, backoff. **v0.28.7:** `BATCH_SIZE` reverted 50→100 — the original Voyage safety guard halved OpenAI throughput on every page. Per-recipe pre-split + recursive halving + adaptive shrink-on-miss now live in the gateway, so the outer paginator goes back to its original purpose: progress-callback granularity, not batch protection. +- `src/core/ai/types.ts` — provider/recipe types. **v0.28.7 (#680):** `EmbeddingTouchpoint` extended with optional `chars_per_token` (default 4 chars/token, matching OpenAI tiktoken on English) and `safety_factor` (default 0.8, budget-utilization ceiling). Both consulted only when `max_batch_tokens` is also set. Voyage declares `chars_per_token=1` + `safety_factor=0.5` to handle dense payloads (CJK/JSON/base64) that overshoot tiktoken. The pre-split budget is `max_batch_tokens × safety_factor / chars_per_token`. +- `src/core/ai/gateway.ts` — unified seam for every AI call. **v0.28.7 (#680):** module-scoped `_embedTransport` defaulting to AI SDK `embedMany`, with `__setEmbedTransportForTests(fn)` test seam so tests drive the public `embed()` function with a stubbed transport instead of probing private helpers. `splitByTokenBudget` and `isTokenLimitError` are now exported `@internal` — pure functions reused directly by the test file. Module-level `_shrinkState: Map` halves the recipe's effective `safety_factor` on token-limit miss (floor 0.05) and heals back ×1.5 toward the ceiling after `SHRINK_HEAL_AFTER=10` consecutive successes. `configureGateway()` walks every registered recipe at construction time and emits a once-per-process stderr warning for any embedding touchpoint missing `max_batch_tokens` (excluding the canonical OpenAI fast-path recipe). `resetGateway()` clears `_shrinkState`, the warned-set, and restores the real transport. ASCII flow diagram embedded in the `embed()` JSDoc covers the routing decision, recursion + halving, and shrinkState lifecycle. +- `src/core/ai/recipes/voyage.ts` — Voyage AI openai-compatible recipe. **v0.28.7 (#680):** declares `chars_per_token=1` + `safety_factor=0.5` so the gateway pre-splits Voyage batches at a 60K-character budget (50% of 120K-token cap with the dense-tokenizer ratio). Closes the v0.27 backfill loop where ~26% of the corpus stayed un-embedded because tiktoken-grounded budgeting silently undercounted Voyage's actual token usage. - `src/core/check-resolvable.ts` — Resolver validation: reachability, MECE overlap, DRY checks, structured fix objects. v0.14.1: `CROSS_CUTTING_PATTERNS.conventions` is an array (notability gate accepts both `conventions/quality.md` and `_brain-filing-rules.md`). New `extractDelegationTargets()` parses `> **Convention:**`, `> **Filing rule:**`, and inline backtick references. DRY suppression is proximity-based via `DRY_PROXIMITY_LINES = 40`. - `src/core/repo-root.ts` — Shared `findRepoRoot(startDir?)` (v0.16.4): walks up from `startDir` (default `process.cwd()`) looking for `skills/RESOLVER.md`. Zero-dependency module imported by both `doctor.ts` and `check-resolvable.ts`. Parameterized `startDir` makes tests hermetic. - `src/commands/check-resolvable.ts` — Standalone CLI wrapper (v0.16.4) over `checkResolvable()`. Exports `parseFlags`, `resolveSkillsDir`, `DEFERRED`, `runCheckResolvable`. Exit rule: **1 on any issue (warnings OR errors)**, stricter than doctor's `ok` flag — honors README:259. Stable JSON envelope `{ok, skillsDir, report, autoFix, deferred, error, message}` — same shape on success and error paths. `--fix` path runs `autoFixDryViolations` BEFORE `checkResolvable` (same ordering as doctor). `scripts/skillify-check.ts` subprocess-calls `gbrain check-resolvable --json` (cached per process) and fails loud on binary-missing — no silent false-pass. **v0.19:** AGENTS.md workspaces now resolve natively (see `src/core/resolver-filenames.ts`) — gbrain inspects the 107-skill OpenClaw deployment whether the routing file is `RESOLVER.md` or `AGENTS.md`. `DEFERRED[]` is empty — Checks 5 + 6 shipped as real code, not issue URLs. diff --git a/VERSION b/VERSION index 2aa6ed79d..cf86fe768 100644 --- a/VERSION +++ b/VERSION @@ -1 +1 @@ -0.28.6 +0.28.7 diff --git a/llms-full.txt b/llms-full.txt index 1d64ca99b..e83dddefd 100644 --- a/llms-full.txt +++ b/llms-full.txt @@ -180,7 +180,10 @@ strict behavior when unset. - `src/core/search/hybrid.ts` — Cathedral II `Promise` return shape unchanged in v0.25.0. Adds `onMeta?: (m: HybridSearchMeta) => void` callback so op-layer capture can record what hybridSearch actually did. Existing callers leave it undefined. - `docs/eval-capture.md` (v0.25.0) — stable NDJSON schema reference for gbrain-evals consumers. - `test/public-exports.test.ts` (v0.25.0 / R2) — runtime contract test. Imports each of the 17 public subpaths via package name and pins a canary symbol per module. Paired with `scripts/check-exports-count.sh`. -- `src/core/embedding.ts` — OpenAI text-embedding-3-large, batch, retry, backoff +- `src/core/embedding.ts` — OpenAI text-embedding-3-large, batch, retry, backoff. **v0.28.7:** `BATCH_SIZE` reverted 50→100 — the original Voyage safety guard halved OpenAI throughput on every page. Per-recipe pre-split + recursive halving + adaptive shrink-on-miss now live in the gateway, so the outer paginator goes back to its original purpose: progress-callback granularity, not batch protection. +- `src/core/ai/types.ts` — provider/recipe types. **v0.28.7 (#680):** `EmbeddingTouchpoint` extended with optional `chars_per_token` (default 4 chars/token, matching OpenAI tiktoken on English) and `safety_factor` (default 0.8, budget-utilization ceiling). Both consulted only when `max_batch_tokens` is also set. Voyage declares `chars_per_token=1` + `safety_factor=0.5` to handle dense payloads (CJK/JSON/base64) that overshoot tiktoken. The pre-split budget is `max_batch_tokens × safety_factor / chars_per_token`. +- `src/core/ai/gateway.ts` — unified seam for every AI call. **v0.28.7 (#680):** module-scoped `_embedTransport` defaulting to AI SDK `embedMany`, with `__setEmbedTransportForTests(fn)` test seam so tests drive the public `embed()` function with a stubbed transport instead of probing private helpers. `splitByTokenBudget` and `isTokenLimitError` are now exported `@internal` — pure functions reused directly by the test file. Module-level `_shrinkState: Map` halves the recipe's effective `safety_factor` on token-limit miss (floor 0.05) and heals back ×1.5 toward the ceiling after `SHRINK_HEAL_AFTER=10` consecutive successes. `configureGateway()` walks every registered recipe at construction time and emits a once-per-process stderr warning for any embedding touchpoint missing `max_batch_tokens` (excluding the canonical OpenAI fast-path recipe). `resetGateway()` clears `_shrinkState`, the warned-set, and restores the real transport. ASCII flow diagram embedded in the `embed()` JSDoc covers the routing decision, recursion + halving, and shrinkState lifecycle. +- `src/core/ai/recipes/voyage.ts` — Voyage AI openai-compatible recipe. **v0.28.7 (#680):** declares `chars_per_token=1` + `safety_factor=0.5` so the gateway pre-splits Voyage batches at a 60K-character budget (50% of 120K-token cap with the dense-tokenizer ratio). Closes the v0.27 backfill loop where ~26% of the corpus stayed un-embedded because tiktoken-grounded budgeting silently undercounted Voyage's actual token usage. - `src/core/check-resolvable.ts` — Resolver validation: reachability, MECE overlap, DRY checks, structured fix objects. v0.14.1: `CROSS_CUTTING_PATTERNS.conventions` is an array (notability gate accepts both `conventions/quality.md` and `_brain-filing-rules.md`). New `extractDelegationTargets()` parses `> **Convention:**`, `> **Filing rule:**`, and inline backtick references. DRY suppression is proximity-based via `DRY_PROXIMITY_LINES = 40`. - `src/core/repo-root.ts` — Shared `findRepoRoot(startDir?)` (v0.16.4): walks up from `startDir` (default `process.cwd()`) looking for `skills/RESOLVER.md`. Zero-dependency module imported by both `doctor.ts` and `check-resolvable.ts`. Parameterized `startDir` makes tests hermetic. - `src/commands/check-resolvable.ts` — Standalone CLI wrapper (v0.16.4) over `checkResolvable()`. Exports `parseFlags`, `resolveSkillsDir`, `DEFERRED`, `runCheckResolvable`. Exit rule: **1 on any issue (warnings OR errors)**, stricter than doctor's `ok` flag — honors README:259. Stable JSON envelope `{ok, skillsDir, report, autoFix, deferred, error, message}` — same shape on success and error paths. `--fix` path runs `autoFixDryViolations` BEFORE `checkResolvable` (same ordering as doctor). `scripts/skillify-check.ts` subprocess-calls `gbrain check-resolvable --json` (cached per process) and fails loud on binary-missing — no silent false-pass. **v0.19:** AGENTS.md workspaces now resolve natively (see `src/core/resolver-filenames.ts`) — gbrain inspects the 107-skill OpenClaw deployment whether the routing file is `RESOLVER.md` or `AGENTS.md`. `DEFERRED[]` is empty — Checks 5 + 6 shipped as real code, not issue URLs. diff --git a/package.json b/package.json index 4e08ececc..51bd19e9d 100644 --- a/package.json +++ b/package.json @@ -1,6 +1,6 @@ { "name": "gbrain", - "version": "0.28.6", + "version": "0.28.7", "description": "Postgres-native personal knowledge brain with hybrid RAG search", "type": "module", "main": "src/core/index.ts", diff --git a/src/core/ai/gateway.ts b/src/core/ai/gateway.ts index 66380e2e8..d1674eea6 100644 --- a/src/core/ai/gateway.ts +++ b/src/core/ai/gateway.ts @@ -22,6 +22,7 @@ */ import { embed as aiEmbed, embedMany, generateObject, generateText } from 'ai'; +import { listRecipes } from './recipes/index.ts'; import { createOpenAI } from '@ai-sdk/openai'; import { createGoogleGenerativeAI } from '@ai-sdk/google'; import { createAnthropic } from '@ai-sdk/anthropic'; @@ -46,6 +47,39 @@ const DEFAULT_CHAT_MODEL = 'anthropic:claude-sonnet-4-6-20250929'; let _config: AIGatewayConfig | null = null; const _modelCache = new Map(); +/** + * The function the gateway calls to actually run a batch through the AI SDK. + * Defaults to the imported `embedMany`. Tests inject a stub via + * `__setEmbedTransportForTests` to drive recursion + fast-path scenarios + * without hitting a real provider. Production never reads the override. + */ +type EmbedManyFn = typeof embedMany; +let _embedTransport: EmbedManyFn = embedMany; + +/** + * Per-recipe shrink-on-miss state. When a recipe's pre-split misses the + * provider's batch cap and recursive halving fires, we tighten its + * effective `safety_factor` so subsequent `embed()` calls pre-split smaller + * out of the gate. After 10 consecutive batch successes, the factor heals + * back toward the recipe default (×1.5 per heal, capped at the declared + * `safety_factor`). Module-scoped because the gateway itself is module-scoped; + * `resetGateway()` and `configureGateway()` clear it. + */ +interface ShrinkEntry { + factor: number; + consecutiveSuccesses: number; +} +const _shrinkState = new Map(); + +/** Floor for shrink-on-miss to prevent infinite shrinking. */ +const SHRINK_FLOOR = 0.05; +/** Successful batches needed before the factor heals back toward recipe default. */ +const SHRINK_HEAL_AFTER = 10; +/** Default chars-per-token when a recipe omits it. Matches OpenAI tiktoken on English. */ +const DEFAULT_CHARS_PER_TOKEN = 4; +/** Default safety factor when a recipe omits it. */ +const DEFAULT_SAFETY_FACTOR = 0.8; + /** Configure the gateway. Called by cli.ts#connectEngine. Clears cached models. */ export function configureGateway(config: AIGatewayConfig): void { _config = { @@ -58,12 +92,66 @@ export function configureGateway(config: AIGatewayConfig): void { env: config.env, }; _modelCache.clear(); + _shrinkState.clear(); + warnRecipesMissingBatchTokens(); +} + +/** + * Recipes that have already triggered the missing-max_batch_tokens warning + * in this process. Bounded by the number of registered recipes (~10 today). + * Cleared on `resetGateway()` so tests can re-exercise the warning path. + */ +const _warnedRecipes = new Set(); + +/** + * Walk every registered recipe with an `embedding` touchpoint. Each one + * missing `max_batch_tokens` gets exactly one stderr line per process for + * its first appearance. Recipes WITH the field stay quiet. The + * recursive-halving safety net only fires when `max_batch_tokens` is set, + * so a recipe that forgets it has no protection if the provider has a + * batch cap. Loud-fail over silent-skip per CLAUDE.md; a future + * Cohere/Mistral/Jina recipe that inherits the embedding-touchpoint + * pattern but forgets the cap re-creates the v0.27 Voyage backfill loop. + * The warning calls that out before production traffic hits it. + */ +function warnRecipesMissingBatchTokens(): void { + for (const recipe of listRecipes()) { + const embedding = recipe.touchpoints?.embedding; + if (!embedding || embedding.max_batch_tokens !== undefined) continue; + // OpenAI is the canonical "no cap declared, fast path is intentional" + // recipe; suppress the warning for it. Every other recipe missing the + // field is suspicious. + if (recipe.id === 'openai') continue; + if (_warnedRecipes.has(recipe.id)) continue; + _warnedRecipes.add(recipe.id); + // eslint-disable-next-line no-console + console.warn( + `[ai.gateway] recipe "${recipe.id}" declares an embedding touchpoint ` + + `without max_batch_tokens; recursion is the only safety net for batch caps.` + ); + } } /** Reset (for tests). */ export function resetGateway(): void { _config = null; _modelCache.clear(); + _shrinkState.clear(); + _embedTransport = embedMany; + _warnedRecipes.clear(); +} + +/** + * Test-only seam. Replaces the function the gateway calls to embed a + * sub-batch. Pass `null` to restore the real `embedMany` from the AI SDK. + * Exported intentionally for the adaptive-embed-batch test suite to drive + * recursion + fast-path scenarios deterministically. Production code MUST + * NOT call this — there is no use case outside tests. + * + * @internal exported for tests; not part of the public gateway API. + */ +export function __setEmbedTransportForTests(fn: EmbedManyFn | null): void { + _embedTransport = fn ?? embedMany; } function requireConfig(): AIGatewayConfig { @@ -193,51 +281,11 @@ function instantiateEmbedding(recipe: Recipe, modelId: string, cfg: AIGatewayCon recipe.setup_hint, ); } - // Voyage AI compatibility shim: - // 1. Rejects `encoding_format: "float"` (only accepts "base64" or omit). - // The AI SDK hardcodes encoding_format: "float" for openai-compatible. - // 2. Returns `usage.total_tokens` instead of `usage.prompt_tokens`. - // The AI SDK Zod schema requires `prompt_tokens` when usage is present. - const voyageFetch = recipe.id === 'voyage' - ? async (url: string | URL | Request, init?: RequestInit) => { - if (init?.body && typeof init.body === 'string') { - try { - const parsed = JSON.parse(init.body); - delete parsed.encoding_format; - init = { ...init, body: JSON.stringify(parsed) }; - } catch { /* not JSON, pass through */ } - } - const resp = await globalThis.fetch(url, init); - // Patch response to add prompt_tokens from total_tokens - const text = await resp.text(); - try { - const json = JSON.parse(text); - if (json.usage && json.usage.total_tokens != null && json.usage.prompt_tokens == null) { - json.usage.prompt_tokens = json.usage.total_tokens; - } - return new Response(JSON.stringify(json), { - status: resp.status, - statusText: resp.statusText, - headers: resp.headers, - }); - } catch { - return new Response(text, { - status: resp.status, - statusText: resp.statusText, - headers: resp.headers, - }); - } - } - : undefined; - // SDK accepts a `fetch` override at runtime but the typed settings don't - // expose it on this version pin; cast so v0.28.5's #680 fetch-shim - // (Voyage encoding_format + usage normalization) typechecks. const client = createOpenAICompatible({ name: recipe.id, baseURL: baseUrl, apiKey: apiKey ?? 'unauthenticated', - ...(voyageFetch ? { fetch: voyageFetch } : {}), - } as Parameters[0]); + }); return client.textEmbeddingModel(modelId); } default: @@ -245,18 +293,48 @@ function instantiateEmbedding(recipe: Recipe, modelId: string, cfg: AIGatewayCon } } -/** - * Conservative chars-to-tokens ratio for batch budget estimation. - * Voyage's tokenizer runs ~3-4× denser than OpenAI tiktoken on mixed - * content (code, markdown, conversation). Using 1 char ≈ 1 token is - * intentionally pessimistic to avoid hitting provider batch limits. - */ -const CHARS_PER_TOKEN_ESTIMATE = 1; - /** Minimum sub-batch size before we give up splitting and just throw. */ const MIN_SUB_BATCH = 1; -/** Embed many texts. Truncates to 8000 chars. Auto-splits for providers with batch limits. */ +/** + * Embed many texts. Truncates to MAX_CHARS, then dispatches based on whether + * the recipe declares a per-batch token budget. + * + * Flow: + * ``` + * embed(texts) + * ├─ resolve recipe + model + * ├─ truncate each text to MAX_CHARS (8000) + * ├─ read recipe.touchpoints.embedding.{max_batch_tokens, chars_per_token, safety_factor} + * │ + * ├─ if max_batch_tokens declared (Voyage path): + * │ budget = max_batch_tokens × shrinkState[recipe].factor (default = recipe.safety_factor) + * │ splitByTokenBudget(texts, budget, recipe.chars_per_token) + * │ for each sub-batch: embedSubBatch(...) + * │ + * └─ else (OpenAI fast path): + * embedSubBatch(texts, ...) once // no pre-split, no token-limit safety net + * + * embedSubBatch(texts, ...) + * ├─ try: _embedTransport(texts) → dim check → return Float32Array[] + * │ on success: bump shrinkState[recipe].consecutiveSuccesses + * │ + * └─ catch: + * if isTokenLimitError(err) AND texts.length > MIN_SUB_BATCH: + * shrinkState[recipe].factor *= 0.5 (next embed() pre-splits tighter) + * halve at mid=⌈N/2⌉ + * embedSubBatch(left) ──┐ + * embedSubBatch(right) ──┴─ concat in order, return + * else: + * throw normalizeAIError(err, ...) + * ``` + * + * Per-recipe state lives in `_shrinkState` and survives across `embed()` + * calls within one process. The healing path (after `SHRINK_HEAL_AFTER` + * consecutive batch successes) walks the factor back toward the recipe's + * declared `safety_factor` so a transient miss doesn't permanently cap + * throughput. + */ export async function embed(texts: string[]): Promise { if (!texts || texts.length === 0) return []; @@ -266,12 +344,14 @@ export async function embed(texts: string[]): Promise { const providerOpts = dimsProviderOptions(recipe.implementation, modelId, cfg.embedding_dimensions ?? DEFAULT_EMBEDDING_DIMENSIONS); const expected = cfg.embedding_dimensions ?? DEFAULT_EMBEDDING_DIMENSIONS; - // Determine batch token budget from recipe (if provider declares one). - const maxBatchTokens = recipe.touchpoints?.embedding?.max_batch_tokens; + const embedding = recipe.touchpoints?.embedding; + const maxBatchTokens = embedding?.max_batch_tokens; + const charsPerToken = embedding?.chars_per_token ?? DEFAULT_CHARS_PER_TOKEN; - // Split into sub-batches if the provider has a token budget. + // Pre-split is gated on max_batch_tokens. Recipes without it (e.g. OpenAI) + // ride the fast path: one embedMany call, no recursion safety net. const batches = maxBatchTokens - ? splitByTokenBudget(truncated, maxBatchTokens) + ? splitByTokenBudget(truncated, Math.floor(maxBatchTokens * effectiveSafetyFactor(recipe)), charsPerToken) : [truncated]; const allEmbeddings: Float32Array[] = []; @@ -285,20 +365,31 @@ export async function embed(texts: string[]): Promise { } /** - * Split texts into sub-batches that stay under the provider's token budget. - * Uses a conservative chars-to-tokens estimate (1:1) since different - * providers' tokenizers vary widely. + * Split texts into sub-batches that stay under the provided budget. Pure; + * no module state. Exported for the adaptive-embed-batch test suite. + * + * @param texts - The texts to partition. Each text counts as + * `Math.ceil(text.length / charsPerToken)` tokens for budget purposes. + * @param budgetTokens - The token ceiling for each sub-batch. Caller is + * responsible for applying any safety-factor shrink before passing in. + * @param charsPerToken - Provider-specific character density. Defaults to + * `DEFAULT_CHARS_PER_TOKEN` (4) when omitted, matching OpenAI tiktoken. + * + * @internal exported for tests; not part of the public gateway API. */ -function splitByTokenBudget(texts: string[], maxTokens: number): string[][] { - // Use 80% of declared limit as safety margin for tokenizer variance. - const budget = Math.floor(maxTokens * 0.8); +export function splitByTokenBudget( + texts: string[], + budgetTokens: number, + charsPerToken: number = DEFAULT_CHARS_PER_TOKEN, +): string[][] { + const ratio = charsPerToken > 0 ? charsPerToken : DEFAULT_CHARS_PER_TOKEN; const batches: string[][] = []; let current: string[] = []; let currentTokens = 0; for (const text of texts) { - const estTokens = Math.ceil(text.length / CHARS_PER_TOKEN_ESTIMATE); - if (current.length > 0 && currentTokens + estTokens > budget) { + const estTokens = Math.ceil(text.length / ratio); + if (current.length > 0 && currentTokens + estTokens > budgetTokens) { batches.push(current); current = []; currentTokens = 0; @@ -311,8 +402,12 @@ function splitByTokenBudget(texts: string[], maxTokens: number): string[][] { return batches; } -/** Returns true if the error looks like a provider batch-token-limit error. */ -function isTokenLimitError(err: unknown): boolean { +/** + * Returns true if the error looks like a provider batch-token-limit error. + * + * @internal exported for tests; not part of the public gateway API. + */ +export function isTokenLimitError(err: unknown): boolean { const msg = err instanceof Error ? err.message : String(err); return ( /max.*allowed.*tokens.*batch/i.test(msg) || @@ -321,6 +416,54 @@ function isTokenLimitError(err: unknown): boolean { ); } +/** + * Resolve the recipe's effective safety factor (declared default, optionally + * shrunk by prior misses in this process). + */ +function effectiveSafetyFactor(recipe: Recipe): number { + const declared = recipe.touchpoints?.embedding?.safety_factor ?? DEFAULT_SAFETY_FACTOR; + const entry = _shrinkState.get(recipe.id); + return entry?.factor ?? declared; +} + +/** Tighten the recipe's effective safety factor on a token-limit miss. */ +function shrinkOnMiss(recipe: Recipe): void { + const declared = recipe.touchpoints?.embedding?.safety_factor ?? DEFAULT_SAFETY_FACTOR; + const current = _shrinkState.get(recipe.id)?.factor ?? declared; + const next = Math.max(SHRINK_FLOOR, current * 0.5); + _shrinkState.set(recipe.id, { factor: next, consecutiveSuccesses: 0 }); +} + +/** Bump the win counter; heal toward declared default after enough wins. */ +function recordSubBatchSuccess(recipe: Recipe): void { + const declared = recipe.touchpoints?.embedding?.safety_factor ?? DEFAULT_SAFETY_FACTOR; + const entry = _shrinkState.get(recipe.id); + if (!entry || entry.factor >= declared) { + // Either no shrink active, or already at/above the declared ceiling — nothing to heal. + if (entry) { + _shrinkState.set(recipe.id, { factor: entry.factor, consecutiveSuccesses: 0 }); + } + return; + } + const wins = entry.consecutiveSuccesses + 1; + if (wins >= SHRINK_HEAL_AFTER) { + const healed = Math.min(declared, entry.factor * 1.5); + _shrinkState.set(recipe.id, { factor: healed, consecutiveSuccesses: 0 }); + } else { + _shrinkState.set(recipe.id, { factor: entry.factor, consecutiveSuccesses: wins }); + } +} + +/** + * Read the current shrink state for a recipe. Test-only seam. + * + * @internal exported for tests; not part of the public gateway API. + */ +export function __getShrinkStateForTests(recipeId: string): ShrinkEntry | undefined { + const entry = _shrinkState.get(recipeId); + return entry ? { ...entry } : undefined; +} + /** * Embed a single sub-batch with automatic halving on token-limit errors. * If the batch is already at MIN_SUB_BATCH and still fails, throws. @@ -334,7 +477,7 @@ async function embedSubBatch( modelId: string, ): Promise { try { - const result = await embedMany({ + const result = await _embedTransport({ model, values: texts, providerOptions: providerOpts, @@ -342,26 +485,20 @@ async function embedSubBatch( const first = result.embeddings?.[0]; if (first && Array.isArray(first) && first.length !== expectedDims) { - // v0.28.5 (#672): the previous fix-hint pointed at `gbrain migrate - // --embedding-model … --embedding-dimensions …` but `gbrain migrate` - // only handles engine migration, not embedding reconfiguration. The - // canonical fix is the manual ALTER recipe in - // docs/embedding-migrations.md, surfaced inline so the user doesn't - // need to context-switch. throw new AIConfigError( `Embedding dim mismatch: model ${modelId} returned ${first.length} but schema expects ${expectedDims}.`, - `Either change models to one that returns ${expectedDims}-d embeddings, ` + - `or migrate the existing brain to ${first.length}-d (destructive — see docs/embedding-migrations.md). ` + - `Quick recipe: DROP INDEX idx_chunks_embedding; ALTER COLUMN embedding TYPE vector(${first.length}); ` + - `UPDATE content_chunks SET embedding = NULL; gbrain config set embedding_dimensions ${first.length}; ` + - `gbrain embed --stale.`, + `Run \`gbrain migrate --embedding-model ${getEmbeddingModel()} --embedding-dimensions ${first.length}\` or change models.`, ); } + recordSubBatchSuccess(recipe); return result.embeddings.map((e: number[]) => new Float32Array(e)); } catch (err) { - // On token-limit error, try splitting the batch in half and retrying. + // On token-limit error, tighten the recipe's effective safety factor + // (so the next embed() pre-splits smaller) and recursively halve THIS + // batch to make forward progress without dropping work. if (isTokenLimitError(err) && texts.length > MIN_SUB_BATCH) { + shrinkOnMiss(recipe); const mid = Math.ceil(texts.length / 2); const left = await embedSubBatch(texts.slice(0, mid), model, providerOpts, expectedDims, recipe, modelId); const right = await embedSubBatch(texts.slice(mid), model, providerOpts, expectedDims, recipe, modelId); diff --git a/src/core/ai/recipes/voyage.ts b/src/core/ai/recipes/voyage.ts index fb85e0586..b4620400c 100644 --- a/src/core/ai/recipes/voyage.ts +++ b/src/core/ai/recipes/voyage.ts @@ -3,10 +3,6 @@ import type { Recipe } from '../types.ts'; /** * Voyage AI exposes an OpenAI-compatible /embeddings endpoint. * Base URL: https://api.voyageai.com/v1 - * - * Voyage 4 family (Jan 2026): shared embedding space across all v4 variants, - * flexible dims (256/512/1024/2048), 32K context, MoE architecture (large). - * You can index with voyage-4-large and query with voyage-4-lite — no reindex. */ export const voyage: Recipe = { id: 'voyage', @@ -20,17 +16,18 @@ export const voyage: Recipe = { }, touchpoints: { embedding: { - models: [ - 'voyage-4-large', 'voyage-4', 'voyage-4-lite', 'voyage-4-nano', - 'voyage-3.5', 'voyage-3-large', 'voyage-3', 'voyage-3-lite', - 'voyage-code-3', 'voyage-finance-2', 'voyage-law-2', - ], + models: ['voyage-3-large', 'voyage-3', 'voyage-3-lite'], default_dims: 1024, cost_per_1m_tokens_usd: 0.18, - price_last_verified: '2026-05-06', - // Voyage enforces 120K tokens per batch. Use conservative limit - // because Voyage's tokenizer runs 3-4× denser than tiktoken. + price_last_verified: '2026-04-20', + // Voyage enforces 120K tokens per batch. Voyage's tokenizer runs + // ~3-4× denser than OpenAI tiktoken on mixed content (code/JSON/CJK), + // so the per-recipe pre-split uses 1 char ≈ 1 token at 0.5 utilization + // (60K char budget). Recursive halving in the gateway is the runtime + // safety net when dense payloads still overshoot. max_batch_tokens: 120_000, + chars_per_token: 1, + safety_factor: 0.5, }, }, setup_hint: 'Get an API key at https://dash.voyageai.com/api-keys, then `export VOYAGE_API_KEY=...`', diff --git a/src/core/ai/types.ts b/src/core/ai/types.ts index b5bc256e5..ec791e869 100644 --- a/src/core/ai/types.ts +++ b/src/core/ai/types.ts @@ -31,12 +31,28 @@ export interface EmbeddingTouchpoint { price_last_verified?: string; // ISO date /** * Maximum tokens per batch for this provider's embedding endpoint. - * When set, the gateway auto-splits large batches to stay under this - * limit. Uses a conservative chars-to-tokens ratio (1 char ≈ 1 token) - * since tokenizer density varies widely across providers (e.g. Voyage - * tokenizes ~3-4× denser than OpenAI tiktoken on mixed content). + * When set, the gateway pre-splits batches at + * `max_batch_tokens × safety_factor / chars_per_token` characters and + * recursively halves on token-limit errors at runtime. When unset, the + * gateway makes a single embedMany() call with no safety net (OpenAI fast + * path). */ max_batch_tokens?: number; + /** + * Expected character density for this provider's tokenizer (chars per + * token). OpenAI tiktoken averages ~4 on English text; Voyage averages + * ~1 on mixed content (code/JSON/CJK). Defaults to 4 if omitted. + * Only consulted when `max_batch_tokens` is also set. + */ + chars_per_token?: number; + /** + * Budget-utilization ceiling in (0, 1]. The gateway pre-splits at + * `safety_factor × max_batch_tokens` to leave headroom for tokenizer + * variance. Defaults to 0.8. Voyage-style providers with dense payloads + * should pin this lower (e.g. 0.5). Only consulted when + * `max_batch_tokens` is also set. + */ + safety_factor?: number; } export interface ExpansionTouchpoint { diff --git a/src/core/embedding.ts b/src/core/embedding.ts index 43c1de7f9..82239cef7 100644 --- a/src/core/embedding.ts +++ b/src/core/embedding.ts @@ -26,12 +26,12 @@ export interface EmbedBatchOptions { } /** - * Embed a batch of texts via the gateway. Sub-batches of 50 so upstream - * progress callbacks fire incrementally on large imports. The gateway - * handles truncation, retries, provider dispatch, and adaptive batch - * splitting for providers with token-per-batch limits (e.g. Voyage). + * Embed a batch of texts via the gateway. Sub-batches of 100 so upstream + * progress callbacks fire incrementally on large imports. The gateway owns + * adaptive batch splitting and per-recipe token-budget logic; this paginator + * is purely about progress-callback granularity. */ -const BATCH_SIZE = 50; +const BATCH_SIZE = 100; export async function embedBatch( texts: string[], options: EmbedBatchOptions = {}, diff --git a/test/ai/adaptive-embed-batch.test.ts b/test/ai/adaptive-embed-batch.test.ts index 4a3eeaa51..6cd8d06ab 100644 --- a/test/ai/adaptive-embed-batch.test.ts +++ b/test/ai/adaptive-embed-batch.test.ts @@ -1,112 +1,390 @@ /** - * Tests for adaptive embed batch splitting (fix/adaptive-embed-batch-sizing). + * Tests for the adaptive embed batch system (PR #680, ships v0.28.7). * - * Validates that splitByTokenBudget correctly partitions texts to stay - * under provider token limits, and that embedSubBatch retries with - * halved batches on token-limit errors. + * Coverage matrix (per the eng-review plan): + * + * 1. Pure helpers exported from gateway.ts: + * - splitByTokenBudget pure-function semantics + chars_per_token threading + * - isTokenLimitError regex coverage + * + * 2. Recursion through public embed() with the AI-SDK transport stubbed. + * We do NOT call private functions; the test seam is the + * __setEmbedTransportForTests hook on the gateway. + * + * 3. Order preservation across recursive halving (left/right concat). + * + * 4. Terminal MIN_SUB_BATCH=1 — single text whose transport always fails + * must throw normalizeAIError, not loop forever. + * + * 5. OpenAI fast path (D3) — recipe with no max_batch_tokens calls the + * transport exactly once with no pre-split. + * + * 6. Shrink-on-miss adaptive cache (D8-A) — first miss halves the factor; + * after SHRINK_HEAL_AFTER successes the factor heals back toward the + * recipe-declared safety_factor. + * + * 7. Startup warning (D9-B) — gateway construction warns once per recipe + * with an embedding touchpoint missing max_batch_tokens (excluding the + * OpenAI canonical fast-path recipe). */ -import { describe, expect, test } from 'bun:test'; +import { afterEach, beforeEach, describe, expect, mock, test } from 'bun:test'; +import { + configureGateway, + resetGateway, + embed, + splitByTokenBudget, + isTokenLimitError, + __setEmbedTransportForTests, + __getShrinkStateForTests, +} from '../../src/core/ai/gateway.ts'; +import { AIConfigError, AITransientError } from '../../src/core/ai/errors.ts'; -// We test the logic by importing the gateway module and exercising the -// exported embed() function through its internal helpers. Since -// splitByTokenBudget and isTokenLimitError are module-private, we test -// them indirectly through behavior. +// --------- Test helpers --------- -// Direct unit tests for splitByTokenBudget logic (extracted for clarity). -describe('splitByTokenBudget logic', () => { - // Re-implement the algorithm locally for unit testing since it's private. - function splitByTokenBudget(texts: string[], maxTokens: number): string[][] { - const budget = Math.floor(maxTokens * 0.8); - const batches: string[][] = []; - let current: string[] = []; - let currentTokens = 0; +/** + * Build an embedding-shape return for an arbitrary number of values. Each + * embedding is `dims` floats, all set to a sentinel index so tests can + * assert order preservation. + */ +function fakeEmbeddings(values: string[], dims: number): { embeddings: number[][] } { + return { + embeddings: values.map((_, i) => + // First slot encodes the input index so we can verify ordering. + Array.from({ length: dims }, (_, j) => (j === 0 ? i : 0.1)), + ), + }; +} - for (const text of texts) { - const estTokens = Math.ceil(text.length / 1); // 1 char ≈ 1 token - if (current.length > 0 && currentTokens + estTokens > budget) { - batches.push(current); - current = []; - currentTokens = 0; - } - current.push(text); - currentTokens += estTokens; - } - if (current.length > 0) batches.push(current); - return batches; - } +const VOYAGE_TOKEN_LIMIT_ERROR = new Error( + "Request to model 'voyage-3-large' failed. The max allowed tokens per submitted batch is 120000.", +); +function configureVoyage(): void { + configureGateway({ + embedding_model: 'voyage:voyage-3-large', + embedding_dimensions: 1024, + env: { VOYAGE_API_KEY: 'sk-fake' }, + }); +} + +function configureOpenAI(): void { + configureGateway({ + embedding_model: 'openai:text-embedding-3-large', + embedding_dimensions: 1536, + env: { OPENAI_API_KEY: 'sk-fake' }, + }); +} + +// --------- 1. Pure helpers --------- + +describe('splitByTokenBudget (pure helper)', () => { test('single small text stays in one batch', () => { - const result = splitByTokenBudget(['hello'], 120_000); - expect(result).toHaveLength(1); - expect(result[0]).toEqual(['hello']); + const result = splitByTokenBudget(['hello'], 120_000, 1); + expect(result).toEqual([['hello']]); }); test('texts fitting within budget stay in one batch', () => { const texts = Array.from({ length: 10 }, () => 'a'.repeat(1000)); - // 10 * 1000 = 10K chars, well under 120K * 0.8 = 96K budget - const result = splitByTokenBudget(texts, 120_000); + const result = splitByTokenBudget(texts, 96_000, 1); expect(result).toHaveLength(1); expect(result[0]).toHaveLength(10); }); test('texts exceeding budget are split into multiple batches', () => { - // Each text is 50K chars. Budget = 120K * 0.8 = 96K. - // First text fits (50K < 96K). Second would make it 100K > 96K → new batch. + // chars_per_token=1, so each 50K-char text counts as 50K tokens. + // Budget 96K → first text fits, second pushes over → new batch. const texts = ['a'.repeat(50_000), 'b'.repeat(50_000), 'c'.repeat(50_000)]; - const result = splitByTokenBudget(texts, 120_000); + const result = splitByTokenBudget(texts, 96_000, 1); expect(result).toHaveLength(3); - expect(result[0]).toHaveLength(1); - expect(result[1]).toHaveLength(1); - expect(result[2]).toHaveLength(1); + expect(result.map(b => b.length)).toEqual([1, 1, 1]); }); - test('packs multiple texts into one batch when under budget', () => { - // Each text is 30K chars. Budget = 96K. Two fit (60K < 96K), three don't (90K < 96K still fits!). - // Actually 30K * 3 = 90K < 96K, so three fit. - const texts = ['a'.repeat(30_000), 'b'.repeat(30_000), 'c'.repeat(30_000), 'd'.repeat(30_000)]; - const result = splitByTokenBudget(texts, 120_000); - // 3 fit in first batch (90K < 96K), 4th starts new batch - expect(result).toHaveLength(2); - expect(result[0]).toHaveLength(3); - expect(result[1]).toHaveLength(1); + test('chars_per_token=4 (OpenAI density) packs 4× more chars per batch', () => { + // Each 50K-char text = 12.5K tokens at chars_per_token=4. Budget 96K + // tokens → 7 fit; the 8th would overflow into a new batch. + const texts = Array.from({ length: 10 }, (_, i) => `${i}`.repeat(50_000)); + const result = splitByTokenBudget(texts, 96_000, 4); + expect(result[0].length).toBe(7); + expect(result[1].length).toBe(3); + }); + + test('default chars_per_token (4) when ratio omitted', () => { + // Same payload as above without the explicit ratio. + const texts = Array.from({ length: 10 }, (_, i) => `${i}`.repeat(50_000)); + const explicit = splitByTokenBudget(texts, 96_000, 4); + const implicit = splitByTokenBudget(texts, 96_000); + expect(implicit).toEqual(explicit); }); test('empty input returns empty array', () => { - const result = splitByTokenBudget([], 120_000); - expect(result).toHaveLength(0); + expect(splitByTokenBudget([], 120_000, 1)).toEqual([]); }); - test('single text larger than budget still goes in a batch', () => { - // A single huge text can't be split further by this function. - const texts = ['a'.repeat(200_000)]; - const result = splitByTokenBudget(texts, 120_000); + test('single text larger than budget still goes in a batch (split helper does not subdivide)', () => { + const result = splitByTokenBudget(['a'.repeat(200_000)], 120_000, 1); expect(result).toHaveLength(1); expect(result[0]).toHaveLength(1); }); + + test('zero or negative chars_per_token falls back to default', () => { + const texts = ['a'.repeat(40_000)]; + expect(splitByTokenBudget(texts, 96_000, 0)).toEqual(splitByTokenBudget(texts, 96_000, 4)); + expect(splitByTokenBudget(texts, 96_000, -1)).toEqual(splitByTokenBudget(texts, 96_000, 4)); + }); }); -describe('isTokenLimitError pattern matching', () => { - function isTokenLimitError(err: unknown): boolean { - const msg = err instanceof Error ? err.message : String(err); - return ( - /max.*allowed.*tokens.*batch/i.test(msg) || - /batch.*too.*many.*tokens/i.test(msg) || - /token.*limit.*exceeded/i.test(msg) - ); - } - +describe('isTokenLimitError (pure helper)', () => { test('matches Voyage error format', () => { - const msg = 'Request to model \'voyage-4-large\' failed. The max allowed tokens per submitted batch is 120000.'; - expect(isTokenLimitError(new Error(msg))).toBe(true); + expect(isTokenLimitError(VOYAGE_TOKEN_LIMIT_ERROR)).toBe(true); + }); + + test('matches "token limit exceeded" variant', () => { + expect(isTokenLimitError(new Error('Token limit exceeded for batch request'))).toBe(true); + }); + + test('matches "batch too many tokens" variant', () => { + expect(isTokenLimitError(new Error('Batch contains too many tokens'))).toBe(true); }); test('does not match unrelated errors', () => { expect(isTokenLimitError(new Error('Connection refused'))).toBe(false); expect(isTokenLimitError(new Error('Invalid API key'))).toBe(false); + expect(isTokenLimitError(new Error('429 rate limited'))).toBe(false); }); - test('matches token limit exceeded variant', () => { - expect(isTokenLimitError(new Error('Token limit exceeded for batch request'))).toBe(true); + test('handles non-Error throwables', () => { + expect(isTokenLimitError('Token limit exceeded')).toBe(true); + expect(isTokenLimitError({ message: 'some other thing' })).toBe(false); + expect(isTokenLimitError(null)).toBe(false); + expect(isTokenLimitError(undefined)).toBe(false); + }); +}); + +// --------- 2-4. Recursion via embed() with stubbed transport --------- + +describe('embed() recursion via stubbed transport', () => { + beforeEach(() => resetGateway()); + afterEach(() => __setEmbedTransportForTests(null)); + + test('halves on token-limit error and concatenates left+right in order', async () => { + configureVoyage(); + + const stub = mock(async ({ values }: { values: string[] }) => { + // First call: full batch fails. Halved calls: succeed. + if (values.length === 50) throw VOYAGE_TOKEN_LIMIT_ERROR; + return fakeEmbeddings(values, 1024); + }); + __setEmbedTransportForTests(stub as any); + + // Build 50 texts that each fit comfortably under any pre-split budget + // (1 char ≈ 1 token in voyage's recipe; 0.5 × 120K = 60K char budget). + const texts = Array.from({ length: 50 }, (_, i) => `t${i}`); + const result = await embed(texts); + + // Stub fired 3 times: 1 fail (length 50) + 2 success (length 25 each). + expect(stub).toHaveBeenCalledTimes(3); + const callLengths = stub.mock.calls.map(([arg]) => (arg as { values: string[] }).values.length); + expect(callLengths.sort((a, b) => a - b)).toEqual([25, 25, 50]); + expect(result).toHaveLength(50); + }); + + test('preserves input order across halving boundaries', async () => { + configureVoyage(); + + const stub = mock(async ({ values }: { values: string[] }) => { + if (values.length === 10) throw VOYAGE_TOKEN_LIMIT_ERROR; + return fakeEmbeddings(values, 1024); + }); + __setEmbedTransportForTests(stub as any); + + const texts = Array.from({ length: 10 }, (_, i) => String.fromCharCode(97 + i)); // a..j + const result = await embed(texts); + + expect(result).toHaveLength(10); + // The fakeEmbeddings helper encodes the within-call index in slot 0; + // halved calls each receive sub-arrays of length 5, so slot 0 reads + // [0,1,2,3,4,0,1,2,3,4] — that's the contract that proves order + // preservation despite the embeddings being concatenated from two calls. + const slotZero = result.map(v => v[0]); + expect(slotZero).toEqual([0, 1, 2, 3, 4, 0, 1, 2, 3, 4]); + }); + + test('terminal case: single text always fails → normalizes and throws (no infinite loop)', async () => { + configureVoyage(); + + const stub = mock(async () => { throw VOYAGE_TOKEN_LIMIT_ERROR; }); + __setEmbedTransportForTests(stub as any); + + let caught: unknown = null; + try { + await embed(['just one text']); + } catch (e) { + caught = e; + } + expect(caught).not.toBeNull(); + expect(caught instanceof AIConfigError || caught instanceof AITransientError).toBe(true); + // Stub fires once for the single-element batch; cannot halve further so + // the recursion gives up at MIN_SUB_BATCH=1 and rethrows. + expect(stub).toHaveBeenCalledTimes(1); + }); +}); + +// --------- 5. OpenAI fast path (D3) --------- + +describe('embed() OpenAI fast path (no max_batch_tokens)', () => { + beforeEach(() => resetGateway()); + afterEach(() => __setEmbedTransportForTests(null)); + + test('recipe without max_batch_tokens calls transport exactly once with no partition', async () => { + configureOpenAI(); + + const stub = mock(async ({ values }: { values: string[] }) => fakeEmbeddings(values, 1536)); + __setEmbedTransportForTests(stub as any); + + const texts = Array.from({ length: 100 }, (_, i) => `text-${i}`); + const result = await embed(texts); + + expect(stub).toHaveBeenCalledTimes(1); + const callValues = (stub.mock.calls[0][0] as { values: string[] }).values; + expect(callValues).toEqual(texts); + expect(result).toHaveLength(100); + }); + + test('OpenAI fast path is unaffected by Voyage shrink state', async () => { + // Configure Voyage first and trigger a shrink… + configureVoyage(); + const voyageStub = mock(async ({ values }: { values: string[] }) => { + if (values.length === 4) throw VOYAGE_TOKEN_LIMIT_ERROR; + return fakeEmbeddings(values, 1024); + }); + __setEmbedTransportForTests(voyageStub as any); + await embed(['a', 'b', 'c', 'd']); + expect(__getShrinkStateForTests('voyage')?.factor).toBe(0.25); + + // …then reconfigure to OpenAI. The shrink state belongs to the prior + // gateway's lifecycle and must not leak. + configureOpenAI(); + const openaiStub = mock(async ({ values }: { values: string[] }) => fakeEmbeddings(values, 1536)); + __setEmbedTransportForTests(openaiStub as any); + await embed(['x', 'y']); + expect(openaiStub).toHaveBeenCalledTimes(1); + expect(__getShrinkStateForTests('voyage')).toBeUndefined(); + }); +}); + +// --------- 6. Shrink-on-miss adaptive cache (D8-A) --------- + +describe('shrink-on-miss adaptive cache', () => { + beforeEach(() => resetGateway()); + afterEach(() => __setEmbedTransportForTests(null)); + + test('first token-limit miss halves the recipe safety factor', async () => { + configureVoyage(); + expect(__getShrinkStateForTests('voyage')).toBeUndefined(); + + const stub = mock(async ({ values }: { values: string[] }) => { + if (values.length === 4) throw VOYAGE_TOKEN_LIMIT_ERROR; + return fakeEmbeddings(values, 1024); + }); + __setEmbedTransportForTests(stub as any); + + await embed(['a', 'b', 'c', 'd']); + // Voyage declares safety_factor=0.5; after one miss → 0.5 × 0.5 = 0.25. + expect(__getShrinkStateForTests('voyage')?.factor).toBe(0.25); + }); + + test('factor floors at SHRINK_FLOOR (0.05) under repeated misses', async () => { + configureVoyage(); + const stub = mock(async ({ values }: { values: string[] }) => { + // Always throw on >1 to keep recursion going until MIN_SUB_BATCH=1 + // succeeds. That gives many shrink events per embed() call. + if (values.length > 1) throw VOYAGE_TOKEN_LIMIT_ERROR; + return fakeEmbeddings(values, 1024); + }); + __setEmbedTransportForTests(stub as any); + + // 16 texts will recurse 4 levels deep, generating multiple shrink events. + await embed(Array.from({ length: 16 }, (_, i) => `t${i}`)); + const factor = __getShrinkStateForTests('voyage')?.factor ?? -1; + expect(factor).toBeGreaterThanOrEqual(0.05); + }); + + test('factor heals back toward declared safety_factor after enough wins', async () => { + configureVoyage(); + const stub = mock(async ({ values }: { values: string[] }) => { + // Once: fail at length 2, succeed everywhere else. Subsequent calls + // all succeed. + if (stub.mock.calls.length === 1 && values.length === 2) throw VOYAGE_TOKEN_LIMIT_ERROR; + return fakeEmbeddings(values, 1024); + }); + __setEmbedTransportForTests(stub as any); + + await embed(['a', 'b']); // 1 fail + 2 successes (length 1 each) → factor 0.25, wins 2 + const afterMiss = __getShrinkStateForTests('voyage')?.factor; + expect(afterMiss).toBe(0.25); + + // Drive 10 more successful calls. SHRINK_HEAL_AFTER=10; on the 10th win + // the factor multiplies by 1.5 (capped at the declared 0.5 ceiling). + for (let i = 0; i < 8; i++) { + await embed(['solo']); + } + const healed = __getShrinkStateForTests('voyage')?.factor ?? 0; + // 0.25 × 1.5 = 0.375. Still below the recipe ceiling of 0.5; the next + // round of 10 wins would bump it to min(0.5, 0.375 × 1.5) = 0.5. + expect(healed).toBeCloseTo(0.375, 5); + }); + + test('healing path cannot exceed the recipe-declared safety_factor', async () => { + configureVoyage(); + const stub = mock(async ({ values }: { values: string[] }) => { + if (stub.mock.calls.length === 1) throw VOYAGE_TOKEN_LIMIT_ERROR; + return fakeEmbeddings(values, 1024); + }); + __setEmbedTransportForTests(stub as any); + + // Trigger one shrink, then drive enough wins to fully heal. + await embed(['one', 'two']); + for (let i = 0; i < 30; i++) { + await embed(['solo']); + } + const factor = __getShrinkStateForTests('voyage')?.factor ?? 0; + // Declared safety_factor is 0.5; healing must clamp at that ceiling. + expect(factor).toBeLessThanOrEqual(0.5); + expect(factor).toBeGreaterThan(0); + }); +}); + +// --------- 7. Startup warning (D9-B) --------- + +describe('startup warning for recipes missing max_batch_tokens', () => { + beforeEach(() => resetGateway()); + + test('first configureGateway call warns about each missing-cap recipe; subsequent calls suppressed', () => { + const warnings: string[] = []; + const original = console.warn; + console.warn = (msg: string) => warnings.push(String(msg)); + try { + configureOpenAI(); + const firstCallCount = warnings.length; + // Reconfigure: the warning should NOT re-fire for the same recipes + // within one process (we already told the operator). + configureOpenAI(); + expect(warnings.length).toBe(firstCallCount); + } finally { + console.warn = original; + } + + // The warning text should match the documented contract. + const contractMatch = warnings.filter(w => + w.includes('[ai.gateway]') && w.includes('declares an embedding touchpoint'), + ); + expect(contractMatch.length).toBeGreaterThan(0); + + // Voyage declares max_batch_tokens → suppressed. OpenAI is the + // canonical fast-path recipe → also suppressed by id. Both must be + // absent from the warnings. + expect(warnings.find(w => w.includes('"voyage"'))).toBeUndefined(); + expect(warnings.find(w => w.includes('"openai"'))).toBeUndefined(); }); });