From cc90471a96fd00b4f5f4004ea4bc5008e489a8f0 Mon Sep 17 00:00:00 2001 From: spiky02plateau <155588579+spiky02plateau@users.noreply.github.com> Date: Mon, 6 Jul 2026 11:58:47 +0200 Subject: [PATCH] perf(contextual-retrieval): bound per-chunk synopsis concurrency Replace the strictly sequential per-chunk synopsis loop with a bounded sliding worker pool (existing runSlidingPool helper). Results land in chunk order via index-addressed writes; code chunks still bypass the wrapper; embedding remains one page-level batch after all synopses. New knob GBRAIN_CONTEXTUAL_CHUNK_CONCURRENCY, default 4, clamped to [1,16]; 1 reproduces the prior sequential behavior exactly. Each chunk task still acquires/releases the global synopsis rate-lease, which remains the cross-worker governor; the lease id now travels from acquire to release instead of shared mutable state, and lease waits are abort-responsive. At 20-45s per synopsis call, a 120-chunk transcript page previously needed 60-90+ min wall time and routinely outlived job timeouts. --- src/core/contextual-retrieval-service.ts | 243 +++++++++---- .../handlers/contextual-reindex-per-chunk.ts | 41 ++- .../contextual-retrieval-service-pure.test.ts | 344 +++++++++++++++++- 3 files changed, 533 insertions(+), 95 deletions(-) diff --git a/src/core/contextual-retrieval-service.ts b/src/core/contextual-retrieval-service.ts index 0e22b7746..b4fe5dc72 100644 --- a/src/core/contextual-retrieval-service.ts +++ b/src/core/contextual-retrieval-service.ts @@ -45,7 +45,6 @@ import { embedBatch } from './embedding.ts'; import { resolveContextualRetrievalMode } from './contextual-retrieval-resolver.ts'; import { buildContextualPrefix, - extractFirstTwoSentences, modeRequiresHaiku, modeRequiresWrapper, sanitizeTitle, @@ -56,10 +55,8 @@ import { SYNOPSIS_PROMPT_VERSION, type GeneratePerChunkSynopsisResult, } from './page-summary.ts'; -import { - logSynopsisFailure, - type SynopsisFailureKind, -} from './audit-synopsis.ts'; +import type { SynopsisFailureKind } from './audit-synopsis.ts'; +import { runSlidingPool } from './worker-pool.ts'; import type { BrainEngine } from './engine.ts'; import type { ChunkInput, CRMode, Page } from './types.ts'; import type { SourceRow } from './sources-ops.ts'; @@ -72,6 +69,24 @@ import type { SourceRow } from './sources-ops.ts'; * corpus_generation hash. */ export const TITLE_WRAPPER_VERSION = 1; +const DEFAULT_HAIKU_MODEL = 'anthropic:claude-haiku-4-5-20251001'; +export const DEFAULT_CONTEXTUAL_CHUNK_CONCURRENCY = 4; +export const MAX_CONTEXTUAL_CHUNK_CONCURRENCY = 16; + +export function resolveContextualChunkConcurrency( + env: Record = process.env, +): number { + const raw = env.GBRAIN_CONTEXTUAL_CHUNK_CONCURRENCY; + if (raw === undefined || raw.trim() === '') return DEFAULT_CONTEXTUAL_CHUNK_CONCURRENCY; + const n = Number(raw); + if (!Number.isFinite(n)) return DEFAULT_CONTEXTUAL_CHUNK_CONCURRENCY; + return clampContextualChunkConcurrency(n); +} + +function clampContextualChunkConcurrency(n: number): number { + if (!Number.isFinite(n)) return DEFAULT_CONTEXTUAL_CHUNK_CONCURRENCY; + return Math.max(1, Math.min(MAX_CONTEXTUAL_CHUNK_CONCURRENCY, Math.trunc(n))); +} /** * Embedding model placeholder. The actual model name lands here from @@ -196,12 +211,16 @@ export interface ReembedPageArgs { * src/core/minions/rate-leases.ts here; inline callers (import-file, * reindex command) pass undefined and rely on gateway-level retry. */ - acquireSynopsisLease?: () => Promise; - releaseSynopsisLease?: () => Promise; + acquireSynopsisLease?: () => Promise; + releaseSynopsisLease?: (lease?: unknown) => Promise; + /** + * Intra-page per-chunk synopsis concurrency. 1 preserves the legacy + * sequential loop exactly; higher values only parallelize Haiku synopsis + * calls. Embedding remains one batch after all synopses succeed. + */ + chunkConcurrency?: number; } -const DEFAULT_HAIKU_MODEL = 'anthropic:claude-haiku-4-5-20251001'; - /** * Re-embed one page through the active CR mode. Implements the D26 P0-2 * two-phase build pattern. @@ -415,82 +434,41 @@ async function tryBuildPhase1(opts: { } // per_chunk_synopsis path. Read source text via fallback chain, - // generate synopsis per chunk sequentially within this page (D10), + // generate synopsis per chunk through a bounded sliding pool, then // batch embed at the end (D27 P2-2). const sourceText = readSourceTextWithFallback(page, chunks); - const wrappedTexts: string[] = []; + const wrappedTexts: string[] = new Array(chunks.length); + const chunkConcurrency = clampContextualChunkConcurrency( + args.chunkConcurrency ?? resolveContextualChunkConcurrency(), + ); - for (let i = 0; i < chunks.length; i++) { - const c = chunks[i]; - - // Code chunks always bypass the wrapper (D20-T4) — pass through. - if (c.chunk_source === 'fenced_code') { - wrappedTexts.push(c.chunk_text); - continue; - } - - // Acquire rate-lease per chunk (D26 P0-3). Inline callers pass no - // hooks; only the Minion handler wires through rate-leases.ts. - if (args.acquireSynopsisLease) { - await args.acquireSynopsisLease(); - } - - let synopsisResult: GeneratePerChunkSynopsisResult; - try { - synopsisResult = await generatePerChunkSynopsis({ - documentText: sourceText, - chunkText: c.chunk_text, - pageTitle: page.title, - pageSlug: args.pageSlug, - sourceId: args.sourceId, - chunkIndex: c.chunk_index, - model: haikuModel, - abortSignal: args.abortSignal, + const poolResult = await runSlidingPool({ + items: chunks, + workers: chunkConcurrency, + signal: args.abortSignal, + onError: 'abort', + failureLabel: (c) => String(c.chunk_index), + onItem: async (c, i) => { + wrappedTexts[i] = await buildWrappedChunkText({ + chunk: c, + sourceText, + safeTitle, + page, + args, + haikuModel, }); - } finally { - if (args.releaseSynopsisLease) { - try { - await args.releaseSynopsisLease(); - } catch { - // Lease release failure shouldn't abort the page; surfacing it - // would race with the synopsis result. Audit-only. - } - } - } + }, + }); - if (synopsisResult.kind === 'success') { - const prefix = buildContextualPrefix(safeTitle, synopsisResult.synopsis); - wrappedTexts.push( - wrapChunkForEmbedding(c.chunk_text, prefix, c.chunk_source), - ); - continue; + if (poolResult.failures.length > 0) { + const failure = [...poolResult.failures].sort((a, b) => a.idx - b.idx)[0].error; + if (failure instanceof ChunkSynopsisPhase1Error) { + return failure.result; } - - // Failure classification per D27 P1-2: - // refusal | empty | malformed → page-level fall-back to title-only - // auth_failure → permanent (won't fix with retry) - // rate_limit | timeout | network | provider_5xx → transient - // source_missing → walked into fallback already; would be 'malformed' - // from generatePerChunkSynopsis if we ever propagated it here - if ( - synopsisResult.kind === 'refusal' || - synopsisResult.kind === 'empty' || - synopsisResult.kind === 'malformed' - ) { - return { kind: 'page_level_fallback_requested', cause: synopsisResult.kind }; - } - if (synopsisResult.kind === 'auth_failure') { - return { - kind: 'permanent', - cause: synopsisResult.kind, - detail: synopsisResult.detail ?? 'auth failure', - }; - } - return { - kind: 'transient', - cause: synopsisResult.kind, - detail: synopsisResult.detail ?? 'transient', - }; + throw failure; + } + if (poolResult.aborted || args.abortSignal?.aborted) { + return { kind: 'transient', cause: 'timeout', detail: 'aborted' }; } // All chunks synthesized successfully. Single batch embed (D27 P2-2). @@ -511,6 +489,113 @@ async function tryBuildPhase1(opts: { } } +class ChunkSynopsisPhase1Error extends Error { + constructor(readonly result: Exclude) { + super(`chunk synopsis failed: ${result.kind}`); + this.name = 'ChunkSynopsisPhase1Error'; + } +} + +async function buildWrappedChunkText(opts: { + chunk: ChunkInput; + sourceText: string; + safeTitle: string; + page: Page; + args: ReembedPageArgs; + haikuModel: string; +}): Promise { + const { chunk: c, sourceText, safeTitle, page, args, haikuModel } = opts; + + // Code chunks always bypass the wrapper (D20-T4) — pass through. + if (c.chunk_source === 'fenced_code') { + return c.chunk_text; + } + + // Acquire rate-lease per chunk (D26 P0-3). Inline callers pass no + // hooks; only the Minion handler wires through rate-leases.ts. + let lease: unknown; + let leaseAcquired = false; + let synopsisResult: GeneratePerChunkSynopsisResult; + try { + if (args.acquireSynopsisLease) { + try { + lease = await args.acquireSynopsisLease(); + } catch (err) { + if (args.abortSignal?.aborted || isAbortError(err)) { + throw new ChunkSynopsisPhase1Error({ + kind: 'transient', + cause: 'timeout', + detail: 'aborted', + }); + } + throw err; + } + leaseAcquired = true; + } + synopsisResult = await generatePerChunkSynopsis({ + documentText: sourceText, + chunkText: c.chunk_text, + pageTitle: page.title, + pageSlug: args.pageSlug, + sourceId: args.sourceId, + chunkIndex: c.chunk_index, + model: haikuModel, + abortSignal: args.abortSignal, + }); + } finally { + if (leaseAcquired && args.releaseSynopsisLease) { + try { + await args.releaseSynopsisLease(lease); + } catch { + // Lease release failure shouldn't abort the page; surfacing it + // would race with the synopsis result. Audit-only. + } + } + } + + if (synopsisResult.kind === 'success') { + const prefix = buildContextualPrefix(safeTitle, synopsisResult.synopsis); + return wrapChunkForEmbedding(c.chunk_text, prefix, c.chunk_source); + } + + // Failure classification per D27 P1-2: + // refusal | empty | malformed → page-level fall-back to title-only + // auth_failure → permanent (won't fix with retry) + // rate_limit | timeout | network | provider_5xx → transient + // source_missing → walked into fallback already; would be 'malformed' + // from generatePerChunkSynopsis if we ever propagated it here + if ( + synopsisResult.kind === 'refusal' || + synopsisResult.kind === 'empty' || + synopsisResult.kind === 'malformed' + ) { + throw new ChunkSynopsisPhase1Error({ + kind: 'page_level_fallback_requested', + cause: synopsisResult.kind, + }); + } + if (synopsisResult.kind === 'auth_failure') { + throw new ChunkSynopsisPhase1Error({ + kind: 'permanent', + cause: synopsisResult.kind, + detail: synopsisResult.detail ?? 'auth failure', + }); + } + throw new ChunkSynopsisPhase1Error({ + kind: 'transient', + cause: synopsisResult.kind, + detail: synopsisResult.detail ?? 'transient', + }); +} + +function isAbortError(err: unknown): boolean { + return ( + typeof err === 'object' && + err !== null && + (err as { name?: unknown }).name === 'AbortError' + ); +} + /** * Source-text fallback chain per D11: * 1. read page.source_path from disk (truest "document") diff --git a/src/core/minions/handlers/contextual-reindex-per-chunk.ts b/src/core/minions/handlers/contextual-reindex-per-chunk.ts index 08b61ff4a..9d03b80f6 100644 --- a/src/core/minions/handlers/contextual-reindex-per-chunk.ts +++ b/src/core/minions/handlers/contextual-reindex-per-chunk.ts @@ -40,6 +40,7 @@ import { UnrecoverableError } from '../types.ts'; import type { BrainEngine } from '../../engine.ts'; import { reembedPageWithContextualRetrieval, + resolveContextualChunkConcurrency, type ReembedPageResult, } from '../../contextual-retrieval-service.ts'; import { @@ -132,7 +133,7 @@ export function makeContextualReindexHandler(opts: MakeContextualReindexHandlerO // call inside the service acquires/releases a lease against the // shared key across all worker processes. const maxConcurrent = resolveMaxConcurrent(); - let currentLeaseId: number | null = null; + const chunkConcurrency = resolveContextualChunkConcurrency(); const result: ReembedPageResult = await reembedPageWithContextualRetrieval({ engine, @@ -141,32 +142,32 @@ export function makeContextualReindexHandler(opts: MakeContextualReindexHandlerO globalMode, killSwitchDisabled, abortSignal: ctx.signal, + chunkConcurrency, acquireSynopsisLease: async () => { // Poll-acquire with brief backoff. The service's per-chunk loop - // is sequential within a page; this guards against the cross- - // worker pile-up. + // is bounded within a page; this guards against the cross-worker + // pile-up and remains the global rate governor. let attempts = 0; const maxAttempts = 60; // ~1 min max wait per chunk before giving up while (attempts < maxAttempts) { + if (ctx.signal.aborted) throw abortError(); const res = await acquireLease(engine, RATE_LEASE_KEY, ctx.id, maxConcurrent, { ttlMs: 60_000, }); if (res.acquired && res.leaseId != null) { - currentLeaseId = res.leaseId; - return; + return res.leaseId; } attempts++; - await new Promise((r) => setTimeout(r, 1000)); + await sleepWithAbort(1000, ctx.signal); } throw new Error( `Failed to acquire ${RATE_LEASE_KEY} lease after ${maxAttempts} attempts; ` + `Haiku rate limit pile-up too deep.`, ); }, - releaseSynopsisLease: async () => { - if (currentLeaseId != null) { - await releaseLease(engine, currentLeaseId); - currentLeaseId = null; + releaseSynopsisLease: async (lease) => { + if (typeof lease === 'number') { + await releaseLease(engine, lease); } }, }); @@ -218,6 +219,26 @@ async function tryLoadPageAcrossSources( return null; } +function sleepWithAbort(ms: number, signal: AbortSignal): Promise { + return new Promise((resolve, reject) => { + if (signal.aborted) { + reject(abortError()); + return; + } + const timer = setTimeout(resolve, ms); + signal.addEventListener('abort', () => { + clearTimeout(timer); + reject(abortError()); + }, { once: true }); + }); +} + +function abortError(): Error { + const err = new Error('aborted'); + err.name = 'AbortError'; + return err; +} + function classifyResult( pageSlug: string, result: ReembedPageResult, diff --git a/test/contextual-retrieval-service-pure.test.ts b/test/contextual-retrieval-service-pure.test.ts index 6918d41b3..3ba8d56d1 100644 --- a/test/contextual-retrieval-service-pure.test.ts +++ b/test/contextual-retrieval-service-pure.test.ts @@ -1,20 +1,38 @@ /** * Pure-function tests for src/core/contextual-retrieval-service.ts. * - * The full service test (PHASE 1 + PHASE 2 happy path, refusal restart, - * transient error propagation) needs a real PGLite + gateway stub seam. - * That lands in test/e2e/contextual-retrieval.test.ts. This file pins - * the service's pure helpers: corpus_generation hash composition + the - * expectedMode helper used by the T9 reindex sweep predicate. + * This file pins the service's pure helpers plus hermetic service behavior + * driven through fake engine + gateway seams. Full PGLite coverage lives in + * test/e2e/contextual-retrieval-pglite.test.ts. */ -import { describe, test, expect } from 'bun:test'; +import { afterEach, describe, test, expect } from 'bun:test'; import { computeCorpusGeneration, computeSourceTextHash, expectedModeForPageSourceOnly, + reembedPageWithContextualRetrieval, + resolveContextualChunkConcurrency, TITLE_WRAPPER_VERSION, } from '../src/core/contextual-retrieval-service.ts'; +import { + __setChatTransportForTests, + __setEmbedTransportForTests, + configureGateway, + resetGateway, + type ChatOpts, + type ChatResult, +} from '../src/core/ai/gateway.ts'; +import type { ChunkInput } from '../src/core/types.ts'; +import { withEnv } from './helpers/with-env.ts'; + +const TEST_DIMS = 1536; + +afterEach(() => { + __setChatTransportForTests(null); + __setEmbedTransportForTests(null); + resetGateway(); +}); describe('computeCorpusGeneration', () => { test('returns 16-char hex hash', () => { @@ -138,3 +156,317 @@ describe('expectedModeForPageSourceOnly (T9 reindex sweep helper)', () => { } }); }); + +describe('resolveContextualChunkConcurrency', () => { + test('defaults to 4 and reads the process env', async () => { + await withEnv({ GBRAIN_CONTEXTUAL_CHUNK_CONCURRENCY: undefined }, async () => { + expect(resolveContextualChunkConcurrency()).toBe(4); + }); + await withEnv({ GBRAIN_CONTEXTUAL_CHUNK_CONCURRENCY: '7' }, async () => { + expect(resolveContextualChunkConcurrency()).toBe(7); + }); + }); + + test('clamps to [1, 16] and ignores invalid values', () => { + expect(resolveContextualChunkConcurrency({ + GBRAIN_CONTEXTUAL_CHUNK_CONCURRENCY: '0', + })).toBe(1); + expect(resolveContextualChunkConcurrency({ + GBRAIN_CONTEXTUAL_CHUNK_CONCURRENCY: '-3', + })).toBe(1); + expect(resolveContextualChunkConcurrency({ + GBRAIN_CONTEXTUAL_CHUNK_CONCURRENCY: '99', + })).toBe(16); + expect(resolveContextualChunkConcurrency({ + GBRAIN_CONTEXTUAL_CHUNK_CONCURRENCY: '1.9', + })).toBe(1); + expect(resolveContextualChunkConcurrency({ + GBRAIN_CONTEXTUAL_CHUNK_CONCURRENCY: 'not-a-number', + })).toBe(4); + }); +}); + +describe('per-chunk synopsis concurrency', () => { + test('concurrency > 1 preserves chunk-order embed input', async () => { + const chunks = makeChunks(['alpha', 'beta', 'gamma', 'delta']); + const delays: Record = { alpha: 30, beta: 5, gamma: 20, delta: 1 }; + const sequential = await runWithChatStub({ + chunks, + concurrency: 1, + delayForChunk: (chunk) => delays[chunk] ?? 1, + }); + const parallel = await runWithChatStub({ + chunks, + concurrency: 4, + delayForChunk: (chunk) => delays[chunk] ?? 1, + }); + + expect(parallel.result.kind).toBe('success'); + expect(parallel.embedInputs).toEqual(sequential.embedInputs); + expect(parallel.embeddedChunks.map((c) => c.chunk_text)).toEqual( + chunks.map((c) => c.chunk_text), + ); + }); + + test('concurrency is bounded', async () => { + let active = 0; + let maxActive = 0; + let leaseActive = 0; + let maxLeaseActive = 0; + let acquired = 0; + let released = 0; + const chunks = makeChunks(Array.from({ length: 8 }, (_, i) => `chunk-${i}`)); + const out = await runWithChatStub({ + chunks, + concurrency: 3, + acquireSynopsisLease: async () => { + acquired++; + leaseActive++; + maxLeaseActive = Math.max(maxLeaseActive, leaseActive); + return acquired; + }, + releaseSynopsisLease: async () => { + released++; + leaseActive--; + }, + chat: async (opts) => { + active++; + maxActive = Math.max(maxActive, active); + try { + await delay(20, opts.abortSignal); + return chatSuccess(`Synopsis for ${extractChunk(opts)}`); + } finally { + active--; + } + }, + }); + + expect(out.result.kind).toBe('success'); + expect(maxActive).toBeGreaterThan(1); + expect(maxActive).toBeLessThanOrEqual(3); + expect(maxLeaseActive).toBeLessThanOrEqual(3); + expect(acquired).toBe(8); + expect(released).toBe(8); + expect(leaseActive).toBe(0); + }); + + test('one chunk failure aborts queued work and falls back at page level', async () => { + let started = 0; + const chunks = makeChunks(Array.from({ length: 9 }, (_, i) => `chunk-${i}`)); + const out = await runWithChatStub({ + chunks, + concurrency: 3, + chat: async (opts) => { + started++; + const chunk = extractChunk(opts); + if (chunk === 'chunk-0') return chatSuccess(''); + await delay(30, opts.abortSignal); + return chatSuccess(`Synopsis for ${chunk}`); + }, + }); + + expect(out.result.kind).toBe('page_fallback'); + expect(started).toBeLessThanOrEqual(3); + }); + + test('fenced code chunks bypass synopsis calls and leases', async () => { + let chatCalls = 0; + let leaseCalls = 0; + const chunks: ChunkInput[] = [ + { chunk_index: 0, chunk_text: 'intro', chunk_source: 'compiled_truth' }, + { chunk_index: 1, chunk_text: 'const x = 1;', chunk_source: 'fenced_code' }, + { chunk_index: 2, chunk_text: 'outro', chunk_source: 'compiled_truth' }, + ]; + + const out = await runWithChatStub({ + chunks, + concurrency: 3, + acquireSynopsisLease: async () => { + leaseCalls++; + }, + releaseSynopsisLease: async () => {}, + chat: async (opts) => { + chatCalls++; + return chatSuccess(`Synopsis for ${extractChunk(opts)}`); + }, + }); + + expect(out.result.kind).toBe('success'); + expect(chatCalls).toBe(2); + expect(leaseCalls).toBe(2); + expect(out.embedInputs[1]).toBe('const x = 1;'); + }); + + test('abortSignal cancels in-flight and queued synopsis work promptly', async () => { + const controller = new AbortController(); + let started = 0; + const chunks = makeChunks(Array.from({ length: 20 }, (_, i) => `chunk-${i}`)); + const startedAt = Date.now(); + const promise = runWithChatStub({ + chunks, + concurrency: 4, + abortSignal: controller.signal, + chat: async (opts) => { + started++; + await delay(1000, opts.abortSignal); + return chatSuccess(`Synopsis for ${extractChunk(opts)}`); + }, + }); + setTimeout(() => controller.abort(), 20); + + const out = await promise; + expect(out.result.kind).toBe('transient_error'); + if (out.result.kind === 'transient_error') { + expect(out.result.cause).toBe('timeout'); + } + expect(started).toBeLessThanOrEqual(4); + expect(Date.now() - startedAt).toBeLessThan(300); + }); +}); + +function makeChunks(texts: string[]): ChunkInput[] { + return texts.map((text, i) => ({ + chunk_index: i, + chunk_text: text, + chunk_source: 'compiled_truth', + })); +} + +async function runWithChatStub(opts: { + chunks: ChunkInput[]; + concurrency: number; + abortSignal?: AbortSignal; + delayForChunk?: (chunk: string) => number; + chat?: (opts: ChatOpts) => Promise; + acquireSynopsisLease?: () => Promise; + releaseSynopsisLease?: (lease?: unknown) => Promise; +}) { + configureGateway({ + embedding_model: 'openai:text-embedding-3-large', + embedding_dimensions: TEST_DIMS, + env: { OPENAI_API_KEY: 'sk-test' }, + }); + + const embedInputs: string[][] = []; + __setEmbedTransportForTests(async ({ values }: any) => { + embedInputs.push([...values]); + return { + embeddings: values.map((_: string, i: number) => + Array.from({ length: TEST_DIMS }, () => 0.001 + i * 0.001), + ), + usage: { tokens: 0 }, + } as any; + }); + + __setChatTransportForTests(opts.chat ?? (async (chatOpts) => { + const chunk = extractChunk(chatOpts); + await delay(opts.delayForChunk?.(chunk) ?? 1, chatOpts.abortSignal); + return chatSuccess(`Synopsis for ${chunk}`); + })); + + const engine = makeServiceEngine(opts.chunks); + const result = await reembedPageWithContextualRetrieval({ + engine, + pageSlug: 'wiki/concepts/concurrency-test', + sourceId: 'default', + globalMode: 'per_chunk_synopsis', + chunkConcurrency: opts.concurrency, + abortSignal: opts.abortSignal, + ...(opts.acquireSynopsisLease && { acquireSynopsisLease: opts.acquireSynopsisLease }), + ...(opts.releaseSynopsisLease && { releaseSynopsisLease: opts.releaseSynopsisLease }), + }); + + return { + result, + embedInputs: embedInputs.flat(), + embeddedChunks: engine.embeddedChunks as ChunkInput[], + }; +} + +function makeServiceEngine(chunks: ChunkInput[]) { + const engine: any = { + embeddedChunks: [] as ChunkInput[], + async getPage() { + return { + id: 1, + slug: 'wiki/concepts/concurrency-test', + source_id: 'default', + type: 'concept', + title: 'Concurrency Test', + compiled_truth: chunks.map((c) => c.chunk_text).join('\n\n'), + timeline: '', + frontmatter: {}, + created_at: new Date('2026-01-01T00:00:00Z'), + updated_at: new Date('2026-01-01T00:00:00Z'), + deleted_at: null, + }; + }, + async executeRaw() { + return [{ + id: 'default', + name: 'Default', + local_path: null, + last_commit: null, + last_sync_at: null, + config: {}, + created_at: new Date('2026-01-01T00:00:00Z'), + contextual_retrieval_mode: null, + trust_frontmatter_overrides: false, + }]; + }, + async getChunks() { + return chunks; + }, + async transaction(fn: (tx: any) => Promise) { + await fn({ + upsertChunks: async (_slug: string, embedded: ChunkInput[]) => { + engine.embeddedChunks = embedded; + }, + updatePageContextualRetrievalState: async () => {}, + }); + }, + async updatePageContextualRetrievalState() {}, + }; + return engine; +} + +function extractChunk(opts: ChatOpts): string { + const content = String(opts.messages[0]?.content ?? ''); + return content.match(/\n([\s\S]*?)\n<\/chunk>/)?.[1] ?? ''; +} + +function chatSuccess(text: string): ChatResult { + return { + text, + blocks: [], + stopReason: 'end', + usage: { + input_tokens: 0, + output_tokens: 0, + cache_read_tokens: 0, + cache_creation_tokens: 0, + }, + model: 'stub:chat', + providerId: 'stub', + }; +} + +function delay(ms: number, signal?: AbortSignal): Promise { + return new Promise((resolve, reject) => { + if (signal?.aborted) { + reject(abortError()); + return; + } + const timer = setTimeout(resolve, ms); + signal?.addEventListener('abort', () => { + clearTimeout(timer); + reject(abortError()); + }, { once: true }); + }); +} + +function abortError(): Error { + const err = new Error('aborted'); + err.name = 'AbortError'; + return err; +}