From 522cfb032a94e268cd347e6d5769dcf956de7e07 Mon Sep 17 00:00:00 2001 From: Garry Tan Date: Sat, 1 Aug 2026 08:34:30 +0800 Subject: [PATCH] =?UTF-8?q?fix(chunker):=20measured=20hard-split=20budgets?= =?UTF-8?q?=20+=20token-aware=20capByChars=20=E2=80=94=20the=20follow-up?= =?UTF-8?q?=20invited=20in=20#3477's=20merge=20review=20(#3564)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Co-Authored-By: YMYD --- src/core/chunkers/code.ts | 166 ++++++++--------- src/core/chunkers/recursive.ts | 100 ++++++++-- src/core/chunkers/token-estimate.ts | 121 ++++++++++++ test/chunkers/embed-estimate-home.test.ts | 216 ++++++++++++++++++++++ 4 files changed, 502 insertions(+), 101 deletions(-) create mode 100644 src/core/chunkers/token-estimate.ts create mode 100644 test/chunkers/embed-estimate-home.test.ts diff --git a/src/core/chunkers/code.ts b/src/core/chunkers/code.ts index f9cdb93ca..b2dbc4949 100644 --- a/src/core/chunkers/code.ts +++ b/src/core/chunkers/code.ts @@ -20,7 +20,14 @@ import { chunkText as recursiveChunk } from './recursive.ts'; import { buildQualifiedName } from './qualified-names.ts'; -import { CJK_SLUG_CHARS, CJK_RANGES_REGEX } from '../cjk.ts'; +import { estimateTokens, estimateEmbedTokens, estimateEmbedTokensCeiling, DEFAULT_MAX_CHUNK_TOKENS } from './token-estimate.ts'; +import { safeSplitIndex } from '../text-safe.ts'; + +// Both estimators moved to token-estimate.ts (#3477 follow-up) so +// recursive.ts can share them without an import cycle. Re-exported here: +// commands/sync.ts, commands/reindex-code.ts, and tests import them from +// this module. +export { estimateTokens, estimateEmbedTokens } from './token-estimate.ts'; // Embed the tree-sitter runtime + per-language grammars as files. // `with { type: 'file' }` returns a path (string) at runtime. Bun bundles @@ -559,7 +566,6 @@ export function parseWithTimeout( } const DEFAULT_CHUNKER_TIMEOUT_MS = 30_000; -const DEFAULT_MAX_CHUNK_TOKENS = 2000; function resolveChunkerTimeoutMs(): number { const raw = process.env.GBRAIN_CHUNKER_TIMEOUT_MS; @@ -851,9 +857,19 @@ function capOversizedChunks( continue; } // Strip the structured header ("[Lang] path:N-M symbol\n\n") so the splitter - // works on the raw body; buildChunk re-adds a header to each piece. - const body = c.text.replace(/^\[[^\]]+\] [^\n]+\n\n/, ''); - for (const piece of splitToTokenBudget(body, cap, opts)) { + // works on the raw body; buildChunk re-adds a header to each piece. The + // re-added header costs tokens too — budget for it, or every piece split + // to exactly `cap` re-emerges a header's-worth over it (measured: a 2,000 + // cap emitted 2,011-token fence chunks when the body alone was capped). + // The reservation must be an UPPER bound on the header's contribution: + // estimateEmbedTokens is super-additive across a mixed-script join, so the + // header's standalone cl100k figure under-counts ~2.5x once the body + // contains CJK and the weighted branch takes over (see + // estimateEmbedTokensCeiling). + const headerMatch = c.text.match(/^\[[^\]]+\] [^\n]+\n\n/); + const body = headerMatch ? c.text.slice(headerMatch[0].length) : c.text; + const bodyCap = Math.max(1, cap - (headerMatch ? estimateEmbedTokensCeiling(headerMatch[0]) : 0)); + for (const piece of splitToTokenBudget(body, bodyCap, opts)) { if (!piece.trim()) continue; out.push(buildChunk({ body: piece, @@ -880,44 +896,68 @@ function splitToTokenBudget(text: string, cap: number, opts: CodeChunkOptions): chunkSize: opts.fallbackChunkSizeWords ?? 300, chunkOverlap: opts.fallbackOverlapWords ?? 50, }).map((p) => p.text); - for (const piece of pieces) { - if (estimateEmbedTokens(piece) <= cap) { - out.push(piece); - continue; + // Hard-split budget is derived from each piece's own measured density + // (chars per estimated token) scaled to the cap, not a fixed chars-per- + // token guess: the previous 3.5 chars/token ASCII assumption undercuts + // URL-dense JSON (~2.6 chars/token measured), leaving 2,070–2,095-token + // slices past a 2,000 cap. Slices are re-measured and re-derived (density + // varies within a piece), so the cap holds by construction. + const hardSplit = (p: string): void => { + const est = estimateEmbedTokens(p); + if (est <= cap) { + out.push(p); + return; } - // Hard-split slice size. Pure-ASCII pieces: ~3.5 chars/token is a - // conservative cl100k estimate for source text. CJK-containing pieces: - // the weighted estimate can reach 1 token/char, so budget 1 char/token - // to keep every slice under cap by construction. - const charBudget = Math.max(1, Math.floor(cap * (CJK_RANGES_REGEX.test(piece) ? 1 : 3.5))); - for (let i = 0; i < piece.length; i += charBudget) out.push(piece.slice(i, i + charBudget)); - } + const charBudget = Math.max(1, Math.floor((p.length * cap) / est)); + if (charBudget >= p.length) { + out.push(p); // 1-char floor on a tiny cap — nothing left to split + return; + } + // Even out the slice width instead of striding by charBudget and shedding + // `p.length mod charBudget` as a standalone piece at EVERY recursion + // level: buildChunk re-headers each remainder into its own embedding row, + // so a 14.4K fence emitted 5 chunks of 50-86 chars (and, deeper in the + // recursion, 3-char slivers) alongside its real content. Evening is free — + // the piece count is ceil(length / charBudget) either way, so the same + // content is spread over the same number of chunks — and width <= + // charBudget by construction, so the token budget still holds. + // + // The width is re-derived from what REMAINS on every step rather than + // fixed up front, because safeSplitIndex can back a cut off by up to two + // units and a fixed width lets that drift accumulate into a tail runt + // (measured on an all-astral blob: 4-unit chunks trailing 724-unit ones). + let i = 0; + while (i < p.length) { + const remaining = p.length - i; + const partsLeft = Math.ceil(remaining / charBudget); + // `i === 0` cannot recurse on the whole piece — charBudget < p.length is + // checked above, so partsLeft >= 2 on the first step. The guard keeps a + // degenerate budget from looping instead of terminating. + if (partsLeft <= 1) { + if (i === 0) out.push(p); + else hardSplit(p.slice(i)); + return; + } + // The budget is derived from measured density, so it has arbitrary + // parity: a raw slice at `i + width` orphans a UTF-16 surrogate half. + // safeSplitIndex backs the cut off a pair (#2011 — a lone surrogate is + // rejected by Postgres inside a ::jsonb cast and aborts the whole batch). + const end = safeSplitIndex(p, i + Math.ceil(remaining / partsLeft)); + if (end <= i) { + // Degenerate width (a 1-char budget backing off a surrogate pair) — + // emit rather than drop, and never re-enter on the same string. + if (i === 0) out.push(p); + else hardSplit(p.slice(i)); + return; + } + hardSplit(p.slice(i, end)); + i = end; + } + }; + for (const piece of pieces) hardSplit(piece); return out; } -const CJK_CHARS_G = new RegExp(`[${CJK_SLUG_CHARS}]`, 'g'); - -/** - * Embedding-safe token estimate for the oversize cap. estimateTokens - * (cl100k) matches embedding-family tokenizers closely on pure-ASCII source - * (measured identical on English prose and JSON vs Qwen3-Embedding), but - * UNDERCOUNTS mixed CJK+ASCII chunks — measured −31% on URL-dense Korean - * text vs the Qwen3 embedding tokenizer, which is exactly the shape that - * overflows strict embedding backends (#2826). For chunks containing CJK, - * take the max of cl100k and a per-char-class overestimate (CJK 1.0/char, - * other non-whitespace 0.75/char, whitespace 0.1/char). CJK-DOMINANT text - * is unaffected too: cl100k already counts it above the weighted form, so - * max() returns the same value as today. Only mixed-script chunks — the - * measured divergence class — estimate higher. - */ -export function estimateEmbedTokens(text: string): number { - const cjk = (text.match(CJK_CHARS_G) || []).length; - if (cjk === 0) return estimateTokens(text); - const ws = (text.match(/\s/g) || []).length; - const weighted = Math.ceil(cjk + (text.length - cjk - ws) * 0.75 + ws * 0.1); - return Math.max(estimateTokens(text), weighted); -} - // ---------- Internals ---------- function fallbackChunks( @@ -1245,54 +1285,6 @@ function sanitize(name: string): string { return name.replace(/[\n\r\t]+/g, ' ').replace(/\s+/g, ' ').trim(); } -// v0.19.0 (Layer 5): accurate token count via @dqbd/tiktoken cl100k_base, -// the same encoder text-embedding-3-large uses. The old len/4 heuristic was -// 2-3x off for code. Lazy-init so dev and compiled-binary both only pay -// the init cost once. Falls back to the heuristic if the encoder fails -// to load (vanishingly unlikely but keeps the chunker available). -let tiktokenEncoder: { encode: (s: string) => Uint32Array; free: () => void } | null = null; -let tiktokenInitialized = false; - -// v0.20.0 Cathedral II Layer 8 (D1) — exported so commands/sync.ts can -// estimate embed cost before a --all sync blows a surprise OpenAI bill. -// Same cl100k_base tokenizer the embedding path actually uses, so cost -// estimates match actual billing within tokenizer noise. -export function estimateTokens(text: string): number { - if (!text) return 0; - if (!tiktokenInitialized) { - try { - // eslint-disable-next-line @typescript-eslint/no-var-requires - const m = require('@dqbd/tiktoken'); - tiktokenEncoder = m.get_encoding('cl100k_base'); - } catch { - tiktokenEncoder = null; - } - tiktokenInitialized = true; - } - if (tiktokenEncoder) { - try { - return tiktokenEncoder.encode(text).length; - } catch { - // Code legitimately contains tiktoken special-token strings (e.g. CLIP/GPT - // tokenizers embed the literal "<|endoftext|>"). The default encode() uses - // disallowed_special='all' and THROWS on those, crashing reindex-code on - // valid source files. For a token COUNT we don't need special-token - // semantics: re-encode treating them as ordinary text (never throws), - // heuristic only if even that fails. - try { - return ( - tiktokenEncoder as unknown as { - encode: (s: string, allowed: string[], disallowed: string[]) => Uint32Array; - } - ).encode(text, [], []).length; - } catch { - return Math.max(1, Math.ceil(text.length / 4)); - } - } - } - return Math.max(1, Math.ceil(text.length / 4)); -} - // v0.20.0 Cathedral II Layer 4: display name derived from the language // manifest. Single source of truth — adding a new language via // registerLanguage() automatically exposes its displayName to chunk diff --git a/src/core/chunkers/recursive.ts b/src/core/chunkers/recursive.ts index a9aa54e88..703acb6ac 100644 --- a/src/core/chunkers/recursive.ts +++ b/src/core/chunkers/recursive.ts @@ -13,11 +13,16 @@ * v0.32.7: maxChars hard cap (default 6000) sliding-window safety belt * guarantees no chunk overflows OpenAI's 8192-token embedding limit even * on pathological CJK / whitespace-less text. + * #3477 follow-up: the belt also bounds ESTIMATED embedding tokens + * (DEFAULT_MAX_CHUNK_TOKENS, shared with the code chunker's oversize cap) — + * a char-only cap cannot bound tokens for CJK/dense text (#3037, #2826). * * Lossless invariant: non-overlapping portions reassemble to original. */ import { countCJKAwareWords, CJK_SENTENCE_DELIMITERS, CJK_CLAUSE_DELIMITERS } from '../cjk.ts'; +import { estimateEmbedTokens, DEFAULT_MAX_CHUNK_TOKENS } from './token-estimate.ts'; +import { safeSplitIndex } from '../text-safe.ts'; /** * Markdown chunker version. Folded into the per-page chunker_version column @@ -109,29 +114,96 @@ export function chunkText(text: string, opts?: ChunkOptions): TextChunk[] { } /** - * Hard-cap a chunk's char length via a sliding window. Returns the input - * unchanged when it's already ≤ maxChars. + * Hard-cap a chunk via a sliding window — by char length AND by estimated + * embedding tokens. Returns the input unchanged when it fits both budgets. * - * Overlap is min(500, maxChars/10) so successive windows preserve semantic + * The char budget (maxChars, default 6000) is the historical belt; the token + * budget (DEFAULT_MAX_CHUNK_TOKENS, shared with the code chunker's oversize + * cap) is the constraint embedders actually enforce. A char-only cap cannot + * bound tokens: 6000 CJK-dense chars run 3-6k tokens, past strict embedder + * contexts (nomic-embed-text 2048), so those chunks fail on every embed + * sweep, silently, forever (#3037) — and URL-dense CJK markdown emits + * over-limit chunks well under maxChars (#2826). When the text over-runs the + * token budget, the window is derived from its own measured density — + * floor(length × budget / estimate) — and every slice is re-checked (local + * density can exceed the whole-text average), re-deriving on the slice until + * each piece fits. ASCII prose is unaffected: 6000 chars measure ~1.5-1.7k + * cl100k tokens, under the budget, so the window stays maxChars. + * + * Overlap is min(500, window/10) so successive windows preserve semantic * continuity across the cut. * - * v0.32.7. BMP-only safe (does not split astral surrogate pairs in practice - * because declared CJK ranges are all BMP; widening to astral Han support - * is a v0.33+ follow-up that requires Array.from-style codepoint iteration). + * v0.32.7. Surrogate-safe: the window is derived from measured density and so + * has arbitrary parity, which a raw slice would use to cut an astral pair in + * half — every boundary goes through safeSplitIndex. (The former "BMP-only + * safe" note rested on maxChars=6000 and stride=5500 both being even; + * deriving the window from density retired that guarantee.) */ -function capByChars(text: string, maxChars: number): string[] { - if (text.length <= maxChars) return text.length > 0 ? [text] : []; - const overlap = Math.min(500, Math.floor(maxChars / 10)); - const stride = Math.max(1, maxChars - overlap); +function capByChars(text: string, maxChars: number, knownEst?: number): string[] { + if (text.length === 0) return []; + const est = knownEst ?? probeEmbedTokens(text); + const window = est <= DEFAULT_MAX_CHUNK_TOKENS + ? maxChars + : Math.max(1, Math.min(maxChars, Math.floor((text.length * DEFAULT_MAX_CHUNK_TOKENS) / est))); + if (text.length <= window) { + // Emitting the text whole is the one path that skips the per-slice + // re-check below, so a PROBED estimate has to be confirmed exactly first: + // a sparse ASCII head can under-read a dense CJK tail. + if (knownEst !== undefined || text.length <= DENSITY_PROBE_CHARS) return [text]; + const exact = estimateEmbedTokens(text); + return exact <= DEFAULT_MAX_CHUNK_TOKENS ? [text] : capByChars(text, maxChars, exact); + } + // The stride keeps its nominal window-minus-overlap value. Evening the + // windows out (as the header-budget hard split does) is WRONG here: that + // splitter partitions, this one overlaps, so shrinking the stride to land + // the last window flush against the end collapses successive windows into + // near-duplicates — measured on scripts/test-weights.json, two 6,047-char + // chunks differing by 47 chars. A short final window is the cheaper end of + // that trade and is the behavior this loop has always had. + const overlap = Math.min(500, Math.floor(window / 10)); + const stride = Math.max(1, window - overlap); const out: string[] = []; - for (let i = 0; i < text.length; i += stride) { - const slice = text.slice(i, i + maxChars).trim(); - if (slice.length > 0) out.push(slice); - if (i + maxChars >= text.length) break; + let i = 0; + while (i < text.length) { + const end = safeSplitIndex(text, Math.min(text.length, i + window)); + const slice = text.slice(i, end).trim(); + if (slice.length > 0) { + const sliceEst = estimateEmbedTokens(slice); + if (sliceEst > DEFAULT_MAX_CHUNK_TOKENS) { + // Denser than the text average — re-derive locally, reusing the exact + // figure just measured (it also guarantees window < slice.length, so + // the recursion strictly shrinks). + out.push(...capByChars(slice, maxChars, sliceEst)); + } else { + out.push(slice); + } + } + if (end >= text.length) break; + const next = safeSplitIndex(text, Math.min(text.length, i + stride)); + i = next > i ? next : i + 1; } return out; } +/** + * Chars measured to derive the window. estimateEmbedTokens is SUPERLINEAR on + * CJK — measured on this repo's encoder: 2K chars 11ms, 6K 99ms, 20K 1,138ms — + * and capByChars runs on every chunk, so measuring the whole text up front + * dominates the chunker (the 20K-char whitespace-less CJK cap test went from + * an O(1) length compare to a 6.7s run, past bun's 5s per-test limit, on a + * cold encoder). The window only needs an approximate density: every emitted + * slice is re-measured exactly, denser-than-average slices recurse on that + * exact figure, and the one path that emits without a re-check confirms + * exactly first — so the cap holds regardless of what the probe reads. + */ +const DENSITY_PROBE_CHARS = 2000; + +function probeEmbedTokens(text: string): number { + if (text.length <= DENSITY_PROBE_CHARS) return estimateEmbedTokens(text); + const head = text.slice(0, safeSplitIndex(text, DENSITY_PROBE_CHARS)); + return Math.ceil((estimateEmbedTokens(head) * text.length) / head.length); +} + function recursiveSplit(text: string, level: number, target: number): string[] { if (level >= DELIMITERS.length) { // Level 4: split on whitespace diff --git a/src/core/chunkers/token-estimate.ts b/src/core/chunkers/token-estimate.ts new file mode 100644 index 000000000..7c1642b71 --- /dev/null +++ b/src/core/chunkers/token-estimate.ts @@ -0,0 +1,121 @@ +/** + * Embedding-token estimation — shared by both chunkers (#3477 follow-up). + * + * Moved out of code.ts so recursive.ts can measure with the same estimator: + * code.ts imports recursive.ts, so recursive.ts could never import these from + * code.ts without a cycle. The natural home suggested in #3477's merge review + * was cjk.ts, but cjk.ts is a check:fuzz-purity bundle target and tiktoken's + * loader pulls node:fs into the bundle — so the estimators live here, one + * layer above cjk.ts (whose char classes they reuse) and below both chunkers. + */ + +import { CJK_SLUG_CHARS } from '../cjk.ts'; + +/** + * Default hard budget for any emitted chunk's estimated embedding tokens. + * Shared by capOversizedChunks (code.ts, #1675) and capByChars (recursive.ts). + * 2000 keeps a margin under the smallest common strict embedder contexts + * (nomic-embed-text 2048 — #3037; llama-server -ub 2048 — #2826). + */ +export const DEFAULT_MAX_CHUNK_TOKENS = 2000; + +// v0.19.0 (Layer 5): accurate token count via @dqbd/tiktoken cl100k_base, +// the same encoder text-embedding-3-large uses. The old len/4 heuristic was +// 2-3x off for code. Lazy-init so dev and compiled-binary both only pay +// the init cost once. Falls back to the heuristic if the encoder fails +// to load (vanishingly unlikely but keeps the chunker available). +let tiktokenEncoder: { encode: (s: string) => Uint32Array; free: () => void } | null = null; +let tiktokenInitialized = false; + +// v0.20.0 Cathedral II Layer 8 (D1) — re-exported from code.ts so +// commands/sync.ts can estimate embed cost before a --all sync blows a +// surprise OpenAI bill. Same cl100k_base tokenizer the embedding path +// actually uses, so cost estimates match actual billing within tokenizer +// noise. +export function estimateTokens(text: string): number { + if (!text) return 0; + if (!tiktokenInitialized) { + try { + // eslint-disable-next-line @typescript-eslint/no-var-requires + const m = require('@dqbd/tiktoken'); + tiktokenEncoder = m.get_encoding('cl100k_base'); + } catch { + tiktokenEncoder = null; + } + tiktokenInitialized = true; + } + if (tiktokenEncoder) { + try { + return tiktokenEncoder.encode(text).length; + } catch { + // Code legitimately contains tiktoken special-token strings (e.g. CLIP/GPT + // tokenizers embed the literal "<|endoftext|>"). The default encode() uses + // disallowed_special='all' and THROWS on those, crashing reindex-code on + // valid source files. For a token COUNT we don't need special-token + // semantics: re-encode treating them as ordinary text (never throws), + // heuristic only if even that fails. + try { + return ( + tiktokenEncoder as unknown as { + encode: (s: string, allowed: string[], disallowed: string[]) => Uint32Array; + } + ).encode(text, [], []).length; + } catch { + return Math.max(1, Math.ceil(text.length / 4)); + } + } + } + return Math.max(1, Math.ceil(text.length / 4)); +} + +const CJK_CHARS_G = new RegExp(`[${CJK_SLUG_CHARS}]`, 'g'); + +/** + * Embedding-safe token estimate for the oversize caps. estimateTokens + * (cl100k) matches embedding-family tokenizers closely on pure-ASCII source + * (measured identical on English prose and JSON vs Qwen3-Embedding), but + * UNDERCOUNTS mixed CJK+ASCII chunks — measured −31% on URL-dense Korean + * text vs the Qwen3 embedding tokenizer, which is exactly the shape that + * overflows strict embedding backends (#2826). For chunks containing CJK, + * take the max of cl100k and a per-char-class overestimate (CJK 1.0/char, + * other non-whitespace 0.75/char, whitespace 0.1/char). CJK-DOMINANT text + * is unaffected too: cl100k already counts it above the weighted form, so + * max() returns the same value as today. Only mixed-script chunks — the + * measured divergence class — estimate higher. + */ +export function estimateEmbedTokens(text: string): number { + const cjk = (text.match(CJK_CHARS_G) || []).length; + if (cjk === 0) return estimateTokens(text); + return Math.max(estimateTokens(text), weightedTokens(text, cjk)); +} + +/** The per-char-class overestimate half of estimateEmbedTokens. Linear (two + * regex scans), unlike the cl100k encoder — see estimateEmbedTokensCeiling. */ +function weightedTokens(text: string, cjk: number): number { + const ws = (text.match(/\s/g) || []).length; + return Math.ceil(cjk + (text.length - cjk - ws) * 0.75 + ws * 0.1); +} + +/** + * Upper bound on what `text` contributes to `estimateEmbedTokens(text + rest)` + * for ANY `rest` — i.e. the figure to RESERVE when a fragment will be glued + * onto a body whose script mix is not yet known. + * + * estimateEmbedTokens is super-additive across a mixed-script join. It only + * switches to the weighted branch when the text it is handed contains CJK, so + * a pure-ASCII fragment measured ALONE costs cl100k (a 59-char structured + * chunk header = 17 tokens) while the SAME fragment inside a chunk whose body + * contains CJK costs ~0.75/char on the weighted branch (~42 tokens). Reserving + * the standalone figure under-counts ~2.5x, and capOversizedChunks then emits + * body-capped pieces that re-emerge over the cap once the header is re-added + * (measured: 2,006- and 2,023-token chunks on src/core/migrate.ts against a + * 2,000 cap, where the pre-#3564 chunker emitted none). + * + * Taking max(cl100k, weighted) unconditionally is a true bound because the + * weighted form is additive per char class, so weighted(a + b) <= + * weighted(a) + weighted(b), and cl100k does not gain tokens across the + * header's trailing blank line. + */ +export function estimateEmbedTokensCeiling(text: string): number { + return Math.max(estimateTokens(text), weightedTokens(text, (text.match(CJK_CHARS_G) || []).length)); +} diff --git a/test/chunkers/embed-estimate-home.test.ts b/test/chunkers/embed-estimate-home.test.ts new file mode 100644 index 000000000..c0dad5534 --- /dev/null +++ b/test/chunkers/embed-estimate-home.test.ts @@ -0,0 +1,216 @@ +/** + * #3477 follow-up — the two items flagged in its merge review: + * + * (1) splitToTokenBudget's hard-split budget is derived from each piece's + * own measured density (chars per estimated token) instead of a fixed + * 3.5 chars/token guess. URL-dense ASCII JSON runs ~2.6 chars/token, so + * the old budget let 2,070–2,299-token slices past a 2,000 cap. + * + * (3) estimateTokens/estimateEmbedTokens moved below both chunkers + * (token-estimate.ts — cjk.ts itself is a check:fuzz-purity target and + * tiktoken's loader pulls node:fs; code.ts imports recursive.ts, so + * recursive.ts could never reuse them without a cycle), letting + * capByChars bound estimated embedding tokens too — + * the fix for the #3037 shape (CJK-dense chunks under maxChars=6000 + * but over the embedder context, permanently unembeddable, silently) + * and #2826's markdown reproduction (URL-dense Korean at defaults + * emitting ~4,200-char / ~2,200-token chunks). + */ + +import { describe, test, expect } from 'bun:test'; +import { + estimateTokens as estimateTokensViaCode, + estimateEmbedTokens as estimateEmbedTokensViaCode, + chunkCodeText, +} from '../../src/core/chunkers/code.ts'; +import { estimateTokens, estimateEmbedTokens, DEFAULT_MAX_CHUNK_TOKENS } from '../../src/core/chunkers/token-estimate.ts'; +import { chunkText } from '../../src/core/chunkers/recursive.ts'; + +/** A string survives a UTF-8 round trip only if it is well-formed UTF-16 — + * i.e. no orphaned surrogate half. Postgres rejects a lone surrogate inside a + * `::jsonb` cast and aborts the whole batch (#2011). */ +function isWellFormedUtf16(s: string): boolean { + return Buffer.from(s, 'utf8').toString('utf8') === s; +} + +/** ASCII structured header + CJK-dense body: the shape where the re-added + * header's token cost is SUPER-additive (see the header-reservation test). */ +function mixedScriptTypeScript(lines: number): string { + const body = Array.from({ length: lines }, (_, i) => + ` // 설정 항목 ${i}: 환경 변수와 기본값을 병합해 최종 구성을 만든다 (참조 config/${i})\n` + + ` const option_${i} = resolveOption('key_${i}', defaults.key_${i}, { 우선순위: ${i} });`, + ).join('\n'); + return `export function loadEverything(defaults: Defaults) {\n${body}\n return { ok: true };\n}\n`; +} + +/** URL-dense ASCII JSON — ~2.6 chars/token, the (1) leak shape. */ +function urlDenseAsciiJson(targetChars: number): string { + const entries: string[] = []; + let i = 0; + let len = 0; + while (len < targetChars) { + const hex = ((i * 48271) % 65521).toString(16) + ((i * 69621) % 233280).toString(16) + ((i * 16807) % 104729).toString(16); + const row = + ` "row_${i}": { "href": "https://api.example.com/v3/resources/${hex}?sig=ab${i}cd&expires=17${i}&scope=read%2Fwrite", "etag": "W/\\"x${i}y\\"", "n": ${i} }`; + entries.push(row); + len += row.length; + i++; + } + return `{\n${entries.join(',\n')}\n}`; +} + +/** URL-dense Korean rollup lines — #2826's markdown reproduction shape. */ +function urlDenseKoreanMarkdown(lines: number): string { + return Array.from({ length: lines }, (_, i) => + `- 항목 ${i}: 검증용 한국어 설명 문장이 이어집니다 · 링크: https://docs.example.com/pages/${String(i).padStart(32, '0')}?v=abcdef0123456789&ref=sample`, + ).join('\n'); +} + +describe('estimator home (cjk.ts) — the (3) move', () => { + test('code.ts re-exports are the same functions (import sites unchanged)', () => { + expect(estimateTokensViaCode).toBe(estimateTokens); + expect(estimateEmbedTokensViaCode).toBe(estimateEmbedTokens); + }); +}); + +describe('splitToTokenBudget — measured hard-split budget, the (1) leak', () => { + test('URL-dense ASCII json fence stays under the default cap — headers included, no slack (previously 2,379-token max)', async () => { + const src = urlDenseAsciiJson(14_400); + const chunks = await chunkCodeText(src, 'fence.json'); + expect(chunks.length).toBeGreaterThan(1); + for (const c of chunks) { + // STRICT: the emitted chunk (structured header + body) fits the cap. + // The splitter reserves the header's tokens from the body budget, so + // no "body capped, header pushed it over" residue survives. + expect(estimateEmbedTokens(c.text)).toBeLessThanOrEqual(DEFAULT_MAX_CHUNK_TOKENS); + } + // Content preserved — first and last rows survive the re-split. + const joined = chunks.map((c) => c.text).join('\n'); + expect(joined).toContain('"row_0"'); + expect(joined).toContain('scope=read%2Fwrite'); + }); +}); + +describe('capByChars — token-aware belt, the (3) payoff', () => { + test('URL-dense Korean markdown at defaults stays under the token budget (previously ~2,200-token chunks)', () => { + const chunks = chunkText(urlDenseKoreanMarkdown(120)); + expect(chunks.length).toBeGreaterThan(0); + for (const c of chunks) { + expect(estimateEmbedTokens(c.text)).toBeLessThanOrEqual(DEFAULT_MAX_CHUNK_TOKENS); + } + }); + + test('low-density bilingual table (the #3037 shape) splits under the budget instead of shipping over-context chunks', () => { + // #3037's failing chunk: mostly-ASCII with a CJK minority (their repro: + // 6001 chars, 942 CJK). Density sits BELOW CJK_DENSITY_THRESHOLD, so the + // word pipeline counts whitespace tokens and happily builds multi- + // thousand-char chunks; the old belt only checked chars (6000), so these + // shipped at token counts past strict embedder contexts. + const table = Array.from({ length: 120 }, (_, i) => + `ITEM-${String(i).padStart(6, '0')} | 环境配置说明 段落${i} | https://wiki.example.com/pages/${String(i).padStart(20, '0')}?rev=${i}&lang=zh | flags=prod,readonly,audit`, + ).join('\n'); + const chunks = chunkText(table); + expect(chunks.length).toBeGreaterThan(1); + for (const c of chunks) { + expect(estimateEmbedTokens(c.text)).toBeLessThanOrEqual(DEFAULT_MAX_CHUNK_TOKENS); + } + }); + + // DROPPED: a 'mixed-density input — every slice re-checked' test used to sit + // here. It passed against the pre-cap behavior too, so it proved nothing. + // Measured why: its dense run was whitespace-less CJK, which countCJKAwareWords + // scores per character, so the word pipeline had already cut it to <=300-char + // pieces before the belt was ever consulted (baseline: 20 dense chunks, max + // 300 chars, cap never fires). The discriminating shape for the belt is the + // LOW-density one — CJK below the density threshold, where the word pipeline + // counts whitespace tokens and builds multi-thousand-char chunks — and that + // is the #3037 bilingual-table test above. + + test('ASCII prose under both budgets passes through untouched (single chunk, verbatim)', () => { + const prose = 'plain english prose that fits comfortably inside every budget. '.repeat(20).trim(); + const chunks = chunkText(prose); + expect(chunks.length).toBe(1); + expect(chunks[0]!.text).toBe(prose); + }); +}); + +describe('header reservation is an UPPER bound — estimateEmbedTokens is super-additive', () => { + test('mixed-script source: the re-added ASCII header never pushes a piece over the cap', async () => { + // estimateEmbedTokens takes max(cl100k, per-char-class weighted) and only + // switches on the weighted branch when the text contains CJK. An ASCII + // header measured ALONE therefore costs cl100k (a 59-char header = 17 + // tokens), but once it is glued onto a body containing CJK the whole + // chunk measures on the weighted branch, where those same 59 chars cost + // ~0.75/char (~42 tokens). Reserving the standalone figure under-counts + // ~2.5x and the re-headered piece re-emerges over the cap: measured 2,006 + // and 2,023-token chunks on src/core/migrate.ts against a 2,000 cap, + // where the pre-fix chunker emitted none (maxEst 1,439). + const src = mixedScriptTypeScript(220); + const chunks = await chunkCodeText(src, 'src/config/loader.ts'); + expect(chunks.length).toBeGreaterThan(1); + for (const c of chunks) { + expect(estimateEmbedTokens(c.text)).toBeLessThanOrEqual(DEFAULT_MAX_CHUNK_TOKENS); + } + }); +}); + +describe('astral surrogate pairs survive the derived-window splits', () => { + // The hard-split budget and the capByChars window are now DERIVED from + // measured density, so they land on arbitrary parity. A raw .slice() at + // such an offset orphans a UTF-16 surrogate half. The repo already ships + // src/core/text-safe.ts:safeSplitIndex for exactly this: #2011 — + // `extract --stale` died at ~1,550 pages because excerpt() raw-sliced a + // window boundary through an emoji, and a lone surrogate is rejected by + // Postgres inside a ::jsonb cast, aborting the WHOLE batch. + const ASTRAL = '\u{20000}\u{20001}\u{20002}\u{1F600}'; + + test('markdown chunker emits well-formed UTF-16 on an astral-only document', () => { + const chunks = chunkText(ASTRAL.repeat(1500)); + expect(chunks.length).toBeGreaterThan(1); + const corrupted = chunks.filter((c) => !isWellFormedUtf16(c.text)); + expect(corrupted.length).toBe(0); + }); + + test('code chunker emits well-formed UTF-16 on a whitespace-less astral blob', async () => { + const chunks = await chunkCodeText(JSON.stringify({ k: ASTRAL.repeat(3000) }), 'blob.json'); + expect(chunks.length).toBeGreaterThan(1); + const corrupted = chunks.filter((c) => !isWellFormedUtf16(c.text)); + expect(corrupted.length).toBe(0); + }); + + test('astral content is preserved across the split, not dropped', () => { + const doc = ASTRAL.repeat(1500); + const rejoined = chunkText(doc).map((c) => c.text).join(''); + expect(rejoined).toContain(ASTRAL.repeat(4)); + }); +}); + +describe('hard split leaves no runt chunks', () => { + // A fixed `i += charBudget` stride sheds `length mod charBudget` chars as a + // standalone piece at EVERY recursion level, and buildChunk re-headers each + // one into its own embedding row — near-empty fragments that still match + // queries. Evening the slice width out costs nothing: the piece COUNT is + // ceil(length / charBudget) either way, so the same content is redistributed + // over the same number of chunks; only the smallest piece changes (measured + // on a 45.6K blob: min 589 -> 1900 chars, same 24 pieces). + const SLIVER_CHARS = 200; + + test('URL-dense ASCII json fence sheds no sliver (pre-fix: 3-56 char fragments)', async () => { + const chunks = await chunkCodeText(urlDenseAsciiJson(14_400), 'fence.json'); + expect(chunks.length).toBeGreaterThan(1); + const bodies = chunks.map((c) => c.text.replace(/^\[[^\]]+\] [^\n]+\n\n/, '')); + expect(bodies.filter((b) => b.length < SLIVER_CHARS)).toEqual([]); + }); + + test('whitespace-less blob — the pure hard-split path — sheds no sliver', async () => { + // No whitespace to break on, so recursiveChunk cannot help and every + // boundary comes from the hard splitter: the shape where the remainder + // stride was most visible (measured 32/54/56-char chunks pre-fix). + const blob = JSON.stringify({ d: 'https://example.com/a/b/c?q=1&r=2#frag-'.repeat(1200) }); + const chunks = await chunkCodeText(blob, 'blob.json'); + expect(chunks.length).toBeGreaterThan(1); + const bodies = chunks.map((c) => c.text.replace(/^\[[^\]]+\] [^\n]+\n\n/, '')); + expect(bodies.filter((b) => b.length < SLIVER_CHARS)).toEqual([]); + for (const c of chunks) expect(estimateEmbedTokens(c.text)).toBeLessThanOrEqual(DEFAULT_MAX_CHUNK_TOKENS); + }); +});