Compare commits

..
Author SHA1 Message Date
88287775e5 fix(chunker): estimated-token hard cap — URL-dense/CJK-fallback chunks overflow strict embedding-server token limits
Takeover of #2847 (rebased onto current master). Fixes #2826.

- cjk.ts: estimateEmbeddingTokens() — conservative per-char-class token
  estimate (CJK 1.0, other 0.75, whitespace 0.1 per code unit).
- recursive.ts (MARKDOWN_CHUNKER_VERSION 3→4): countWords floored at
  ceil(nonWhitespaceChars/6); capByEstimatedTokens() final pass with
  ChunkOptions.maxTokens (default 1500).
- code.ts (CHUNKER_VERSION 4→5): capCodeChunks() applies the same cap to
  AST-path chunks that splitLargeNode can't subdivide.

Co-authored-by: paul-0320 <paul-0320@users.noreply.github.com>

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
2026-07-21 14:18:40 -07:00
20 changed files with 465 additions and 520 deletions
+5 -10
View File
@@ -170,14 +170,10 @@ export async function runImport(
// v0.22.13 (PR #490 Q2): shared parseWorkers helper rejects bad input
// (--workers 0, -3, "foo") with a loud error instead of silently falling
// through to 1. Mirrors sync.ts's flag handling.
const { parseWorkers, autoConcurrency } = await import('../core/sync-concurrency.ts');
// #1207: undefined (no --workers flag) defers to autoConcurrency below —
// the shared sync/import policy (PGLite → 1, >100 files → 4) — instead of
// hardcoding serial. Large Postgres imports stop paying one embedding
// round-trip per file in sequence.
let workerCount: number | undefined;
const { parseWorkers } = await import('../core/sync-concurrency.ts');
let workerCount: number;
try {
workerCount = parseWorkers(workersArg ?? undefined);
workerCount = parseWorkers(workersArg ?? undefined) ?? 1;
} catch (e) {
console.error(e instanceof Error ? e.message : String(e));
process.exit(1);
@@ -256,9 +252,8 @@ export async function runImport(
}
const files = resumeFilter(allFiles, dir, completed);
// Determine actual worker count. Explicit --workers wins; otherwise the
// shared autoConcurrency policy decides from engine kind + file count.
const actualWorkers = autoConcurrency(engine, files.length, workerCount);
// Determine actual worker count
const actualWorkers = workerCount > 1 ? workerCount : 1;
if (actualWorkers > 1) {
console.log(`Using ${actualWorkers} parallel workers`);
}
+6 -24
View File
@@ -1513,21 +1513,12 @@ export async function embed(texts: string[], opts?: EmbedOpts): Promise<Float32A
const embedding = recipe.touchpoints?.embedding;
const maxBatchTokens = embedding?.max_batch_tokens;
const maxBatchCount = embedding?.max_batch_count;
const charsPerToken = embedding?.chars_per_token ?? DEFAULT_CHARS_PER_TOKEN;
// Pre-split is gated on max_batch_tokens / max_batch_count. Recipes with
// neither (e.g. OpenAI) ride the fast path: one embedMany call, no
// recursion safety net.
const batches = (maxBatchTokens || maxBatchCount)
? splitByTokenBudget(
truncated,
maxBatchTokens
? Math.floor(maxBatchTokens * effectiveSafetyFactor(recipe))
: Number.MAX_SAFE_INTEGER,
charsPerToken,
maxBatchCount,
)
// 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, Math.floor(maxBatchTokens * effectiveSafetyFactor(recipe)), charsPerToken)
: [truncated];
const allEmbeddings: Float32Array[] = [];
@@ -1577,9 +1568,6 @@ export async function embed(texts: string[], opts?: EmbedOpts): Promise<Float32A
* 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.
* @param maxBatchCount - #1199: optional cap on INPUTS per sub-batch, for
* providers that reject batches by count (DashScope: 10). When omitted,
* only the token budget governs.
*
* @internal exported for tests; not part of the public gateway API.
*/
@@ -1587,17 +1575,15 @@ export function splitByTokenBudget(
texts: string[],
budgetTokens: number,
charsPerToken: number = DEFAULT_CHARS_PER_TOKEN,
maxBatchCount?: number,
): string[][] {
const ratio = charsPerToken > 0 ? charsPerToken : DEFAULT_CHARS_PER_TOKEN;
const maxCount = maxBatchCount !== undefined && maxBatchCount > 0 ? maxBatchCount : Infinity;
const batches: string[][] = [];
let current: string[] = [];
let currentTokens = 0;
for (const text of texts) {
const estTokens = Math.ceil(text.length / ratio);
if (current.length > 0 && (currentTokens + estTokens > budgetTokens || current.length >= maxCount)) {
if (current.length > 0 && currentTokens + estTokens > budgetTokens) {
batches.push(current);
current = [];
currentTokens = 0;
@@ -1623,11 +1609,7 @@ export function isTokenLimitError(err: unknown): boolean {
/token.*limit.*exceeded/i.test(msg) ||
// OpenAI embeddings: "Invalid 'input': maximum request size is 300000 tokens per request."
/maximum request size.*tokens/i.test(msg) ||
/max.*tokens.*per.*request/i.test(msg) ||
// DashScope: "batch size is invalid, it should not be larger than 10." (#1199)
// Count-cap error, but recursive halving shrinks count too, so the same
// safety net converges.
/batch size is invalid/i.test(msg)
/max.*tokens.*per.*request/i.test(msg)
);
}
-4
View File
@@ -31,10 +31,6 @@ export const dashscope: Recipe = {
// path. Conservative declaration so the gateway pre-splits before
// hitting whatever undocumented server-side limit exists.
max_batch_tokens: 8192,
// #1199: DashScope hard-caps embeddings at 10 inputs per request
// ("batch size is invalid, it should not be larger than 10"). The
// token budget alone admits far more than 10 short chunks per batch.
max_batch_count: 10,
// text-embedding-v3 mixes English + CJK heavily; the tokenizer is
// closer to Voyage density than OpenAI tiktoken for CJK-dominant
// content. Conservative chars_per_token=2 leaves headroom.
-9
View File
@@ -16,15 +16,6 @@ export const google: Recipe = {
dims_options: [768, 1536, 3072],
cost_per_1m_tokens_usd: 0.15,
price_last_verified: '2026-04-20',
// #970: Gemini's documented limits are per-INPUT (2048 tokens,
// silently truncated beyond) and per-REQUEST count (batchEmbedContents
// caps at 100 inputs). There is no separate per-request token cap, so
// the token budget is derived: 100 inputs × 2048 tokens. The count cap
// binds first for typical chunk sizes. Do NOT copy the 2048 per-input
// limit into max_batch_tokens — that would over-split 50×.
max_batch_tokens: 204_800,
chars_per_token: 4,
max_batch_count: 100,
},
expansion: {
models: ['gemini-2.0-flash', 'gemini-2.0-flash-lite'],
+1 -4
View File
@@ -58,8 +58,5 @@ export function getRecipe(id: string): Recipe | undefined {
}
export function listRecipes(): Recipe[] {
// Read the map (not ALL) so there is one source of truth — getRecipe,
// model-resolver, and listRecipes all see the same registry, and tests
// can inject a synthetic recipe via RECIPES to exercise registry walks.
return [...RECIPES.values()];
return [...ALL];
}
-10
View File
@@ -46,16 +46,6 @@ export interface EmbeddingTouchpoint {
* Only consulted when `max_batch_tokens` is also set.
*/
chars_per_token?: number;
/**
* #1199: maximum number of INPUTS per embedding request, for providers
* that hard-cap batch size by count rather than (or in addition to)
* tokens — DashScope text-embedding-v3 rejects batches > 10 with
* `InvalidParameter`, Gemini batchEmbedContents caps at 100 requests.
* When set, the gateway's pre-split flushes a sub-batch at this count
* even if the token budget still has room. Independent of
* `max_batch_tokens`; either alone triggers the pre-split.
*/
max_batch_count?: number;
/**
* Budget-utilization ceiling in (0, 1]. The gateway pre-splits at
* `safety_factor × max_batch_tokens` to leave headroom for tokenizer
+39 -3
View File
@@ -18,8 +18,9 @@
* at runtime.
*/
import { chunkText as recursiveChunk } from './recursive.ts';
import { chunkText as recursiveChunk, capByEstimatedTokens, DEFAULT_MAX_EST_TOKENS } from './recursive.ts';
import { buildQualifiedName } from './qualified-names.ts';
import { estimateEmbeddingTokens } from '../cjk.ts';
// Embed the tree-sitter runtime + per-language grammars as files.
// `with { type: 'file' }` returns a path (string) at runtime. Bun bundles
@@ -111,7 +112,15 @@ import G_ZIG from '../../assets/wasm/grammars/tree-sitter-zig.wasm' with { type:
// chunks get the new columns populated. Without this, the v28 backfill
// gives every existing chunk a search_vector but subsequent Layer 5 AST
// work would silently no-op.
export const CHUNKER_VERSION = 4;
//
// v5: estimated-token hard cap on AST-path chunks (capCodeChunks). A node
// splitLargeNode can't subdivide (giant single-statement function, huge
// literal) previously shipped WHOLE regardless of size and could overflow
// strict per-request embedding-token limits (local llama-server crashes
// past ~2,050 tokens, measured). Mirrors the markdown
// chunker's v4 cap; fallback-path chunks are already capped inside
// recursiveChunk.
export const CHUNKER_VERSION = 5;
// Lazy-loaded tree-sitter module (v0.22.x API: Parser is default export)
let Parser: typeof import('web-tree-sitter') | null = null;
@@ -708,7 +717,7 @@ export async function chunkCodeTextFull(
if (chunks.length === 0) {
return { chunks: fallbackChunks(source, filePath, language, opts), edges: rawEdges };
}
return { chunks: mergeSmallSiblings(chunks, chunkTarget), edges: rawEdges };
return { chunks: capCodeChunks(mergeSmallSiblings(chunks, chunkTarget)), edges: rawEdges };
} catch {
return { chunks: fallbackChunks(source, filePath, language, opts), edges: [] };
} finally {
@@ -791,6 +800,33 @@ function mergeSmallSiblings(chunks: CodeChunk[], chunkTarget: number): CodeChunk
return merged;
}
/**
* v5 final safety pass for AST-path chunks: split any chunk whose
* ESTIMATED embedding tokens (conservative per-char-class heuristic,
* cjk.ts) exceed DEFAULT_MAX_EST_TOKENS. Reaches chunks the AST logic
* can't subdivide — splitLargeNode returns [] for nodes with < 2 body
* children (giant single-statement functions, huge literals), which
* previously shipped whole at any size.
*
* Split pieces inherit the source chunk's metadata verbatim; start/end
* lines become approximate for pieces after the first. Acceptable —
* these chunks exist for embedding + retrieval, and the alternative was
* an embedding request the server rejects (or worse, crashes on).
*/
function capCodeChunks(chunks: CodeChunk[]): CodeChunk[] {
if (chunks.every((c) => estimateEmbeddingTokens(c.text) <= DEFAULT_MAX_EST_TOKENS)) {
return chunks;
}
const out: CodeChunk[] = [];
for (const c of chunks) {
const pieces = capByEstimatedTokens(c.text, DEFAULT_MAX_EST_TOKENS);
for (const piece of pieces) {
out.push({ ...c, text: piece, index: out.length, metadata: { ...c.metadata } });
}
}
return out;
}
function buildMergedChunk(group: CodeChunk[], index: number): CodeChunk {
const first = group[0]!;
const last = group[group.length - 1]!;
+117 -6
View File
@@ -17,7 +17,13 @@
* Lossless invariant: non-overlapping portions reassemble to original.
*/
import { countCJKAwareWords, CJK_SENTENCE_DELIMITERS, CJK_CLAUSE_DELIMITERS } from '../cjk.ts';
import {
countCJKAwareWords,
CJK_SENTENCE_DELIMITERS,
CJK_CLAUSE_DELIMITERS,
charEmbedTokenWeight,
estimateEmbeddingTokens,
} from '../cjk.ts';
/**
* Markdown chunker version. Folded into the per-page chunker_version column
@@ -33,8 +39,20 @@ import { countCJKAwareWords, CJK_SENTENCE_DELIMITERS, CJK_CLAUSE_DELIMITERS } fr
* re-embed (not re-chunk) so existing pages pick up the wrapper on the
* post-upgrade reembed sweep. See
* `src/core/contextual-retrieval-service.ts`.
*
* v4: estimated-token hard cap + whitespace-word undercount fix. The word
* pipeline counted a 150-char URL as ONE whitespace word, so URL/phone/
* email-dense docs (CJK density < 0.30 → whitespace fallback) produced
* 3-4K-char chunks that overflow strict per-request embedding-token
* limits (measured: local llama-server crashes past ~2,050 tokens; URL
* soup tokenizes at ~1.6 chars/token). Two changes:
* 1. countWords() floors the count at ceil(nonWhitespaceChars/6) so a
* URL counts roughly per-character, not as one word.
* 2. capByEstimatedTokens() final pass guarantees every chunk fits
* `maxTokens` (default 1500) under a conservative per-char-class
* token estimate, regardless of how word counting misjudged it.
*/
export const MARKDOWN_CHUNKER_VERSION = 3;
export const MARKDOWN_CHUNKER_VERSION = 4;
const DELIMITERS: string[][] = [
['\n\n'], // L0: paragraphs
@@ -48,8 +66,20 @@ export interface ChunkOptions {
chunkSize?: number; // target words per chunk (default 300)
chunkOverlap?: number; // overlap words (default 50)
maxChars?: number; // hard cap on any chunk's char length (default 6000)
/**
* v4: hard cap on any chunk's ESTIMATED embedding tokens (default 1500).
* Estimate = conservative per-char-class weights (see cjk.ts
* estimateEmbeddingTokens) — deliberately high, so the real tokenizer
* count stays below this value. Default leaves headroom for the
* contextual-retrieval wrapper (≤ ~630 chars) under a ~2,050-token
* per-request embedding server limit.
*/
maxTokens?: number;
}
/** v4 default for ChunkOptions.maxTokens — see the field doc above. */
export const DEFAULT_MAX_EST_TOKENS = 1500;
export interface TextChunk {
text: string;
index: number;
@@ -73,6 +103,7 @@ export function chunkText(text: string, opts?: ChunkOptions): TextChunk[] {
const chunkSize = opts?.chunkSize || 300;
const chunkOverlap = opts?.chunkOverlap || 50;
const maxChars = opts?.maxChars || 6000;
const maxTokens = opts?.maxTokens || DEFAULT_MAX_EST_TOKENS;
if (!text || text.trim().length === 0) return [];
@@ -89,8 +120,9 @@ export function chunkText(text: string, opts?: ChunkOptions): TextChunk[] {
const wordCount = countWords(stripped);
if (wordCount <= chunkSize) {
// Single-chunk path: still apply the maxChars cap.
const capped = capByChars(stripped.trim(), maxChars);
// Single-chunk path: still apply the maxChars + maxTokens caps.
const capped = capByChars(stripped.trim(), maxChars)
.flatMap((t) => capByEstimatedTokens(t, maxTokens));
return capped.map((t, i) => ({ text: t, index: i }));
}
@@ -101,9 +133,14 @@ export function chunkText(text: string, opts?: ChunkOptions): TextChunk[] {
// v0.32.7: hard char cap. Catches pathological CJK + whitespace-less text
// that the word-level pipeline can't bound (a single Chinese paragraph can
// exceed 8192 OpenAI embedding tokens at any word count).
// v4: estimated-token cap on top — the char cap alone passes token-dense
// content (URL soup at ~1.6 chars/token) that overflows strict embedding
// server limits.
const capped: string[] = [];
for (const chunk of withOverlap) {
capped.push(...capByChars(chunk.trim(), maxChars));
for (const piece of capByChars(chunk.trim(), maxChars)) {
capped.push(...capByEstimatedTokens(piece, maxTokens));
}
}
return capped.map((t, i) => ({ text: t, index: i }));
}
@@ -132,6 +169,68 @@ function capByChars(text: string, maxChars: number): string[] {
return out;
}
/**
* How far back (in chars) the token cap looks for a friendly cut point
* before falling back to a hard cut. 300 covers typical rollup/list line
* lengths so forced splits land at line starts, not mid-URL.
*/
const TOKEN_CAP_CUT_LOOKBACK = 300;
/**
* v4: hard-cap a chunk's ESTIMATED embedding tokens. Final safety pass —
* runs after capByChars on every chunk, so no upstream miscounting
* (whitespace-word fallback, overlap inflation, char-cap survivors) can
* emit a chunk past `maxTokens`.
*
* Cut placement prefers, within the last TOKEN_CAP_CUT_LOOKBACK chars of
* the window: a newline, then any whitespace, then a hard cut. This keeps
* forced splits off mid-line/mid-URL positions for list-shaped content
* and inside code fences. No overlap is added (pieces stay lossless
* modulo the trims the char cap already applies).
*
* @internal exported for the code chunker (code.ts) and tests.
*/
export function capByEstimatedTokens(text: string, maxTokens: number): string[] {
if (text.length === 0) return [];
if (estimateEmbeddingTokens(text) <= maxTokens) return [text];
const out: string[] = [];
let start = 0;
while (start < text.length) {
// Greedily extend the window until the next char would break the cap.
// Always take at least one char so the loop makes forward progress.
let est = 0;
let end = start;
while (end < text.length) {
const w = charEmbedTokenWeight(text.charCodeAt(end));
if (est + w > maxTokens && end > start) break;
est += w;
end++;
}
if (end < text.length) {
const windowStart = Math.max(start + 1, end - TOKEN_CAP_CUT_LOOKBACK);
let cut = text.lastIndexOf('\n', end - 1);
if (cut < windowStart) {
cut = -1;
for (let i = end - 1; i >= windowStart; i--) {
const code = text.charCodeAt(i);
if (code === 0x20 || (code >= 0x09 && code <= 0x0d)) {
cut = i;
break;
}
}
}
if (cut >= windowStart) end = cut + 1;
}
const slice = text.slice(start, end).trim();
if (slice.length > 0) out.push(slice);
start = end;
}
return out;
}
function recursiveSplit(text: string, level: number, target: number): string[] {
if (level >= DELIMITERS.length) {
// Level 4: split on whitespace
@@ -317,7 +416,19 @@ function extractTrailingContext(text: string, targetWords: number): string {
* Delegated to src/core/cjk.ts so the slugify whitelist, expansion
* detection, and PGLite keyword fallback all agree on what "CJK enough"
* means.
*
* v4: floored at ceil(nonWhitespaceChars/6). The whitespace fallback
* counts a 150-char URL as ONE word, so URL/phone/email-dense docs
* (whose ASCII mass pushes CJK density below the 0.30 threshold) were
* sized at a fraction of their real bulk and merged into 3-4K-char
* chunks. The floor makes long whitespace-less runs count roughly
* per-character while leaving normal Latin prose untouched (average
* English word ≈ 5 chars < 6, so the whitespace count still wins).
* Kept local to the chunker — search/expansion.ts keeps the original
* countCJKAwareWords semantics for its query-length check.
*/
function countWords(text: string): number {
return countCJKAwareWords(text);
const cjkAware = countCJKAwareWords(text);
const nonWhitespace = text.replace(/\s/g, '').length;
return Math.max(cjkAware, Math.ceil(nonWhitespace / 6));
}
+61
View File
@@ -65,3 +65,64 @@ export function countCJKAwareWords(s: string): number {
export function escapeLikePattern(s: string): string {
return s.replace(/\\/g, '\\\\').replace(/%/g, '\\%').replace(/_/g, '\\_');
}
/**
* Conservative per-char-class embedding-token weights (markdown chunker v4).
*
* Why this exists: the chunker's "word" counting drastically UNDER-counts
* whitespace-less ASCII runs (a 150-char URL = 1 whitespace word), so
* word-based size targets can emit chunks that overflow an embedding
* server's per-request token limit. Measured on a local Qwen3-embedding
* llama-server stack:
* - URL/phone/email-dense text tokenizes at ~1.6 chars/token
* - base64-ish / minified blobs approach ~1.3 chars/token (worst case)
* - Korean prose tokenizes NO WORSE than 1 char/token in practice
*
* Weights are deliberately HIGH (tokens are overestimated) so any cap
* based on this estimate is safe against real tokenizers:
* - CJK char → 1.0 token (real CJK prose is cheaper)
* - other non-space → 0.75 token (≈1.33 chars/token, covers base64)
* - whitespace → 0.1 token (mostly folds into neighbor tokens)
*/
export const EMBED_TOKEN_WEIGHT_CJK = 1.0;
export const EMBED_TOKEN_WEIGHT_OTHER = 0.75;
export const EMBED_TOKEN_WEIGHT_WS = 0.1;
/** BMP CJK check by UTF-16 code unit — same ranges as CJK_SLUG_CHARS. */
export function isCJKCodeUnit(code: number): boolean {
return (
(code >= 0x4e00 && code <= 0x9fff) || // Han
(code >= 0x3040 && code <= 0x309f) || // Hiragana
(code >= 0x30a0 && code <= 0x30ff) || // Katakana
(code >= 0xac00 && code <= 0xd7af) // Hangul Syllables
);
}
/**
* Per-code-unit token weight. Unrecognized whitespace (exotic Unicode
* spaces) intentionally falls into OTHER — that only overestimates.
*/
export function charEmbedTokenWeight(code: number): number {
if (isCJKCodeUnit(code)) return EMBED_TOKEN_WEIGHT_CJK;
if (
code === 0x20 || (code >= 0x09 && code <= 0x0d) ||
code === 0xa0 || code === 0x3000
) {
return EMBED_TOKEN_WEIGHT_WS;
}
return EMBED_TOKEN_WEIGHT_OTHER;
}
/**
* Tokenizer-free embedding-token estimate (conservative overestimate).
* See weight docs above. Astral chars count as 2 OTHER code units —
* another overestimate, which is the safe direction.
*/
export function estimateEmbeddingTokens(s: string): number {
if (s.length === 0) return 0;
let est = 0;
for (let i = 0; i < s.length; i++) {
est += charEmbedTokenWeight(s.charCodeAt(i));
}
return Math.ceil(est);
}
+6 -56
View File
@@ -79,34 +79,15 @@ export interface EmbedBatchOptions {
* and amplify rate-limit pressure.
*/
maxRetries?: number;
/**
* #1818: bounded parallelism across BATCH_SIZE sub-batches. Defaults to
* `GBRAIN_EMBED_BATCH_CONCURRENCY` env, else 4. Results are
* index-addressed so output order always matches input order. Set 1 to
* force the pre-v0.42 serial dispatch.
*/
concurrency?: number;
}
/**
* 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
* owns progress-callback granularity and (#1818) bounded parallel dispatch
* of the sub-batches — the embed-stale.ts worker-pool pattern, scoped down.
* is purely about progress-callback granularity.
*/
const BATCH_SIZE = 100;
const DEFAULT_EMBED_BATCH_CONCURRENCY = 4;
function resolveEmbedBatchConcurrency(options: EmbedBatchOptions): number {
if (options.concurrency !== undefined) {
return Math.max(1, Math.floor(options.concurrency));
}
const env = Number(process.env.GBRAIN_EMBED_BATCH_CONCURRENCY);
if (Number.isFinite(env) && env >= 1) return Math.floor(env);
return DEFAULT_EMBED_BATCH_CONCURRENCY;
}
export async function embedBatch(
texts: string[],
options: EmbedBatchOptions = {},
@@ -122,44 +103,13 @@ export async function embedBatch(
if (texts.length <= BATCH_SIZE && !options.onBatchComplete) {
return gatewayEmbed(texts, gwOpts);
}
// #1818: dispatch sub-batches through a bounded worker pool instead of a
// serial loop. Results are written into a preallocated index-addressed
// array so output order matches input order regardless of completion
// order; onBatchComplete reports a monotonic completed-embedding count.
const slices: Array<{ start: number; texts: string[] }> = [];
const results: Float32Array[] = [];
for (let i = 0; i < texts.length; i += BATCH_SIZE) {
slices.push({ start: i, texts: texts.slice(i, i + BATCH_SIZE) });
const slice = texts.slice(i, i + BATCH_SIZE);
const out = await gatewayEmbed(slice, gwOpts);
results.push(...out);
options.onBatchComplete?.(results.length, texts.length);
}
const results = new Array<Float32Array>(texts.length);
let next = 0;
let done = 0;
const numWorkers = Math.min(resolveEmbedBatchConcurrency(options), slices.length);
// Once any sub-batch fails, `failed` stops the surviving workers from
// dispatching FURTHER slices — the whole call is rejecting anyway, so
// continuing would burn real provider spend in the background and fire
// onBatchComplete after the caller already saw the failure (worst with
// embedBatchWithBackoff, whose 429 backoff assumes nothing is in flight).
// In-flight sibling calls still run to completion (bounded by numWorkers-1).
let failed = false;
const worker = async (): Promise<void> => {
while (!failed && next < slices.length) {
// NOTE: no local aborted-check here — an aborted signal makes the next
// gatewayEmbed call throw (SDK-side), which rejects the pool. Returning
// silently instead would resolve with holes in `results`.
const slice = slices[next++];
let out: Float32Array[];
try {
out = await gatewayEmbed(slice.texts, gwOpts);
} catch (err) {
failed = true;
throw err;
}
for (let j = 0; j < out.length; j++) results[slice.start + j] = out[j];
done += out.length;
if (!failed) options.onBatchComplete?.(done, texts.length);
}
};
await Promise.all(Array.from({ length: numWorkers }, () => worker()));
return results;
}
+5 -109
View File
@@ -39,8 +39,6 @@ import {
__getShrinkStateForTests,
} from '../../src/core/ai/gateway.ts';
import { AIConfigError, AITransientError } from '../../src/core/ai/errors.ts';
import { RECIPES } from '../../src/core/ai/recipes/index.ts';
import type { Recipe } from '../../src/core/ai/types.ts';
// The last test in this file leaves the gateway configured with a remote
// provider + fake key and a REAL embed transport. Without a final reset,
@@ -95,14 +93,6 @@ function configureGoogle(): void {
});
}
function configureDashscope(): void {
configureGateway({
embedding_model: 'dashscope:text-embedding-v3',
embedding_dimensions: 1024,
env: { DASHSCOPE_API_KEY: 'sk-fake' },
});
}
// --------- 1. Pure helpers ---------
describe('splitByTokenBudget (pure helper)', () => {
@@ -159,27 +149,6 @@ describe('splitByTokenBudget (pure helper)', () => {
expect(splitByTokenBudget(texts, 96_000, 0)).toEqual(splitByTokenBudget(texts, 96_000, 4));
expect(splitByTokenBudget(texts, 96_000, -1)).toEqual(splitByTokenBudget(texts, 96_000, 4));
});
// #1199: count cap for providers that reject batches by input count.
test('max_batch_count flushes even when token budget has room', () => {
const texts = Array.from({ length: 25 }, (_, i) => `t${i}`);
const result = splitByTokenBudget(texts, 1_000_000, 4, 10);
expect(result.map(b => b.length)).toEqual([10, 10, 5]);
expect(result.flat()).toEqual(texts);
});
test('token budget still governs alongside max_batch_count', () => {
const texts = ['a'.repeat(50_000), 'b'.repeat(50_000), 'c'.repeat(50_000)];
const result = splitByTokenBudget(texts, 96_000, 1, 10);
expect(result).toHaveLength(3);
});
test('undefined / zero / negative max_batch_count is ignored', () => {
const texts = Array.from({ length: 25 }, () => 'x');
expect(splitByTokenBudget(texts, 1_000_000, 4, undefined)).toHaveLength(1);
expect(splitByTokenBudget(texts, 1_000_000, 4, 0)).toHaveLength(1);
expect(splitByTokenBudget(texts, 1_000_000, 4, -5)).toHaveLength(1);
});
});
describe('isTokenLimitError (pure helper)', () => {
@@ -210,12 +179,6 @@ describe('isTokenLimitError (pure helper)', () => {
expect(isTokenLimitError(new Error('Exceeded 300000 max tokens per request'))).toBe(true);
});
test('matches DashScope batch-count error (#1199)', () => {
expect(isTokenLimitError(new Error(
'InvalidParameter: batch size is invalid, it should not be larger than 10.',
))).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);
@@ -424,92 +387,26 @@ describe('shrink-on-miss adaptive cache', () => {
});
});
// --------- 8. Pre-split count cap through public embed() (#1199 / #970) ---------
describe('embed() pre-split honors max_batch_count', () => {
beforeEach(() => resetGateway());
afterEach(() => __setEmbedTransportForTests(null));
test('dashscope never dispatches more than 10 inputs per call (#1199)', async () => {
configureDashscope();
const stub = mock(async ({ values }: { values: string[] }) => fakeEmbeddings(values, 1024));
__setEmbedTransportForTests(stub as any);
// 25 short texts fit trivially in the 8192-token budget; without the
// count cap they'd ship as ONE batch and DashScope would reject it.
const texts = Array.from({ length: 25 }, (_, i) => `short-${i}`);
const result = await embed(texts);
expect(result).toHaveLength(25);
const callLengths = stub.mock.calls.map(([arg]) => (arg as { values: string[] }).values.length);
expect(Math.max(...callLengths)).toBeLessThanOrEqual(10);
expect(callLengths.reduce((a, b) => a + b, 0)).toBe(25);
// Order preserved across sub-batches.
expect((stub.mock.calls[0][0] as { values: string[] }).values[0]).toBe('short-0');
});
test('google pre-splits at 100 inputs per batchEmbedContents call (#970)', async () => {
configureGoogle();
const stub = mock(async ({ values }: { values: string[] }) => fakeEmbeddings(values, 768));
__setEmbedTransportForTests(stub as any);
const texts = Array.from({ length: 250 }, (_, i) => `g${i}`);
const result = await embed(texts);
expect(result).toHaveLength(250);
const callLengths = stub.mock.calls.map(([arg]) => (arg as { values: string[] }).values.length);
expect(callLengths).toEqual([100, 100, 50]);
});
});
// --------- 7. Startup warning (D9-B) ---------
describe('startup warning for recipes missing max_batch_tokens', () => {
beforeEach(() => resetGateway());
// #970 closed google's missing cap, so no registered recipe is capless
// anymore. Inject a synthetic capless recipe to keep the warning path
// covered for the NEXT recipe that forgets the field.
const caplessRecipe: Recipe = {
id: 'capless-test',
name: 'Capless Test Provider',
tier: 'openai-compat',
implementation: 'openai-compatible',
base_url_default: 'https://example.invalid/v1',
auth_env: { required: [] },
touchpoints: {
embedding: { models: ['capless-embed-1'], default_dims: 768 },
},
};
function configureCapless(): void {
configureGateway({
embedding_model: 'capless-test:capless-embed-1',
embedding_dimensions: 768,
env: {},
});
}
test('configured missing-cap recipe warns once; unrelated recipes stay quiet', () => {
const warnings: string[] = [];
const original = console.warn;
console.warn = (msg: string) => warnings.push(String(msg));
RECIPES.set(caplessRecipe.id, caplessRecipe);
try {
configureOpenAI();
expect(warnings.length).toBe(0);
// #970 regression: google now declares max_batch_tokens → quiet.
configureGoogle();
expect(warnings.length).toBe(0);
configureCapless();
const firstCallCount = warnings.length;
// Reconfigure: the warning should NOT re-fire for the same recipes
// within one process (we already told the operator).
configureCapless();
configureGoogle();
expect(warnings.length).toBe(firstCallCount);
} finally {
console.warn = original;
RECIPES.delete(caplessRecipe.id);
}
// The warning text should match the documented contract.
@@ -518,12 +415,11 @@ describe('startup warning for recipes missing max_batch_tokens', () => {
);
expect(contractMatch.length).toBe(1);
// Voyage + google declare max_batch_tokens → suppressed. OpenAI is the
// canonical fast-path recipe → also suppressed by id. All must be
// absent from the warnings; only the synthetic capless recipe fires.
// 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();
expect(warnings.find(w => w.includes('"google"'))).toBeUndefined();
expect(warnings.find(w => w.includes('"capless-test"'))).toBeDefined();
expect(warnings.find(w => w.includes('"google"'))).toBeDefined();
});
});
+13 -13
View File
@@ -52,7 +52,16 @@ describe('v0.32 #779: no_batch_cap suppresses the missing-max_batch_tokens warni
}
});
test('configureGateway does NOT warn for google now that it declares batch caps (#970)', () => {
test('configureGateway warns for google only when google embedding is configured', () => {
warnSpy.mockClear();
resetGateway();
configureGateway({ env: {} });
let messages = warnSpy.mock.calls.map(c => String(c[0] ?? ''));
expect(
messages.some(m => m.includes('"google"') && m.includes('without max_batch_tokens')),
'google should not warn while OpenAI default is configured',
).toBe(false);
warnSpy.mockClear();
resetGateway();
configureGateway({
@@ -60,20 +69,11 @@ describe('v0.32 #779: no_batch_cap suppresses the missing-max_batch_tokens warni
embedding_dimensions: 768,
env: { GOOGLE_GENERATIVE_AI_API_KEY: 'fake' },
});
const messages = warnSpy.mock.calls.map(c => String(c[0] ?? ''));
messages = warnSpy.mock.calls.map(c => String(c[0] ?? ''));
expect(
messages.some(m => m.includes('"google"') && m.includes('without max_batch_tokens')),
'google declares max_batch_tokens/max_batch_count since #970 — no warning',
).toBe(false);
});
test('google recipe declares its derived batch caps (#970)', () => {
const e = getRecipe('google')!.touchpoints.embedding!;
// Count cap is the REAL Gemini limit (batchEmbedContents: 100 inputs);
// the token budget is derived (100 × 2048 per-input tokens), NOT the
// 2048 per-input limit — copying that verbatim would over-split 50×.
expect(e.max_batch_count).toBe(100);
expect(e.max_batch_tokens).toBe(204_800);
'google should warn when configured because it has fixed-cap models',
).toBe(true);
});
test('every recipe with empty models[] declares user_provided_models OR has openai-fast-path', () => {
-5
View File
@@ -55,11 +55,6 @@ describe('recipe: dashscope', () => {
expect(r.touchpoints.embedding!.chars_per_token).toBeGreaterThan(0);
});
test('declares max_batch_count: 10 — DashScope rejects larger batches (#1199)', () => {
const r = getRecipe('dashscope')!;
expect(r.touchpoints.embedding!.max_batch_count).toBe(10);
});
test('dimsProviderOptions threads dimensions for text-embedding-v3 (Matryoshka)', async () => {
// Codex finding #1: DashScope text-embedding-v3 is Matryoshka 64-1024.
// Without `dimensions` on the wire, user-selected non-default dims are
+6 -3
View File
@@ -15,19 +15,22 @@ import { describe, test, expect } from 'bun:test';
import { CHUNKER_VERSION } from '../src/core/chunkers/code.ts';
describe('Layer 12 — CHUNKER_VERSION constant', () => {
test('bumped to 4 for Cathedral II', () => {
test('bumped to 5 for the estimated-token hard cap', () => {
// v3: v0.19.0 Chonkie parity (tokenizer + small-sibling merge).
// v4: v0.20.0 Cathedral II (qualified names + parent scope + doc_comment
// + fence extraction + chunk-grain FTS). Folded into content_hash
// so any bump forces clean re-chunks on next sync.
expect(CHUNKER_VERSION).toBe(4);
// v5: estimated-token hard cap on AST-path chunks (capCodeChunks) so
// un-subdividable giant nodes can't overflow strict embedding
// server token limits.
expect(CHUNKER_VERSION).toBe(5);
});
test('is stable across imports (not recomputed at call time)', async () => {
const a = (await import('../src/core/chunkers/code.ts')).CHUNKER_VERSION;
const b = (await import('../src/core/chunkers/code.ts')).CHUNKER_VERSION;
expect(a).toBe(b);
expect(a).toBe(4);
expect(a).toBe(5);
});
});
+2 -2
View File
@@ -10,8 +10,8 @@ import { describe, test, expect } from 'bun:test';
import { chunkCodeText, detectCodeLanguage, CHUNKER_VERSION } from '../../src/core/chunkers/code.ts';
describe('CHUNKER_VERSION', () => {
test('v0.20.0 Cathedral II Layer 12 bumped to 4', () => {
expect(CHUNKER_VERSION).toBe(4);
test('v5: estimated-token hard cap on AST-path chunks', () => {
expect(CHUNKER_VERSION).toBe(5);
});
});
+6 -5
View File
@@ -135,13 +135,14 @@ describe('Recursive Text Chunker', () => {
});
describe('CJK chunking (v0.32.7)', () => {
test('MARKDOWN_CHUNKER_VERSION is 3', async () => {
test('MARKDOWN_CHUNKER_VERSION is 4', async () => {
// v0.40.3.0: bumped 2→3 to signal the post-upgrade reembed sweep that
// contextual retrieval wrapping is now applied at embed time. Chunk
// boundaries themselves are unchanged; the bump forces re-embed for
// pages where chunker_version < 3.
// contextual retrieval wrapping is now applied at embed time.
// v4: estimated-token hard cap + whitespace-word undercount floor
// (URL-dense docs produced chunks past strict embedding server token
// limits). Boundary change → forces re-chunk for chunker_version < 4.
const mod = await import('../../src/core/chunkers/recursive.ts');
expect(mod.MARKDOWN_CHUNKER_VERSION).toBe(3);
expect(mod.MARKDOWN_CHUNKER_VERSION).toBe(4);
});
test('long pure-Chinese paragraph splits into multiple chunks', () => {
+195
View File
@@ -0,0 +1,195 @@
/**
* Markdown chunker v4 / code chunker v5 estimated-token hard cap
* regression tests.
*
* Reproduces a field failure: a local llama-server embedding backend
* (`-ub 2048`) crashes deterministically (trace/BPT trap EOF at the
* client) when a single chunk exceeds ~2,050 real tokens. Two content
* shapes triggered it:
*
* 1. Korean docs carrying one long source URL per line.
* The URLs' ASCII mass pushes CJK density below 0.30, flipping
* countCJKAwareWords to whitespace counting, where a 150-char URL
* counts as ONE word chunks ballooned to 3-4K chars 2,000+
* real tokens (URL soup tokenizes at ~1.6 chars/token).
*
* 2. Large JSON code blocks (~7K chars) that the word pipeline
* undercounts the same way (few whitespace tokens).
*
* The fix: every emitted chunk must satisfy
* estimateEmbeddingTokens(chunk) <= maxTokens (default 1500)
* where the estimate deliberately OVERSTATES real tokenizer counts.
*/
import { describe, test, expect } from 'bun:test';
import { chunkText, capByEstimatedTokens, DEFAULT_MAX_EST_TOKENS } from '../../src/core/chunkers/recursive.ts';
import { chunkCodeText } from '../../src/core/chunkers/code.ts';
import { estimateEmbeddingTokens } from '../../src/core/cjk.ts';
/** Synthesize the failing shape: Korean rollup lines each ending in a long Notion URL. */
function urlDenseKoreanRollup(lines: number): string {
const out: string[] = ['# 링크가 줄마다 붙는 한국어 예시 문서', ''];
for (let i = 0; i < lines; i++) {
const hex32 = (i * 2654435761 >>> 0).toString(16).padStart(8, '0').repeat(4);
out.push(
`- **항목 ${i}**: 이 줄은 청커 동작 검증을 위한 의미 없는 한국어 예시 문장입니다 · 전화 000-0000-${String(1000 + i)} · ` +
`이메일 user${i}@example.com · 링크: https://docs.example.com/pages/${hex32}?v=abcdef0123456789&ref=sample`,
);
}
return out.join('\n');
}
/** Synthesize a large pretty-printed JSON block with CJK values. */
function bigJsonBlock(targetChars: number): string {
const entries: string[] = [];
let i = 0;
let len = 0;
while (len < targetChars) {
const row =
` "item_${i}": { "name": "예시-${i}", "url": "https://example.com/api/v2/items/${i}?token=abc${i}def", "qty": ${i % 100}, "memo": "한국어 값이 섞인 예시 데이터" }`;
entries.push(row);
len += row.length;
i++;
}
return `{\n${entries.join(',\n')}\n}`;
}
describe('v4 estimated-token cap — URL-dense Korean doc (field-failure shape)', () => {
test('every chunk stays under the estimated-token cap', () => {
const md = urlDenseKoreanRollup(60);
const chunks = chunkText(md);
expect(chunks.length).toBeGreaterThan(0);
for (const c of chunks) {
expect(estimateEmbeddingTokens(c.text)).toBeLessThanOrEqual(DEFAULT_MAX_EST_TOKENS);
}
});
test('no chunk reaches the measured 3K-char danger zone for URL soup', () => {
const md = urlDenseKoreanRollup(60);
const chunks = chunkText(md);
// 1500 est tokens at the OTHER weight (0.75/char) bounds chunks to
// ~2,000 chars for pure ASCII — well under the ~3,300 chars where
// URL-dense content crosses ~2,050 real tokens (1.6 chars/token).
for (const c of chunks) {
expect(c.text.length).toBeLessThanOrEqual(2600);
}
});
test('content is preserved (no lines dropped by the cap)', () => {
const md = urlDenseKoreanRollup(60);
const chunks = chunkText(md);
const joined = chunks.map((c) => c.text).join('\n');
// Spot-check first / middle / last rollup lines survive chunking.
for (const marker of ['항목 0', '항목 30', '항목 59']) {
expect(joined).toContain(marker);
}
});
});
describe('v4 estimated-token cap — large JSON blocks', () => {
test('7K-char pretty JSON through the prose path stays under the cap', () => {
const md = `설정 파일 원문 보존:\n\n\`\`\`\n${bigJsonBlock(7000)}\n\`\`\`\n`;
const chunks = chunkText(md);
expect(chunks.length).toBeGreaterThan(1);
for (const c of chunks) {
expect(estimateEmbeddingTokens(c.text)).toBeLessThanOrEqual(DEFAULT_MAX_EST_TOKENS);
}
});
test('7K-char minified JSON (single whitespace-less token) stays under the cap', () => {
const minified = bigJsonBlock(7000).replace(/\n\s*/g, '');
const chunks = chunkText(minified);
expect(chunks.length).toBeGreaterThan(1);
for (const c of chunks) {
expect(estimateEmbeddingTokens(c.text)).toBeLessThanOrEqual(DEFAULT_MAX_EST_TOKENS);
}
});
test('json fence via the code chunker stays under the cap (+header slack)', async () => {
const chunks = await chunkCodeText(bigJsonBlock(7000), 'fence.json');
expect(chunks.length).toBeGreaterThan(0);
for (const c of chunks) {
// buildChunk prepends a short "[JSON] fence.json:…" header AFTER the
// body-level cap; allow ~60 est tokens of header slack. Real-token
// safety margin (2,050 overestimated 1,500) absorbs this easily.
expect(estimateEmbeddingTokens(c.text)).toBeLessThanOrEqual(DEFAULT_MAX_EST_TOKENS + 60);
}
});
});
describe('v4 word-count floor — behavior preserved for normal content', () => {
test('Latin prose chunking is unchanged by the floor (avg word < 6 chars)', () => {
const prose = Array.from({ length: 120 }, (_, i) =>
`This is sentence number ${i} and it talks about ordinary things in plain words.`,
).join(' ');
const chunks = chunkText(prose);
// Historical behavior: ~1,560 whitespace words → multiple ~300-word chunks.
expect(chunks.length).toBeGreaterThan(3);
for (const c of chunks) {
const words = c.text.split(/\s+/).length;
expect(words).toBeLessThanOrEqual(300 * 1.5 + 50); // merge cap + overlap
}
});
test('Korean prose (CJK-dense, no URLs) never triggers the token cap', () => {
const prose = Array.from({ length: 80 }, (_, i) =>
`이 문장은 순수 한국어 산문의 청킹 동작을 확인하기 위한 ${i}번째 예시 문장입니다.`,
).join(' ');
const chunks = chunkText(prose);
expect(chunks.length).toBeGreaterThan(1);
for (const c of chunks) {
// CJK-dense chunks are char-counted (≈450 max) — nowhere near 1500.
expect(estimateEmbeddingTokens(c.text)).toBeLessThanOrEqual(700);
}
});
});
describe('capByEstimatedTokens unit behavior', () => {
test('returns input unchanged when under the cap', () => {
expect(capByEstimatedTokens('short text', 1500)).toEqual(['short text']);
expect(capByEstimatedTokens('', 1500)).toEqual([]);
});
test('prefers newline cut points within the lookback window', () => {
const line = 'x'.repeat(100);
const text = Array.from({ length: 40 }, () => line).join('\n');
const pieces = capByEstimatedTokens(text, 1000);
expect(pieces.length).toBeGreaterThan(1);
for (const p of pieces) {
// Every piece should be whole lines (multiples of the 100-char line).
for (const l of p.split('\n')) {
expect(l).toBe(line);
}
}
});
test('makes forward progress on whitespace-less input (hard cut)', () => {
const blob = 'a'.repeat(10_000);
const pieces = capByEstimatedTokens(blob, 1000);
expect(pieces.length).toBeGreaterThan(1);
expect(pieces.join('')).toBe(blob);
for (const p of pieces) {
expect(estimateEmbeddingTokens(p)).toBeLessThanOrEqual(1000);
}
});
});
describe('estimateEmbeddingTokens — weight sanity', () => {
test('overestimates URL-dense ASCII (0.75/char ≥ measured ~0.63/char)', () => {
const url = 'https://docs.example.com/pages/a1b2c3d4e5f6a1b2c3d4e5f6a1b2c3d4?v=abc&ref=sample';
const est = estimateEmbeddingTokens(url);
expect(est).toBeGreaterThanOrEqual(Math.floor(url.length * 0.7));
});
test('counts CJK at 1 token/char', () => {
expect(estimateEmbeddingTokens('가나다라마')).toBe(5);
});
test('whitespace is nearly free', () => {
expect(estimateEmbeddingTokens(' \n\t ')).toBeLessThanOrEqual(1);
});
test('empty string is 0', () => {
expect(estimateEmbeddingTokens('')).toBe(0);
});
});
-161
View File
@@ -1,161 +0,0 @@
/**
* #1818: embedBatch dispatches its 100-input sub-batches through a bounded
* worker pool (the embed-stale.ts concurrency pattern) instead of a serial
* `for` loop. This file pins:
*
* - output order matches input order regardless of completion order
* (index-addressed results)
* - parallelism actually happens (max in-flight > 1) and stays bounded
* (max in-flight <= configured concurrency)
* - concurrency: 1 restores the serial pre-#1818 dispatch
* - GBRAIN_EMBED_BATCH_CONCURRENCY env is honored when the option is unset
* - onBatchComplete reports a monotonic completed count ending at total
*
* Transport is stubbed via the gateway's __setEmbedTransportForTests seam
* (same pattern as test/ai/adaptive-embed-batch.test.ts). OpenAI recipe =
* fast path (no pre-split), so each embedBatch sub-batch is exactly one
* transport call.
*/
import { afterAll, afterEach, beforeEach, describe, expect, test } from 'bun:test';
import {
configureGateway,
resetGateway,
__setEmbedTransportForTests,
} from '../src/core/ai/gateway.ts';
import { embedBatch } from '../src/core/embedding.ts';
import { withEnv } from './helpers/with-env.ts';
const DIMS = 1536;
function configureOpenAI(): void {
configureGateway({
embedding_model: 'openai:text-embedding-3-large',
embedding_dimensions: DIMS,
env: { OPENAI_API_KEY: 'sk-fake' },
});
}
/**
* Install a transport whose returned embedding encodes the GLOBAL input
* index in dim 0 (texts are `t<N>`), so order can be asserted end-to-end.
* Tracks the max number of concurrently in-flight transport calls.
*/
function installTrackingTransport(delayMs = 5): { maxInFlight: () => number } {
let inFlight = 0;
let maxInFlight = 0;
__setEmbedTransportForTests((async ({ values }: { values: string[] }) => {
inFlight++;
maxInFlight = Math.max(maxInFlight, inFlight);
await new Promise(r => setTimeout(r, delayMs));
inFlight--;
return {
embeddings: values.map(v => {
const idx = Number(v.slice(1));
return Array.from({ length: DIMS }, (_, j) => (j === 0 ? idx : 0.1));
}),
};
}) as any);
return { maxInFlight: () => maxInFlight };
}
const texts = Array.from({ length: 250 }, (_, i) => `t${i}`);
afterAll(() => resetGateway());
describe('embedBatch bounded parallelism (#1818)', () => {
beforeEach(() => {
resetGateway();
configureOpenAI();
});
afterEach(() => {
__setEmbedTransportForTests(null);
});
test('default pool dispatches sub-batches in parallel, order preserved', async () => {
const tracker = installTrackingTransport();
const result = await embedBatch(texts, { onBatchComplete: () => {} });
expect(result).toHaveLength(250);
for (let i = 0; i < 250; i++) {
expect(result[i][0]).toBe(i);
}
// 250 texts → 3 sub-batches; default concurrency 4 → all 3 in flight.
expect(tracker.maxInFlight()).toBeGreaterThan(1);
expect(tracker.maxInFlight()).toBeLessThanOrEqual(4);
});
test('concurrency: 1 keeps the serial dispatch', async () => {
const tracker = installTrackingTransport();
const result = await embedBatch(texts, { concurrency: 1, onBatchComplete: () => {} });
expect(result).toHaveLength(250);
expect(tracker.maxInFlight()).toBe(1);
});
test('GBRAIN_EMBED_BATCH_CONCURRENCY env bounds the pool when option unset', async () => {
const tracker = installTrackingTransport();
await withEnv({ GBRAIN_EMBED_BATCH_CONCURRENCY: '2' }, async () => {
await embedBatch(texts, { onBatchComplete: () => {} });
});
expect(tracker.maxInFlight()).toBeGreaterThan(1);
expect(tracker.maxInFlight()).toBeLessThanOrEqual(2);
});
test('onBatchComplete reports a monotonic count ending at total', async () => {
installTrackingTransport();
const seen: number[] = [];
await embedBatch(texts, {
onBatchComplete: (done, total) => {
expect(total).toBe(250);
seen.push(done);
},
});
expect(seen).toHaveLength(3); // 100 + 100 + 50 sub-batches
for (let i = 1; i < seen.length; i++) {
expect(seen[i]).toBeGreaterThan(seen[i - 1]);
}
expect(seen[seen.length - 1]).toBe(250);
});
test('a failing sub-batch rejects the whole call', async () => {
let call = 0;
__setEmbedTransportForTests((async ({ values }: { values: string[] }) => {
call++;
if (call === 2) throw new Error('boom');
await new Promise(r => setTimeout(r, 2));
return { embeddings: values.map(() => Array.from({ length: DIMS }, () => 0.1)) };
}) as any);
await expect(embedBatch(texts, { onBatchComplete: () => {} })).rejects.toThrow();
});
test('after a failure, surviving workers stop dispatching new slices', async () => {
// 1000 texts → 10 slices, concurrency 2. First call fails immediately;
// without the `failed` flag the second worker would keep draining all
// 10 slices in the background AFTER embedBatch already rejected —
// burning provider spend and firing onBatchComplete post-rejection.
let calls = 0;
const completions: number[] = [];
__setEmbedTransportForTests((async ({ values }: { values: string[] }) => {
calls++;
if (calls === 1) throw new Error('boom');
await new Promise(r => setTimeout(r, 5));
return { embeddings: values.map(() => Array.from({ length: DIMS }, () => 0.1)) };
}) as any);
const many = Array.from({ length: 1000 }, (_, i) => `t${i}`);
await expect(
embedBatch(many, { concurrency: 2, onBatchComplete: d => completions.push(d) }),
).rejects.toThrow('boom');
const callsAtRejection = calls;
await new Promise(r => setTimeout(r, 50)); // would-be background drain window
expect(calls).toBe(callsAtRejection); // no new dispatch after rejection
expect(calls).toBeLessThanOrEqual(2); // only the in-flight sibling ran
expect(completions).toHaveLength(0); // no progress reported after failure
});
test('single small batch without callback stays on the one-call fast path', async () => {
const tracker = installTrackingTransport(1);
const result = await embedBatch(['t0', 't1', 't2']);
expect(result).toHaveLength(3);
expect(result[1][0]).toBe(1);
expect(tracker.maxInFlight()).toBe(1);
});
});
+3 -27
View File
@@ -19,7 +19,7 @@
* overwrites this preload.
*/
import { configureGateway, getEmbeddingDimensions } from '../../src/core/ai/gateway.ts';
import { afterEach, beforeEach } from 'bun:test';
import { beforeEach } from 'bun:test';
const LEGACY_CONFIG = {
embedding_model: 'openai:text-embedding-3-large',
@@ -52,7 +52,7 @@ applyLegacy();
// 2. file-local beforeAll → may overwrite to ZE/1280
// Since beforeAll runs once per file BEFORE the first beforeEach,
// file-local beforeAll wins for that file's tests. ✓
function applyLegacyIfEmpty() {
beforeEach(() => {
try {
// Only re-apply if the gateway was reset (or never configured).
// Tests that explicitly configured a different model in their
@@ -62,28 +62,4 @@ function applyLegacyIfEmpty() {
} catch {
applyLegacy();
}
}
beforeEach(applyLegacyIfEmpty);
// PR #3130 shard-order fix: beforeEach alone leaves ONE window open — a file
// whose LAST afterEach calls resetGateway() poisons the NEXT file's
// beforeAll, which runs BEFORE any beforeEach fires. A beforeAll there that
// does engine.initSchema() then sizes the embedding column from the gateway
// DEFAULTS (zembed-1/1280d) instead of the pinned legacy 1536, and every
// 1536-d Float32Array fixture in that file dies with
// "expected 1280 dimensions, not 1536". Which file pair collides is a
// function of shard composition, so adding/removing ANY test file can
// surface it (that is exactly how it bit shard 9).
//
// Preload hooks are registered before any file-local hooks, and bun runs
// after-hooks inside-out (file-local afterEach first, then this one), so
// this repairs the empty slot immediately after the poisoning reset —
// before the next file's beforeAll can observe it.
//
// Known remaining window: a file whose afterAll() resets the gateway (no
// hook runs between its afterAll and the next file's beforeAll). Files
// that reset in afterAll and can precede a schema-creating file should
// re-apply their own config, or the victim file should configureGateway()
// explicitly in its beforeAll.
afterEach(applyLegacyIfEmpty);
});
-69
View File
@@ -1,69 +0,0 @@
/**
* #1207: `gbrain import` without `--workers` used to hardcode workerCount=1,
* so a large Postgres import paid one serial embedding round-trip per file.
* runImport now routes the default through the shared autoConcurrency policy
* (PGLite 1, >100 files on Postgres DEFAULT_PARALLEL_WORKERS), while an
* explicit `--workers N` still wins.
*
* The engine here is a minimal postgres-kind stub with no database_url in
* config runImport's parallel branch then falls back to serial processing
* (its PR #490 guard) but the WORKER-COUNT DECISION (the thing #1207 fixes)
* is still observable via the "Using N parallel workers" log line. Per-file
* imports fail against the stub engine and are swallowed by runImport's
* per-file catch; that's fine this test pins the policy, not the import.
*/
import { afterEach, beforeEach, describe, expect, test } from 'bun:test';
import { mkdtempSync, writeFileSync, mkdirSync, rmSync, realpathSync } from 'fs';
import { tmpdir } from 'os';
import { join } from 'path';
import { withEnv } from './helpers/with-env.ts';
import { runImport } from '../src/commands/import.ts';
const fakePostgresEngine = {
kind: 'postgres',
executeRaw: async () => [],
logIngest: async () => {},
setConfig: async () => {},
getConfig: async () => null,
} as any;
let workspace: string;
let brainDir: string;
let logs: string[];
const realLog = console.log;
beforeEach(() => {
workspace = mkdtempSync(join(tmpdir(), 'gbrain-import-workers-home-'));
mkdirSync(join(workspace, '.gbrain'), { recursive: true });
brainDir = realpathSync(mkdtempSync(join(tmpdir(), 'gbrain-import-workers-brain-')));
// 101 files: one past AUTO_CONCURRENCY_FILE_THRESHOLD (100).
for (let i = 0; i < 101; i++) {
writeFileSync(join(brainDir, `page-${i}.md`), `# Page ${i}\n\nbody ${i}\n`);
}
logs = [];
console.log = (msg?: unknown) => logs.push(String(msg));
});
afterEach(() => {
console.log = realLog;
rmSync(workspace, { recursive: true, force: true });
rmSync(brainDir, { recursive: true, force: true });
});
describe('import default worker count (#1207)', () => {
test('no --workers flag → autoConcurrency picks 4 for >100 files on Postgres', async () => {
await withEnv({ GBRAIN_HOME: join(workspace, '.gbrain'), GBRAIN_SOURCE: undefined }, async () => {
await runImport(fakePostgresEngine, [brainDir, '--no-embed'], { sourceId: 'default' });
});
expect(logs.some(l => l.includes('Using 4 parallel workers'))).toBe(true);
});
test('explicit --workers 2 still wins over the auto policy', async () => {
await withEnv({ GBRAIN_HOME: join(workspace, '.gbrain'), GBRAIN_SOURCE: undefined }, async () => {
await runImport(fakePostgresEngine, [brainDir, '--no-embed', '--workers', '2'], { sourceId: 'default' });
});
expect(logs.some(l => l.includes('Using 2 parallel workers'))).toBe(true);
expect(logs.some(l => l.includes('Using 4 parallel workers'))).toBe(false);
});
});