From d859d2470b0dfe54cdf1935e543a47c67048d215 Mon Sep 17 00:00:00 2001 From: dvcdsys Date: Tue, 18 Aug 2026 12:28:32 +0100 Subject: [PATCH 1/6] feat(voyage): exact token counting via the model's own BPE tokenizer MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The voyage provider had no tokenizer, so it guessed token cost as len(bytes)/2 and lived with the consequences: batches sized against a guess, a TPM throttle metering phantom tokens, and over-long inputs cut into byte windows whose vectors were then averaged. This adds a pure-Go, count-only implementation of the model's actual pipeline and wires it in behind an interface, so the guessing path becomes the fallback rather than the only option. WHAT THE HEURISTIC ACTUALLY COST, measured against Voyage's own usage.total_tokens (412 single-input requests) and against the HuggingFace Rust tokenizer as oracle (20k real chunks from a 45-repo corpus): - overestimates by 1.94x on average, so the TPM token bucket subtracts roughly twice the tokens actually spent. At the observed indexing rate (128 chunks/s, mean 221 tokens/chunk) real usage was ~1.7M TPM while the limiter believed 3.3M and throttled against a 3M budget — the run was self-limited to about half the tier's real capacity. - AND still undercounts 0.5% of chunks, worst case -41%, which is the direction that ships an over-limit POST. The safety the constant was chosen for is not actually delivered. - the "~1.4 bytes/token worst case" in the old comment is optimistic; measured worst is ~1.18. WHY A HAND-ROLLED TOKENIZER RATHER THAN A LIBRARY. Both pure-Go candidates were evaluated against the same oracle on the same corpus: - sugarme/tokenizer cannot load this tokenizer.json at all. It panics in regexp.MustCompile: the Qwen2 pre-tokenizer pattern contains \s+(?!\S), and Go's RE2 has no lookahead. - the ollama tree's pure-Go tokenizer knows about that and rewrites the pattern, but its compensation only covers space indentation. Tabs take a different branch, so it mismatches HF on 8.19% of real chunks (30.2% of chunks containing a tab), 1471 of them undercounts, worst -16.7% on tab-indented JSON. It also skips the declared NFC normalizer, so decomposed input costs an extra token. bpecount instead implements the seven alternation branches by hand with Perl leftmost-first semantics, which is what makes the whitespace precedence come out right. Validated at 0 mismatches against the HF oracle across 70,412 inputs (412 Voyage-verified, 20k real chunks, 50k fuzz). It needs no new module — golang.org/x/text was already an indirect dependency — and counts 27.5k chunks/s single-threaded, which is orders of magnitude above any indexing rate on the 2-vCPU box this targets. Only the merges array is read; the vocabulary is not needed for a count. New pre-tokens are merged through a linked list plus a rank heap rather than the naive pair rescan, because minified and generated files are real: a 30 KB run of one character took 9s the naive way and takes ~16ms this way. WHAT CHANGES FOR CALLERS. tokenizer.Budget is one interface, not a mandatory counter plus an optional splitter, so the chunker takes a single constructor argument. ExactCounts() is how a caller learns whether it may size to the limit or must keep the old margin. The doc comment records the cost asymmetry that makes two methods worth having: both do the same single pass over the same memo, but CountTokens allocates nothing and runs per chunk (1.9M times on the reference corpus) while SplitPoints builds an offsets slice and runs only when an input does not fit (5 times on that same corpus). SplitPoints is exact, not iterative. Because Split runs with Isolated behaviour and BPE is applied per pre-token, merges never cross a pre-token boundary, so token counts are additive over pre-tokens and a single left-to-right pass yields cut points whose pieces provably sum to the whole. Only a single pre-token larger than the budget needs a search, and it gets a binary search on bytes inside that one pre-token. Batch cap goes 80K -> 115K when counts are exact: the 40K of headroom existed to absorb the heuristic's undercount against Voyage's 120K hard limit, not to absorb anything on Voyage's side. An operator's explicit MaxTokensPerRequest still wins over both. MEASURED EFFECT, so nobody expects the wrong thing. Simulating both packers over 200k real chunks: 1724 requests today vs 1568 with exact counts — 1.10x, not the 2x the 1.94x inflation suggests. Batches are bound by the 128-input cap of voyage-code-*, not by tokens (mean chunk is 221 tokens, so 128 inputs is ~28K against an 80K cap). The throughput win is in the TPM throttle, not in packing; the packing win is real but small. DELIBERATELY NOT IN THIS CHANGE. The chunker does not yet consume Budget — splitOversizeInput still byte-windows on the provider side, and only the provider-side counting is switched over. That integration is the follow-up this interface exists for, and it is where averaging over byte windows finally goes away. On the reference corpus 5 chunks in 1.9M exceeded the old 30 KB threshold, so the damage is bounded and self-healing: those files re-embed correctly the next time they change, and can be found with SELECT file_path FROM vector_contents WHERE LENGTH(content) > 30000. TESTING. The golden counts need the real 7 MB tokenizer.json, which is not in the repo, so those tests skip on a clean checkout — CI coverage comes from a synthetic merge table that exercises the splitter, the merge loop and the additivity property SplitPoints depends on. Packaging that file (embed vs fetch-and-cache; only ~3.5 MB of it is load-bearing) is still open, which is why the path is an operator-set config field for now and an unreadable path degrades to the estimate with a warning rather than failing to start. Co-Authored-By: Claude Opus 5 --- .../embeddings/provider/voyage/factory.go | 7 + .../embeddings/provider/voyage/voyage.go | 122 ++++- .../embeddings/provider/voyage/voyage_test.go | 37 +- .../internal/tokenizer/bpecount/bpecount.go | 481 ++++++++++++++++++ .../tokenizer/bpecount/bpecount_test.go | 233 +++++++++ server/internal/tokenizer/budget.go | 49 ++ 6 files changed, 920 insertions(+), 9 deletions(-) create mode 100644 server/internal/tokenizer/bpecount/bpecount.go create mode 100644 server/internal/tokenizer/bpecount/bpecount_test.go create mode 100644 server/internal/tokenizer/budget.go diff --git a/server/internal/embeddings/provider/voyage/factory.go b/server/internal/embeddings/provider/voyage/factory.go index 64d12c57..926e4869 100644 --- a/server/internal/embeddings/provider/voyage/factory.go +++ b/server/internal/embeddings/provider/voyage/factory.go @@ -33,6 +33,13 @@ func (factory) SchemaJSON() []byte { Description: "int8 is dequantized to float32 on the server side.", }, {Name: "truncation", Label: "Truncate over-length input", Kind: "bool", Default: true}, + { + Name: "tokenizer_path", Label: "Tokenizer file", Kind: "string", + Description: "Absolute path to the model's tokenizer.json (huggingface.co/voyageai/). " + + "Set it and token counts become exact: batches pack to the real limit instead of a " + + "byte guess that overestimates ~2x, and over-long inputs split on token boundaries " + + "instead of byte windows. Empty falls back to the estimate.", + }, {Name: "api_key_env", Label: "API key env var", Kind: "secret-env", Required: true, Default: defaultAPIKeyEnv}, }, } diff --git a/server/internal/embeddings/provider/voyage/voyage.go b/server/internal/embeddings/provider/voyage/voyage.go index 04318f85..14f0bb96 100644 --- a/server/internal/embeddings/provider/voyage/voyage.go +++ b/server/internal/embeddings/provider/voyage/voyage.go @@ -36,6 +36,7 @@ import ( "golang.org/x/time/rate" "github.com/dvcdsys/code-index/server/internal/embeddings/provider" + "github.com/dvcdsys/code-index/server/internal/tokenizer/bpecount" ) // voyageBatchTooLargeRegex matches Voyage's per-batch token-limit @@ -159,6 +160,15 @@ type Config struct { // all in-flight + recent requests). 0 = no throttling. RateLimitTPM int `json:"rate_limit_tpm,omitempty"` + // TokenizerPath points at the model's tokenizer.json (the file + // Voyage publishes at huggingface.co/voyageai/). When set and + // loadable, token counts become EXACT and the per-batch cap rises to + // exactTokensPerBatch — the 40K of headroom the byte heuristic needed + // is headroom against the heuristic, not against Voyage. When empty or + // unreadable the provider logs once and falls back to estimateTokens, + // so a missing file degrades throughput, never correctness. + TokenizerPath string `json:"tokenizer_path,omitempty"` + // MaxInputsPerRequest overrides defaultMaxBatchSize. 0 = use // the default (128, safe for voyage-code-*). Operators running // only voyage-3* may bump this to 1000 for fewer round-trips. @@ -266,6 +276,10 @@ type Provider struct { // budget is a sliding minute and bursting saves nothing. reqLimiter *rate.Limiter + // counter is the model's real tokenizer, or nil when no tokenizer.json + // was configured or it failed to load. Safe for concurrent use. + counter *bpecount.Counter + // tokenLimiter caps tokens-per-minute when cfg.RateLimitTPM > 0. // Burst is set to maxTokensPerBatch so a single full-budget POST // can pass even when the bucket is otherwise empty (we'd just @@ -294,6 +308,19 @@ func New(cfg Config, secrets provider.SecretLookup, logger *slog.Logger) *Provid secrets: secrets, http: &http.Client{Timeout: 60 * time.Second}, } + if cfg.TokenizerPath != "" { + c, err := bpecount.Load(cfg.TokenizerPath) + if err != nil { + // Not fatal: the byte heuristic still works. Loud because the + // operator asked for exact counts and is not getting them. + logger.Warn("voyage: tokenizer load failed, falling back to byte estimate", + "path", cfg.TokenizerPath, "err", err) + } else { + p.counter = c + logger.Info("voyage: exact token counting enabled", "path", cfg.TokenizerPath) + } + } + // Convert RPM/TPM to per-second token-bucket rates. burst on the // request bucket is 1 (one request worth of "credit"); burst on // the token bucket equals one full POST so we don't deadlock a @@ -302,7 +329,7 @@ func New(cfg Config, secrets provider.SecretLookup, logger *slog.Logger) *Provid p.reqLimiter = rate.NewLimiter(rate.Limit(float64(cfg.RateLimitRPM)/60.0), 1) } if cfg.RateLimitTPM > 0 { - p.tokenLimiter = rate.NewLimiter(rate.Limit(float64(cfg.RateLimitTPM)/60.0), cfg.maxTokensPerBatch()) + p.tokenLimiter = rate.NewLimiter(rate.Limit(float64(cfg.RateLimitTPM)/60.0), p.maxTokensPerBatch()) } return p } @@ -428,7 +455,7 @@ func (p *Provider) embedAndAverage(ctx context.Context, texts []string, inputTyp } // Phase 2: batch + POST as before, on the expanded slice. - batches := planBatches(expanded, p.cfg.maxBatchSize(), p.cfg.maxTokensPerBatch()) + batches := planBatches(expanded, p.cfg.maxBatchSize(), p.maxTokensPerBatch(), p.CountTokens) if len(batches) > 1 { p.logger.Info("voyage: splitting batch", "model", p.cfg.Model, @@ -561,7 +588,10 @@ func (p *Provider) embedWithAdaptiveSplit(ctx context.Context, texts []string, i // operator can override them via the admin form when their tier or // chosen model allows a higher cap (e.g. voyage-3-large at 1000 // inputs/POST instead of 128). -func planBatches(texts []string, maxInputs, maxTokens int) [][]string { +func planBatches(texts []string, maxInputs, maxTokens int, count func(string) int) [][]string { + if count == nil { + count = estimateTokens + } if len(texts) == 0 { return nil } @@ -569,7 +599,7 @@ func planBatches(texts []string, maxInputs, maxTokens int) [][]string { var current []string currentTokens := 0 for _, t := range texts { - est := estimateTokens(t) + est := count(t) // Close the current batch when adding this text would exceed // either limit (and the batch already has something to send). if len(current) > 0 && (len(current) >= maxInputs || currentTokens+est > maxTokens) { @@ -586,9 +616,12 @@ func planBatches(texts []string, maxInputs, maxTokens int) [][]string { return batches } -// estimateTokens returns a conservative upper bound on the token cost -// of one text, in Voyage's tokenizer. Uses byte-length divided by a -// chars-per-token heuristic; see bytesPerToken doc for rationale. +// estimateTokens is the FALLBACK used only when no tokenizer.json is +// loaded. Measured against Voyage's own usage.total_tokens on 20k real +// chunks it overestimates by 1.94x on average — which wastes round-trips +// — while still undercounting 0.5% of chunks, worst case -41%. That is +// the wrong error in both directions, and it is why loading the real +// tokenizer is worth the 7 MB: see Provider.countTokens. func estimateTokens(s string) int { return len(s) / bytesPerToken } @@ -784,3 +817,78 @@ func dequantize(raw json.RawMessage, dtype string) ([]float32, error) { func (p *Provider) apiKey() (string, bool) { return provider.ResolveAPIKey(p.secrets, p.cfg.APIKeyEnv) } + +// ---------- tokenizer.Budget ---------- +// +// Implemented on Provider so the chunker can be handed the live provider and +// stay ignorant of which model is active: only the provider knows whether +// tokens come from a real BPE table, from llama-server's /tokenize, or from a +// byte guess. + +// exactTokensPerBatch is the per-POST cap once counts are exact. +// +// The 80K default exists to survive the byte heuristic's ~43% undercount +// against Voyage's 120K hard limit. With the real tokenizer the count is the +// count — measured against usage.total_tokens it is never below what Voyage +// bills — so the headroom collapses to a margin for Voyage-side accounting +// drift rather than for our own error. +const exactTokensPerBatch = 115_000 + +// maxInputTokens is voyage-code-3's per-input context window. The 32K applies +// to voyage-code-* and voyage-3*; smaller models would need a table here, but +// undershooting only costs an unnecessary split. +const maxInputTokens = 32_000 + +// maxTokensPerBatch is the provider-level cap: an explicit operator override +// wins, then the exact-counting cap, then the conservative byte-heuristic one. +func (p *Provider) maxTokensPerBatch() int { + if p.cfg.MaxTokensPerRequest > 0 { + return p.cfg.MaxTokensPerRequest + } + if p.counter != nil { + return exactTokensPerBatch + } + return defaultMaxTokensPerBatch +} + +// MaxInputTokens reports the model's context window for a single input. +func (p *Provider) MaxInputTokens() int { return maxInputTokens } + +// ExactCounts reports whether CountTokens/SplitPoints are exact rather than +// estimated. False means no tokenizer.json was loaded. +func (p *Provider) ExactCounts() bool { return p.counter != nil } + +// CountTokens returns the token cost of s. Allocation-free on the exact path; +// this is the hot one — it runs for every chunk that gets embedded. +func (p *Provider) CountTokens(s string) int { + if p.counter != nil { + return p.counter.Count(s) + } + return estimateTokens(s) +} + +// SplitPoints returns byte offsets at which s must be cut so no piece exceeds +// budget tokens, and s's total token count. +// +// Exact when a tokenizer is loaded: cuts land on pre-token boundaries, where +// BPE merges never reach across, so the pieces provably add up to the whole. +// Without a tokenizer it degrades to rune-aligned byte windows — the old +// behaviour, kept only so a caller that ignores ExactCounts still gets +// something it can send. Check ExactCounts before trusting these. +func (p *Provider) SplitPoints(s string, budget int) ([]int, int) { + if p.counter != nil { + return p.counter.SplitPoints(s, budget) + } + var offsets []int + maxBytes := budget * bytesPerToken + if maxBytes <= 0 { + return nil, estimateTokens(s) + } + for off := maxBytes; off < len(s); off += maxBytes { + for off > 0 && !utf8.RuneStart(s[off]) { + off-- + } + offsets = append(offsets, off) + } + return offsets, estimateTokens(s) +} diff --git a/server/internal/embeddings/provider/voyage/voyage_test.go b/server/internal/embeddings/provider/voyage/voyage_test.go index 542e0a3c..94869308 100644 --- a/server/internal/embeddings/provider/voyage/voyage_test.go +++ b/server/internal/embeddings/provider/voyage/voyage_test.go @@ -5,6 +5,7 @@ import ( "encoding/base64" "encoding/json" "fmt" + "github.com/dvcdsys/code-index/server/internal/tokenizer" "io" "net/http" "net/http/httptest" @@ -187,7 +188,7 @@ func TestPlanBatches_SplitsByTokenBudget(t *testing.T) { small := "tiny" texts := []string{big, small, small, small, small, small} - batches := planBatches(texts, defaultMaxBatchSize, defaultMaxTokensPerBatch) + batches := planBatches(texts, defaultMaxBatchSize, defaultMaxTokensPerBatch, nil) if len(batches) < 2 { t.Fatalf("expected at least 2 batches, got %d", len(batches)) } @@ -213,7 +214,7 @@ func TestPlanBatches_RespectsCountCap(t *testing.T) { for i := range texts { texts[i] = "chunk" } - batches := planBatches(texts, defaultMaxBatchSize, defaultMaxTokensPerBatch) + batches := planBatches(texts, defaultMaxBatchSize, defaultMaxTokensPerBatch, nil) if len(batches) != 2 { t.Fatalf("expected 2 batches (128 + 72), got %d", len(batches)) } @@ -652,3 +653,35 @@ func TestInt8Dequantize_Base64(t *testing.T) { t.Errorf("base64 int8 dequantized values out of range: %v", v) } } + +// TestProviderSatisfiesBudget pins the provider to the interface the chunker +// consumes. A compile-time assertion rather than a runtime test: the whole +// point of the interface is that the chunker never imports this package. +func TestProviderSatisfiesBudget(t *testing.T) { + var _ tokenizer.Budget = (*Provider)(nil) +} + +// TestFallbackWithoutTokenizer covers the degraded path: no tokenizer.json +// means estimates, the conservative batch cap, and ExactCounts()==false so a +// caller can widen its margins instead of trusting the number. +func TestFallbackWithoutTokenizer(t *testing.T) { + p := &Provider{cfg: Config{Model: "voyage-code-3"}} + if p.ExactCounts() { + t.Error("ExactCounts must be false without a tokenizer") + } + if got := p.maxTokensPerBatch(); got != defaultMaxTokensPerBatch { + t.Errorf("batch cap = %d, want the conservative %d", got, defaultMaxTokensPerBatch) + } + if got, want := p.CountTokens("hello world"), len("hello world")/bytesPerToken; got != want { + t.Errorf("CountTokens = %d, want the byte estimate %d", got, want) + } +} + +// TestOperatorOverrideWinsOverExactCap — an explicit MaxTokensPerRequest is +// the operator's call and must not be silently raised by exact counting. +func TestOperatorOverrideWinsOverExactCap(t *testing.T) { + p := &Provider{cfg: Config{Model: "voyage-code-3", MaxTokensPerRequest: 42_000}} + if got := p.maxTokensPerBatch(); got != 42_000 { + t.Errorf("batch cap = %d, want the operator's 42000", got) + } +} diff --git a/server/internal/tokenizer/bpecount/bpecount.go b/server/internal/tokenizer/bpecount/bpecount.go new file mode 100644 index 00000000..381ef416 --- /dev/null +++ b/server/internal/tokenizer/bpecount/bpecount.go @@ -0,0 +1,481 @@ +// Package bpecount is a minimal, count-only byte-level BPE tokenizer for +// GPT-2/Qwen2-style tokenizer.json files (voyage-code-3, Qwen2, GPT-4o…). +// +// It reproduces the HuggingFace pipeline: +// +// normalizer = NFC +// pretokenize = Split(Qwen2 regex, Isolated) + ByteLevel(add_prefix_space=false) +// model = BPE (greedy lowest-rank merge) +// +// The Split regex contains `\s+(?!\S)`, a negative lookahead Go's RE2 +// cannot express, so the splitter is hand-rolled rather than compiled. +// Only a COUNT is produced — no ids, no offsets. +package bpecount + +import ( + "container/heap" + "encoding/json" + "fmt" + "os" + "strings" + "sync" + "unicode" + "unicode/utf8" + + "golang.org/x/text/unicode/norm" +) + +// Counter holds the merge table and a memo of pre-token → token count. +type Counter struct { + merges map[string]int32 // "left right" -> rank + + mu sync.RWMutex + memo map[string]int +} + +type tokJSON struct { + Model struct { + Type string `json:"type"` + Merges json.RawMessage `json:"merges"` + } `json:"model"` +} + +// Load reads a tokenizer.json and keeps only what a count needs: the merges. +func Load(path string) (*Counter, error) { + b, err := os.ReadFile(path) + if err != nil { + return nil, err + } + return LoadBytes(b) +} + +func LoadBytes(b []byte) (*Counter, error) { + var tj tokJSON + if err := json.Unmarshal(b, &tj); err != nil { + return nil, err + } + if tj.Model.Type != "BPE" { + return nil, fmt.Errorf("bpecount: unsupported model type %q", tj.Model.Type) + } + // merges is either ["a b", ...] (v1) or [["a","b"], ...] (v2). + var flat []string + m := make(map[string]int32) + if err := json.Unmarshal(tj.Model.Merges, &flat); err == nil { + for i, s := range flat { + m[s] = int32(i) + } + } else { + var pairs [][]string + if err := json.Unmarshal(tj.Model.Merges, &pairs); err != nil { + return nil, fmt.Errorf("bpecount: merges: %w", err) + } + for i, p := range pairs { + if len(p) == 2 { + m[p[0]+" "+p[1]] = int32(i) + } + } + } + return &Counter{merges: m, memo: make(map[string]int, 1<<16)}, nil +} + +// ---------- byte-level alphabet (GPT-2 bytes_to_unicode) ---------- + +var byteRune [256]rune + +func init() { + for b := 0; b < 256; b++ { + r := rune(b) + switch { + case r == 0xad: + r = 0x143 + case r <= 0x20: + r += 0x100 + case r >= 0x7f && r <= 0xa0: + r += 0xa2 + } + byteRune[b] = r + } +} + +// ---------- hand-rolled splitter ---------- +// +// Qwen2 pattern, alternation tried left to right (Perl leftmost-first): +// +// (?i:'s|'t|'re|'ve|'m|'ll|'d) +// [^\r\n\p{L}\p{N}]?\p{L}+ +// \p{N} +// ?[^\s\p{L}\p{N}]+[\r\n]* +// \s*[\r\n]+ +// \s+(?!\S) +// \s+ +// +// Every rune is covered by some branch, so Split(Isolated) yields no gaps. + +func isL(r rune) bool { return unicode.IsLetter(r) } +func isN(r rune) bool { return unicode.IsNumber(r) } +func isWS(r rune) bool { return unicode.IsSpace(r) } +func isNL(r rune) bool { return r == '\r' || r == '\n' } + +var contractions = []string{"s", "t", "re", "ve", "m", "ll", "d"} + +// nextToken returns the byte length of the pre-token starting at s[0]. +func nextToken(s string) int { + r0, w0 := decode(s, 0) + + // A: contraction + if r0 == '\'' && len(s) > w0 { + low := strings.ToLower(s) + for _, c := range contractions { + if strings.HasPrefix(low[w0:], c) { + return w0 + len(c) + } + } + } + + // B: [^\r\n\p{L}\p{N}]? \p{L}+ + { + i := 0 + if !isNL(r0) && !isL(r0) && !isN(r0) { + i = w0 + } + j := i + for j < len(s) { + r, w := decode(s, j) + if !isL(r) { + break + } + j += w + } + if j > i { // at least one letter followed + return j + } + } + + // C: single \p{N} + if isN(r0) { + return w0 + } + + // D: " ?" [^\s\p{L}\p{N}]+ [\r\n]* + { + i := 0 + if r0 == ' ' { + i = w0 + } + j := i + for j < len(s) { + r, w := decode(s, j) + if isWS(r) || isL(r) || isN(r) { + break + } + j += w + } + if j > i { + for j < len(s) { + r, w := decode(s, j) + if !isNL(r) { + break + } + j += w + } + return j + } + } + + // E/F/G: whitespace run. + if isWS(r0) { + // maximal whitespace run + end := 0 + lastNL := -1 + for end < len(s) { + r, w := decode(s, end) + if !isWS(r) { + break + } + if isNL(r) { + lastNL = end + w + } + end += w + } + // E: \s*[\r\n]+ — run truncated after its LAST \r or \n. + if lastNL >= 0 { + return lastNL + } + // F: \s+(?!\S) — whole run at EOF, else run minus its last rune. + if end == len(s) { + return end + } + _, lw := decodeLast(s[:end]) + if end-lw > 0 { + return end - lw + } + // G: \s+ (single whitespace rune followed by a non-space) + return end + } + + // Unreachable for well-formed input; make progress anyway. + return w0 +} + +func decode(s string, i int) (rune, int) { + if s[i] < utf8.RuneSelf { + return rune(s[i]), 1 + } + return utf8.DecodeRuneInString(s[i:]) +} + +func decodeLast(s string) (rune, int) { + return utf8.DecodeLastRuneInString(s) +} + +// ---------- BPE ---------- + +func (c *Counter) bpeLen(piece string) int { + c.mu.RLock() + n, ok := c.memo[piece] + c.mu.RUnlock() + if ok { + return n + } + n = c.bpe(piece) + c.mu.Lock() + if len(c.memo) < 1<<20 { + c.memo[piece] = n + } + c.mu.Unlock() + return n +} + +// node is one symbol in the doubly-linked list the merge loop walks. +type node struct { + prev, next int + s string + alive bool +} + +// cand is a candidate merge sitting in the priority queue. +type cand struct { + rank int32 + l, r int + // len of the two symbols when the candidate was pushed; a stale entry + // (one side already merged into something longer) is detected by comparing. + ll, rl int +} + +type candHeap []cand + +func (h candHeap) Len() int { return len(h) } +func (h candHeap) Less(i, j int) bool { + if h[i].rank != h[j].rank { + return h[i].rank < h[j].rank + } + return h[i].l < h[j].l // ties: leftmost first, matching HF +} +func (h candHeap) Swap(i, j int) { h[i], h[j] = h[j], h[i] } +func (h *candHeap) Push(x any) { *h = append(*h, x.(cand)) } +func (h *candHeap) Pop() any { + old := *h + n := len(old) + v := old[n-1] + *h = old[:n-1] + return v +} + +// bpe applies the greedy lowest-rank-first merge with a linked list + heap, +// so a pathological single pre-token (a 30 KB run of '-' or one enormous +// identifier) stays near-linear instead of the O(n^2) rescan a naive loop does. +func (c *Counter) bpe(piece string) int { + nodes := make([]node, 0, len(piece)) + for _, r := range piece { + i := len(nodes) + nodes = append(nodes, node{prev: i - 1, next: i + 1, s: string(r), alive: true}) + } + n := len(nodes) + if n < 2 { + return n + } + nodes[n-1].next = -1 + + h := make(candHeap, 0, n) + push := func(l, r int) { + if l < 0 || r < 0 || r >= n { + return + } + if rk, ok := c.merges[nodes[l].s+" "+nodes[r].s]; ok { + h = append(h, cand{rank: rk, l: l, r: r, ll: len(nodes[l].s), rl: len(nodes[r].s)}) + } + } + for i := 0; i+1 < n; i++ { + push(i, i+1) + } + heap.Init(&h) + + live := n + for h.Len() > 0 { + cd := heap.Pop(&h).(cand) + l, r := cd.l, cd.r + // Reject stale entries: either side merged away or grew since push. + if !nodes[l].alive || !nodes[r].alive || nodes[l].next != r || + len(nodes[l].s) != cd.ll || len(nodes[r].s) != cd.rl { + continue + } + nodes[l].s += nodes[r].s + nodes[r].alive = false + nodes[l].next = nodes[r].next + if nodes[r].next >= 0 { + nodes[nodes[r].next].prev = l + } + live-- + if live == 1 { + return 1 + } + if p := nodes[l].prev; p >= 0 { + if rk, ok := c.merges[nodes[p].s+" "+nodes[l].s]; ok { + heap.Push(&h, cand{rank: rk, l: p, r: l, ll: len(nodes[p].s), rl: len(nodes[l].s)}) + } + } + if nx := nodes[l].next; nx >= 0 { + if rk, ok := c.merges[nodes[l].s+" "+nodes[nx].s]; ok { + heap.Push(&h, cand{rank: rk, l: l, r: nx, ll: len(nodes[l].s), rl: len(nodes[nx].s)}) + } + } + } + return live +} + +// Count returns the number of tokens voyage/HF would produce for text. +func (c *Counter) Count(text string) int { + if text == "" { + return 0 + } + s := norm.NFC.String(text) + total := 0 + var sb strings.Builder + for len(s) > 0 { + n := nextToken(s) + if n <= 0 { + n = 1 + } + piece := s[:n] + s = s[n:] + // ByteLevel map + sb.Reset() + sb.Grow(len(piece) * 2) + for i := 0; i < len(piece); i++ { + sb.WriteRune(byteRune[piece[i]]) + } + total += c.bpeLen(sb.String()) + } + return total +} + +// SplitPoints returns the byte offsets at which text must be cut so that no +// piece exceeds budget tokens, plus the total token count of the whole text. +// +// The cuts are exact, not estimated, and they need no search. BPE merges never +// cross a pre-token boundary in this pipeline (Split runs with Isolated +// behaviour, and ByteLevel+BPE are applied per pre-token), so a text's token +// count is the SUM of its pre-tokens' counts. Cutting on a pre-token boundary +// therefore leaves both sides tokenising exactly as they did inside the whole: +// the parts always add up to the total, with no drift to correct for. +// +// One pass, left to right, reusing the same memo Count uses — so a split costs +// what a count costs, plus the offsets slice. +// +// A single pre-token larger than budget cannot be honoured on a boundary (its +// merges DO interact internally). Rather than silently emit an over-budget +// piece, splitInside falls back to a binary search on bytes within that one +// pre-token. That path is for base64 blobs and minified lines with no +// whitespace; on a 45-repo corpus it fires for 5 chunks in 1.9M. +// +// Offsets are cut points only: nil means the text already fits. +func (c *Counter) SplitPoints(text string, budget int) (offsets []int, total int) { + if text == "" { + return nil, 0 + } + if budget <= 0 { + return nil, c.Count(text) + } + s := norm.NFC.String(text) + + acc := 0 // tokens accumulated in the current piece + pos := 0 // byte offset into s + var sb strings.Builder + for pos < len(s) { + n := nextToken(s[pos:]) + if n <= 0 { + n = 1 + } + piece := s[pos : pos+n] + + sb.Reset() + sb.Grow(len(piece) * 2) + for i := 0; i < len(piece); i++ { + sb.WriteRune(byteRune[piece[i]]) + } + tk := c.bpeLen(sb.String()) + + switch { + case tk > budget: + // Does not fit even alone. Close the current piece, then cut + // inside this pre-token. + if acc > 0 { + offsets = append(offsets, pos) + acc = 0 + } + inner := c.splitInside(s[pos:pos+n], budget) + for _, off := range inner { + offsets = append(offsets, pos+off) + } + // Tail of the pre-token starts a fresh piece; its token count is + // unknown without re-counting, so charge it conservatively as a + // full budget minus nothing and let the next boundary close it. + acc = 0 + case acc+tk > budget: + offsets = append(offsets, pos) + acc = tk + default: + acc += tk + } + total += tk + pos += n + } + return offsets, total +} + +// splitInside cuts one over-budget pre-token by binary search on its bytes. +// Inside a pre-token counts are not additive, so every candidate cut is +// re-counted — but the search converges in a handful of probes because +// bytes-per-token is near-constant within a homogeneous run. +func (c *Counter) splitInside(piece string, budget int) []int { + var cuts []int + start := 0 + for start < len(piece) { + if c.Count(piece[start:]) <= budget { + break + } + lo, hi := start+1, len(piece) + best := start + 1 + for lo <= hi { + mid := (lo + hi) / 2 + for mid > start && mid < len(piece) && !utf8.RuneStart(piece[mid]) { + mid-- + } + if mid <= start { + break + } + if c.Count(piece[start:mid]) <= budget { + best = mid + lo = mid + 1 + } else { + hi = mid - 1 + } + } + if best >= len(piece) { + break + } + cuts = append(cuts, best) + start = best + } + return cuts +} diff --git a/server/internal/tokenizer/bpecount/bpecount_test.go b/server/internal/tokenizer/bpecount/bpecount_test.go new file mode 100644 index 00000000..546ec1f0 --- /dev/null +++ b/server/internal/tokenizer/bpecount/bpecount_test.go @@ -0,0 +1,233 @@ +package bpecount + +import ( + "os" + "testing" +) + +// tokenizerPath is the real voyage-code-3 tokenizer.json. The tests that need +// it skip when it is absent so a checkout without the 7 MB file still builds +// and tests clean. +const tokenizerPath = "../../../../loadtests/bench/voyage-code-3.tokenizer.json" + +func load(t *testing.T) *Counter { + t.Helper() + if _, err := os.Stat(tokenizerPath); err != nil { + t.Skip("tokenizer.json not present") + } + c, err := Load(tokenizerPath) + if err != nil { + t.Fatalf("load: %v", err) + } + return c +} + +// TestCountMatchesReference pins the counts that were verified against +// Voyage's own usage.total_tokens and the HuggingFace Rust tokenizer. The tab +// cases are the ones a RE2 rewrite of the pre-tokenizer regex gets wrong: the +// `\s+(?!\S)` lookahead cannot be expressed, and a naive rewrite absorbs a +// leading tab into the punctuation branch that may only absorb a space. +func TestCountMatchesReference(t *testing.T) { + c := load(t) + for _, tc := range []struct { + in string + want int + }{ + {"func main() {\n\tfmt.Println(\"hi\")\n}\n", 10}, + {"\t\t\"a\"", 4}, + {"\t\t\t\"end\": {", 6}, + {"a\t\t-b", 4}, + {"class A:\n def g(self):\n x = 1\n", 14}, + {"hello world", 2}, + {"#ifdef USE_THREADS", 3}, + {"", 0}, + } { + if got := c.Count(tc.in); got != tc.want { + t.Errorf("Count(%q) = %d, want %d", tc.in, got, tc.want) + } + } +} + +// TestNFCNormalisation covers the second gap in the ollama tokenizer: the +// tokenizer.json declares an NFC normalizer, and skipping it makes decomposed +// input cost an extra token. +func TestNFCNormalisation(t *testing.T) { + c := load(t) + nfc := "caf\u00e9" // é as one code point + nfd := "cafe\u0301" // e + combining acute + if a, b := c.Count(nfc), c.Count(nfd); a != b { + t.Errorf("NFC %d != NFD %d — normaliser not applied", a, b) + } +} + +// TestSplitPointsAreExact is the property the splitter exists for: because BPE +// merges never cross a pre-token boundary, the pieces must add up to the whole +// and none may exceed the budget. +func TestSplitPointsAreExact(t *testing.T) { + c := load(t) + src := "" + for i := 0; i < 400; i++ { + src += "func handler(w http.ResponseWriter, r *http.Request) {\n\tdefer r.Body.Close()\n}\n" + } + const budget = 500 + + offsets, total := c.SplitPoints(src, budget) + if total != c.Count(src) { + t.Fatalf("SplitPoints total %d != Count %d", total, c.Count(src)) + } + if len(offsets) == 0 { + t.Fatalf("expected cuts for %d tokens at budget %d", total, budget) + } + + sum, prev := 0, 0 + for _, off := range append(offsets, len(src)) { + n := c.Count(src[prev:off]) + if n > budget { + t.Errorf("piece [%d:%d] is %d tokens, over budget %d", prev, off, n, budget) + } + sum += n + prev = off + } + if sum != total { + t.Errorf("pieces sum to %d, whole is %d — merges leaked across a cut", sum, total) + } +} + +// TestSplitPointsFitsAlready — a text under budget must not be cut. +func TestSplitPointsFitsAlready(t *testing.T) { + c := load(t) + offsets, total := c.SplitPoints("package main\n", 1000) + if offsets != nil { + t.Errorf("expected no cuts, got %v", offsets) + } + if total == 0 { + t.Error("total should be counted even when no cut is needed") + } +} + +// TestSplitInsidePreToken covers the one case boundaries cannot serve: a +// single pre-token bigger than the budget (base64 blobs, minified lines). +func TestSplitPointsOversizePreToken(t *testing.T) { + c := load(t) + blob := "" + for i := 0; i < 4000; i++ { + blob += "aB3" + } + const budget = 100 + offsets, _ := c.SplitPoints(blob, budget) + if len(offsets) == 0 { + t.Fatal("expected the blob to be cut") + } + prev := 0 + for _, off := range append(offsets, len(blob)) { + if n := c.Count(blob[prev:off]); n > budget { + t.Errorf("piece [%d:%d] is %d tokens, over budget %d", prev, off, n, budget) + } + prev = off + } +} + +// --- CI coverage without the 7 MB file --- +// +// The golden-count tests above need the real voyage-code-3 tokenizer.json, +// which is not in the repo, so they skip on a clean checkout. The mechanics — +// pre-token splitting, merge application, the additivity SplitPoints relies on +// — do not need that vocabulary. A hand-built merge table exercises them, so +// CI still fails if the splitter or the merge loop regresses. +func syntheticCounter(t *testing.T) *Counter { + t.Helper() + // Merges are ranked: "a b" collapses first, then "ab c". + tj := `{"model":{"type":"BPE","merges":["a b","ab c","f u","fu n"]}}` + c, err := LoadBytes([]byte(tj)) + if err != nil { + t.Fatalf("LoadBytes: %v", err) + } + return c +} + +// TestSyntheticMergesApply — with only "a b" known, "abc" costs one merge plus +// the leftover byte; the second merge then folds that leftover in. +func TestSyntheticMergesApply(t *testing.T) { + c := syntheticCounter(t) + if got, want := c.Count("abc"), 1; got != want { + t.Errorf(`Count("abc") = %d, want %d (a+b -> ab, ab+c -> abc)`, got, want) + } + if got, want := c.Count("ab"), 1; got != want { + t.Errorf(`Count("ab") = %d, want %d`, got, want) + } + if got, want := c.Count("acb"), 3; got != want { + t.Errorf(`Count("acb") = %d, want %d (no merge applies)`, got, want) + } +} + +// TestSyntheticAdditivity is the property the whole splitter rests on: BPE +// never merges across a pre-token boundary, so counts add up. If a future +// change made merges span boundaries, cuts would silently produce over-budget +// pieces — this catches it without needing the real vocabulary. +func TestSyntheticAdditivity(t *testing.T) { + c := syntheticCounter(t) + const text = "abc abc\n\tabc fun fun" + whole := c.Count(text) + + offsets, total := c.SplitPoints(text, 3) + if total != whole { + t.Fatalf("SplitPoints total %d != Count %d", total, whole) + } + sum, prev := 0, 0 + for _, off := range append(offsets, len(text)) { + n := c.Count(text[prev:off]) + if n > 3 { + t.Errorf("piece %q is %d tokens, over budget 3", text[prev:off], n) + } + sum += n + prev = off + } + if sum != whole { + t.Errorf("pieces sum to %d, whole is %d", sum, whole) + } +} + +// TestPreTokenBoundaries pins the hand-rolled splitter against the branches of +// the Qwen2 pattern that a RE2 rewrite gets wrong — whitespace runs, and a tab +// that must NOT be absorbed into the punctuation branch. +func TestPreTokenBoundaries(t *testing.T) { + for _, tc := range []struct { + in string + want []string + }{ + {"a b", []string{"a", " b"}}, + {"a b", []string{"a", " ", " b"}}, + // A tab is a legal single-character prefix for the letter branch + // ([^\r\n\p{L}\p{N}]?\p{L}+), so it attaches to what follows. + {"x\n\ty", []string{"x", "\n", "\ty"}}, + // The lookahead branch \s+(?!\S) matches a whitespace run only when + // nothing non-space follows it. The first tab qualifies (a tab + // follows); the second does not (a quote follows) and falls through + // to plain \s+. Hence two separate pre-tokens, not one run — this is + // precisely what a RE2 rewrite of the pattern gets wrong. + {"\t\t\"a\"", []string{"\t", "\t", "\"a", "\""}}, + {"it's", []string{"it", "'s"}}, + {"a1", []string{"a", "1"}}, + } { + var got []string + s := tc.in + for len(s) > 0 { + n := nextToken(s) + if n <= 0 { + t.Fatalf("nextToken(%q) returned %d", s, n) + } + got = append(got, s[:n]) + s = s[n:] + } + if len(got) != len(tc.want) { + t.Errorf("split(%q) = %q, want %q", tc.in, got, tc.want) + continue + } + for i := range got { + if got[i] != tc.want[i] { + t.Errorf("split(%q) = %q, want %q", tc.in, got, tc.want) + break + } + } + } +} diff --git a/server/internal/tokenizer/budget.go b/server/internal/tokenizer/budget.go new file mode 100644 index 00000000..90d6406c --- /dev/null +++ b/server/internal/tokenizer/budget.go @@ -0,0 +1,49 @@ +// Package tokenizer carries the model-token knowledge the chunker needs and +// the embedding providers own. +// +// The chunker decides WHERE to cut; only the provider knows WHAT the model +// counts. Keeping the two apart is what lets a model change without teaching +// the chunker about byte-level BPE, SentencePiece, or llama-server's +// /tokenize endpoint. +package tokenizer + +// Budget is the whole surface the chunker sees. One interface rather than a +// mandatory one plus an optional splitter: a single constructor argument, and +// callers that need both never have to type-assert. +// +// Cost note, because the two methods look interchangeable and are not: +// CountTokens and SplitPoints do the same single left-to-right pass over the +// same memo, so neither is algorithmically cheaper. What differs is +// allocation. CountTokens returns an int and allocates nothing; SplitPoints +// must build a slice of offsets — for a 60 KB input that is on the order of +// 15k entries. CountTokens runs on every chunk (1.9M of them on the reference +// corpus) while SplitPoints runs only on inputs that exceed the model's +// context (5 of that same 1.9M). Call CountTokens on the hot path and reach +// for SplitPoints only once a text is known not to fit. +// +// A third shortcut avoids both: byte-level BPE cannot emit a token covering +// less than one byte, so len(text) <= budget PROVES the text fits, with no +// tokenisation at all. Use it before calling anything here. +type Budget interface { + // MaxInputTokens is the model's context window for a single input. + MaxInputTokens() int + + // ExactCounts reports whether CountTokens and SplitPoints are exact. + // False means the provider has no tokenizer and is estimating from + // byte length: counts may be wrong in both directions and split points + // are byte windows, not token boundaries. Callers that need a + // guarantee must widen their safety margin when this is false — + // silently trusting an estimate is what the byte-window splitter used + // to do, and it produced averaged vectors nobody could see was wrong. + ExactCounts() bool + + // CountTokens returns the number of tokens the model will charge for. + CountTokens(s string) int + + // SplitPoints returns byte offsets at which s must be cut so no piece + // exceeds budget tokens, plus the total token count of s. Offsets, not + // substrings, so the caller keeps ownership of the metadata that hangs + // off those positions — line numbers, symbol names, byte ranges. + // A nil offsets slice means s already fits. + SplitPoints(s string, budget int) (offsets []int, total int) +} From a3a81526e8a01d313cd87af570dbbd0d1ce31f06 Mon Sep 17 00:00:00 2001 From: dvcdsys Date: Tue, 18 Aug 2026 12:38:29 +0100 Subject: [PATCH 2/6] feat(chunker): size chunks in tokens, not in a byte stand-in MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The chunker's size limit was 1500*3 bytes — a token target expressed in the only unit it had. That ratio holds for dense ASCII code and misses everywhere else: a comment block in Cyrillic or CJK costs two to three bytes per character, so a byte-capped chunk carries a third of the tokens intended, while minified JavaScript packs several times more. The chunker now asks the active embedding provider what a chunk actually costs, when that provider can answer. Wiring: embeddings.Service.TokenBudget() hands out the live provider when it implements tokenizer.Budget, nil otherwise. The indexer asks PER FILE rather than caching it, because a provider swap between files is legitimate while mixing two models' limits inside one file's chunk set is not. CIX_MAX_CHUNK_TOKENS finally does something: it was parsed into config and only ever used to size llama's context. The bound is applied ABOVE the chunking paths, not inside the tree-sitter one. That placement is the point: minified files and files with no grammar fall through to chunkFallback's sliding window, and those are precisely the inputs that blow the model's input limit. A budget that only covered the tree-sitter path would have missed exactly the cases it exists for — the first draft did sit inside that path, and the two tests that now cover minified JS and a grammarless file both failed against it. With a budget in hand the byte cap is deliberately raised out of the way (innerMax = MaxInt32) so semantic units arrive whole and are cut once, in tokens. Without one, nothing changes: nil budget, or a provider that reports ExactCounts() == false, keeps the byte path byte for byte. An estimating provider is routed to the byte path on purpose — its numbers are the same guess, and dressing them as a token budget would hide that from the caller. splitChunkTokens also closes a hole the byte splitter had. That loop requires len(currentLines) > 1, so a single line longer than the cap was never split: on the 45-repo reference corpus that produced a 65 KB chunk, which the voyage provider then cut into byte windows and averaged the vectors of — a vector representing neither half, with nothing in the logs to say so. A line that cannot fit alone is now cut on real token boundaries from Budget.SplitPoints, and the tail seeds the next chunk so short following lines can still join it. Attribution rule is unchanged: only the first piece keeps SymbolName and ChunkType, the rest become `block`, so one long function does not produce N symbol rows all claiming to be it. Co-Authored-By: Claude Opus 5 --- server/cmd/cix-server/main.go | 1 + server/internal/chunker/chunker.go | 168 +++++++++++++++++- .../internal/chunker/chunker_tokens_test.go | 145 +++++++++++++++ server/internal/embeddings/service.go | 20 +++ server/internal/indexer/indexer.go | 22 ++- 5 files changed, 352 insertions(+), 4 deletions(-) create mode 100644 server/internal/chunker/chunker_tokens_test.go diff --git a/server/cmd/cix-server/main.go b/server/cmd/cix-server/main.go index a2b8e7c6..54c23cc8 100644 --- a/server/cmd/cix-server/main.go +++ b/server/cmd/cix-server/main.go @@ -437,6 +437,7 @@ func run() (restart bool, err error) { idx := indexer.New(database, vsHolder, embedSvc, logger) idx.SetEmbedIncludePath(cfg.EmbedIncludePath) + idx.SetMaxChunkTokens(cfg.MaxChunkTokens) // Record the active embedding model on every indexed project so the // dashboard can highlight stale vectors when the runtime provider / // model changes. Wire it as a live lookup so a runtime provider diff --git a/server/internal/chunker/chunker.go b/server/internal/chunker/chunker.go index 6391ba1b..7371f91e 100644 --- a/server/internal/chunker/chunker.go +++ b/server/internal/chunker/chunker.go @@ -10,7 +10,9 @@ package chunker import ( + "github.com/dvcdsys/code-index/server/internal/tokenizer" "log/slog" + "math" "path/filepath" "strings" "sync" @@ -457,15 +459,77 @@ type Reference struct { // falls back to sliding-window chunking for unsupported languages. The maxSize // parameter controls per-chunk character limit; pass 0 to use the default. func ChunkFile(filePath, content, language string, maxSize int) ([]Chunk, []Reference, error) { + return ChunkFileTokens(filePath, content, language, maxSize, nil, 0) +} + +// ChunkFileTokens is ChunkFile with a token budget. +// +// maxSize (bytes) has always been a stand-in for a token limit: the default +// 4500 is "1500 tokens x 3 bytes", a ratio that holds for dense ASCII code and +// falls apart everywhere else — Cyrillic or CJK comments cost two to three +// bytes per character, so a byte-sized chunk carries far fewer tokens than +// intended, while minified JavaScript packs far more. +// +// When budget is non-nil and reports exact counts, the size decision is made +// in tokens instead, and an over-budget chunk is cut on real token boundaries +// (via Budget.SplitPoints) rather than on a byte count. maxTokens <= 0 uses +// defaultMaxChunkTokens. A nil or estimating budget keeps the byte path +// unchanged, so nothing about existing behaviour depends on a provider +// having a tokenizer. +func ChunkFileTokens(filePath, content, language string, maxSize int, budget tokenizer.Budget, maxTokens int) ([]Chunk, []Reference, error) { if maxSize <= 0 { maxSize = maxChunkSize } - chunks, refs, err := chunkWithTreesitter(filePath, content, language, maxSize) + if budget != nil && !budget.ExactCounts() { + // An estimate is what the byte path already is; do not pretend + // otherwise by routing through the token splitter. + budget = nil + } + if budget != nil { + if maxTokens <= 0 { + maxTokens = defaultMaxChunkTokens + } + if lim := budget.MaxInputTokens(); lim > 0 && maxTokens > lim { + maxTokens = lim + } + } + // With a token budget the byte cap must not fire first: it is the very + // bias being removed (a byte limit cuts Cyrillic or CJK three times + // sooner than ASCII for the same token cost). Let the inner path emit + // whole semantic units and bound them in tokens afterwards. + innerMax := maxSize + if budget != nil { + innerMax = math.MaxInt32 + } + + chunks, refs, err := chunkWithTreesitter(filePath, content, language, innerMax) if err != nil { // Fallback: sliding window, no references. - return chunkFallback(filePath, content, language), nil, nil + return boundTokens(chunkFallback(filePath, content, language), budget, maxTokens), nil, nil + } + return boundTokens(chunks, budget, maxTokens), refs, nil +} + +// boundTokens enforces the token budget over chunks from ANY path — the +// tree-sitter one, the bash regex extractor, or the sliding-window fallback. +// Applying it here rather than inside the tree-sitter path is deliberate: +// minified JavaScript and files with no grammar are exactly the inputs that +// reach the fallback, and they are also the ones most likely to blow the +// model's input limit. A budget that only covered the happy path would miss +// them. +func boundTokens(chunks []Chunk, budget tokenizer.Budget, maxTokens int) []Chunk { + if budget == nil || maxTokens <= 0 { + return chunks + } + out := make([]Chunk, 0, len(chunks)) + for _, c := range chunks { + if budget.CountTokens(c.Content) > maxTokens { + out = append(out, splitChunkTokens(c, budget, maxTokens)...) + continue + } + out = append(out, c) } - return chunks, refs, nil + return out } // chunkFallback returns reasonable chunks for content that the tree-sitter @@ -1060,3 +1124,101 @@ func sortRanges(ranges [][2]int) { } } } + +// defaultMaxChunkTokens is the token equivalent of maxChunkSize. The byte +// default was written as 1500*3 — a 1500-token target at three bytes each — +// so the token target is that same 1500, now expressed in the unit that +// actually matters. CIX_MAX_CHUNK_TOKENS overrides it. +const defaultMaxChunkTokens = 1500 + +// splitChunkTokens cuts an over-budget chunk into pieces of <= maxTokens. +// +// It is the token-aware sibling of splitChunk and keeps its attribution rule: +// only the first piece inherits SymbolName/ChunkType, the rest become `block`, +// so one long function does not produce N rows all claiming to be that symbol. +// +// Two differences that matter: +// +// - The running size is counted, not estimated, so a chunk of Cyrillic +// comments is no longer cut three times sooner than an equivalent chunk of +// ASCII. +// - A single line longer than the budget is cut INSIDE the line, on token +// boundaries from Budget.SplitPoints. The byte splitter could not do this +// (its loop requires len(currentLines) > 1), which is how a 65 KB minified +// line reached the embedder as one chunk and got byte-windowed and +// vector-averaged downstream. +func splitChunkTokens(chunk Chunk, budget tokenizer.Budget, maxTokens int) []Chunk { + lines := splitLines(chunk.Content) + var subChunks []Chunk + + emit := func(content string, startLine, endLine int) { + if content == "" { + return + } + c := Chunk{ + Content: content, + FilePath: chunk.FilePath, + StartLine: startLine, + EndLine: endLine, + Language: chunk.Language, + ParentName: chunk.ParentName, + } + if len(subChunks) == 0 { + c.ChunkType = chunk.ChunkType + c.SymbolName = chunk.SymbolName + c.SymbolSignature = chunk.SymbolSignature + } else { + c.ChunkType = "block" + } + subChunks = append(subChunks, c) + } + + var currentLines []string + currentStart := chunk.StartLine + currentTokens := 0 + + flush := func(endLine int) { + if len(currentLines) == 0 { + return + } + emit(joinLines(currentLines), currentStart, endLine) + currentLines = nil + currentTokens = 0 + } + + for i, line := range lines { + lineNo := chunk.StartLine + i + n := budget.CountTokens(line) + + // A line that cannot fit on its own: close what we have, then cut + // the line itself on token boundaries. + if n > maxTokens { + flush(lineNo - 1) + offsets, _ := budget.SplitPoints(line, maxTokens) + prev := 0 + for _, off := range offsets { + emit(line[prev:off], lineNo, lineNo) + prev = off + } + // Tail of the line seeds the next piece so short following + // lines can still join it. + currentLines = []string{line[prev:]} + currentStart = lineNo + currentTokens = budget.CountTokens(line[prev:]) + continue + } + + if len(currentLines) > 0 && currentTokens+n > maxTokens { + flush(lineNo - 1) + currentStart = lineNo + } + currentLines = append(currentLines, line) + currentTokens += n + } + flush(chunk.StartLine + len(lines) - 1) + + if len(subChunks) == 0 { + return []Chunk{chunk} + } + return subChunks +} diff --git a/server/internal/chunker/chunker_tokens_test.go b/server/internal/chunker/chunker_tokens_test.go new file mode 100644 index 00000000..65c5440d --- /dev/null +++ b/server/internal/chunker/chunker_tokens_test.go @@ -0,0 +1,145 @@ +package chunker + +import ( + "strings" + "testing" + + "github.com/dvcdsys/code-index/server/internal/tokenizer" +) + +// fakeBudget counts one token per whitespace-separated word and cuts on word +// boundaries. Deterministic and independent of any vocabulary, so these tests +// assert the CHUNKER's behaviour rather than a tokenizer's — the real +// tokenizer has its own tests. +type fakeBudget struct{ maxInput int } + +func (f fakeBudget) MaxInputTokens() int { return f.maxInput } +func (f fakeBudget) ExactCounts() bool { return true } + +func (f fakeBudget) CountTokens(s string) int { return len(strings.Fields(s)) } + +func (f fakeBudget) SplitPoints(s string, budget int) ([]int, int) { + var offsets []int + count, since := 0, 0 + inWord := false + for i := 0; i < len(s); i++ { + isSpace := s[i] == ' ' || s[i] == '\t' || s[i] == '\n' + if !isSpace && !inWord { + inWord = true + count++ + since++ + if since > budget { + offsets = append(offsets, i) + since = 1 + } + } else if isSpace { + inWord = false + } + } + return offsets, count +} + +var _ tokenizer.Budget = fakeBudget{} + +// TestTokenBudgetBoundsEveryChunk is the property the integration exists for: +// with a budget in hand, no emitted chunk may exceed it. +func TestTokenBudgetBoundsEveryChunk(t *testing.T) { + var sb strings.Builder + sb.WriteString("func run() {\n") + for i := 0; i < 300; i++ { + sb.WriteString("\tdo something with several words on this line\n") + } + sb.WriteString("}\n") + + b := fakeBudget{maxInput: 4096} + chunks, _, err := ChunkFileTokens("x.go", sb.String(), "go", 0, b, 50) + if err != nil { + t.Fatalf("chunk: %v", err) + } + if len(chunks) < 2 { + t.Fatalf("expected the body to be split, got %d chunk(s)", len(chunks)) + } + for i, c := range chunks { + if n := b.CountTokens(c.Content); n > 50 { + t.Errorf("chunk %d is %d tokens, over budget 50", i, n) + } + } +} + +// TestLongSingleLineIsSplit covers the hole the byte splitter had: its loop +// requires more than one line, so a minified file arrived at the embedder as +// one enormous chunk. On the reference corpus that produced a 65 KB chunk +// whose vector was an average of byte windows. +func TestLongSingleLineIsSplit(t *testing.T) { + line := strings.TrimSpace(strings.Repeat("token ", 5000)) + b := fakeBudget{maxInput: 4096} + + chunks, _, err := ChunkFileTokens("min.js", line, "javascript", 0, b, 100) + if err != nil { + t.Fatalf("chunk: %v", err) + } + for i, c := range chunks { + if n := b.CountTokens(c.Content); n > 100 { + t.Errorf("chunk %d is %d tokens, over budget 100 — long line not cut", i, n) + } + } + if len(chunks) < 2 { + t.Fatalf("expected the single line to be cut, got %d chunk(s)", len(chunks)) + } +} + +// TestBudgetCappedByModelContext — a chunk target above the model's own input +// window is meaningless; the smaller of the two must win. +func TestBudgetCappedByModelContext(t *testing.T) { + src := strings.TrimSpace(strings.Repeat("word ", 400)) + b := fakeBudget{maxInput: 40} + + chunks, _, err := ChunkFileTokens("x.txt", src, "text", 0, b, 10000) + if err != nil { + t.Fatalf("chunk: %v", err) + } + for i, c := range chunks { + if n := b.CountTokens(c.Content); n > 40 { + t.Errorf("chunk %d is %d tokens, over the model's %d-token window", i, n, 40) + } + } +} + +// TestEstimatingBudgetKeepsBytePath — a provider without a real tokenizer must +// not be routed through the token splitter: its numbers are the same guess the +// byte path already makes, and pretending otherwise hides that from the caller. +func TestEstimatingBudgetKeepsBytePath(t *testing.T) { + src := strings.Repeat("x := 1\n", 2000) + got, _, err := ChunkFileTokens("x.go", src, "go", 0, estimatingBudget{}, 10) + if err != nil { + t.Fatalf("chunk: %v", err) + } + want, _, err := ChunkFile("x.go", src, "go", 0) + if err != nil { + t.Fatalf("chunk: %v", err) + } + if len(got) != len(want) { + t.Errorf("estimating budget changed chunking: %d chunks vs %d on the byte path", + len(got), len(want)) + } +} + +type estimatingBudget struct{ fakeBudget } + +func (estimatingBudget) ExactCounts() bool { return false } + +// TestNilBudgetUnchanged pins that the default path is untouched. +func TestNilBudgetUnchanged(t *testing.T) { + src := strings.Repeat("func f() { return 1 }\n", 500) + a, _, err := ChunkFileTokens("x.go", src, "go", 0, nil, 0) + if err != nil { + t.Fatalf("chunk: %v", err) + } + b, _, err := ChunkFile("x.go", src, "go", 0) + if err != nil { + t.Fatalf("chunk: %v", err) + } + if len(a) != len(b) { + t.Errorf("nil budget diverged from ChunkFile: %d vs %d chunks", len(a), len(b)) + } +} diff --git a/server/internal/embeddings/service.go b/server/internal/embeddings/service.go index cbb12c70..8422d370 100644 --- a/server/internal/embeddings/service.go +++ b/server/internal/embeddings/service.go @@ -20,6 +20,7 @@ import ( // provider purely by kind string — these imports are the wiring. _ "github.com/dvcdsys/code-index/server/internal/embeddings/provider/openai" _ "github.com/dvcdsys/code-index/server/internal/embeddings/provider/voyage" + "github.com/dvcdsys/code-index/server/internal/tokenizer" ) // Service is the public embeddings API used by handlers and the indexer. @@ -459,6 +460,25 @@ func (s *Service) Status() provider.Status { // CurrentKind reports the kind of the active provider, or "" when // disabled / not yet built. Used by /status and admin endpoints. +// TokenBudget returns the active provider as a token budget when it can +// count tokens, and nil otherwise. The chunker uses it to size chunks in the +// model's own unit; nil keeps it on the byte heuristic. +// +// Snapshotted under the read lock like CurrentKind, because a provider swap +// mid-file would otherwise mix two models' limits inside one chunk set. +func (s *Service) TokenBudget() tokenizer.Budget { + s.mu.RLock() + defer s.mu.RUnlock() + if s.current == nil { + return nil + } + b, ok := s.current.(tokenizer.Budget) + if !ok { + return nil + } + return b +} + func (s *Service) CurrentKind() string { if s == nil || s.disabled { return "" diff --git a/server/internal/indexer/indexer.go b/server/internal/indexer/indexer.go index 03a17ad5..eb4eeddb 100644 --- a/server/internal/indexer/indexer.go +++ b/server/internal/indexer/indexer.go @@ -23,6 +23,7 @@ import ( "github.com/dvcdsys/code-index/server/internal/embeddings" "github.com/dvcdsys/code-index/server/internal/langdetect" "github.com/dvcdsys/code-index/server/internal/symbolindex" + "github.com/dvcdsys/code-index/server/internal/tokenizer" "github.com/dvcdsys/code-index/server/internal/vectorstore" ) @@ -141,6 +142,10 @@ type Service struct { // reindexed under the new format. embedIncludePath bool + // maxChunkTokens is the per-chunk token target (CIX_MAX_CHUNK_TOKENS). + // 0 means the chunker's own default. + maxChunkTokens int + // embeddingModel is the active embedding model identifier persisted on // projects.indexed_with_model at FinishIndexing. Set via // SetEmbeddingModel from main; empty string keeps the column NULL so @@ -208,6 +213,12 @@ func (s *Service) SetEmbedIncludePath(v bool) { s.embedIncludePath = v } +// SetMaxChunkTokens sets the per-chunk token target used when the active +// embedding provider can count tokens exactly. +func (s *Service) SetMaxChunkTokens(n int) { + s.maxChunkTokens = n +} + // SetEmbeddingModel records the model identifier the indexer will write to // projects.indexed_with_model at FinishIndexing. Called from main once the // runtime config is resolved; empty string disables the write (the column @@ -705,7 +716,16 @@ func (s *Service) ProcessFilesStreaming( language = "text" } - chunks, refs, err := chunker.ChunkFile(fp.Path, fp.Content, language, 0) + // The token budget comes from the LIVE provider, asked per file: a + // provider swap between files is legitimate, mixing two models' + // limits inside one file's chunks is not. + var budget tokenizer.Budget + if tb, ok := s.emb.(interface { + TokenBudget() tokenizer.Budget + }); ok { + budget = tb.TokenBudget() + } + chunks, refs, err := chunker.ChunkFileTokens(fp.Path, fp.Content, language, 0, budget, s.maxChunkTokens) if err != nil { s.logger.Warn("indexer: chunk file failed", "path", fp.Path, "err", err) progressSend(progress, ProgressEvent{ From e343209e52a5a71d92f81ad9b6f6dcd9c3086cd7 Mon Sep 17 00:00:00 2001 From: dvcdsys Date: Tue, 18 Aug 2026 12:49:28 +0100 Subject: [PATCH 3/6] fix(chunker): take cut points from the tokenizer, not from summed line counts MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The first token splitter counted each line, accumulated, and re-joined the lines it had collected. Joining reinserts the newlines, and a newline plus the next line's indentation is its own pre-token, so the joined text costs more than the sum the loop was tracking. On real files a 1500-token budget produced chunks of up to 1546 tokens — verified by counting the largest chunks of two freshly reindexed repos with the model's own tokenizer. Cut positions now come from Budget.SplitPoints over the whole content, which is exact by construction, and the pieces are SUBSTRINGS of that content rather than reassembled text. The class of error disappears instead of being compensated for. Two consequences worth spelling out, both learned from tests that failed: Snapping a cut back to a line start (so a piece never begins mid-line and its recorded line range does not lie) shrinks the piece BEFORE the cut and grows the one after it. Reusing the remaining offsets SplitPoints had returned would hand the next piece the tokens this one gave up and push it over budget. So each piece recomputes its cut from the actual new start; the loop is not an inefficiency, it is the correctness. The content-preservation invariant is asserted on splitChunkTokens directly rather than through ChunkFileTokens, because the sliding-window fallback deliberately overlaps its windows for recall — whole-pipeline output is not expected to concatenate back to the source, and a test that assumed otherwise was wrong about the pipeline, not about the splitter. Re-verified on the file that produced the reference corpus's 65 KB chunk: 48 pieces, largest exactly 1500 tokens, none over. The byte path leaves that file as 6 chunks whose largest is 65,553 tokens — twice voyage-code-3's 32K context, which with truncation enabled means the tail was silently dropped before embedding. Co-Authored-By: Claude Opus 5 --- server/internal/chunker/chunker.go | 142 ++++++++---------- .../internal/chunker/chunker_tokens_test.go | 72 +++++++++ 2 files changed, 137 insertions(+), 77 deletions(-) diff --git a/server/internal/chunker/chunker.go b/server/internal/chunker/chunker.go index 7371f91e..8a2b938f 100644 --- a/server/internal/chunker/chunker.go +++ b/server/internal/chunker/chunker.go @@ -1133,92 +1133,80 @@ const defaultMaxChunkTokens = 1500 // splitChunkTokens cuts an over-budget chunk into pieces of <= maxTokens. // -// It is the token-aware sibling of splitChunk and keeps its attribution rule: -// only the first piece inherits SymbolName/ChunkType, the rest become `block`, -// so one long function does not produce N rows all claiming to be that symbol. +// It keeps splitChunk's attribution rule: only the first piece inherits +// SymbolName/ChunkType, the rest become `block`, so one long function does not +// produce N symbol rows all claiming to be it. // -// Two differences that matter: +// The cut positions come from Budget.SplitPoints over the WHOLE content rather +// than from summing per-line counts. Per-line summing is off by the separators +// — joining lines reinserts newlines, and a newline plus the next line's +// indentation forms its own pre-token — so a budget of 1500 produced chunks of +// up to 1546 tokens on real files, roughly one extra token per line boundary. +// SplitPoints is exact by construction, so the bound actually holds. // -// - The running size is counted, not estimated, so a chunk of Cyrillic -// comments is no longer cut three times sooner than an equivalent chunk of -// ASCII. -// - A single line longer than the budget is cut INSIDE the line, on token -// boundaries from Budget.SplitPoints. The byte splitter could not do this -// (its loop requires len(currentLines) > 1), which is how a 65 KB minified -// line reached the embedder as one chunk and got byte-windowed and -// vector-averaged downstream. +// Each exact cut is then pulled BACK to the nearest line start, because a chunk +// that begins mid-line reads badly in search results and its line range lies. +// Moving a cut backwards only ever shrinks the piece before it, so the budget +// survives the adjustment. A line longer than the whole budget has no earlier +// boundary to snap to; there the exact cut stands and the line is split +// internally — which is the case the byte splitter could not handle at all. func splitChunkTokens(chunk Chunk, budget tokenizer.Budget, maxTokens int) []Chunk { - lines := splitLines(chunk.Content) - var subChunks []Chunk - - emit := func(content string, startLine, endLine int) { - if content == "" { - return - } - c := Chunk{ - Content: content, - FilePath: chunk.FilePath, - StartLine: startLine, - EndLine: endLine, - Language: chunk.Language, - ParentName: chunk.ParentName, - } - if len(subChunks) == 0 { - c.ChunkType = chunk.ChunkType - c.SymbolName = chunk.SymbolName - c.SymbolSignature = chunk.SymbolSignature - } else { - c.ChunkType = "block" + var out []Chunk + rest := chunk.Content + line := chunk.StartLine + + for rest != "" { + cuts, _ := budget.SplitPoints(rest, maxTokens) + if len(cuts) == 0 { + out = append(out, mkPiece(chunk, rest, line, len(out) == 0)) + break } - subChunks = append(subChunks, c) - } - var currentLines []string - currentStart := chunk.StartLine - currentTokens := 0 - - flush := func(endLine int) { - if len(currentLines) == 0 { - return + at := cuts[0] + // Pull the cut back to a line start so a piece never begins + // mid-line: search results and the stored line range both lie + // otherwise. Moving backwards only shrinks this piece, so it stays + // inside the budget. A line wider than the whole budget has no + // earlier boundary — there the exact cut stands and the line is + // split internally, which is the case the byte splitter could not + // handle at all. + if nl := strings.LastIndexByte(rest[:at], '\n'); nl >= 0 { + at = nl + 1 } - emit(joinLines(currentLines), currentStart, endLine) - currentLines = nil - currentTokens = 0 - } - - for i, line := range lines { - lineNo := chunk.StartLine + i - n := budget.CountTokens(line) - - // A line that cannot fit on its own: close what we have, then cut - // the line itself on token boundaries. - if n > maxTokens { - flush(lineNo - 1) - offsets, _ := budget.SplitPoints(line, maxTokens) - prev := 0 - for _, off := range offsets { - emit(line[prev:off], lineNo, lineNo) - prev = off - } - // Tail of the line seeds the next piece so short following - // lines can still join it. - currentLines = []string{line[prev:]} - currentStart = lineNo - currentTokens = budget.CountTokens(line[prev:]) - continue + if at == 0 { + at = cuts[0] } - if len(currentLines) > 0 && currentTokens+n > maxTokens { - flush(lineNo - 1) - currentStart = lineNo - } - currentLines = append(currentLines, line) - currentTokens += n + piece := rest[:at] + out = append(out, mkPiece(chunk, piece, line, len(out) == 0)) + line += strings.Count(piece, "\n") + rest = rest[at:] + // Recomputing the next cut from the NEW start is the whole point of + // the loop: snapping back moved the boundary, so the remaining cuts + // SplitPoints returned no longer apply — reusing them would hand the + // next piece the words this one gave up and push it over budget. } - flush(chunk.StartLine + len(lines) - 1) + return out +} - if len(subChunks) == 0 { - return []Chunk{chunk} +// mkPiece builds one output chunk, preserving splitChunk's attribution rule: +// only the first piece keeps SymbolName/ChunkType, so a long function does not +// produce N symbol rows all claiming to be it. +func mkPiece(src Chunk, content string, startLine int, first bool) Chunk { + c := Chunk{ + Content: content, + FilePath: src.FilePath, + StartLine: startLine, + EndLine: startLine + strings.Count(strings.TrimSuffix(content, "\n"), "\n"), + Language: src.Language, + ParentName: src.ParentName, } - return subChunks + if first { + c.ChunkType = src.ChunkType + c.SymbolName = src.SymbolName + c.SymbolSignature = src.SymbolSignature + } else { + c.ChunkType = "block" + } + return c } diff --git a/server/internal/chunker/chunker_tokens_test.go b/server/internal/chunker/chunker_tokens_test.go index 65c5440d..769630fe 100644 --- a/server/internal/chunker/chunker_tokens_test.go +++ b/server/internal/chunker/chunker_tokens_test.go @@ -143,3 +143,75 @@ func TestNilBudgetUnchanged(t *testing.T) { t.Errorf("nil budget diverged from ChunkFile: %d vs %d chunks", len(a), len(b)) } } + +// TestTokenSplitPreservesContent pins the invariant that makes the token +// splitter exact: its pieces are SUBSTRINGS of the chunk it was given, so +// concatenating them reproduces it byte for byte — nothing lost, nothing +// duplicated. +// +// The first implementation instead re-joined lines it had counted separately, +// and the newlines it reinserted cost tokens the running total never saw. A +// 1500-token budget produced 1546-token chunks on real files. Slicing the +// original removes that class of error rather than compensating for it. +// +// Asserted on splitChunkTokens directly, not through ChunkFileTokens: the +// sliding-window fallback deliberately overlaps its windows for recall, so +// whole-pipeline output is not expected to concatenate back. +func TestTokenSplitPreservesContent(t *testing.T) { + var sb strings.Builder + for i := 0; i < 200; i++ { + sb.WriteString("some line with a handful of words in it\n") + } + src := Chunk{ + Content: sb.String(), + FilePath: "x.txt", + StartLine: 1, + EndLine: 200, + ChunkType: "function", + SymbolName: strPtr("run"), + } + b := fakeBudget{maxInput: 4096} + + pieces := splitChunkTokens(src, b, 40) + var rebuilt strings.Builder + for i, c := range pieces { + rebuilt.WriteString(c.Content) + if n := b.CountTokens(c.Content); n > 40 { + t.Errorf("piece %d is %d tokens, over budget 40", i, n) + } + } + if rebuilt.String() != src.Content { + t.Errorf("concatenated pieces differ from the source (%d bytes vs %d)", + rebuilt.Len(), len(src.Content)) + } + if pieces[0].SymbolName == nil || *pieces[0].SymbolName != "run" || pieces[0].ChunkType != "function" { + t.Error("first piece must inherit the symbol") + } + for i, c := range pieces[1:] { + if c.SymbolName != nil || c.ChunkType != "block" { + t.Errorf("piece %d must be an anonymous block, got type %q", i+1, c.ChunkType) + } + } +} + +// TestTokenSplitLineNumbers — a piece that starts mid-file must report the +// line it actually starts on, or `cix search` sends the reader to the wrong +// place. +func TestTokenSplitLineNumbers(t *testing.T) { + var sb strings.Builder + for i := 0; i < 100; i++ { + sb.WriteString("word word word word word\n") + } + b := fakeBudget{maxInput: 4096} + + chunks := splitChunkTokens(Chunk{Content: sb.String(), FilePath: "x.txt", StartLine: 1}, b, 20) + line := 1 + for i, c := range chunks { + if c.StartLine != line { + t.Errorf("chunk %d starts at line %d, expected %d", i, c.StartLine, line) + } + line += strings.Count(c.Content, "\n") + } +} + +func strPtr(s string) *string { return &s } From 7906b4abf21717854c4bf908c36eb1e8a7b8fb00 Mon Sep 17 00:00:00 2001 From: dvcdsys Date: Tue, 18 Aug 2026 13:03:07 +0100 Subject: [PATCH 4/6] test(chunker): make the double adversarial and check the properties on a real corpus MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Three bugs shipped into this branch before being caught, and the reason each one survived is the same: the only thing verifying the token layer was my own model of it. The tokenizer next door had zero bugs across 70,412 inputs because it could be checked against HuggingFace's implementation and against Voyage's billing. Chunk splitting has no such oracle — where to cut a chunk is our decision, not a spec — so the oracle has to be built out of invariants and inputs nobody wrote for the test. fakeBudget was the direct cause of the worst one. It counted whitespace-separated words and let newlines be free, which made the sum of per-line counts equal the count of the joined text — exactly the assumption the implementation got wrong. The double agreed with the bug, the tests passed, and real files came out 3% over budget. It now charges a token per newline, like the real tokenizer does, so summing parts no longer equals the whole unless the pieces are substrings — which is the property the splitter must have. The corpus tests run the chunker over real files from the local fixture: 396 files across 45 repositories, 5,199 chunks, largest exactly 1,500 against a 1,500 budget; 132 files large enough to need splitting, all reconstructing byte for byte with correct line numbers. Those files are what found the original bugs — a 65 KB single-line Zig literal, minified JavaScript that has no grammar and so takes the fallback path — and hand-written cases had not. The fixture is tens of gigabytes and is not in the repository, so the tests skip unless CIX_TEST_CORPUS_DIR and CIX_TEST_TOKENIZER are set. A clean checkout and CI stay green; anyone with a fixture gets the coverage. The file header says how to point them at one. Co-Authored-By: Claude Opus 5 --- .../internal/chunker/chunker_tokens_test.go | 69 ++++-- .../internal/chunker/corpus_property_test.go | 208 ++++++++++++++++++ 2 files changed, 263 insertions(+), 14 deletions(-) create mode 100644 server/internal/chunker/corpus_property_test.go diff --git a/server/internal/chunker/chunker_tokens_test.go b/server/internal/chunker/chunker_tokens_test.go index 769630fe..f73e1a37 100644 --- a/server/internal/chunker/chunker_tokens_test.go +++ b/server/internal/chunker/chunker_tokens_test.go @@ -7,36 +7,77 @@ import ( "github.com/dvcdsys/code-index/server/internal/tokenizer" ) -// fakeBudget counts one token per whitespace-separated word and cuts on word -// boundaries. Deterministic and independent of any vocabulary, so these tests -// assert the CHUNKER's behaviour rather than a tokenizer's — the real -// tokenizer has its own tests. +// fakeBudget is deliberately ADVERSARIAL: a newline costs a token, exactly +// like it does in the real tokenizer, where a line break plus the next line's +// indentation forms its own pre-token. +// +// The first version of this double counted whitespace-separated words and let +// newlines be free. That made the sum of per-line counts equal the count of +// the joined text — which is precisely the assumption the implementation got +// wrong, so the double agreed with the bug and the tests passed while real +// files came out 3% over budget. A test double that cannot express the +// failure mode cannot catch it. +// +// Counting rule: one token per word start, one per newline. Sum over pieces +// therefore does NOT equal the count of the concatenation unless the pieces +// are substrings — which is the property the splitter must have. type fakeBudget struct{ maxInput int } func (f fakeBudget) MaxInputTokens() int { return f.maxInput } func (f fakeBudget) ExactCounts() bool { return true } -func (f fakeBudget) CountTokens(s string) int { return len(strings.Fields(s)) } +func (f fakeBudget) CountTokens(s string) int { + n, inWord := 0, false + for i := 0; i < len(s); i++ { + switch c := s[i]; { + case c == '\n': + n++ + inWord = false + case c == ' ' || c == '\t' || c == '\r': + inWord = false + default: + if !inWord { + inWord = true + n++ + } + } + } + return n +} +// SplitPoints cuts before the token that would overflow the budget, so every +// piece it produces costs at most budget under CountTokens above. func (f fakeBudget) SplitPoints(s string, budget int) ([]int, int) { var offsets []int - count, since := 0, 0 + total, since := 0, 0 inWord := false + cut := func(at int) { + offsets = append(offsets, at) + since = 1 + } for i := 0; i < len(s); i++ { - isSpace := s[i] == ' ' || s[i] == '\t' || s[i] == '\n' - if !isSpace && !inWord { - inWord = true - count++ + switch c := s[i]; { + case c == '\n': + total++ since++ if since > budget { - offsets = append(offsets, i) - since = 1 + cut(i) } - } else if isSpace { inWord = false + case c == ' ' || c == '\t' || c == '\r': + inWord = false + default: + if !inWord { + inWord = true + total++ + since++ + if since > budget { + cut(i) + } + } } } - return offsets, count + return offsets, total } var _ tokenizer.Budget = fakeBudget{} diff --git a/server/internal/chunker/corpus_property_test.go b/server/internal/chunker/corpus_property_test.go new file mode 100644 index 00000000..f414af6a --- /dev/null +++ b/server/internal/chunker/corpus_property_test.go @@ -0,0 +1,208 @@ +package chunker + +import ( + "math/rand" + "os" + "path/filepath" + "sort" + "strings" + "testing" + + "github.com/dvcdsys/code-index/server/internal/tokenizer" + "github.com/dvcdsys/code-index/server/internal/tokenizer/bpecount" +) + +// Property test over a real corpus. +// +// Chunk splitting under a token budget has no external oracle: unlike the +// tokenizer, which can be checked against HuggingFace's implementation and +// against Voyage's own billing, "where should a chunk be cut" is our decision +// and there is nothing to compare it to. What can be checked is that the +// properties we chose actually hold on inputs we did not write — and the +// bugs this file exists to catch were all found by real files rather than by +// hand-written cases: +// +// - a 65 KB single line (a Zig integer literal) that the byte splitter left +// whole, at 65,553 tokens against a 32K context; +// - minified JavaScript, which has no grammar and therefore reaches the +// sliding-window fallback rather than the tree-sitter path; +// - per-line token counting, which was 3% under the truth because joining +// lines reinserts newlines that cost tokens. +// +// The corpus is not in the repository — it is a local fixture of cloned +// repositories, tens of gigabytes. Point the test at one: +// +// CIX_TEST_CORPUS_DIR=…/loadtests/data/repos/repos \ +// CIX_TEST_TOKENIZER=…/voyage-code-3.tokenizer.json \ +// go test ./internal/chunker/ -run Corpus +// +// Without those it skips, so a clean checkout and CI stay green. + +type realBudget struct{ tk *bpecount.Counter } + +func (realBudget) MaxInputTokens() int { return 32000 } +func (realBudget) ExactCounts() bool { return true } +func (b realBudget) CountTokens(s string) int { return b.tk.Count(s) } +func (b realBudget) SplitPoints(s string, n int) ([]int, int) { return b.tk.SplitPoints(s, n) } + +var _ tokenizer.Budget = realBudget{} + +var extLang = map[string]string{ + ".go": "go", ".py": "python", ".ts": "typescript", ".tsx": "tsx", + ".js": "javascript", ".jsx": "javascript", ".java": "java", ".rs": "rust", + ".c": "c", ".h": "c", ".cpp": "cpp", ".rb": "ruby", ".php": "php", + ".kt": "kotlin", ".swift": "swift", ".ex": "elixir", ".zig": "zig", + ".lua": "lua", ".sh": "bash", ".md": "markdown", ".json": "json", +} + +// sampleCorpus walks the fixture and returns up to n files, deterministically +// shuffled so a failure is reproducible. +func sampleCorpus(t *testing.T, root string, n int) []string { + t.Helper() + var files []string + err := filepath.WalkDir(root, func(path string, d os.DirEntry, err error) error { + if err != nil { + return nil //nolint:nilerr // an unreadable entry is not this test's problem + } + if d.IsDir() { + if d.Name() == ".git" { + return filepath.SkipDir + } + return nil + } + if _, ok := extLang[strings.ToLower(filepath.Ext(path))]; ok { + files = append(files, path) + } + return nil + }) + if err != nil { + t.Fatalf("walk corpus: %v", err) + } + sort.Strings(files) + rng := rand.New(rand.NewSource(20260818)) + rng.Shuffle(len(files), func(i, j int) { files[i], files[j] = files[j], files[i] }) + if len(files) > n { + files = files[:n] + } + return files +} + +func corpusBudget(t *testing.T) (realBudget, string) { + t.Helper() + dir := os.Getenv("CIX_TEST_CORPUS_DIR") + tok := os.Getenv("CIX_TEST_TOKENIZER") + if dir == "" || tok == "" { + t.Skip("set CIX_TEST_CORPUS_DIR and CIX_TEST_TOKENIZER to run corpus property tests") + } + tk, err := bpecount.Load(tok) + if err != nil { + t.Fatalf("load tokenizer: %v", err) + } + return realBudget{tk}, dir +} + +// TestCorpusChunksRespectBudget is the property that matters to the API: no +// chunk may cost more tokens than the budget, whatever path produced it. +func TestCorpusChunksRespectBudget(t *testing.T) { + b, dir := corpusBudget(t) + const budget = 1500 + + files := sampleCorpus(t, dir, 400) + if len(files) == 0 { + t.Skip("corpus contains no recognised source files") + } + + var checked, chunks, worst int + for _, f := range files { + src, err := os.ReadFile(f) + if err != nil || len(src) == 0 { + continue + } + lang := extLang[strings.ToLower(filepath.Ext(f))] + got, _, err := ChunkFileTokens(f, string(src), lang, 0, b, budget) + if err != nil { + t.Errorf("%s: %v", f, err) + continue + } + checked++ + chunks += len(got) + for i, c := range got { + n := b.CountTokens(c.Content) + if n > worst { + worst = n + } + if n > budget { + t.Errorf("%s chunk %d: %d tokens, over budget %d", f, i, n, budget) + } + } + } + t.Logf("%d files, %d chunks, largest %d tokens (budget %d)", checked, chunks, worst, budget) +} + +// TestCorpusSplitPreservesContent asserts the splitter loses and duplicates +// nothing, on chunks taken from real files rather than constructed ones. Run +// against splitChunkTokens directly: the sliding-window fallback overlaps its +// windows by design, so whole-pipeline output is not expected to concatenate +// back to the source. +func TestCorpusSplitPreservesContent(t *testing.T) { + b, dir := corpusBudget(t) + const budget = 300 + + var split int + for _, f := range sampleCorpus(t, dir, 200) { + src, err := os.ReadFile(f) + if err != nil || len(src) == 0 { + continue + } + whole := Chunk{ + Content: string(src), + FilePath: f, + StartLine: 1, + ChunkType: "file", + } + if b.CountTokens(whole.Content) <= budget { + continue + } + pieces := splitChunkTokens(whole, b, budget) + split++ + + var rebuilt strings.Builder + for i, p := range pieces { + rebuilt.WriteString(p.Content) + if n := b.CountTokens(p.Content); n > budget { + t.Errorf("%s piece %d: %d tokens, over budget %d", f, i, n, budget) + } + } + if rebuilt.String() != whole.Content { + t.Errorf("%s: pieces do not reconstruct the file (%d bytes vs %d)", + f, rebuilt.Len(), len(whole.Content)) + } + } + t.Logf("%d files exceeded the budget and were split", split) +} + +// TestCorpusLineNumbers — a chunk's StartLine must point at the line its text +// actually begins on, or search results send the reader to the wrong place. +func TestCorpusLineNumbers(t *testing.T) { + b, dir := corpusBudget(t) + const budget = 300 + + for _, f := range sampleCorpus(t, dir, 150) { + src, err := os.ReadFile(f) + if err != nil || len(src) == 0 { + continue + } + whole := Chunk{Content: string(src), FilePath: f, StartLine: 1} + if b.CountTokens(whole.Content) <= budget { + continue + } + line := 1 + for i, p := range splitChunkTokens(whole, b, budget) { + if p.StartLine != line { + t.Errorf("%s piece %d starts at line %d, expected %d", f, i, p.StartLine, line) + break + } + line += strings.Count(p.Content, "\n") + } + } +} From c598cb5536a70c24420bc2fbdeac0a7428f83808 Mon Sep 17 00:00:00 2001 From: dvcdsys Date: Tue, 18 Aug 2026 14:20:29 +0100 Subject: [PATCH 5/6] =?UTF-8?q?fix:=20review=20findings=20=E2=80=94=20offs?= =?UTF-8?q?et=20mapping,=20tokenizer=20validation,=20model=20limits?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Nine findings from review, grouped by what they were actually about. CORRECTNESS — offsets into the wrong string (blocking). SplitPoints normalised its input and returned offsets into that NFC copy while callers slice the original. NFC changes lengths in both directions: decomposed input shrinks (e + U+0301 becomes one code point) so cuts landed mid-rune, and the composition exclusions at U+0958..U+095F DEcompose under NFC, making the normalised form longer, so an offset could exceed len(raw) and panic the indexer on the slice expression. Reproduced both ways before fixing. Input that is already NFC — nearly all source — takes a fast path where the two coincide; anything else is cut one piece at a time, each cut mapped back to a raw offset on a normalisation boundary, rounding down so the budget survives the mapping. Every previous fixture was ASCII, which is why nothing caught it. CORRECTNESS — an over-budget pre-token left its tail uncounted (blocking in spirit; latent for today's only caller). After splitInside the accumulator was reset to zero although the tail after the last inner cut opens the next piece, so following pre-tokens could stack a full budget on top of it and produce pieces of nearly twice the budget. The chunker escaped it by consuming only the first cut and recomputing, but the doc invites consuming the whole slice. CORRECTNESS — an unvalidated tokenizer could be confidently wrong (blocking). LoadBytes checked only model.type == "BPE" while the Qwen2 pre-tokenizer regex and the NFC normalizer are hardcoded here. A GPT-2 or o200k tokenizer.json parses fine, declares BPE, and would have produced plausible wrong counts with ExactCounts() reporting true — which then sizes chunks and packs batches. The declared normalizer and pre-tokenizer are now compared against what this package implements, and a mismatch fails the load so the caller stays on its estimate, which is wrong but knows it. CORRECTNESS — MaxInputTokens was a constant 32000, but the provider's own model enum offers voyage-code-2 at 16K. It is a per-model table now, with unknown models falling back to the smaller window: undershooting costs a needless split, overshooting costs silently truncated input. CORRECTNESS — token-sized chunks reopened the averaging path. The chunker lifts the byte cap when it has a budget, but the provider still cut any input over 30 KB into byte windows and AVERAGED their vectors — the invisible quality loss this branch exists to remove. At CIX_MAX_CHUNK_TOKENS=20000, legal and well inside the window, that would have been routine. With a tokenizer the provider now asks the real question (does this exceed the model's context) and cuts on token boundaries; byte windows remain only where there is no tokenizer and therefore no better answer. CORRECTNESS — the fallback path was still byte-biased. innerMax reached only the tree-sitter path, so files with no grammar kept the fixed 4000-byte sliding window, and boundTokens could not repair it: it splits chunks that are too big and cannot grow ones that are too small. Multi-byte text — which correlates with "no grammar" more than one would like — produced chunks worth a third of the budget on exactly the path the description claimed was covered. ROBUSTNESS — Service.TokenBudget was the only Service method without the nil/disabled guard, and a typed-nil *Service satisfies the capability interface the indexer asserts on, so the first indexed file would panic in RLock. The indexer's inline anonymous assertion is now the named TokenBudgetSource, hoisted out of the per-file loop (the per-file CALL stays — a provider swap between files is legitimate): renaming TokenBudget is now a compile error rather than a silent return to byte chunking. EFFICIENCY — splitChunkTokens called SplitPoints per piece and used only the first cut, rescanning the whole remainder each time: on the 66 KB single-line case that is ~990 suffix scans instead of ~44. The cut list is now reused while it stays valid and recomputed only when snapping to a line start actually moves a boundary — which on single-line content never happens. boundTokens also skips tokenising entirely when len(content) <= budget: a byte-level BPE token cannot cover less than a byte, so that comparison PROVES the chunk fits, and most chunks are small. Corpus property tests dropped from 5.4s/7.6s to 3.8s/2.2s. EFFICIENCY — the contraction branch lowercased the entire remaining suffix for every apostrophe, though only two bytes can match. A 512 KB file with 5,000 quotes moved over a gigabyte through the allocator on the indexing hot path. Also: the dead estimating branch of Provider.SplitPoints (unreachable, and it looped forever when the budget was smaller than one leading multi-byte rune) is gone; the batch log reported the Config cap while packing used the Provider cap; CIX_MAX_CHUNK_TOKENS's default now comes from chunker.DefaultMaxChunkTokens instead of a second copy of 1500; and the /loadtests/ gitignore entry ships here rather than sitting in a working tree, since committed tests reference that path. Tests: NFD and composition-exclusion cases (both previously absent — every fixture was ASCII), a foreign-pipeline rejection case, an uncounted-tail case, a fallback-fills-the-budget case, and TestSplitPointsOversizePreToken fixed to actually enter splitInside — "aB3" repeated pre-tokenizes into 2-byte pieces that never exceed any budget, so the path it named was never executed. Co-Authored-By: Claude Opus 5 --- .gitignore | 3 + server/internal/chunker/chunker.go | 133 ++++++++++---- .../internal/chunker/chunker_tokens_test.go | 35 ++++ server/internal/config/config.go | 3 +- .../embeddings/provider/voyage/voyage.go | 96 +++++++--- server/internal/embeddings/service.go | 10 +- server/internal/indexer/indexer.go | 24 ++- .../internal/tokenizer/bpecount/bpecount.go | 164 ++++++++++++++++-- .../tokenizer/bpecount/bpecount_test.go | 109 +++++++++++- 9 files changed, 489 insertions(+), 88 deletions(-) diff --git a/.gitignore b/.gitignore index e9f580c5..f9dcd5a7 100644 --- a/.gitignore +++ b/.gitignore @@ -93,3 +93,6 @@ server/internal/httpapi/dashboard/dist/* portainer_mcp/ portainer-mcp* tools.yaml + +# Local load-test corpus + throwaway server instance (never committed) +/loadtests/ diff --git a/server/internal/chunker/chunker.go b/server/internal/chunker/chunker.go index 8a2b938f..7cae9c3c 100644 --- a/server/internal/chunker/chunker.go +++ b/server/internal/chunker/chunker.go @@ -487,7 +487,7 @@ func ChunkFileTokens(filePath, content, language string, maxSize int, budget tok } if budget != nil { if maxTokens <= 0 { - maxTokens = defaultMaxChunkTokens + maxTokens = DefaultMaxChunkTokens } if lim := budget.MaxInputTokens(); lim > 0 && maxTokens > lim { maxTokens = lim @@ -505,7 +505,10 @@ func ChunkFileTokens(filePath, content, language string, maxSize int, budget tok chunks, refs, err := chunkWithTreesitter(filePath, content, language, innerMax) if err != nil { // Fallback: sliding window, no references. - return boundTokens(chunkFallback(filePath, content, language), budget, maxTokens), nil, nil + return chunkFallbackTokens(filePath, content, language, budget, maxTokens), nil, nil + } + if len(chunks) == 0 { + return chunkFallbackTokens(filePath, content, language, budget, maxTokens), nil, nil } return boundTokens(chunks, budget, maxTokens), refs, nil } @@ -523,6 +526,14 @@ func boundTokens(chunks []Chunk, budget tokenizer.Budget, maxTokens int) []Chunk } out := make([]Chunk, 0, len(chunks)) for _, c := range chunks { + // A byte-level BPE token can never cover less than one byte, so a + // chunk shorter than the budget provably fits and needs no counting + // at all. Most chunks are, so this skips the tokenizer on the hot + // path rather than paying for an answer already known. + if len(c.Content) <= maxTokens { + out = append(out, c) + continue + } if budget.CountTokens(c.Content) > maxTokens { out = append(out, splitChunkTokens(c, budget, maxTokens)...) continue @@ -549,6 +560,44 @@ func chunkFallback(filePath, content, language string) []Chunk { return chunkSlidingWindow(filePath, content, language) } +// chunkFallbackTokens is chunkFallback with a token budget: the sliding window +// walks token boundaries instead of a fixed byte count. +// +// The byte window is 4000 bytes regardless of what those bytes contain, so a +// file of Cyrillic or CJK prose — two to three bytes per character — produced +// windows worth a third of the intended tokens, and this is the path such +// files take, because "no grammar" and "not ASCII" go together often enough to +// matter. boundTokens alone could not fix it: it splits chunks that are too +// large and has no way to grow ones that are too small. +func chunkFallbackTokens(filePath, content, language string, budget tokenizer.Budget, maxTokens int) []Chunk { + if budget == nil || maxTokens <= 0 { + return chunkFallback(filePath, content, language) + } + if language == "bash" { + if c := bashRegexChunks(filePath, content); len(c) > 0 { + return boundTokens(c, budget, maxTokens) + } + } + if len(content) == 0 { + return nil + } + // One chunk, then let the token splitter cut it — same code path, same + // guarantees (pieces are substrings, none over budget, line numbers + // tracked) as every other over-budget chunk in this package. + whole := Chunk{ + Content: content, + ChunkType: "block", + FilePath: filePath, + StartLine: 1, + EndLine: countNewlines(content) + 1, + Language: language, + } + if budget.CountTokens(content) <= maxTokens { + return []Chunk{whole} + } + return splitChunkTokens(whole, budget, maxTokens) +} + // --------------------------------------------------------------------------- // Tree-sitter path // --------------------------------------------------------------------------- @@ -1125,11 +1174,15 @@ func sortRanges(ranges [][2]int) { } } -// defaultMaxChunkTokens is the token equivalent of maxChunkSize. The byte +// DefaultMaxChunkTokens is the token equivalent of maxChunkSize. The byte // default was written as 1500*3 — a 1500-token target at three bytes each — // so the token target is that same 1500, now expressed in the unit that -// actually matters. CIX_MAX_CHUNK_TOKENS overrides it. -const defaultMaxChunkTokens = 1500 +// actually matters. +// +// Exported so config.go can use it as the CIX_MAX_CHUNK_TOKENS default rather +// than repeating the number: two copies of a chunk-size default drifting apart +// is the exact failure this change removes for maxChunkSize. +const DefaultMaxChunkTokens = 1500 // splitChunkTokens cuts an over-budget chunk into pieces of <= maxTokens. // @@ -1151,40 +1204,60 @@ const defaultMaxChunkTokens = 1500 // boundary to snap to; there the exact cut stands and the line is split // internally — which is the case the byte splitter could not handle at all. func splitChunkTokens(chunk Chunk, budget tokenizer.Budget, maxTokens int) []Chunk { + content := chunk.Content var out []Chunk - rest := chunk.Content - line := chunk.StartLine - - for rest != "" { - cuts, _ := budget.SplitPoints(rest, maxTokens) - if len(cuts) == 0 { - out = append(out, mkPiece(chunk, rest, line, len(out) == 0)) - break + pos, line := 0, chunk.StartLine + var cuts []int // absolute offsets into content; nil means "recompute" + + for pos < len(content) { + if cuts == nil { + rel, _ := budget.SplitPoints(content[pos:], maxTokens) + if len(rel) == 0 { + out = append(out, mkPiece(chunk, content[pos:], line, len(out) == 0)) + break + } + cuts = make([]int, 0, len(rel)) + for _, r := range rel { + cuts = append(cuts, pos+r) + } } at := cuts[0] - // Pull the cut back to a line start so a piece never begins - // mid-line: search results and the stored line range both lie - // otherwise. Moving backwards only shrinks this piece, so it stays - // inside the budget. A line wider than the whole budget has no - // earlier boundary — there the exact cut stands and the line is - // split internally, which is the case the byte splitter could not - // handle at all. - if nl := strings.LastIndexByte(rest[:at], '\n'); nl >= 0 { - at = nl + 1 + // Pull the cut back to a line start so a piece never begins mid-line: + // search results and the stored line range both lie otherwise. Moving + // backwards only shrinks this piece, so it stays inside the budget. A + // line wider than the whole budget has no earlier boundary — there the + // exact cut stands and the line is split internally, which is the case + // the byte splitter could not handle at all. + snapped := at + if nl := strings.LastIndexByte(content[pos:at], '\n'); nl >= 0 { + snapped = pos + nl + 1 } - if at == 0 { - at = cuts[0] + if snapped <= pos { + snapped = at } - piece := rest[:at] + piece := content[pos:snapped] out = append(out, mkPiece(chunk, piece, line, len(out) == 0)) line += strings.Count(piece, "\n") - rest = rest[at:] - // Recomputing the next cut from the NEW start is the whole point of - // the loop: snapping back moved the boundary, so the remaining cuts - // SplitPoints returned no longer apply — reusing them would hand the - // next piece the words this one gave up and push it over budget. + pos = snapped + + if snapped == at { + // The boundary landed where the tokenizer put it, so the cuts + // after it are still valid and can be consumed without another + // pass. This is what keeps a 66 KB single line — where the + // newline snap never fires — from costing one full scan per + // piece. + cuts = cuts[1:] + if len(cuts) == 0 { + cuts = nil + } + continue + } + // Snapping moved the boundary: every later cut was measured from a + // start that no longer exists, and reusing them would hand the next + // piece the tokens this one gave up. Recompute. + cuts = nil } return out } diff --git a/server/internal/chunker/chunker_tokens_test.go b/server/internal/chunker/chunker_tokens_test.go index f73e1a37..e0dec211 100644 --- a/server/internal/chunker/chunker_tokens_test.go +++ b/server/internal/chunker/chunker_tokens_test.go @@ -256,3 +256,38 @@ func TestTokenSplitLineNumbers(t *testing.T) { } func strPtr(s string) *string { return &s } + +// TestFallbackFillsTheBudget covers the path a file with no grammar takes. +// The byte-sized sliding window cut every 4000 bytes regardless of content, so +// multi-byte text produced windows worth a fraction of the intended tokens — +// and boundTokens could not repair that, since it only splits chunks that are +// too big and cannot merge ones that are too small. +func TestFallbackFillsTheBudget(t *testing.T) { + // Two-bytes-per-character text, well past one byte window. + src := strings.Repeat("привіт світ це коментар українською\n", 400) + b := fakeBudget{maxInput: 4096} + const budget = 200 + + chunks, _, err := ChunkFileTokens("notes.unknownlang", src, "unknownlang", 0, b, budget) + if err != nil { + t.Fatalf("chunk: %v", err) + } + if len(chunks) < 2 { + t.Fatalf("expected several chunks, got %d", len(chunks)) + } + var under int + for i, c := range chunks { + n := b.CountTokens(c.Content) + if n > budget { + t.Errorf("chunk %d is %d tokens, over budget %d", i, n, budget) + } + // The last chunk is a remainder and may legitimately be short. + if i < len(chunks)-1 && n < budget/2 { + under++ + } + } + if under > 0 { + t.Errorf("%d of %d chunks are under half the budget — the window is still byte-sized", + under, len(chunks)) + } +} diff --git a/server/internal/config/config.go b/server/internal/config/config.go index 4743a947..2a346c59 100644 --- a/server/internal/config/config.go +++ b/server/internal/config/config.go @@ -5,6 +5,7 @@ package config import ( "fmt" + "github.com/dvcdsys/code-index/server/internal/chunker" "net" "os" "path/filepath" @@ -352,7 +353,7 @@ func Load() (*Config, error) { } c.ChunkMaxConcurrent = chunkConc - maxChunk, err := getenvInt("CIX_MAX_CHUNK_TOKENS", 1500) + maxChunk, err := getenvInt("CIX_MAX_CHUNK_TOKENS", chunker.DefaultMaxChunkTokens) if err != nil { return nil, err } diff --git a/server/internal/embeddings/provider/voyage/voyage.go b/server/internal/embeddings/provider/voyage/voyage.go index 14f0bb96..3c952f99 100644 --- a/server/internal/embeddings/provider/voyage/voyage.go +++ b/server/internal/embeddings/provider/voyage/voyage.go @@ -199,7 +199,13 @@ func (c *Config) maxBatchSize() int { return defaultMaxBatchSize } -// maxTokensPerBatch returns the effective per-POST token cap. +// maxTokensPerBatch returns the cap implied by config alone — the operator's +// override, or the conservative byte-heuristic default. +// +// Callers on the hot path want (*Provider).maxTokensPerBatch instead, which +// also knows whether a tokenizer is loaded. Two same-named methods one on +// Config and one on Provider is how the batch log came to report 80K while +// packing used 115K, so this one is only for the Provider method to build on. func (c *Config) maxTokensPerBatch() int { if c.MaxTokensPerRequest > 0 { return c.MaxTokensPerRequest @@ -429,6 +435,38 @@ func (p *Provider) EmbedDocuments(ctx context.Context, texts []string) ([][]floa // such chunk, but oversize chunks are rare on well-chunked // indexes — the indexer should already be cutting at function / // class boundaries. +// splitForInput cuts one input down to what the model can read. +// +// With a tokenizer, the question "does this fit" has an exact answer, so the +// byte cap is not consulted at all: an input under the model's context window +// goes through whole, however many bytes it is, and one over it is cut on real +// token boundaries. Without a tokenizer we are back to guessing, and the byte +// cap is the guess. +// +// This matters because the chunker now sizes in tokens. Its bound and the +// provider's byte cap are different units: at CIX_MAX_CHUNK_TOKENS=20000 — +// legal, well inside the 32K window — chunks of 40-80 KB are ordinary, and +// every one of them used to be byte-windowed here and have its window vectors +// averaged into a single vector representing neither half. The averaging path +// now only runs where it is genuinely needed: no tokenizer, no exact answer. +func (p *Provider) splitForInput(text string, maxBytes int) []string { + if p.counter == nil { + return splitOversizeInput(text, maxBytes) + } + limit := p.MaxInputTokens() + offsets, total := p.counter.SplitPoints(text, limit) + if total <= limit || len(offsets) == 0 { + return []string{text} + } + out := make([]string, 0, len(offsets)+1) + prev := 0 + for _, off := range offsets { + out = append(out, text[prev:off]) + prev = off + } + return append(out, text[prev:]) +} + func (p *Provider) embedAndAverage(ctx context.Context, texts []string, inputType string) ([][]float32, error) { maxIn := p.cfg.maxInputBytes() @@ -438,7 +476,7 @@ func (p *Provider) embedAndAverage(ctx context.Context, texts []string, inputTyp var expanded []string totalSplits := 0 for i, t := range texts { - windows := splitOversizeInput(t, maxIn) + windows := p.splitForInput(t, maxIn) spans[i] = span{start: len(expanded), length: len(windows)} expanded = append(expanded, windows...) if len(windows) > 1 { @@ -462,7 +500,7 @@ func (p *Provider) embedAndAverage(ctx context.Context, texts []string, inputTyp "total_inputs", len(expanded), "sub_batches", len(batches), "limit_inputs", p.cfg.maxBatchSize(), - "limit_tokens", p.cfg.maxTokensPerBatch(), + "limit_tokens", p.maxTokensPerBatch(), ) } allVecs := make([][]float32, 0, len(expanded)) @@ -834,25 +872,43 @@ func (p *Provider) apiKey() (string, bool) { // drift rather than for our own error. const exactTokensPerBatch = 115_000 -// maxInputTokens is voyage-code-3's per-input context window. The 32K applies -// to voyage-code-* and voyage-3*; smaller models would need a table here, but -// undershooting only costs an unnecessary split. -const maxInputTokens = 32_000 +// modelContextTokens is the per-input context window, per model. It is a table +// rather than a constant because the factory's own enum offers voyage-code-2, +// whose window is 16K — half of what the rest of the list takes. Treating that +// as 32K would let the chunker build inputs the model cannot read, and with +// truncation enabled Voyage would silently drop the tail. +// +// Unknown models fall back to the conservative 16K: undershooting costs an +// unnecessary split, overshooting costs silent data loss. +var modelContextTokens = map[string]int{ + "voyage-code-3": 32_000, + "voyage-3-large": 32_000, + "voyage-3": 32_000, + "voyage-3-lite": 32_000, + "voyage-code-2": 16_000, +} + +const fallbackContextTokens = 16_000 // maxTokensPerBatch is the provider-level cap: an explicit operator override // wins, then the exact-counting cap, then the conservative byte-heuristic one. func (p *Provider) maxTokensPerBatch() int { if p.cfg.MaxTokensPerRequest > 0 { - return p.cfg.MaxTokensPerRequest + return p.cfg.maxTokensPerBatch() } if p.counter != nil { return exactTokensPerBatch } - return defaultMaxTokensPerBatch + return p.cfg.maxTokensPerBatch() } // MaxInputTokens reports the model's context window for a single input. -func (p *Provider) MaxInputTokens() int { return maxInputTokens } +func (p *Provider) MaxInputTokens() int { + if n, ok := modelContextTokens[p.cfg.Model]; ok { + return n + } + return fallbackContextTokens +} // ExactCounts reports whether CountTokens/SplitPoints are exact rather than // estimated. False means no tokenizer.json was loaded. @@ -876,19 +932,13 @@ func (p *Provider) CountTokens(s string) int { // behaviour, kept only so a caller that ignores ExactCounts still gets // something it can send. Check ExactCounts before trusting these. func (p *Provider) SplitPoints(s string, budget int) ([]int, int) { - if p.counter != nil { - return p.counter.SplitPoints(s, budget) - } - var offsets []int - maxBytes := budget * bytesPerToken - if maxBytes <= 0 { + if p.counter == nil { + // No tokenizer: there are no token boundaries to report. Returning + // byte windows here would be the old behaviour wearing the new + // interface's clothes, and callers check ExactCounts() precisely so + // they can avoid it. splitForInput still byte-windows internally + // where that is genuinely all we have. return nil, estimateTokens(s) } - for off := maxBytes; off < len(s); off += maxBytes { - for off > 0 && !utf8.RuneStart(s[off]) { - off-- - } - offsets = append(offsets, off) - } - return offsets, estimateTokens(s) + return p.counter.SplitPoints(s, budget) } diff --git a/server/internal/embeddings/service.go b/server/internal/embeddings/service.go index 8422d370..b1b18c3d 100644 --- a/server/internal/embeddings/service.go +++ b/server/internal/embeddings/service.go @@ -458,8 +458,6 @@ func (s *Service) Status() provider.Status { return st } -// CurrentKind reports the kind of the active provider, or "" when -// disabled / not yet built. Used by /status and admin endpoints. // TokenBudget returns the active provider as a token budget when it can // count tokens, and nil otherwise. The chunker uses it to size chunks in the // model's own unit; nil keeps it on the byte heuristic. @@ -467,6 +465,12 @@ func (s *Service) Status() provider.Status { // Snapshotted under the read lock like CurrentKind, because a provider swap // mid-file would otherwise mix two models' limits inside one chunk set. func (s *Service) TokenBudget() tokenizer.Budget { + // A typed-nil *Service still satisfies the capability interface the + // indexer asserts on, so the guard is not decoration: without it the + // first indexed file panics inside RLock. + if s == nil || s.disabled { + return nil + } s.mu.RLock() defer s.mu.RUnlock() if s.current == nil { @@ -479,6 +483,8 @@ func (s *Service) TokenBudget() tokenizer.Budget { return b } +// CurrentKind reports the kind of the active provider, or "" when +// disabled / not yet built. Used by /status and admin endpoints. func (s *Service) CurrentKind() string { if s == nil || s.disabled { return "" diff --git a/server/internal/indexer/indexer.go b/server/internal/indexer/indexer.go index eb4eeddb..4608f01d 100644 --- a/server/internal/indexer/indexer.go +++ b/server/internal/indexer/indexer.go @@ -113,6 +113,17 @@ type TokenAwareEmbedder interface { TokenizeAndEmbed(ctx context.Context, texts []string) ([][]float32, error) } +// TokenBudgetSource is the capability of telling the chunker what a chunk +// costs in the active model's tokens. Named rather than asserted inline so a +// rename of TokenBudget is a compile error somewhere instead of a silent +// return to byte-sized chunking everywhere. +// +// *embeddings.Service satisfies it; test fakes generally do not, and get the +// byte path. +type TokenBudgetSource interface { + TokenBudget() tokenizer.Budget +} + // Service owns sessions and wires dependencies for the three-phase protocol. type Service struct { db *sql.DB @@ -676,6 +687,7 @@ func (s *Service) ProcessFilesStreaming( // is CPU-local and cheap, so it stays sequential to keep progress-event // order; the expensive embed work is parallelised in stage 2. prep := make([]*preparedFile, 0, len(files)) + budgetSrc, _ := s.emb.(TokenBudgetSource) for fi, fp := range files { // file_started — emit even for files we'll skip below, so the client // counter advances monotonically and rendering stays aligned with N. @@ -716,14 +728,12 @@ func (s *Service) ProcessFilesStreaming( language = "text" } - // The token budget comes from the LIVE provider, asked per file: a - // provider swap between files is legitimate, mixing two models' - // limits inside one file's chunks is not. + // The budget is re-read per file: a provider swap between files is + // legitimate, mixing two models' limits inside one file's chunks is + // not. The type assertion itself is hoisted out of the loop. var budget tokenizer.Budget - if tb, ok := s.emb.(interface { - TokenBudget() tokenizer.Budget - }); ok { - budget = tb.TokenBudget() + if budgetSrc != nil { + budget = budgetSrc.TokenBudget() } chunks, refs, err := chunker.ChunkFileTokens(fp.Path, fp.Content, language, 0, budget, s.maxChunkTokens) if err != nil { diff --git a/server/internal/tokenizer/bpecount/bpecount.go b/server/internal/tokenizer/bpecount/bpecount.go index 381ef416..3bd9d1d3 100644 --- a/server/internal/tokenizer/bpecount/bpecount.go +++ b/server/internal/tokenizer/bpecount/bpecount.go @@ -38,6 +38,55 @@ type tokJSON struct { Type string `json:"type"` Merges json.RawMessage `json:"merges"` } `json:"model"` + Normalizer struct { + Type string `json:"type"` + } `json:"normalizer"` + PreTokenizer struct { + Type string `json:"type"` + PreTokenizers []struct { + Type string `json:"type"` + Pattern struct { + Regex string `json:"Regex"` + } `json:"pattern"` + } `json:"pretokenizers"` + } `json:"pre_tokenizer"` +} + +// qwen2SplitPattern is the pre-tokenizer regex this package implements by +// hand. It is compared, not compiled: the point of the hand-rolled splitter is +// that Go's RE2 cannot express the `\s+(?!\S)` lookahead in it. +// +// The comparison is the load-time guard against a plausible and silent +// failure: GPT-2 and o200k tokenizer.json files parse fine, declare +// model.type "BPE", and would produce confidently wrong counts against this +// splitter — GPT-2 has no Split stage at all, o200k has a different pattern. +// Refusing them keeps ExactCounts() false and the caller on its estimate, +// which is wrong but knows it is. +const qwen2SplitPattern = `(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\r\n\p{L}\p{N}]?\p{L}+|\p{N}| ?[^\s\p{L}\p{N}]+[\r\n]*|\s*[\r\n]+|\s+(?!\S)|\s+` + +// checkPipeline refuses a tokenizer whose normalizer or pre-tokenizer is not +// the one this package reimplements. Counting a different pipeline with this +// splitter does not fail loudly — it returns plausible numbers that are wrong, +// and the caller then packs batches and sizes chunks against them. +func checkPipeline(tj tokJSON) error { + if tj.Normalizer.Type != "NFC" { + return fmt.Errorf("bpecount: normalizer is %q, this package implements NFC", + tj.Normalizer.Type) + } + if tj.PreTokenizer.Type != "Sequence" || len(tj.PreTokenizer.PreTokenizers) < 2 { + return fmt.Errorf("bpecount: pre_tokenizer is %q, expected Sequence[Split, ByteLevel]", + tj.PreTokenizer.Type) + } + split, byteLevel := tj.PreTokenizer.PreTokenizers[0], tj.PreTokenizer.PreTokenizers[1] + if split.Type != "Split" || byteLevel.Type != "ByteLevel" { + return fmt.Errorf("bpecount: pre_tokenizer is Sequence[%s, %s], expected Sequence[Split, ByteLevel]", + split.Type, byteLevel.Type) + } + if split.Pattern.Regex != qwen2SplitPattern { + return fmt.Errorf("bpecount: Split pattern is not the one implemented here " + + "(a GPT-2 or o200k tokenizer would count wrong rather than fail)") + } + return nil } // Load reads a tokenizer.json and keeps only what a count needs: the merges. @@ -57,6 +106,9 @@ func LoadBytes(b []byte) (*Counter, error) { if tj.Model.Type != "BPE" { return nil, fmt.Errorf("bpecount: unsupported model type %q", tj.Model.Type) } + if err := checkPipeline(tj); err != nil { + return nil, err + } // merges is either ["a b", ...] (v1) or [["a","b"], ...] (v2). var flat []string m := make(map[string]int32) @@ -123,10 +175,20 @@ func nextToken(s string) int { r0, w0 := decode(s, 0) // A: contraction + // + // Only the two bytes after the apostrophe can matter (the longest + // contraction is "re"/"ve"/"ll"), so lowercase just those. Lowercasing the + // whole remaining string here allocated a copy of the suffix for every + // apostrophe in the file: a 512 KB source with 5,000 quotes moved over a + // gigabyte through the allocator, on the indexing hot path. if r0 == '\'' && len(s) > w0 { - low := strings.ToLower(s) + tail := s[w0:] + if len(tail) > 2 { + tail = tail[:2] + } + low := strings.ToLower(tail) for _, c := range contractions { - if strings.HasPrefix(low[w0:], c) { + if strings.HasPrefix(low, c) { return w0 + len(c) } } @@ -372,6 +434,14 @@ func (c *Counter) Count(text string) int { // SplitPoints returns the byte offsets at which text must be cut so that no // piece exceeds budget tokens, plus the total token count of the whole text. // +// Offsets index the CALLER's string, which is not the same thing as indexing +// the normalised copy counting works on. NFC can both shrink the text (a +// decomposed "e"+U+0301 becomes one code point) and grow it (the composition +// exclusions at U+0958..U+095F decompose under NFC), so an offset taken from +// the normalised copy can land mid-rune or past the end of the original — the +// latter panics the caller's slice expression. Already-normalised input, which +// is nearly all source code, takes the fast path where the two coincide. +// // The cuts are exact, not estimated, and they need no search. BPE merges never // cross a pre-token boundary in this pipeline (Split runs with Isolated // behaviour, and ByteLevel+BPE are applied per pre-token), so a text's token @@ -379,14 +449,11 @@ func (c *Counter) Count(text string) int { // therefore leaves both sides tokenising exactly as they did inside the whole: // the parts always add up to the total, with no drift to correct for. // -// One pass, left to right, reusing the same memo Count uses — so a split costs -// what a count costs, plus the offsets slice. -// // A single pre-token larger than budget cannot be honoured on a boundary (its -// merges DO interact internally). Rather than silently emit an over-budget -// piece, splitInside falls back to a binary search on bytes within that one -// pre-token. That path is for base64 blobs and minified lines with no -// whitespace; on a 45-repo corpus it fires for 5 chunks in 1.9M. +// merges DO interact internally), so splitInside falls back to a binary search +// on bytes within that one pre-token. That path is for base64 blobs and +// minified lines with no whitespace; on a 45-repo corpus it fires for 5 chunks +// in 1.9M. // // Offsets are cut points only: nil means the text already fits. func (c *Counter) SplitPoints(text string, budget int) (offsets []int, total int) { @@ -396,10 +463,68 @@ func (c *Counter) SplitPoints(text string, budget int) (offsets []int, total int if budget <= 0 { return nil, c.Count(text) } - s := norm.NFC.String(text) + if norm.NFC.IsNormalString(text) { + return c.splitNormalized(text, budget) + } + return c.splitDenormalized(text, budget) +} + +// splitDenormalized handles input that is not already NFC. +// +// It cuts one piece at a time: the first cut is computed on the normalised +// form, mapped back to a raw offset on a normalisation boundary, and the +// remainder is then processed from its real start. Recomputing per piece +// rather than translating a whole offsets slice is deliberate — mapping a cut +// backwards to a boundary shrinks the piece before it and grows the one after, +// so offsets computed against the old start would no longer hold. The cost is +// one pass per piece, paid only by input that is not already normalised, which +// source code essentially never is. +func (c *Counter) splitDenormalized(text string, budget int) ([]int, int) { + total := c.Count(text) + var offsets []int + base := 0 + for base < len(text) { + rest := text[base:] + cuts, _ := c.splitNormalized(norm.NFC.String(rest), budget) + if len(cuts) == 0 { + break + } + at := rawOffsetOf(rest, cuts[0]) + if at <= 0 || at >= len(rest) { + break + } + offsets = append(offsets, base+at) + base += at + } + return offsets, total +} - acc := 0 // tokens accumulated in the current piece - pos := 0 // byte offset into s +// rawOffsetOf maps a byte offset in NFC(raw) back to a byte offset in raw, +// rounding DOWN to a normalisation boundary — rounding down can only shrink +// the piece that ends there, so the budget survives the rounding. +func rawOffsetOf(raw string, normOff int) int { + rawPos, normPos, lastRaw := 0, 0, 0 + for rawPos < len(raw) { + n := norm.NFC.NextBoundaryInString(raw[rawPos:], true) + if n <= 0 { + break + } + segNorm := len(norm.NFC.String(raw[rawPos : rawPos+n])) + if normPos+segNorm > normOff { + return lastRaw + } + normPos += segNorm + rawPos += n + lastRaw = rawPos + } + return lastRaw +} + +// splitNormalized is SplitPoints for input already known to be NFC, where a +// normalised offset IS a raw offset. +func (c *Counter) splitNormalized(s string, budget int) (offsets []int, total int) { + acc := 0 // tokens accumulated in the current piece + pos := 0 // byte offset into s var sb strings.Builder for pos < len(s) { n := nextToken(s[pos:]) @@ -423,14 +548,19 @@ func (c *Counter) SplitPoints(text string, budget int) (offsets []int, total int offsets = append(offsets, pos) acc = 0 } - inner := c.splitInside(s[pos:pos+n], budget) + inner := c.splitInside(piece, budget) for _, off := range inner { offsets = append(offsets, pos+off) } - // Tail of the pre-token starts a fresh piece; its token count is - // unknown without re-counting, so charge it conservatively as a - // full budget minus nothing and let the next boundary close it. - acc = 0 + // The tail after the last inner cut opens the next piece and + // must be CHARGED for: leaving acc at zero let the following + // pre-tokens add a full budget on top of it, producing pieces of + // up to twice the budget. + tailStart := 0 + if len(inner) > 0 { + tailStart = inner[len(inner)-1] + } + acc = c.Count(piece[tailStart:]) case acc+tk > budget: offsets = append(offsets, pos) acc = tk diff --git a/server/internal/tokenizer/bpecount/bpecount_test.go b/server/internal/tokenizer/bpecount/bpecount_test.go index 546ec1f0..1adf5f1c 100644 --- a/server/internal/tokenizer/bpecount/bpecount_test.go +++ b/server/internal/tokenizer/bpecount/bpecount_test.go @@ -1,7 +1,9 @@ package bpecount import ( + "encoding/json" "os" + "strings" "testing" ) @@ -53,8 +55,8 @@ func TestCountMatchesReference(t *testing.T) { // input cost an extra token. func TestNFCNormalisation(t *testing.T) { c := load(t) - nfc := "caf\u00e9" // é as one code point - nfd := "cafe\u0301" // e + combining acute + nfc := "caf\u00e9" // é as one code point + nfd := "cafe\u0301" // e + combining acute if a, b := c.Count(nfc), c.Count(nfd); a != b { t.Errorf("NFC %d != NFD %d — normaliser not applied", a, b) } @@ -109,10 +111,11 @@ func TestSplitPointsFitsAlready(t *testing.T) { // single pre-token bigger than the budget (base64 blobs, minified lines). func TestSplitPointsOversizePreToken(t *testing.T) { c := load(t) - blob := "" - for i := 0; i < 4000; i++ { - blob += "aB3" - } + // One unbroken run of a single character class. "aB3" repeated would NOT + // do: the pre-tokenizer breaks letters from digits, so it yields 2-byte + // pre-tokens that never exceed the budget and the splitInside path this + // test exists for is never entered. + blob := strings.Repeat("a", 4000) const budget = 100 offsets, _ := c.SplitPoints(blob, budget) if len(offsets) == 0 { @@ -137,8 +140,7 @@ func TestSplitPointsOversizePreToken(t *testing.T) { func syntheticCounter(t *testing.T) *Counter { t.Helper() // Merges are ranked: "a b" collapses first, then "ab c". - tj := `{"model":{"type":"BPE","merges":["a b","ab c","f u","fu n"]}}` - c, err := LoadBytes([]byte(tj)) + c, err := LoadBytes(syntheticJSON("a b", "ab c", "f u", "fu n")) if err != nil { t.Fatalf("LoadBytes: %v", err) } @@ -231,3 +233,94 @@ func TestPreTokenBoundaries(t *testing.T) { } } } + +// TestNFDPiecesRespectBudget covers input that is not already NFC, where the +// normalised copy counting works on has different byte offsets from the +// caller's string. Before the fix, offsets computed against the normalised +// copy were returned as-is: decomposed text made them land mid-rune, and the +// composition exclusions at U+0958..U+095F (which DEcompose under NFC, making +// the normalised form longer) pushed them past the end of the original, so the +// caller's slice expression panicked. Every other fixture in this file is +// ASCII or already-NFC and could not see it. +func TestNFDPiecesRespectBudget(t *testing.T) { + c := load(t) + for _, raw := range []string{ + strings.Repeat("// café comment here\n", 200), + strings.Repeat("x क़ख़ग़ ", 400), + strings.Repeat("Ώ ", 900), + } { + offs, total := c.SplitPoints(raw, 50) + prev := 0 + for _, off := range append(offs, len(raw)) { + if off > len(raw) || off < prev { + t.Fatalf("bad offset %d (len %d, prev %d)", off, len(raw), prev) + } + if n := c.Count(raw[prev:off]); n > 50 { + t.Errorf("piece [%d:%d] is %d tokens, over budget 50", prev, off, n) + } + prev = off + } + if total != c.Count(raw) { + t.Errorf("total %d != Count %d", total, c.Count(raw)) + } + } +} + +// TestSplitInsideChargesTail — an over-budget pre-token used to leave its tail +// uncounted, letting the next piece stack a full budget on top of it. +func TestSplitInsideChargesTail(t *testing.T) { + c, err := LoadBytes(syntheticJSON("a b")) + if err != nil { + t.Fatal(err) + } + in := strings.Repeat("x", 10) + " " + strings.Repeat("y", 4) + offs, _ := c.SplitPoints(in, 5) + prev := 0 + for _, off := range append(offs, len(in)) { + if n := c.Count(in[prev:off]); n > 5 { + t.Errorf("piece %q is %d tokens, over budget 5", in[prev:off], n) + } + prev = off + } +} + +// syntheticJSON builds a tokenizer.json with a hand-picked merge table and the +// real pipeline sections, so LoadBytes's compatibility check sees what it +// expects. Tests that only exercise merging still have to declare the pipeline +// they are pretending to be — which is the point of the check. +func syntheticJSON(merges ...string) []byte { + doc := map[string]any{ + "model": map[string]any{"type": "BPE", "merges": merges}, + "normalizer": map[string]any{"type": "NFC"}, + "pre_tokenizer": map[string]any{ + "type": "Sequence", + "pretokenizers": []any{ + map[string]any{"type": "Split", "pattern": map[string]any{"Regex": qwen2SplitPattern}}, + map[string]any{"type": "ByteLevel"}, + }, + }, + } + b, err := json.Marshal(doc) + if err != nil { + panic(err) + } + return b +} + +// TestRejectsForeignPipeline — a GPT-2 or o200k tokenizer.json parses cleanly +// and declares BPE, but its pre-tokenizer is not the one implemented here. It +// must be refused rather than counted wrongly. +func TestRejectsForeignPipeline(t *testing.T) { + for name, doc := range map[string]string{ + "gpt2 (no Split stage)": `{"model":{"type":"BPE","merges":["a b"]}, + "normalizer":null,"pre_tokenizer":{"type":"ByteLevel"}}`, + "o200k (different pattern)": `{"model":{"type":"BPE","merges":["a b"]}, + "normalizer":{"type":"NFC"},"pre_tokenizer":{"type":"Sequence","pretokenizers":[ + {"type":"Split","pattern":{"Regex":"[^\\r\\n\\p{L}\\p{N}]?[\\p{L}]+"}}, + {"type":"ByteLevel"}]}}`, + } { + if _, err := LoadBytes([]byte(doc)); err == nil { + t.Errorf("%s: expected a load error, got none", name) + } + } +} From 0328c18a9b0c5585403e8789cc531359532c7e44 Mon Sep 17 00:00:00 2001 From: dvcdsys Date: Tue, 18 Aug 2026 15:52:26 +0100 Subject: [PATCH 6/6] fix(chunker,tokenizer): unreachable fallback and a deadlock on multi-byte runs MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Two blockers from the second review, both confirmed before fixing. DEADLOCK. splitInside binary-searched over byte offsets and aligned its midpoint to a rune start by DECREMENTING. When alignment pulled the midpoint below lo, the next lo = mid+1 did not advance, and the (lo, hi) pair repeated forever. Any run of multi-byte runes long enough to exceed the budget as one pre-token reached it — a box-drawing comment separator does, and so does an arrow run or a run of combining marks. The worker did not error or slow down, it stopped. Reproduced at budgets 5 and 50 on four inputs, including already-NFC ones on the fast path, so this was not confined to the new denormalised route. The search now runs over rune-start INDICES enumerated once, so every candidate is a valid boundary by construction and no alignment step exists to misbehave. A single rune that exceeds the budget is emitted rather than refused: not advancing is the failure being removed. All four inputs are tests now, at both budgets, asserting termination AND that no piece is cut mid-rune. Strictly this loop predates the previous commit, but the denormalised path added routes into it. UNREACHABLE FALLBACK. chunkFallbackTokens was dead code. chunkWithTreesitter never returns an error — all six of its decline paths called chunkFallback themselves and returned the byte windows as SUCCESS — so the caller's fallback branch could not run, and no-grammar, minified, parse-failure and empty-AST files kept getting 4000-byte windows regardless of any token budget. That is the finding the previous commit claimed to fix. Those six sites now return errUseFallback and the caller decides, because the caller is the one that knows whether a budget is in play. ChunkFile's byte path is unchanged: it already mapped err to chunkFallback. The test that was supposed to catch this passed by arithmetic. At budget 200 a 4000-byte window of the fixture is ~358 fake tokens, so boundTokens split every window into 200+158 and both halves cleared half-budget without the fallback being token-aware at all. Raised to 800, where a raw window is BELOW half the budget: verified it now fails against the old routing (7 of 8 chunks under half budget) and passes against the new. Also from the review: the split log reported max_input_bytes even when the split was token-bounded; the provider-level token-split branch had no direct test (added, covering both pass-through of a large-but-legal input and cutting one past the window); a doc comment still named the old lowercase constant. The token fallback drops the byte window's 500-byte overlap. That is deliberate and now documented at the site: overlap has to be expressed in tokens to coexist with a token budget — size pieces at budget minus overlap, then extend each start back into its predecessor — which is a change worth making on its own rather than smuggling into a correctness fix. The tree-sitter path has never overlapped, so the two paths are now consistent rather than the fallback being worse than its neighbours. Co-Authored-By: Claude Opus 5 --- server/internal/chunker/chunker.go | 46 +++++++++++----- .../internal/chunker/chunker_tokens_test.go | 8 ++- .../embeddings/provider/voyage/voyage.go | 24 ++++++--- .../embeddings/provider/voyage/voyage_test.go | 45 ++++++++++++++++ .../internal/tokenizer/bpecount/bpecount.go | 53 +++++++++++++------ .../tokenizer/bpecount/bpecount_test.go | 48 +++++++++++++++++ 6 files changed, 190 insertions(+), 34 deletions(-) diff --git a/server/internal/chunker/chunker.go b/server/internal/chunker/chunker.go index 7cae9c3c..5b147971 100644 --- a/server/internal/chunker/chunker.go +++ b/server/internal/chunker/chunker.go @@ -10,6 +10,7 @@ package chunker import ( + "errors" "github.com/dvcdsys/code-index/server/internal/tokenizer" "log/slog" "math" @@ -473,7 +474,7 @@ func ChunkFile(filePath, content, language string, maxSize int) ([]Chunk, []Refe // When budget is non-nil and reports exact counts, the size decision is made // in tokens instead, and an over-budget chunk is cut on real token boundaries // (via Budget.SplitPoints) rather than on a byte count. maxTokens <= 0 uses -// defaultMaxChunkTokens. A nil or estimating budget keeps the byte path +// DefaultMaxChunkTokens. A nil or estimating budget keeps the byte path // unchanged, so nothing about existing behaviour depends on a provider // having a tokenizer. func ChunkFileTokens(filePath, content, language string, maxSize int, budget tokenizer.Budget, maxTokens int) ([]Chunk, []Reference, error) { @@ -503,11 +504,11 @@ func ChunkFileTokens(filePath, content, language string, maxSize int, budget tok } chunks, refs, err := chunkWithTreesitter(filePath, content, language, innerMax) - if err != nil { - // Fallback: sliding window, no references. - return chunkFallbackTokens(filePath, content, language, budget, maxTokens), nil, nil - } - if len(chunks) == 0 { + if err != nil || len(chunks) == 0 { + // Tree-sitter declined (no grammar, parse failure, minified input) or + // produced nothing. Either way the fallback runs here, where the + // budget is known, rather than inside the tree-sitter path where it + // is not. return chunkFallbackTokens(filePath, content, language, budget, maxTokens), nil, nil } return boundTokens(chunks, budget, maxTokens), refs, nil @@ -551,6 +552,17 @@ func boundTokens(chunks []Chunk, budget tokenizer.Budget, maxTokens int) []Chunk // `block` ones, which is much more useful for semantic search. If the // extractor returns nil (no symbols found), we fall through to the universal // sliding-window strategy so the file content is still indexed. +// errUseFallback tells the caller that the tree-sitter path declined this +// file and the fallback chunker must run instead. +// +// It exists because chunkWithTreesitter used to CALL chunkFallback itself and +// return the result as success. That made the caller's fallback branch +// unreachable — which is exactly how a token-aware fallback shipped as dead +// code: no-grammar, minified, parse-failure and empty-AST files all came back +// as byte windows wearing a success return, and no budget ever reached them. +// The decision belongs to whoever knows whether a token budget is in play. +var errUseFallback = errors.New("chunker: tree-sitter declined, use fallback") + func chunkFallback(filePath, content, language string) []Chunk { if language == "bash" { if c := bashRegexChunks(filePath, content); len(c) > 0 { @@ -584,6 +596,16 @@ func chunkFallbackTokens(filePath, content, language string, budget tokenizer.Bu // One chunk, then let the token splitter cut it — same code path, same // guarantees (pieces are substrings, none over budget, line numbers // tracked) as every other over-budget chunk in this package. + // + // Note this drops the byte window's 500-byte overlap. That overlap existed + // so a match spanning a window boundary would still be found in one of the + // two windows, and nothing replaces it here: pieces are contiguous. The + // trade is deliberate for now — overlap has to be expressed in tokens to + // coexist with a token budget (size pieces at budget minus overlap, then + // extend each start back into the previous piece), which is a change worth + // making on its own rather than smuggling into a correctness fix. The + // tree-sitter path has never overlapped, so this makes the two paths + // consistent rather than making the fallback worse than its neighbours. whole := Chunk{ Content: content, ChunkType: "block", @@ -642,11 +664,11 @@ func chunkWithTreesitter(filePath, content, language string, maxSize int) ([]Chu registryMu.RUnlock() if !ok { - return chunkFallback(filePath, content, language), nil, nil + return nil, nil, errUseFallback } if nodeKinds == nil { // Grammar exists but we don't have node definitions → sliding window. - return chunkFallback(filePath, content, language), nil, nil + return nil, nil, errUseFallback } if looksMinified(filePath, content, language) { // Minified/bundled sources are the parser's pathological case: a @@ -654,7 +676,7 @@ func chunkWithTreesitter(filePath, content, language string, maxSize int) ([]Chu // instance to its memory cap, and forces a pool recycle — all to // produce AST chunks with near-zero semantic-search value. Skip // straight to the sliding window. - return chunkFallback(filePath, content, language), nil, nil + return nil, nil, errUseFallback } // Build flat target → kind map. @@ -672,10 +694,10 @@ func chunkWithTreesitter(filePath, content, language string, maxSize int) ([]Chu // fall back to sliding window so the file is still indexed. slog.Warn("chunker: wasm parse failed, falling back to sliding window", "path", filePath, "language", language, "err", err) - return chunkFallback(filePath, content, language), nil, nil + return nil, nil, errUseFallback } if len(nodes) == 0 { - return chunkFallback(filePath, content, language), nil, nil + return nil, nil, errUseFallback } tree := buildFlatTree(nodes) @@ -717,7 +739,7 @@ func chunkWithTreesitter(filePath, content, language string, maxSize int) ([]Chu } if len(finalChunks) == 0 { - return chunkFallback(filePath, content, language), nil, nil + return nil, nil, errUseFallback } return finalChunks, refs, nil } diff --git a/server/internal/chunker/chunker_tokens_test.go b/server/internal/chunker/chunker_tokens_test.go index e0dec211..552ddeb5 100644 --- a/server/internal/chunker/chunker_tokens_test.go +++ b/server/internal/chunker/chunker_tokens_test.go @@ -266,7 +266,13 @@ func TestFallbackFillsTheBudget(t *testing.T) { // Two-bytes-per-character text, well past one byte window. src := strings.Repeat("привіт світ це коментар українською\n", 400) b := fakeBudget{maxInput: 4096} - const budget = 200 + // The budget must exceed one byte window's token worth, or the test + // passes on arithmetic: a 4000-byte window of this text is ~358 fake + // tokens, so at budget 200 boundTokens splits every window into 200+158 + // and both halves clear half-budget without the fallback ever being + // token-aware. At 800 the raw window is BELOW half the budget, so an + // under-filled chunk can only come from a byte-sized window. + const budget = 800 chunks, _, err := ChunkFileTokens("notes.unknownlang", src, "unknownlang", 0, b, budget) if err != nil { diff --git a/server/internal/embeddings/provider/voyage/voyage.go b/server/internal/embeddings/provider/voyage/voyage.go index 3c952f99..c9e73375 100644 --- a/server/internal/embeddings/provider/voyage/voyage.go +++ b/server/internal/embeddings/provider/voyage/voyage.go @@ -484,12 +484,24 @@ func (p *Provider) embedAndAverage(ctx context.Context, texts []string, inputTyp } } if totalSplits > 0 { - p.logger.Info("voyage: oversize inputs split into byte-windows", - "original_inputs", len(texts), - "total_windows", len(expanded), - "split_windows", totalSplits, - "max_input_bytes", maxIn, - ) + // Report the unit the split actually used: with a tokenizer the cut is + // on token boundaries against the model's context, and logging a byte + // cap there sends whoever reads this to the wrong knob. + if p.counter != nil { + p.logger.Info("voyage: oversize inputs split on token boundaries", + "original_inputs", len(texts), + "total_windows", len(expanded), + "split_windows", totalSplits, + "max_input_tokens", p.MaxInputTokens(), + ) + } else { + p.logger.Info("voyage: oversize inputs split into byte-windows", + "original_inputs", len(texts), + "total_windows", len(expanded), + "split_windows", totalSplits, + "max_input_bytes", maxIn, + ) + } } // Phase 2: batch + POST as before, on the expanded slice. diff --git a/server/internal/embeddings/provider/voyage/voyage_test.go b/server/internal/embeddings/provider/voyage/voyage_test.go index 94869308..4eb4dadf 100644 --- a/server/internal/embeddings/provider/voyage/voyage_test.go +++ b/server/internal/embeddings/provider/voyage/voyage_test.go @@ -6,9 +6,11 @@ import ( "encoding/json" "fmt" "github.com/dvcdsys/code-index/server/internal/tokenizer" + "github.com/dvcdsys/code-index/server/internal/tokenizer/bpecount" "io" "net/http" "net/http/httptest" + "os" "strings" "sync/atomic" "testing" @@ -685,3 +687,46 @@ func TestOperatorOverrideWinsOverExactCap(t *testing.T) { t.Errorf("batch cap = %d, want the operator's 42000", got) } } + +// TestSplitForInputUsesTokenBoundaries exercises the provider-level split with +// a tokenizer loaded — the branch that keeps a large-but-legal chunk out of the +// byte-window-and-average path. It needs the real tokenizer.json, so it skips +// on a clean checkout like the other fixture-backed tests. +func TestSplitForInputUsesTokenBoundaries(t *testing.T) { + const tokPath = "../../../../../loadtests/bench/voyage-code-3.tokenizer.json" + if _, err := os.Stat(tokPath); err != nil { + t.Skip("tokenizer.json not present") + } + c, err := bpecount.Load(tokPath) + if err != nil { + t.Fatalf("load tokenizer: %v", err) + } + p := &Provider{cfg: Config{Model: "voyage-code-3"}, counter: c} + + // Comfortably over the old 30 KB byte cap, comfortably under the model's + // 32K-token window: byte-windowing would split and average this, token + // counting must pass it through whole. + big := strings.Repeat("func handler(w http.ResponseWriter) { defer r.Body.Close() }\n", 700) + if n := p.CountTokens(big); n >= p.MaxInputTokens() { + t.Fatalf("fixture is %d tokens, needs to be under %d", n, p.MaxInputTokens()) + } + if got := p.splitForInput(big, 30_000); len(got) != 1 { + t.Errorf("input of %d bytes / %d tokens split into %d windows; a token-sized "+ + "input must pass through whole", len(big), p.CountTokens(big), len(got)) + } + + // Past the window: must split, and every piece must fit. + huge := strings.Repeat("x := compute(alpha, beta, gamma) // annotate the result\n", 40_000) + pieces := p.splitForInput(huge, 30_000) + if len(pieces) < 2 { + t.Fatalf("input of %d tokens was not split", p.CountTokens(huge)) + } + for i, piece := range pieces { + if n := p.CountTokens(piece); n > p.MaxInputTokens() { + t.Errorf("piece %d is %d tokens, over the %d-token window", i, n, p.MaxInputTokens()) + } + } + if strings.Join(pieces, "") != huge { + t.Error("pieces do not reconstruct the input") + } +} diff --git a/server/internal/tokenizer/bpecount/bpecount.go b/server/internal/tokenizer/bpecount/bpecount.go index 3bd9d1d3..64c7387f 100644 --- a/server/internal/tokenizer/bpecount/bpecount.go +++ b/server/internal/tokenizer/bpecount/bpecount.go @@ -578,34 +578,57 @@ func (c *Counter) splitNormalized(s string, budget int) (offsets []int, total in // re-counted — but the search converges in a handful of probes because // bytes-per-token is near-constant within a homogeneous run. func (c *Counter) splitInside(piece string, budget int) []int { + // Candidate cut positions are rune starts, enumerated once. The search + // then runs over INDICES into that list rather than over byte offsets. + // + // The byte-offset version of this loop deadlocked: it aligned a midpoint + // to a rune start by decrementing, and when alignment pulled the midpoint + // below lo, the next lo = mid+1 did not advance, so the (lo, hi) pair + // repeated forever. Any run of multi-byte runes long enough to exceed the + // budget reached it — a box-drawing comment separator is enough, and that + // hung the indexing worker with no error and no progress. Searching over + // rune indices removes the failure rather than guarding it: every + // candidate is a valid boundary by construction, so no alignment step + // exists to misbehave. + starts := make([]int, 0, len(piece)/2+2) + for i := 0; i < len(piece); { + starts = append(starts, i) + _, w := utf8.DecodeRuneInString(piece[i:]) + if w <= 0 { + w = 1 + } + i += w + } + starts = append(starts, len(piece)) + var cuts []int - start := 0 - for start < len(piece) { - if c.Count(piece[start:]) <= budget { + si := 0 + for si < len(starts)-1 { + if c.Count(piece[starts[si]:]) <= budget { break } - lo, hi := start+1, len(piece) - best := start + 1 + // Largest j > si whose prefix still fits. + lo, hi, best := si+1, len(starts)-1, -1 for lo <= hi { mid := (lo + hi) / 2 - for mid > start && mid < len(piece) && !utf8.RuneStart(piece[mid]) { - mid-- - } - if mid <= start { - break - } - if c.Count(piece[start:mid]) <= budget { + if c.Count(piece[starts[si]:starts[mid]]) <= budget { best = mid lo = mid + 1 } else { hi = mid - 1 } } - if best >= len(piece) { + if best < 0 { + // Even one rune exceeds the budget. Emit it anyway: refusing to + // advance is the deadlock this rewrite exists to remove, and a + // budget smaller than a single token is the caller's problem. + best = si + 1 + } + if starts[best] >= len(piece) { break } - cuts = append(cuts, best) - start = best + cuts = append(cuts, starts[best]) + si = best } return cuts } diff --git a/server/internal/tokenizer/bpecount/bpecount_test.go b/server/internal/tokenizer/bpecount/bpecount_test.go index 1adf5f1c..d41a6dfc 100644 --- a/server/internal/tokenizer/bpecount/bpecount_test.go +++ b/server/internal/tokenizer/bpecount/bpecount_test.go @@ -5,6 +5,8 @@ import ( "os" "strings" "testing" + "time" + "unicode/utf8" ) // tokenizerPath is the real voyage-code-3 tokenizer.json. The tests that need @@ -324,3 +326,49 @@ func TestRejectsForeignPipeline(t *testing.T) { } } } + +// TestMultibyteRunsTerminate covers runs of multi-byte runes long enough to +// exceed the budget as a single pre-token — a box-drawing comment separator, +// an arrow run, a run of combining marks. +// +// The byte-offset binary search this replaced aligned its midpoint to a rune +// start by DECREMENTING, so when alignment pulled the midpoint below lo, the +// next lo = mid+1 did not advance and the search spun forever. It took no +// error path and produced no output: the indexing worker simply stopped. All +// three inputs below hung at budgets 5 and 50. +func TestMultibyteRunsTerminate(t *testing.T) { + c := load(t) + inputs := map[string]string{ + "box drawing separator": "// " + strings.Repeat("\u2500", 400) + "\n", + "arrow run": strings.Repeat("\u2192", 400), + "combining marks": strings.Repeat("\u0301", 50), + "composition exclusion": strings.Repeat("\u0958", 2000), + } + for _, budget := range []int{5, 50} { + for name, in := range inputs { + done := make(chan struct{}) + var offs []int + go func(s string, b int) { + offs, _ = c.SplitPoints(s, b) + close(done) + }(in, budget) + + select { + case <-done: + case <-time.After(10 * time.Second): + t.Fatalf("%s at budget %d: SplitPoints did not return", name, budget) + } + + prev := 0 + for _, off := range append(offs, len(in)) { + if off > len(in) || off < prev { + t.Fatalf("%s: offset %d out of range (len %d, prev %d)", name, off, len(in), prev) + } + if !utf8.ValidString(in[prev:off]) { + t.Errorf("%s: piece [%d:%d] is not valid UTF-8 — cut mid-rune", name, prev, off) + } + prev = off + } + } + } +}