Compare commits

...
Author SHA1 Message Date
Garry TanandClaude Fable 5 60fb33c0d9 fix(embed): stop worker pool from dispatching new slices after a sub-batch failure
Review finding on #3130: when one sub-batch rejected, the surviving pool
workers kept draining ALL remaining slices in the background after
embedBatch had already rejected — real provider spend post-failure,
onBatchComplete firing after the caller handled the error, and stacked
429 pressure when embedBatchWithBackoff retried while the failed run was
still draining. A shared failed flag now stops further dispatch (in-flight
sibling calls still settle, bounded by concurrency-1) and suppresses
post-failure progress callbacks. Pinned by a new test: 10 slices /
concurrency 2 / first call fails → no calls after rejection, no
completions reported.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
2026-07-22 11:52:30 -07:00
Garry TanandClaude Fable 5 11ed0871c2 test: fix CI red on #3130 — withEnv for batch-concurrency env + close resetGateway shard-order poison window
Two real failures surfaced by this PR's re-sharding:

1. verify/check-test-isolation: embed-batch-concurrency.test.ts mutated
   process.env directly (R1). Now uses withEnv().

2. test (9) source-health "expected 1280 dimensions, not 1536": a file
   whose last afterEach calls resetGateway() leaves the gateway slot
   empty during the NEXT file's beforeAll (which runs before any
   beforeEach can restore the legacy 1536 pin), so initSchema() sizes
   the embedding column from the zembed-1/1280 defaults and every
   1536-d fixture in that file fails. Which pair collides depends on
   shard composition, so adding test files (as this PR does) can
   surface it anywhere. The legacy-embedding preload now also repairs
   the empty slot in a global afterEach (preload after-hooks run after
   file-local ones), closing the window at the root instead of
   patching one victim file.

Reproduced locally with a poison/afterEach-reset file followed by a
schema-creating file: embedding column typmod 1280 before the fix,
1536 after. check-test-isolation, typecheck, and the affected suites
all pass.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
2026-07-22 11:08:07 -07:00
Garry TanandClaude Fable 5 595eeb7d6f fix(embed): per-request batch caps (google/dashscope) + parallel batch dispatch (#970 #1199 #1207 #1818)
Four embedding-throughput/correctness fixes:

- #970: google recipe now declares max_batch_tokens (204,800 — derived
  from Gemini's real limits: 100 inputs per batchEmbedContents × 2048
  tokens per input) + max_batch_count 100 + chars_per_token, silencing
  the missing-cap startup warning and enabling the gateway pre-split.
  Deliberately NOT the 2048 per-input limit, which would over-split 50x.

- #1199: new optional EmbeddingTouchpoint.max_batch_count enforced in
  splitByTokenBudget (flush at N inputs even when the token budget has
  room); dashscope sets 10 (provider hard-caps embeddings at 10 inputs
  per request). isTokenLimitError also learns DashScope's
  "batch size is invalid" message so recursive halving backstops it.

- #1207: gbrain import without --workers now resolves through the shared
  autoConcurrency policy (PGLite → 1, >100 files on Postgres → 4)
  instead of hardcoding serial; explicit --workers still wins.

- #1818: embedBatch dispatches its 100-input sub-batches through a
  bounded worker pool (default 4; EmbedBatchOptions.concurrency /
  GBRAIN_EMBED_BATCH_CONCURRENCY override) with index-addressed results
  so output order is preserved; single-batch fast path unchanged.

