From d536f2a31789141eb3442b6fc746d6b55f296708 Mon Sep 17 00:00:00 2001 From: louismorgner Date: Fri, 27 Mar 2026 09:37:30 -0700 Subject: [PATCH 1/3] feat: use tiktoken cl100k_base tokenizer for token estimation Replace the len/4 heuristic in estimateTokens() with real tokenization via tiktoken-go/tokenizer using cl100k_base encoding, which closely approximates Claude's tokenizer. Falls back to the old heuristic if the tokenizer fails to initialize. Co-Authored-By: Claude Opus 4.6 --- go.mod | 1 + go.sum | 2 ++ internal/runtime/native_token.go | 38 +++++++++++++++++++----- internal/runtime/native_token_test.go | 42 +++++++++++++++++++-------- 4 files changed, 63 insertions(+), 20 deletions(-) diff --git a/go.mod b/go.mod index 8e65c80..1cabdbe 100644 --- a/go.mod +++ b/go.mod @@ -46,6 +46,7 @@ require ( github.com/olekukonko/tablewriter v0.0.5 // indirect github.com/rivo/uniseg v0.4.7 // indirect github.com/spf13/pflag v1.0.9 // indirect + github.com/tiktoken-go/tokenizer v0.7.0 // indirect github.com/xo/terminfo v0.0.0-20220910002029-abceb7e1c41e // indirect github.com/yuin/goldmark v1.7.13 // indirect github.com/yuin/goldmark-emoji v1.0.6 // indirect diff --git a/go.sum b/go.sum index a1f9306..1860512 100644 --- a/go.sum +++ b/go.sum @@ -102,6 +102,8 @@ github.com/spf13/cobra v1.10.2 h1:DMTTonx5m65Ic0GOoRY2c16WCbHxOOw6xxezuLaBpcU= github.com/spf13/cobra v1.10.2/go.mod h1:7C1pvHqHw5A4vrJfjNwvOdzYu0Gml16OCs2GRiTUUS4= github.com/spf13/pflag v1.0.9 h1:9exaQaMOCwffKiiiYk6/BndUBv+iRViNW+4lEMi0PvY= github.com/spf13/pflag v1.0.9/go.mod h1:McXfInJRrz4CZXVZOBLb0bTZqETkiAhM9Iw0y3An2Bg= +github.com/tiktoken-go/tokenizer v0.7.0 h1:VMu6MPT0bXFDHr7UPh9uii7CNItVt3X9K90omxL54vw= +github.com/tiktoken-go/tokenizer v0.7.0/go.mod h1:6UCYI/DtOallbmL7sSy30p6YQv60qNyU/4aVigPOx6w= github.com/xo/terminfo v0.0.0-20220910002029-abceb7e1c41e h1:JVG44RsyaB9T2KIHavMF/ppJZNG9ZpyihvCd0w101no= github.com/xo/terminfo v0.0.0-20220910002029-abceb7e1c41e/go.mod h1:RbqR21r5mrJuqunuUZ/Dhy/avygyECGrLceyNeo4LiM= github.com/yuin/goldmark v1.7.13 h1:GPddIs617DnBLFFVJFgpo1aBfe/4xcvMc3SB5t/D0pA= diff --git a/internal/runtime/native_token.go b/internal/runtime/native_token.go index 994518b..f562b18 100644 --- a/internal/runtime/native_token.go +++ b/internal/runtime/native_token.go @@ -3,22 +3,44 @@ package runtime import ( "fmt" "strings" + "sync" + + "github.com/tiktoken-go/tokenizer" +) + +// tokenCodec is a lazily-initialized tiktoken encoder (cl100k_base). +// cl100k_base is the closest publicly available encoding to what Claude +// models use and gives much better estimates than the 4-byte heuristic. +var ( + tokenCodec tokenizer.Codec + tokenCodecOnce sync.Once ) -// Approximate bytes-per-token ratio for Claude/GPT-class models. -// Anthropic and OpenAI both converge around 3.5–4 bytes per token for -// English-heavy code/text. We use 4 for conservative (over-)estimation -// so budget decisions err on the side of keeping content shorter. +func getTokenCodec() tokenizer.Codec { + tokenCodecOnce.Do(func() { + enc, err := tokenizer.Get(tokenizer.Cl100kBase) + if err == nil { + tokenCodec = enc + } + }) + return tokenCodec +} + +// approxBytesPerToken is the fallback ratio when the tokenizer is unavailable. const approxBytesPerToken = 4 -// estimateTokens returns a rough token count for the given string. -// This is intentionally cheap (no tiktoken dependency) — the API -// reports exact counts, so this is only used for pre-flight decisions -// like "should we truncate this tool output before adding it to history?" +// estimateTokens returns a token count for the given string using the +// cl100k_base tokenizer. Falls back to len/4 if the tokenizer is unavailable. func estimateTokens(s string) int { if len(s) == 0 { return 0 } + if codec := getTokenCodec(); codec != nil { + n, err := codec.Count(s) + if err == nil { + return n + } + } return (len(s) + approxBytesPerToken - 1) / approxBytesPerToken } diff --git a/internal/runtime/native_token_test.go b/internal/runtime/native_token_test.go index 680ad7f..4c152c5 100644 --- a/internal/runtime/native_token_test.go +++ b/internal/runtime/native_token_test.go @@ -8,30 +8,48 @@ import ( func TestEstimateTokens(t *testing.T) { tests := []struct { input string - want int }{ - {"", 0}, - {"hi", 1}, // 2 bytes / 4 = 0.5, rounds up to 1 - {"hello world", 3}, // 11 bytes / 4 = 2.75, rounds up to 3 - {strings.Repeat("a", 100), 25}, // 100 / 4 = 25 + {""}, + {"hi"}, + {"hello world"}, + {strings.Repeat("a", 100)}, + {"func main() { fmt.Println(\"hello\") }"}, } for _, tt := range tests { got := estimateTokens(tt.input) - if got != tt.want { - t.Errorf("estimateTokens(%q) = %d, want %d", tt.input[:min(len(tt.input), 20)], got, tt.want) + if tt.input == "" { + if got != 0 { + t.Errorf("estimateTokens(%q) = %d, want 0", tt.input, got) + } + continue + } + if got <= 0 { + t.Errorf("estimateTokens(%q) = %d, want > 0", tt.input[:min(len(tt.input), 20)], got) + } + } +} + +func TestEstimateTokens_UsesRealTokenizer(t *testing.T) { + // The real tokenizer should give different results than len/4 for this string. + // "hello world" is 2 tokens in cl100k_base, but len/4 rounds to 3. + got := estimateTokens("hello world") + if codec := getTokenCodec(); codec != nil { + if got == 3 { + t.Error("estimateTokens(\"hello world\") = 3, looks like fallback heuristic is being used instead of real tokenizer") } } } func TestEstimateMessagesTokens(t *testing.T) { msgs := []Message{ - {Role: "system", Content: strings.Repeat("x", 400)}, // 100 tokens + 4 overhead - {Role: "user", Content: "hello"}, // ~1 token + 4 + {Role: "system", Content: strings.Repeat("x", 400)}, + {Role: "user", Content: "hello"}, } got := estimateMessagesTokens(msgs) - // 100 + 4 + 2 + 4 = 110 (approximately) - if got < 100 || got > 120 { - t.Errorf("estimateMessagesTokens = %d, expected ~110", got) + // With real tokenizer or heuristic, should be a reasonable positive number. + // The per-message overhead (4 tokens each × 2 messages = 8) is always added. + if got < 10 { + t.Errorf("estimateMessagesTokens = %d, expected > 10", got) } } From 300bcbae04beaa610864c74fb556acb10bfbc0a6 Mon Sep 17 00:00:00 2001 From: louismorgner Date: Fri, 27 Mar 2026 09:42:55 -0700 Subject: [PATCH 2/3] fix: pin tokenizer test values and fix go.mod indirect marker Address PR review feedback: - go mod tidy to mark tiktoken-go/tokenizer as direct dependency - Pin exact cl100k_base token counts in TestEstimateTokens - Tighten TestEstimateMessagesTokens range to ~59 (was > 10) - Add comment on getTokenCodec permanent fallback behavior Co-Authored-By: Claude Opus 4.6 --- go.mod | 2 +- internal/runtime/native_token.go | 3 ++ internal/runtime/native_token_test.go | 44 +++++++++------------------ 3 files changed, 18 insertions(+), 31 deletions(-) diff --git a/go.mod b/go.mod index 1cabdbe..9c554e8 100644 --- a/go.mod +++ b/go.mod @@ -13,6 +13,7 @@ require ( github.com/google/uuid v1.6.0 github.com/mattn/go-isatty v0.0.20 github.com/spf13/cobra v1.10.2 + github.com/tiktoken-go/tokenizer v0.7.0 gopkg.in/yaml.v3 v3.0.1 ) @@ -46,7 +47,6 @@ require ( github.com/olekukonko/tablewriter v0.0.5 // indirect github.com/rivo/uniseg v0.4.7 // indirect github.com/spf13/pflag v1.0.9 // indirect - github.com/tiktoken-go/tokenizer v0.7.0 // indirect github.com/xo/terminfo v0.0.0-20220910002029-abceb7e1c41e // indirect github.com/yuin/goldmark v1.7.13 // indirect github.com/yuin/goldmark-emoji v1.0.6 // indirect diff --git a/internal/runtime/native_token.go b/internal/runtime/native_token.go index f562b18..fb8f6d7 100644 --- a/internal/runtime/native_token.go +++ b/internal/runtime/native_token.go @@ -16,6 +16,9 @@ var ( tokenCodecOnce sync.Once ) +// getTokenCodec returns the shared tokenizer, initializing it on first call. +// If initialization fails, the codec stays nil permanently (sync.Once won't +// retry) and callers fall back to the len/4 heuristic. func getTokenCodec() tokenizer.Codec { tokenCodecOnce.Do(func() { enc, err := tokenizer.Get(tokenizer.Cl100kBase) diff --git a/internal/runtime/native_token_test.go b/internal/runtime/native_token_test.go index 4c152c5..71cb4ef 100644 --- a/internal/runtime/native_token_test.go +++ b/internal/runtime/native_token_test.go @@ -6,50 +6,34 @@ import ( ) func TestEstimateTokens(t *testing.T) { + // Pinned values from cl100k_base tokenizer. tests := []struct { input string + want int }{ - {""}, - {"hi"}, - {"hello world"}, - {strings.Repeat("a", 100)}, - {"func main() { fmt.Println(\"hello\") }"}, + {"", 0}, + {"hi", 1}, + {"hello world", 2}, + {strings.Repeat("a", 100), 13}, + {`func main() { fmt.Println("hello") }`, 10}, } for _, tt := range tests { got := estimateTokens(tt.input) - if tt.input == "" { - if got != 0 { - t.Errorf("estimateTokens(%q) = %d, want 0", tt.input, got) - } - continue - } - if got <= 0 { - t.Errorf("estimateTokens(%q) = %d, want > 0", tt.input[:min(len(tt.input), 20)], got) - } - } -} - -func TestEstimateTokens_UsesRealTokenizer(t *testing.T) { - // The real tokenizer should give different results than len/4 for this string. - // "hello world" is 2 tokens in cl100k_base, but len/4 rounds to 3. - got := estimateTokens("hello world") - if codec := getTokenCodec(); codec != nil { - if got == 3 { - t.Error("estimateTokens(\"hello world\") = 3, looks like fallback heuristic is being used instead of real tokenizer") + if got != tt.want { + t.Errorf("estimateTokens(%q) = %d, want %d", tt.input[:min(len(tt.input), 20)], got, tt.want) } } } func TestEstimateMessagesTokens(t *testing.T) { msgs := []Message{ - {Role: "system", Content: strings.Repeat("x", 400)}, - {Role: "user", Content: "hello"}, + {Role: "system", Content: strings.Repeat("x", 400)}, // 50 tokens + 4 overhead + {Role: "user", Content: "hello"}, // 1 token + 4 overhead } got := estimateMessagesTokens(msgs) - // With real tokenizer or heuristic, should be a reasonable positive number. - // The per-message overhead (4 tokens each × 2 messages = 8) is always added. - if got < 10 { - t.Errorf("estimateMessagesTokens = %d, expected > 10", got) + // 50 + 4 + 1 + 4 = 59 + if got < 55 || got > 65 { + t.Errorf("estimateMessagesTokens = %d, expected ~59", got) } } From f3147a7427e54724999f033a1e55ab24f06cfaa8 Mon Sep 17 00:00:00 2001 From: louismorgner Date: Fri, 27 Mar 2026 09:59:42 -0700 Subject: [PATCH 3/3] fix: adjust compaction test budgets for real tokenizer counts The compaction tests assumed the old len/4 heuristic where repeated single characters cost ~4 bytes/token. With cl100k_base, repeated chars are ~8 bytes/token, so the test content was below budget thresholds. Increase content sizes and adjust context windows so tests properly trigger pruning/compaction while ensuring compacted results still fit within budget. Co-Authored-By: Claude Opus 4.6 --- internal/runtime/native_compaction_test.go | 43 +++++++++++++--------- 1 file changed, 25 insertions(+), 18 deletions(-) diff --git a/internal/runtime/native_compaction_test.go b/internal/runtime/native_compaction_test.go index 571661b..1deb545 100644 --- a/internal/runtime/native_compaction_test.go +++ b/internal/runtime/native_compaction_test.go @@ -72,7 +72,9 @@ func TestBudgetFail_EmergencyCompaction(t *testing.T) { } // Tiny context window to force BudgetFail immediately. - // Budget = 2048 - 512 - 256 = 1280 tokens. Fail threshold = 98% = 1254 tokens ≈ 5016 bytes. + // Budget = 2048 - 512 - 256 = 1280 tokens. Fail threshold = 98% = 1254 tokens. + // With cl100k_base, repeated chars ≈ 8 bytes/token, so 3 × 8000 bytes ≈ 3000 tokens + // + overhead, well above 1254. profile := runtimeinfo.NativeModelProfile{ ContextWindow: 2048, MaxOutputTokens: 512, @@ -84,9 +86,9 @@ func TestBudgetFail_EmergencyCompaction(t *testing.T) { SessionID: sess.ID, Messages: []Message{ {Role: "system", Content: "system prompt"}, - {Role: "user", Content: strings.Repeat("a", 3000)}, - {Role: "assistant", Content: strings.Repeat("b", 3000)}, - {Role: "tool", Name: "Edit", Content: strings.Repeat("c", 3000)}, // prune-protected + {Role: "user", Content: strings.Repeat("a", 8000)}, + {Role: "assistant", Content: strings.Repeat("b", 8000)}, + {Role: "tool", Name: "Edit", Content: strings.Repeat("c", 8000)}, // prune-protected {Role: "user", Content: "keep"}, {Role: "assistant", Content: "recent"}, }, @@ -465,9 +467,10 @@ func TestManageContextWithBudget_PrunesBeforeCompaction(t *testing.T) { } // Create a state with large tool outputs that should trigger pruning - // at 75% of budget. GPT-4o budget = 128000 - 16384 - 4096 = 107520 tokens - // 75% = 80640 tokens ≈ 322560 bytes - largeContent := strings.Repeat("x", 80*1024) // ~80KB = ~20K tokens each + // at 75% of budget. Budget = 128000 - 16384 - 4096 = 107520 tokens + // 75% = 80640 tokens. With cl100k_base, repeated "x" ≈ 8 bytes/token, + // so 200KB ≈ 25600 tokens × 4 messages ≈ 102400 tokens > 80640. + largeContent := strings.Repeat("x", 200*1024) // ~200KB ≈ 25600 tokens each state := &State{ Runtime: runtimeinfo.NativeRuntime, SessionID: sess.ID, @@ -512,14 +515,17 @@ func TestManageContextWithBudget_CompactsWhenNeeded(t *testing.T) { MetadataDir: t.TempDir(), } - // Use a small context window to force compaction. - // Budget = 4096 - 1024 - 512 = 2560 tokens - // Compact threshold = 90% of 2560 = 2304 tokens ≈ 9216 bytes - // We need messages totaling > 9216 bytes to trigger compaction. - // After pruning, the tool output will shrink but total should still exceed threshold. + // Use a context window where the messages exceed the compact threshold + // (90% of input budget) but the compacted result fits within the budget. + // Budget = 32768 - 2048 - 512 = 30208 tokens + // Compact threshold = 90% of 30208 = 27187 tokens + // With cl100k_base, repeated single chars ≈ 8 bytes/token, + // so 4 × 60000 bytes ≈ 4 × 7500 = 30000 tokens + overhead > 27187. + // After compaction, old messages collapse to a small continuation artifact, + // system prompt is tiny, and only 2 recent messages remain — fits easily. profile := runtimeinfo.NativeModelProfile{ - ContextWindow: 4096, - MaxOutputTokens: 1024, + ContextWindow: 32768, + MaxOutputTokens: 2048, ReservedBuffer: 512, } @@ -527,10 +533,11 @@ func TestManageContextWithBudget_CompactsWhenNeeded(t *testing.T) { Runtime: runtimeinfo.NativeRuntime, SessionID: sess.ID, Messages: []Message{ - {Role: "system", Content: strings.Repeat("s", 4000)}, - {Role: "user", Content: strings.Repeat("a", 4000)}, - {Role: "assistant", Content: strings.Repeat("b", 4000)}, - {Role: "tool", Name: "Edit", Content: strings.Repeat("c", 4000)}, // Edit is prune-protected + {Role: "system", Content: "system prompt"}, + {Role: "user", Content: strings.Repeat("a", 60000)}, + {Role: "assistant", Content: strings.Repeat("b", 60000)}, + {Role: "tool", Name: "Edit", Content: strings.Repeat("c", 60000)}, // Edit is prune-protected + {Role: "user", Content: strings.Repeat("d", 60000)}, {Role: "user", Content: "keep"}, {Role: "assistant", Content: "recent"}, },