diff --git a/CHANGELOG.md b/CHANGELOG.md index 85f87c32..f7bfabd0 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -16,11 +16,33 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] ### Changed +- **BPE estimator cache no longer retains input texts** (`internal/core`). + The 8192-entry sharded LRU previously stored the full original string per + entry for collision checks and copied every input to `[]byte` on each hash. + It now hashes with `hash/maphash` (zero-copy, random per-process seed), + validates hits by `(hash, length)`, and skips caching texts above a 1 MiB + cutoff. This bounds the cache at small fixed-size records regardless of + prompt size — previously, `Compress` estimating the original plus every + intermediate layer output could pin hundreds of MB in the cache. + +### Fixed +- **`RestorationTracker` no longer retains file contents unboundedly.** + `Track` previously did an O(n) linear dedup scan per call (O(n²) over a + session) and kept every tracked file's full content in memory until + `Clear()`. Dedup is now O(1) via a path-indexed map, and total retained + content is capped by a token budget (`DefaultRetainedTokenBudget`, + 200_000 tokens — one default context window, ~800 KB) that evicts the + least recently tracked entries; configurable via + `WithRetainedTokenBudget`. The most recently tracked entry is always kept. - **Version re-baselined to `0.1.0`** across `CITATION.cff` and the embedded `Version` constant, aligning tok with the rest of the hawk-eco ecosystem (`hawk`, `eyrie`, `yaad`, `sight`, `inspect`). No git tag has been cut yet. ### Removed +- Dead `ParallelExecutor` (`internal/filter`): ran every layer on the full + input concurrently and kept the shortest output. No production callers in + tok or hawk, no tests, and it measured output size in bytes rather than + tokens, so it could never be wired into the token-budget pipeline as-is. - Unused internal config structs `QuestionAwareLayerConfig`, `DensityAdaptiveLayerConfig`, `NumericalQuantLayerConfig`, and `DynamicRatioLayerConfig` (and their fields on the internal `LayerConfig`) diff --git a/internal/core/estimator.go b/internal/core/estimator.go index d46389f3..e29d07f0 100644 --- a/internal/core/estimator.go +++ b/internal/core/estimator.go @@ -3,7 +3,7 @@ package core import ( "container/list" "fmt" - "hash/fnv" + "hash/maphash" "strings" "sync" "sync/atomic" @@ -21,12 +21,28 @@ type BPETokenizer struct { } // cacheEntry stores cached token counts with LRU metadata. +// It deliberately does NOT retain the original text: Compress estimates the +// original prompt plus every intermediate layer output, so pinning full texts +// in the LRU could retain hundreds of MB. Hit validation relies on the +// 64-bit maphash plus the text length (see hashText). type cacheEntry struct { - count int - text string // original text for collision detection - elem *list.Element // pointer to list element for O(1) removal + count int + length int // text length, secondary collision check + elem *list.Element // pointer to list element for O(1) removal } +// cacheSeed is generated once per process. The token cache is purely +// in-process, so a random seed is sufficient for correctness and avoids +// fixed-seed hash-flooding concerns. +var cacheSeed = maphash.MakeSeed() + +// maxCacheableTextBytes is the size cutoff above which texts bypass the +// cache entirely. Giant inputs are rare in steady state, dominate the cache's +// memory footprint if admitted, and their Count cost is dominated by the BPE +// encode itself, so caching them buys little. 1 MiB bounds the worst case +// while still covering every realistic prompt, layer output, or file read. +const maxCacheableTextBytes = 1 << 20 // 1 MiB + // lruItem holds the key for list element tracking. type lruItem struct { key uint64 @@ -69,21 +85,31 @@ func (c *tokenCache) getShard(key uint64) *tokenCacheShard { return &c.shards[key%64] } +// hashText returns a 64-bit hash of text without copying it. +// maphash.String reads the string's backing bytes directly (no []byte(text) +// conversion) and is seeded per process. Collision handling: entries are +// additionally validated by text length, so a false hit requires two distinct +// texts with identical length AND identical 64-bit hash. With <=8192 cached +// entries the birthday-bound probability of that is ~1e-12 per process — +// negligible next to the heuristic fallback path — so no stronger check +// (e.g. content prefix) is retained. func hashText(text string) uint64 { - h := fnv.New64a() - _, _ = h.Write([]byte(text)) // #nosec G104 -- hash.Hash.Write never returns an error - return h.Sum64() + return maphash.String(cacheSeed, text) } // get retrieves a cached token count and promotes the entry in the LRU list. // Uses per-shard locking to reduce contention across concurrent goroutines. -// Compares original text to handle FNV-64a hash collisions. +// Validates (hash, length) — the cache never stores the text itself. +// Texts larger than maxCacheableTextBytes always miss. func (c *tokenCache) get(text string) (int, bool) { + if len(text) > maxCacheableTextBytes { + return 0, false + } key := hashText(text) shard := c.getShard(key) shard.mu.Lock() entry, ok := shard.items[key] - if !ok || entry.text != text { + if !ok || entry.length != len(text) { shard.mu.Unlock() return 0, false } @@ -96,12 +122,15 @@ func (c *tokenCache) get(text string) (int, bool) { } func (c *tokenCache) set(text string, count int) { + if len(text) > maxCacheableTextBytes { + return // giant inputs never enter the LRU + } key := hashText(text) shard := c.getShard(key) shard.mu.Lock() defer shard.mu.Unlock() - if entry, ok := shard.items[key]; ok && entry.text == text { + if entry, ok := shard.items[key]; ok && entry.length == len(text) { entry.count = count shard.ll.MoveToFront(entry.elem) return @@ -113,9 +142,9 @@ func (c *tokenCache) set(text string, count int) { elem := shard.ll.PushFront(&lruItem{key: key}) shard.items[key] = &cacheEntry{ - count: count, - text: text, - elem: elem, + count: count, + length: len(text), + elem: elem, } } diff --git a/internal/core/estimator_test.go b/internal/core/estimator_test.go index e797bfc7..73fdef6a 100644 --- a/internal/core/estimator_test.go +++ b/internal/core/estimator_test.go @@ -1,6 +1,11 @@ package core -import "testing" +import ( + "reflect" + "strings" + "sync/atomic" + "testing" +) func TestEstimateTokensExact(t *testing.T) { t.Parallel() @@ -115,3 +120,77 @@ func TestEstimateTokensPreciseShortString(t *testing.T) { t.Logf("short=%q fast=%d precise=%d (difference acceptable; precise is authoritative)", short, fast, precise) } + +// TestTokenCacheHitAcrossCalls verifies the LRU cache returns a cached count +// for identical text across separate calls (and bumps the hit counter). +func TestTokenCacheHitAcrossCalls(t *testing.T) { + t.Parallel() + c := newTokenCache(64) + c.set("hello token cache world", 5) + + got, ok := c.get("hello token cache world") + if !ok { + t.Fatal("expected cache hit for identical text") + } + if got != 5 { + t.Errorf("cached count = %d, want 5", got) + } + if hits := atomic.LoadInt64(&c.hits); hits != 1 { + t.Errorf("hits = %d, want 1", hits) + } +} + +// TestTokenCacheSameLengthDifferentTextMisses verifies hit validation: +// two distinct texts of identical length must not collide. +func TestTokenCacheSameLengthDifferentTextMisses(t *testing.T) { + t.Parallel() + c := newTokenCache(64) + a := "aaaa first text" + b := "bbbb second one" // same length, different content + if len(a) != len(b) { + t.Fatalf("test bug: lengths differ (%d vs %d)", len(a), len(b)) + } + c.set(a, 7) + + if _, ok := c.get(b); ok { + t.Error("different text with same length must be a cache miss") + } + if _, ok := c.get(a); !ok { + t.Error("original text must still hit after miss probe") + } +} + +// TestTokenCacheOversizeTextBypassed verifies texts above the size cutoff +// never enter the LRU, so giant inputs cannot pin cache memory. +func TestTokenCacheOversizeTextBypassed(t *testing.T) { + t.Parallel() + c := newTokenCache(64) + huge := strings.Repeat("x", maxCacheableTextBytes+1) + + c.set(huge, 12345) + if _, ok := c.get(huge); ok { + t.Errorf("text of %d bytes (> cutoff %d) must bypass the cache", + len(huge), maxCacheableTextBytes) + } + + // The boundary itself must still be cacheable. + boundary := strings.Repeat("y", maxCacheableTextBytes) + c.set(boundary, 42) + if got, ok := c.get(boundary); !ok || got != 42 { + t.Errorf("text exactly at cutoff must hit, got (count=%d, ok=%v)", got, ok) + } +} + +// TestCacheEntryDoesNotRetainText pins the memory fix: cacheEntry must not +// carry any string field. Storing the original text meant the 8192-entry LRU +// retained full prompts and every intermediate layer output. +func TestCacheEntryDoesNotRetainText(t *testing.T) { + t.Parallel() + typ := reflect.TypeOf(cacheEntry{}) + for i := 0; i < typ.NumField(); i++ { + if typ.Field(i).Type.Kind() == reflect.String { + t.Errorf("cacheEntry field %q is a string — the cache must not retain text", + typ.Field(i).Name) + } + } +} diff --git a/internal/filter/parallel_executor.go b/internal/filter/parallel_executor.go deleted file mode 100644 index 76aecb41..00000000 --- a/internal/filter/parallel_executor.go +++ /dev/null @@ -1,63 +0,0 @@ -package filter - -import "sync" - -// ParallelExecutor runs independent layers concurrently -type ParallelExecutor struct { - pool *sync.Pool -} - -// NewParallelExecutor creates a parallel execution engine -func NewParallelExecutor() *ParallelExecutor { - return &ParallelExecutor{ - pool: &sync.Pool{ - New: func() interface{} { - return &layerResult{} - }, - }, - } -} - -type layerResult struct { - output string - tokens int -} - -// ExecuteParallel runs layers concurrently and merges results -func (pe *ParallelExecutor) ExecuteParallel(input string, layers []Filter) (string, int) { - if len(layers) == 0 { - return input, 0 - } - - results := make([]*layerResult, len(layers)) - var wg sync.WaitGroup - - for i, layer := range layers { - wg.Add(1) - go func(idx int, l Filter) { - defer wg.Done() - r, _ := pe.pool.Get().(*layerResult) - r.output, r.tokens = l.Apply(input, ModeMinimal) - results[idx] = r - }(i, layer) - } - - wg.Wait() - - // Use best result (highest compression) - best := results[0] - for _, r := range results[1:] { - if len(r.output) < len(best.output) { - best = r - } - } - - // Copy before returning to pool — best points into results, which we Put below. - bestOutput, bestTokens := best.output, best.tokens - - for _, r := range results { - pe.pool.Put(r) - } - - return bestOutput, bestTokens -} diff --git a/restoration.go b/restoration.go index e9227b63..359edf49 100644 --- a/restoration.go +++ b/restoration.go @@ -1,6 +1,7 @@ package tok import ( + "container/list" "sort" "sync" "time" @@ -12,6 +13,14 @@ const ( PriorityWriteEdit = 100 PrioritySkill = 80 PriorityRead = 50 + + // DefaultRetainedTokenBudget caps the total token cost of tracked + // content held in memory between Clear() calls. 200_000 matches the + // default context window: the tracker can never usefully re-inject more + // than one context window of material, so retaining more is pure memory + // waste (~800 KB of text). When the budget is exceeded the least + // recently tracked entries are evicted first. + DefaultRetainedTokenBudget = 200_000 ) // RestorationTracker tracks file operations for selective context re-injection @@ -19,12 +28,20 @@ const ( // writes, skill activations, file reads) can be lost. The tracker records these // operations and, after compaction, selects the highest-priority items to // re-inject within a token budget. +// +// Memory is bounded: deduplication is O(1) via a path-indexed map, and the +// total retained content is capped by the retained token budget (see +// DefaultRetainedTokenBudget and WithRetainedTokenBudget), evicting the least +// recently tracked entries when exceeded. type RestorationTracker struct { - mu sync.Mutex - entries []RestorationEntry - maxBudget int - budgetPct float64 - contextWindow int + mu sync.Mutex + entries *list.List // *RestorationEntry, front = most recently tracked + index map[string]*list.Element + totalTokens int + maxBudget int + budgetPct float64 + contextWindow int + retainedBudget int } // RestorationEntry represents a tracked file operation. @@ -40,9 +57,12 @@ type RestorationEntry struct { // NewRestorationTracker creates a tracker with default settings. func NewRestorationTracker() *RestorationTracker { return &RestorationTracker{ - maxBudget: DefaultMaxBudget, - budgetPct: DefaultBudgetPct, - contextWindow: 200000, // default 200K tokens + entries: list.New(), + index: make(map[string]*list.Element), + maxBudget: DefaultMaxBudget, + budgetPct: DefaultBudgetPct, + contextWindow: 200000, // default 200K tokens + retainedBudget: DefaultRetainedTokenBudget, } } @@ -70,7 +90,20 @@ func (rt *RestorationTracker) WithContextWindow(tokens int) *RestorationTracker return rt } -// Track records a file operation. Deduplicates by path (latest wins). +// WithRetainedTokenBudget caps the total tokens of content retained in memory. +// Once exceeded, the least recently tracked entries are evicted. +func (rt *RestorationTracker) WithRetainedTokenBudget(tokens int) *RestorationTracker { + rt.mu.Lock() + defer rt.mu.Unlock() + rt.retainedBudget = tokens + rt.evictOverBudgetLocked() + return rt +} + +// Track records a file operation. Deduplicates by path in O(1) (latest wins, +// and the entry is refreshed as most recently tracked). When the retained +// token budget is exceeded, the least recently tracked entries are evicted; +// the entry just tracked is always kept, even if it alone exceeds the budget. func (rt *RestorationTracker) Track(path, content, opType string) { rt.mu.Lock() defer rt.mu.Unlock() @@ -92,14 +125,39 @@ func (rt *RestorationTracker) Track(path, content, opType string) { Timestamp: time.Now(), } - // Deduplicate: replace existing entry for same path - for i, e := range rt.entries { - if e.Path == path { - rt.entries[i] = entry - return + // Deduplicate: replace existing entry for same path, refresh recency. + if elem, ok := rt.index[path]; ok { + if old, ok := elem.Value.(*RestorationEntry); ok { + rt.totalTokens -= old.Tokens } + elem.Value = &entry + rt.entries.MoveToFront(elem) + } else { + rt.index[path] = rt.entries.PushFront(&entry) + } + rt.totalTokens += entry.Tokens + + rt.evictOverBudgetLocked() +} + +// evictOverBudgetLocked drops least recently tracked entries until the total +// retained tokens fit the budget. Caller must hold rt.mu. The most recently +// tracked entry is always kept so a single over-budget Track is not a no-op. +func (rt *RestorationTracker) evictOverBudgetLocked() { + for rt.totalTokens > rt.retainedBudget && rt.entries.Len() > 1 { + back := rt.entries.Back() + if back == nil { + break + } + old, ok := back.Value.(*RestorationEntry) + if !ok { + rt.entries.Remove(back) + continue + } + rt.totalTokens -= old.Tokens + delete(rt.index, old.Path) + rt.entries.Remove(back) } - rt.entries = append(rt.entries, entry) } // ComputeBudget calculates the available restoration budget. @@ -131,13 +189,16 @@ func (rt *RestorationTracker) GetRestorations(budget int) []RestorationEntry { rt.mu.Lock() defer rt.mu.Unlock() - if len(rt.entries) == 0 || budget <= 0 { + if rt.entries.Len() == 0 || budget <= 0 { return nil } - // Sort by priority desc, then timestamp desc - sorted := make([]RestorationEntry, len(rt.entries)) - copy(sorted, rt.entries) + sorted := make([]RestorationEntry, 0, rt.entries.Len()) + for e := rt.entries.Front(); e != nil; e = e.Next() { + if entry, ok := e.Value.(*RestorationEntry); ok { + sorted = append(sorted, *entry) + } + } sort.Slice(sorted, func(i, j int) bool { if sorted[i].Priority != sorted[j].Priority { return sorted[i].Priority > sorted[j].Priority @@ -157,12 +218,17 @@ func (rt *RestorationTracker) GetRestorations(budget int) []RestorationEntry { return result } -// Entries returns a snapshot of all tracked entries (for testing/debugging). +// Entries returns a snapshot of all tracked entries, most recently tracked +// first (for testing/debugging). func (rt *RestorationTracker) Entries() []RestorationEntry { rt.mu.Lock() defer rt.mu.Unlock() - result := make([]RestorationEntry, len(rt.entries)) - copy(result, rt.entries) + result := make([]RestorationEntry, 0, rt.entries.Len()) + for e := rt.entries.Front(); e != nil; e = e.Next() { + if entry, ok := e.Value.(*RestorationEntry); ok { + result = append(result, *entry) + } + } return result } @@ -170,12 +236,14 @@ func (rt *RestorationTracker) Entries() []RestorationEntry { func (rt *RestorationTracker) Clear() { rt.mu.Lock() defer rt.mu.Unlock() - rt.entries = nil + rt.entries.Init() + rt.index = make(map[string]*list.Element) + rt.totalTokens = 0 } // Len returns the number of tracked entries. func (rt *RestorationTracker) Len() int { rt.mu.Lock() defer rt.mu.Unlock() - return len(rt.entries) + return rt.entries.Len() } diff --git a/restoration_test.go b/restoration_test.go index 811160e3..81cb8d6b 100644 --- a/restoration_test.go +++ b/restoration_test.go @@ -1,6 +1,7 @@ package tok import ( + "fmt" "strings" "testing" "time" @@ -153,3 +154,109 @@ func TestRestorationTracker_EmptyGetRestorations(t *testing.T) { t.Errorf("expected nil for empty tracker, got %v", got) } } + +func TestRestorationTracker_TrackSamePathRepeatedly(t *testing.T) { + rt := NewRestorationTracker() + for i := 0; i < 500; i++ { + rt.Track("/main.go", fmt.Sprintf("revision %d", i), "write") + } + + if rt.Len() != 1 { + t.Errorf("expected 1 entry after 500 tracks of the same path, got %d", rt.Len()) + } + entries := rt.Entries() + if entries[0].Content != "revision 499" { + t.Errorf("expected latest content 'revision 499', got %q", entries[0].Content) + } +} + +func TestRestorationTracker_RetainedBudgetEvictsOldest(t *testing.T) { + // Identical content shape so all three entries cost the same T tokens. + content := strings.Repeat("word ", 100) + perEntry := EstimateTokens(content) + + rt := NewRestorationTracker().WithRetainedTokenBudget(2 * perEntry) + rt.Track("/a.go", content, "read") + rt.Track("/b.go", content, "read") + rt.Track("/c.go", content, "read") // over budget: oldest (/a.go) evicted + + if rt.Len() != 2 { + t.Fatalf("expected 2 entries after eviction, got %d", rt.Len()) + } + paths := map[string]bool{} + for _, e := range rt.Entries() { + paths[e.Path] = true + } + if paths["/a.go"] { + t.Error("oldest entry /a.go should have been evicted") + } + if !paths["/b.go"] || !paths["/c.go"] { + t.Errorf("recent entries should be retained, got %v", paths) + } + + // Clear still resets everything. + rt.Clear() + if rt.Len() != 0 || len(rt.Entries()) != 0 { + t.Errorf("expected empty tracker after Clear, got len=%d", rt.Len()) + } +} + +func TestRestorationTracker_RetainedBudgetEvictsLeastRecentlyTracked(t *testing.T) { + content := strings.Repeat("token ", 100) + perEntry := EstimateTokens(content) + + rt := NewRestorationTracker().WithRetainedTokenBudget(2 * perEntry) + rt.Track("/a.go", content, "read") + rt.Track("/b.go", content, "read") + rt.Track("/a.go", content, "write") // refresh /a.go recency; /b.go is now oldest + rt.Track("/c.go", content, "read") // over budget: evicts /b.go, not /a.go + + if rt.Len() != 2 { + t.Fatalf("expected 2 entries after eviction, got %d", rt.Len()) + } + paths := map[string]bool{} + for _, e := range rt.Entries() { + paths[e.Path] = true + } + if paths["/b.go"] { + t.Error("/b.go was least recently tracked and should have been evicted") + } + if !paths["/a.go"] || !paths["/c.go"] { + t.Errorf("expected /a.go and /c.go retained, got %v", paths) + } +} + +func TestRestorationTracker_RetainedBudgetKeepsSingleOversizeEntry(t *testing.T) { + huge := strings.Repeat("big ", 10000) // far above any tiny budget + rt := NewRestorationTracker().WithRetainedTokenBudget(10) + rt.Track("/huge.go", huge, "write") + + if rt.Len() != 1 { + t.Fatalf("most recent entry must be kept even when over budget, got %d", rt.Len()) + } + entries := rt.Entries() + if entries[0].Path != "/huge.go" || entries[0].Content != huge { + t.Error("oversize entry should be retained verbatim") + } +} + +func TestRestorationTracker_WithRetainedTokenBudgetShrinksExisting(t *testing.T) { + content := strings.Repeat("data ", 100) + perEntry := EstimateTokens(content) + + rt := NewRestorationTracker() + rt.Track("/a.go", content, "read") + rt.Track("/b.go", content, "read") + rt.Track("/c.go", content, "read") + + // Lowering the budget evicts immediately down to the new cap. + rt.WithRetainedTokenBudget(2 * perEntry) + if rt.Len() != 2 { + t.Errorf("expected 2 entries after budget reduction, got %d", rt.Len()) + } + for _, e := range rt.Entries() { + if e.Path == "/a.go" { + t.Error("oldest entry /a.go should have been evicted by budget reduction") + } + } +}