Also: listRecipes() now reads the exported RECIPES map instead of the
private ALL array (one source of truth; lets tests inject a synthetic
capless recipe to keep the startup-warning path covered now that every
real recipe declares a cap).

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
2026-07-21 14:39:33 -07:00
13 changed files with 501 additions and 39 deletions
+10 -5
View File
@@ -170,10 +170,14 @@ 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 } = await import('../core/sync-concurrency.ts');
let workerCount: number;
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;
try {
workerCount = parseWorkers(workersArg ?? undefined) ?? 1;
workerCount = parseWorkers(workersArg ?? undefined);
} catch (e) {
console.error(e instanceof Error ? e.message : String(e));
process.exit(1);
@@ -252,8 +256,9 @@ export async function runImport(
}
const files = resumeFilter(allFiles, dir, completed);
// Determine actual worker count
const actualWorkers = workerCount > 1 ? workerCount : 1;
// 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);
if (actualWorkers > 1) {
console.log(`Using ${actualWorkers} parallel workers`);
}
+24 -6
View File
@@ -1513,12 +1513,21 @@ 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. 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)
// 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,
)
: [truncated];
const allEmbeddings: Float32Array[] = [];
@@ -1568,6 +1577,9 @@ 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.
*/
@@ -1575,15 +1587,17 @@ 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) {
if (current.length > 0 && (currentTokens + estTokens > budgetTokens || current.length >= maxCount)) {
batches.push(current);
current = [];
currentTokens = 0;
@@ -1609,7 +1623,11 @@ 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)
/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)
);
}
+4
View File
@@ -31,6 +31,10 @@ 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,6 +16,15 @@ 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'],
+4 -1
View File
@@ -58,5 +58,8 @@ export function getRecipe(id: string): Recipe | undefined {
}
export function listRecipes(): Recipe[] {
return [...ALL];
// 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()];
}
+10
View File
@@ -46,6 +46,16 @@ 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
+56 -6
View File
@@ -79,15 +79,34 @@ 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
* is purely about progress-callback granularity.
* owns progress-callback granularity and (#1818) bounded parallel dispatch
* of the sub-batches the embed-stale.ts worker-pool pattern, scoped down.
*/
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 = {},
@@ -103,13 +122,44 @@ export async function embedBatch(
if (texts.length <= BATCH_SIZE && !options.onBatchComplete) {
return gatewayEmbed(texts, gwOpts);
}
const results: Float32Array[] = [];
// #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[] }> = [];
for (let i = 0; i < texts.length; 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);
slices.push({ start: i, texts: texts.slice(i, i + BATCH_SIZE) });
}
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;
}
+109 -5
View File
@@ -39,6 +39,8 @@ 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,
@@ -93,6 +95,14 @@ 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)', () => {
@@ -149,6 +159,27 @@ 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)', () => {
@@ -179,6 +210,12 @@ 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);
@@ -387,26 +424,92 @@ 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).
configureGoogle();
configureCapless();
expect(warnings.length).toBe(firstCallCount);
} finally {
console.warn = original;
RECIPES.delete(caplessRecipe.id);
}
// The warning text should match the documented contract.
@@ -415,11 +518,12 @@ describe('startup warning for recipes missing max_batch_tokens', () => {
);
expect(contractMatch.length).toBe(1);
// Voyage declares max_batch_tokens → suppressed. OpenAI is the
// canonical fast-path recipe → also suppressed by id. Both must be
// absent from the warnings.
// 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.
expect(warnings.find(w => w.includes('"voyage"'))).toBeUndefined();
expect(warnings.find(w => w.includes('"openai"'))).toBeUndefined();
expect(warnings.find(w => w.includes('"google"'))).toBeDefined();
expect(warnings.find(w => w.includes('"google"'))).toBeUndefined();
expect(warnings.find(w => w.includes('"capless-test"'))).toBeDefined();
});
});
+13 -13
View File
@@ -52,16 +52,7 @@ describe('v0.32 #779: no_batch_cap suppresses the missing-max_batch_tokens warni
}
});
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);
test('configureGateway does NOT warn for google now that it declares batch caps (#970)', () => {
warnSpy.mockClear();
resetGateway();
configureGateway({
@@ -69,11 +60,20 @@ describe('v0.32 #779: no_batch_cap suppresses the missing-max_batch_tokens warni
embedding_dimensions: 768,
env: { GOOGLE_GENERATIVE_AI_API_KEY: 'fake' },
});
messages = warnSpy.mock.calls.map(c => String(c[0] ?? ''));
const messages = warnSpy.mock.calls.map(c => String(c[0] ?? ''));
expect(
messages.some(m => m.includes('"google"') && m.includes('without max_batch_tokens')),
'google should warn when configured because it has fixed-cap models',
).toBe(true);
'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);
});
test('every recipe with empty models[] declares user_provided_models OR has openai-fast-path', () => {
+5
View File
@@ -55,6 +55,11 @@ 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
+161
View File
@@ -0,0 +1,161 @@
/**
* #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);
});
});
+27 -3
View File
@@ -19,7 +19,7 @@
* overwrites this preload.
*/
import { configureGateway, getEmbeddingDimensions } from '../../src/core/ai/gateway.ts';
import { beforeEach } from 'bun:test';
import { afterEach, 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. ✓
beforeEach(() => {
function applyLegacyIfEmpty() {
try {
// Only re-apply if the gateway was reset (or never configured).
// Tests that explicitly configured a different model in their
@@ -62,4 +62,28 @@ beforeEach(() => {
} 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
@@ -0,0 +1,69 @@
/**
* #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);
});
});