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
25 changed files with 550 additions and 387 deletions
+1 -12
View File
@@ -54,7 +54,7 @@ export function bigintToStringReplacer(_key: string, value: unknown): unknown {
}
// CLI-only commands that bypass the operation layer
export const CLI_ONLY = new Set(['init', 'reinit-pglite', 'upgrade', 'post-upgrade', 'check-update', 'integrations', 'publish', 'check-backlinks', 'lint', 'report', 'import', 'export', 'files', 'embed', 'serve', 'call', 'config', 'doctor', 'migrate', 'eval', 'sync', 'extract', 'extract-conversation-facts', 'enrich', 'features', 'autopilot', 'graph-query', 'jobs', 'agent', 'apply-migrations', 'skillpack-check', 'skillpack', 'resolvers', 'integrity', 'repair-jsonb', 'orphans', 'sources', 'mounts', 'dream', 'check-resolvable', 'routing-eval', 'skillify', 'smoke-test', 'providers', 'storage', 'repos', 'code-def', 'code-refs', 'reindex', 'reindex-code', 'reindex-frontmatter', 'code-callers', 'code-callees', 'reconcile-links', 'frontmatter', 'auth', 'friction', 'claw-test', 'book-mirror', 'takes', 'think', 'salience', 'anomalies', 'calibration', 'transcripts', 'models', 'remote', 'recall', 'forget', 'edges-backfill', 'facts', 'cache', 'ze-switch', 'founder', 'brainstorm', 'lsd', 'schema', 'capture', 'onboard', 'conversation-parser', 'status', 'connect', 'skillopt', 'quarantine', 'self-upgrade', 'advisor', 'watch', 'reindex-search-vector']);
export const CLI_ONLY = new Set(['init', 'reinit-pglite', 'upgrade', 'post-upgrade', 'check-update', 'integrations', 'publish', 'check-backlinks', 'lint', 'report', 'import', 'export', 'files', 'embed', 'serve', 'call', 'config', 'doctor', 'migrate', 'eval', 'sync', 'extract', 'extract-conversation-facts', 'enrich', 'features', 'autopilot', 'graph-query', 'jobs', 'agent', 'apply-migrations', 'skillpack-check', 'skillpack', 'resolvers', 'integrity', 'repair-jsonb', 'orphans', 'sources', 'mounts', 'dream', 'check-resolvable', 'routing-eval', 'skillify', 'smoke-test', 'providers', 'storage', 'repos', 'code-def', 'code-refs', 'reindex', 'reindex-code', 'reindex-frontmatter', 'code-callers', 'code-callees', 'reconcile-links', 'frontmatter', 'auth', 'friction', 'claw-test', 'book-mirror', 'takes', 'think', 'salience', 'anomalies', 'calibration', 'transcripts', 'models', 'remote', 'recall', 'forget', 'edges-backfill', 'cache', 'ze-switch', 'founder', 'brainstorm', 'lsd', 'schema', 'capture', 'onboard', 'conversation-parser', 'status', 'connect', 'skillopt', 'quarantine', 'self-upgrade', 'advisor', 'watch', 'reindex-search-vector']);
// CLI-only commands whose handlers print their own --help text. These are
// excluded from the generic short-circuit so detailed per-command and
// per-subcommand usage stays reachable.
@@ -991,8 +991,6 @@ const THIN_CLIENT_REFUSED_COMMANDS = new Set([
// hint pointing at the routable MCP tools; per-subcommand splits are
// a v0.31.x follow-up TODO.
'takes', 'sources',
// #1867: fence-backfill edits local .md fences + stamps the local DB.
'facts',
// v0.32 thin-client routing audit (Codex round 2 findings #2, #4):
// - `pages` purge-deleted is admin+localOnly (operations.ts:856-864)
// - `files` list / file_url MCP ops are localOnly (operations.ts:1769-1879)
@@ -1028,7 +1026,6 @@ const THIN_CLIENT_REFUSE_HINTS: Record<string, string> = {
migrate: "migrate runs on the host's local engine. Run on the host machine.",
'apply-migrations': 'schema migrations run on the host. SSH and run there.',
'repair-jsonb': 'repair-jsonb operates on the local DB only.',
facts: 'facts fence-backfill edits local entity-page fences. Run on the host machine.',
integrity: 'integrity scans local files. Run on the host machine.',
serve: 'serve starts a server. Run on the host, not the thin client.',
dream: 'dream runs the autopilot cycle on the host. `gbrain remote ping` queues one. (Native `gbrain dream` thin-client routing planned for v0.31.2.)',
@@ -1858,13 +1855,6 @@ async function handleCliOnly(command: string, args: string[]) {
await runEdgesBackfill(engine, args);
break;
}
case 'facts': {
// #1867 — re-runnable fence-backfill for row_num-NULL legacy fact
// rows (idempotent v0_32_2 phase B, exposed as an operator command).
const { runFactsCommand } = await import('./commands/facts.ts');
await runFactsCommand(engine, args);
break;
}
case 'whoknows': {
// v0.33 (Issue #?): expertise + relationship-proximity routing.
// MCP op `find_experts` (read-scoped) backs the same code path; CLI
@@ -2345,7 +2335,6 @@ TOOLS
check-backlinks <check|fix> [dir] Find/fix missing back-links across brain
lint <dir|file> [--fix] Catch LLM artifacts, placeholder dates, bad frontmatter
orphans [--json] [--count] Find pages with no inbound wikilinks
facts fence-backfill [--dry-run] Fence legacy fact rows (row_num NULL) onto entity pages
salience [--days N] [--kind P] v0.29: pages ranked by emotional + activity salience
anomalies [--since D] [--sigma N] v0.29: cohort-based statistical anomalies (tag, type)
transcripts recent [--days N] v0.29: recent raw .txt transcripts (local-only)
-48
View File
@@ -1,48 +0,0 @@
/**
* gbrain facts — fact-store maintenance surface (#1867).
*
* `fence-backfill` re-runs the v0_32_2 fence-backfill phase on demand.
* Remote `extract_facts` deposits that predate the fence-write backstop
* (and any legacy DB-only insert) leave `row_num IS NULL` rows that the
* cycle extract_facts guard refuses to reconcile past — previously the
* only remedy was the one-shot v0_32_2 migration, which the ledger marks
* complete and never re-runs. The phase is idempotent (only touches
* `row_num IS NULL` rows), so exposing it as a command is safe to re-run
* any time the backlog reappears.
*/
import type { BrainEngine } from '../core/engine.ts';
import { setCliExitVerdict } from '../core/cli-force-exit.ts';
import { phaseBFenceFacts } from './migrations/v0_32_2.ts';
function printHelp(): void {
process.stderr.write(
`Usage: gbrain facts fence-backfill [--dry-run]\n\n` +
`Fence-backfill: appends every legacy fact row (row_num IS NULL) to its\n` +
`entity page's \`## Facts\` fence and stamps row_num + source_markdown_slug\n` +
`back onto the DB row. Idempotent — re-runs only pick up rows still\n` +
`missing a fence assignment. Clears the backlog that makes the cycle's\n` +
`extract_facts phase skip fence→DB reconciliation.\n\n` +
` --dry-run report what would be fenced; no FS or DB writes\n`,
);
}
export async function runFactsCommand(engine: BrainEngine, args: string[]): Promise<void> {
const sub = args[0];
if (!sub || sub === '--help' || sub === '-h') {
printHelp();
return;
}
if (sub !== 'fence-backfill') {
process.stderr.write(`Unknown facts subcommand: ${sub}\n`);
printHelp();
setCliExitVerdict(1);
return;
}
const dryRun = args.includes('--dry-run');
const result = await phaseBFenceFacts(engine, { dryRun });
process.stderr.write(
`fence-backfill: ${result.status}${result.detail ? `${result.detail}` : ''}\n`,
);
if (result.status === 'failed') setCliExitVerdict(1);
}
+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`);
}
+4 -32
View File
@@ -39,7 +39,6 @@ import type { BrainEngine } from '../../core/engine.ts';
import { loadConfig, toEngineConfig } from '../../core/config.ts';
import { createEngine } from '../../core/engine-factory.ts';
import { upsertFactRow, parseFactsFence } from '../../core/facts-fence.ts';
import { resolvePageFilePath } from '../../core/markdown.ts';
let testEngineOverride: BrainEngine | null = null;
export function __setTestEngineOverride(engine: BrainEngine | null): void {
@@ -149,16 +148,9 @@ function isLocalPathDirty(localPath: string): boolean {
}
}
/**
* Exported (not just via `__testing`) because `gbrain facts fence-backfill`
* (#1867) re-runs this phase on demand: remote `extract_facts` deposits that
* predate the fence-write backstop leave row_num-NULL rows the cycle guard
* refuses to reconcile past. The phase is idempotent (only touches
* `row_num IS NULL` rows), so re-running is always safe.
*/
export async function phaseBFenceFacts(
async function phaseBFenceFacts(
engine: BrainEngine | null,
opts: Pick<OrchestratorOpts, 'dryRun'>,
opts: OrchestratorOpts,
): Promise<OrchestratorPhaseResult> {
if (opts.dryRun) {
// Dry-run: report what WOULD happen without touching FS or DB.
@@ -246,11 +238,7 @@ export async function phaseBFenceFacts(
for (const [key, group] of groups) {
const [sourceId, entitySlug] = key.split('\0');
const localPath = localPathById.get(sourceId)!;
// resolvePageFilePath, NOT a bare join — non-default sources fence
// into `<local_path>/.sources/<id>/<slug>.md`, the same path the
// fence-write backstop and put_page write-through compute. A bare
// join here diverges fence and DB for non-default sources (#2044).
const filePath = resolvePageFilePath(localPath, entitySlug, sourceId);
const filePath = join(localPath, `${entitySlug}.md`);
const tmpPath = `${filePath}.tmp`;
try {
@@ -281,21 +269,6 @@ export async function phaseBFenceFacts(
const existingFence = parseFactsFence(body);
const existingKeySet = new Set(existingFence.facts.map(f => `${f.claim}\0${f.source ?? ''}`));
// Seed appended row_nums from MAX(fence max, DB max) — same #2044
// divergence guard as writeFactsToFence. When the fence at the
// resolved path is missing/behind but the DB already holds stamped
// rows for this slug (legacy wrong-path fence writes), fence-max+1
// collides with idx_facts_fence_key on the post-rename UPDATE,
// failing the page and leaving fence and DB disagreeing.
const dbMaxRows = await engine.executeRaw<{ max: number | string | null }>(
`SELECT MAX(row_num) AS max FROM facts
WHERE source_id = $1 AND source_markdown_slug = $2`,
[sourceId, entitySlug],
);
const dbMaxRowNum = Number(dbMaxRows[0]?.max ?? 0) || 0;
const fenceMaxRowNum = existingFence.facts.reduce((m, f) => Math.max(m, f.rowNum), 0);
let nextRowNum = Math.max(fenceMaxRowNum, dbMaxRowNum) + 1;
const assignments: Array<{ id: string; row_num: number }> = [];
for (const row of group) {
const key = `${row.fact}\0${row.source ?? ''}`;
@@ -318,7 +291,6 @@ export async function phaseBFenceFacts(
.toISOString().slice(0, 10)
: undefined;
const { body: updated, rowNum } = upsertFactRow(body, {
rowNum: nextRowNum++,
claim: row.fact,
kind: row.kind,
confidence: row.confidence,
@@ -409,7 +381,7 @@ async function phaseCVerify(
for (const g of groups) {
const localPath = localPathById.get(g.source_id);
if (!localPath) continue;
const filePath = resolvePageFilePath(localPath, g.source_markdown_slug, g.source_id);
const filePath = join(localPath, `${g.source_markdown_slug}.md`);
if (!existsSync(filePath)) {
mismatches.push(`${g.source_markdown_slug} (file missing)`);
continue;
+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
+3 -5
View File
@@ -176,11 +176,9 @@ export async function runExtractFacts(
if (legacyCount > 0) {
result.guardTriggered = true;
result.warnings.push(
`extract_facts: ${legacyCount} legacy fact rows pending fence backfill ` +
`(row_num IS NULL — v0.31 rows or remote extract_facts deposits that ` +
`predate the fence backstop). Run \`gbrain facts fence-backfill\` ` +
`(idempotent, re-runnable) before this phase can safely reconcile ` +
`fence → DB.`,
`extract_facts: ${legacyCount} legacy v0.31 fact rows pending fence backfill. ` +
`Run \`gbrain apply-migrations --yes\` to complete v0_32_2 before this phase ` +
`can safely reconcile fence → DB.`,
);
return result;
}
+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;
}
+6 -7
View File
@@ -1739,14 +1739,13 @@ export interface BrainEngine {
* single-row supersede flow because fence reconciliation is the canonical
* source-of-truth direction, not the consolidator path.
*
* Insertion runs in a single transaction. A collision on the v51
* partial UNIQUE index `(source_id, source_markdown_slug, row_num)`
* skips ONLY that row (ON CONFLICT DO NOTHING, #2044) — the rest of
* the batch still commits, so a redundant deposit against an
* already-indexed fence row is idempotent instead of a hard failure.
* Insertion is atomic per call: all rows commit in a single transaction
* or none commit (the transaction rolls back on any constraint
* violation, e.g. the v51 partial UNIQUE index on
* `(source_id, source_markdown_slug, row_num)`).
*
* Returns the inserted ids in input-order (colliding rows omitted) so
* callers can correlate fence-row → DB-id without a separate lookup.
* Returns the inserted ids in input-order so callers can correlate
* fence-row → DB-id without a separate lookup.
*/
insertFacts(
rows: Array<NewFact & { row_num: number; source_markdown_slug: string }>,
+1 -17
View File
@@ -218,27 +218,11 @@ export async function writeFactsToFence(
}
// 2. Upsert each fact onto the fence in input order. row_num
// monotonically increases, append-only, seeded from the MAX of
// the fence and the DB index (#2044): when fence and DB have
// diverged (e.g. legacy writes that stamped DB rows against a
// fence at a path this code no longer reads), fence-max+1 can
// collide with an existing DB row_num, tripping
// idx_facts_fence_key and rolling back the whole batch.
const dbMaxRows = await engine.executeRaw<{ max: number | string | null }>(
`SELECT MAX(row_num) AS max FROM facts
WHERE source_id = $1 AND source_markdown_slug = $2`,
[target.sourceId, target.slug],
);
const dbMaxRowNum = Number(dbMaxRows[0]?.max ?? 0) || 0;
const fenceMaxRowNum = parseFactsFence(body).facts
.reduce((m, f) => Math.max(m, f.rowNum), 0);
let nextRowNum = Math.max(fenceMaxRowNum, dbMaxRowNum) + 1;
// monotonically increases (max-existing + 1 per call, append-only).
const assignedRowNums: number[] = [];
for (const f of facts) {
const validFromStr = (f.validFrom ?? new Date()).toISOString().slice(0, 10);
const { body: updated, rowNum } = upsertFactRow(body, {
rowNum: nextRowNum++,
claim: f.fact,
kind: (f.kind ?? 'fact') as 'fact' | 'event' | 'preference' | 'commitment' | 'belief',
confidence: f.confidence ?? 1.0,
+8 -18
View File
@@ -4102,13 +4102,11 @@ export class PGLiteEngine implements BrainEngine {
): Promise<{ inserted: number; ids: number[] }> {
if (rows.length === 0) return { inserted: 0, ids: [] };
// Single transaction; per-row INSERTs (not multi-row VALUES) keep the
// embedding-vs-no-embedding branching readable; batch sizes are small
// (5-30 rows per page in practice) so the loop overhead is negligible
// vs the embedding compute cost. #2044: ON CONFLICT DO NOTHING on the
// v51 partial UNIQUE index makes a residual fence/DB row_num collision
// skip that row instead of rolling back the whole batch (parity with
// postgres-engine.ts).
// Single transaction so the v51 partial UNIQUE index can roll back the
// whole batch on constraint violation. Per-row INSERTs (not multi-row
// VALUES) keep the embedding-vs-no-embedding branching readable; batch
// sizes are small (5-30 rows per page in practice) so the loop overhead
// is negligible vs the embedding compute cost.
const ids = await this.db.transaction(async (tx) => {
const out: number[] = [];
for (const input of rows) {
@@ -4151,11 +4149,7 @@ export class PGLiteEngine implements BrainEngine {
$14, $15,
$16, $17, $18, $19,
$20
)
ON CONFLICT (source_id, source_markdown_slug, row_num)
WHERE row_num IS NOT NULL
DO NOTHING
RETURNING id`
) RETURNING id`
: `INSERT INTO facts (
source_id, entity_slug, fact, kind, visibility, notability, context,
valid_from, valid_until, source, source_session, confidence,
@@ -4169,16 +4163,12 @@ export class PGLiteEngine implements BrainEngine {
$15, $16,
$17, $18, $19, $20,
$21
)
ON CONFLICT (source_id, source_markdown_slug, row_num)
WHERE row_num IS NOT NULL
DO NOTHING
RETURNING id`,
) RETURNING id`,
embedStr === null
? [ctx.source_id, entitySlug, input.fact, kind, visibility, notability, context, validFrom, validUntil, input.source, sourceSession, confidence, embeddedAt, input.row_num, input.source_markdown_slug, claimMetric, claimValue, claimUnit, claimPeriod, eventType]
: [ctx.source_id, entitySlug, input.fact, kind, visibility, notability, context, validFrom, validUntil, input.source, sourceSession, confidence, embedStr, embeddedAt, input.row_num, input.source_markdown_slug, claimMetric, claimValue, claimUnit, claimPeriod, eventType],
);
if (ins.rows[0]) out.push(ins.rows[0].id);
out.push(ins.rows[0].id);
}
return out;
});
+6 -12
View File
@@ -4302,12 +4302,10 @@ export class PostgresEngine implements BrainEngine {
// ONCE per process so the cast matches the actual column type
// (halfvec vs vector). The probe is cached after first call.
const castSuffix = await this.resolveFactsEmbeddingCast();
// Single transaction; per-row INSERTs (not multi-row VALUES) keep the
// embedding-vs-no-embedding branching readable; batch sizes are small
// (5-30 rows per page in practice). #2044: ON CONFLICT DO NOTHING on
// the v51 partial UNIQUE index makes a residual fence/DB row_num
// collision skip that row instead of rolling back the whole batch —
// the fence stays system-of-record and reconciliation catches up.
// Single transaction so the v51 partial UNIQUE index can roll back
// the whole batch on constraint violation. Per-row INSERTs (not
// multi-row VALUES) keep the embedding-vs-no-embedding branching
// readable; batch sizes are small (5-30 rows per page in practice).
// No supersede flow in this path — fence reconciliation is the
// canonical source-of-truth direction, not the consolidator path.
const ids = await sql.begin(async (tx) => {
@@ -4348,13 +4346,9 @@ export class PostgresEngine implements BrainEngine {
${input.row_num}, ${input.source_markdown_slug},
${claimMetric}, ${claimValue}, ${claimUnit}, ${claimPeriod},
${eventType}
)
ON CONFLICT (source_id, source_markdown_slug, row_num)
WHERE row_num IS NOT NULL
DO NOTHING
RETURNING id
) RETURNING id
`;
if (ins[0]) out.push(Number(ins[0].id));
out.push(Number(ins[0].id));
}
return out;
});
+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);
});
});
+1 -3
View File
@@ -304,9 +304,7 @@ describe('runExtractFacts — empty-fence guard (Codex R2-#7)', () => {
expect(r.legacyRowsPending).toBe(1);
expect(r.factsInserted).toBe(0);
expect(r.factsDeleted).toBe(0);
// #1867: the remedy hint points at the re-runnable backfill command,
// not the one-shot v0_32_2 migration (which the ledger never re-runs).
expect(r.warnings.some(w => w.includes('gbrain facts fence-backfill'))).toBe(true);
expect(r.warnings.some(w => w.includes('apply-migrations'))).toBe(true);
// Legacy row was NOT touched.
// eslint-disable-next-line @typescript-eslint/no-explicit-any
-133
View File
@@ -1,133 +0,0 @@
/**
* #1867 — `gbrain facts fence-backfill` command tests.
*
* The command re-runs the (idempotent) v0_32_2 phase B on demand so
* row_num-NULL backlogs — remote extract_facts deposits that predate
* the fence-write backstop — can be cleared without re-running the
* one-shot migration. Real PGLite + real tempdir filesystem.
*/
import { describe, test, expect, beforeAll, afterAll, beforeEach } from 'bun:test';
import { mkdtempSync, rmSync, existsSync, readFileSync } from 'node:fs';
import { tmpdir } from 'node:os';
import { join } from 'node:path';
import { PGLiteEngine } from '../src/core/pglite-engine.ts';
import { runFactsCommand } from '../src/commands/facts.ts';
import { phaseBFenceFacts } from '../src/commands/migrations/v0_32_2.ts';
let engine: PGLiteEngine;
let brainDir: string;
beforeAll(async () => {
engine = new PGLiteEngine();
await engine.connect({});
await engine.initSchema();
});
afterAll(async () => {
await engine.disconnect();
try {
if (brainDir) rmSync(brainDir, { recursive: true, force: true });
} catch { /* best-effort */ }
});
beforeEach(async () => {
brainDir = mkdtempSync(join(tmpdir(), 'facts-backfill-cmd-test-'));
// eslint-disable-next-line @typescript-eslint/no-explicit-any
await (engine as any).db.query('DELETE FROM facts');
// eslint-disable-next-line @typescript-eslint/no-explicit-any
await (engine as any).db.query(
`UPDATE sources SET local_path = $1 WHERE id = 'default'`,
[brainDir],
);
});
async function seedLegacyFact(fact: string): Promise<void> {
// The row_num-NULL shape a remote extract_facts deposit leaves behind
// when it lands via the legacy DB-only path.
// eslint-disable-next-line @typescript-eslint/no-explicit-any
await (engine as any).db.query(
`INSERT INTO facts (source_id, entity_slug, fact, kind, visibility, notability,
valid_from, source, confidence)
VALUES ('default', 'people/alice', $1, 'fact', 'private', 'medium',
now(), 'mcp:extract_facts', 1.0)`,
[fact],
);
}
describe('gbrain facts fence-backfill', () => {
test('fences row_num-NULL rows and stamps the DB', async () => {
await seedLegacyFact('Deposited remotely');
await runFactsCommand(engine, ['fence-backfill']);
// The fence exists on disk with the claim.
const filePath = join(brainDir, 'people/alice.md');
expect(existsSync(filePath)).toBe(true);
expect(readFileSync(filePath, 'utf-8')).toContain('Deposited remotely');
// The backlog is cleared: no row_num-NULL rows remain.
// eslint-disable-next-line @typescript-eslint/no-explicit-any
const rows = await (engine as any).db.query(
'SELECT row_num, source_markdown_slug FROM facts',
);
expect(rows.rows).toHaveLength(1);
expect(rows.rows[0].row_num).toBe(1);
expect(rows.rows[0].source_markdown_slug).toBe('people/alice');
});
test('re-run is a no-op (idempotent)', async () => {
await seedLegacyFact('Deposited remotely');
await runFactsCommand(engine, ['fence-backfill']);
await runFactsCommand(engine, ['fence-backfill']);
// eslint-disable-next-line @typescript-eslint/no-explicit-any
const rows = await (engine as any).db.query('SELECT id FROM facts');
expect(rows.rows).toHaveLength(1);
const body = readFileSync(join(brainDir, 'people/alice.md'), 'utf-8');
expect(body.match(/Deposited remotely/g)).toHaveLength(1);
});
test('--dry-run reports without writing', async () => {
await seedLegacyFact('Deposited remotely');
await runFactsCommand(engine, ['fence-backfill', '--dry-run']);
expect(existsSync(join(brainDir, 'people/alice.md'))).toBe(false);
// eslint-disable-next-line @typescript-eslint/no-explicit-any
const rows = await (engine as any).db.query(
'SELECT row_num FROM facts',
);
expect(rows.rows[0].row_num).toBeNull();
});
test('diverged page: appends past the DB row_num max instead of colliding (#2044 class)', async () => {
// The #2044 divergence shape the backfill must survive: the DB already
// holds stamped rows 1..3 for the slug (legacy wrong-path fence write),
// but the fence at the resolved path is missing. Fence-max+1 (= 1) would
// collide with the stamped rows on the post-rename UPDATE, failing the
// page and leaving the renamed fence disagreeing with the DB.
for (let n = 1; n <= 3; n++) {
// eslint-disable-next-line @typescript-eslint/no-explicit-any
await (engine as any).db.query(
`INSERT INTO facts (source_id, entity_slug, fact, kind, visibility, notability,
valid_from, source, confidence, row_num, source_markdown_slug)
VALUES ('default', 'people/alice', $1, 'fact', 'private', 'medium',
now(), 'mcp:extract_facts', 1.0, $2, 'people/alice')`,
[`stamped ${n}`, n],
);
}
await seedLegacyFact('Deposited remotely');
const result = await phaseBFenceFacts(engine, { dryRun: false });
expect(result.status).toBe('complete');
// eslint-disable-next-line @typescript-eslint/no-explicit-any
const rows = await (engine as any).db.query(
`SELECT row_num FROM facts WHERE fact = 'Deposited remotely'`,
);
expect(rows.rows[0].row_num).toBe(4);
const body = readFileSync(join(brainDir, 'people/alice.md'), 'utf-8');
expect(body).toContain('| 4 | Deposited remotely |');
});
});
-43
View File
@@ -292,49 +292,6 @@ describe('lookupSourceLocalPath', () => {
});
});
describe('writeFactsToFence — fence/DB divergence (#2044)', () => {
test('seeds row_num past the DB max when the fence lags the DB', async () => {
// Simulate the divergence class from #2044: DB rows were stamped with
// row_nums against a fence written at a path this code no longer reads
// (e.g. the pre-"Local patch 2026-06-11" wrong-path writes). The page
// on disk has NO fence, but the DB already holds row_num 1..3 for the
// slug. Pre-fix, the next deposit re-assigned row_num=1 from fence
// text alone and the whole insertFacts batch failed on
// idx_facts_fence_key.
for (let n = 1; n <= 3; n++) {
// eslint-disable-next-line @typescript-eslint/no-explicit-any
await (engine as any).db.query(
`INSERT INTO facts (source_id, entity_slug, fact, kind, visibility, notability,
valid_from, source, confidence, row_num, source_markdown_slug)
VALUES ('default', 'people/dana', $1, 'fact', 'private', 'medium',
now(), 'mcp:extract_facts', 1.0, $2, 'people/dana')`,
[`old claim ${n}`, n],
);
}
const result = await writeFactsToFence(
engine,
{ sourceId: 'default', localPath: brainDir, slug: 'people/dana' },
[baseInput({ fact: 'second deposit' })],
);
expect(result.fenceWriteFailed).toBeUndefined();
expect(result.inserted).toBe(1);
// The new row landed PAST the DB max, not at fence-max+1 (= 1).
// eslint-disable-next-line @typescript-eslint/no-explicit-any
const rows = await (engine as any).db.query(
'SELECT row_num FROM facts WHERE id = $1',
[result.ids[0]],
);
expect(rows.rows[0].row_num).toBe(4);
// And the on-disk fence carries the same row_num — fence and DB agree.
const body = readFileSync(join(brainDir, 'people/dana.md'), 'utf-8');
expect(body).toContain('| 4 | second deposit |');
});
});
// Cleanup any leftover tempdirs after the whole suite.
afterAll(() => {
// No-op: each test cleaned up via the beforeEach; this is a safety net.
+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);
});
});
+19 -18
View File
@@ -5,7 +5,7 @@
* - Batch insert N rows persists row_num + source_markdown_slug
* - Empty batch is a no-op
* - Returns ids in input-order
* - v51 partial UNIQUE collision skips only the colliding row (#2044)
* - v51 partial UNIQUE index rolls back the whole batch on a collision
* - deleteFactsForPage scopes by (source_id, source_markdown_slug);
* never touches other pages or pre-v51 NULL-source_markdown_slug rows
* - deleteFactsForPage on an empty page returns deleted:0 (idempotent)
@@ -135,29 +135,30 @@ describe('engine.insertFacts — batch insert', () => {
});
});
test('v51 partial UNIQUE collision skips ONLY the colliding row (#2044)', async () => {
test('v51 partial UNIQUE index rolls back the whole batch on collision', async () => {
// Seed row #1 first.
await engine.insertFacts([fixtureFact(1, { fact: 'seeded' })], { source_id: 'default' });
// Batch-insert rows that include a colliding row_num=1. Pre-#2044 this
// threw and rolled back the whole batch, making a second remote
// extract_facts deposit to an already-fenced page a hard failure. Now
// ON CONFLICT DO NOTHING skips the colliding row and keeps the rest.
const result = await engine.insertFacts(
[
fixtureFact(2, { fact: 'second' }),
fixtureFact(1, { fact: 'collides' }), // row_num=1 on same (source_id, source_markdown_slug)
fixtureFact(3, { fact: 'third' }),
],
{ source_id: 'default' },
);
expect(result.inserted).toBe(2);
expect(result.ids).toHaveLength(2);
// Now try to batch-insert rows that include a colliding row_num=1.
let threw = false;
try {
await engine.insertFacts(
[
fixtureFact(2, { fact: 'second' }),
fixtureFact(1, { fact: 'collides' }), // row_num=1 on same (source_id, source_markdown_slug)
fixtureFact(3, { fact: 'third' }),
],
{ source_id: 'default' },
);
} catch {
threw = true;
}
expect(threw).toBe(true);
// The seeded row survives untouched; the colliding claim is skipped.
// Verify the transaction rolled back — only the seeded row should remain.
// eslint-disable-next-line @typescript-eslint/no-explicit-any
const rows = await (engine as any).db.query('SELECT fact FROM facts ORDER BY id');
expect(rows.rows.map((r: { fact: string }) => r.fact)).toEqual(['seeded', 'second', 'third']);
expect(rows.rows.map((r: { fact: string }) => r.fact)).toEqual(['seeded']);
});
test('different source_markdown_slug values DO NOT collide on the same row_num', async () => {