diff --git a/go.mod b/go.mod index 8e65c80..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 ) 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_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"}, }, diff --git a/internal/runtime/native_token.go b/internal/runtime/native_token.go index 994518b..fb8f6d7 100644 --- a/internal/runtime/native_token.go +++ b/internal/runtime/native_token.go @@ -3,22 +3,47 @@ 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. +// 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) + 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..71cb4ef 100644 --- a/internal/runtime/native_token_test.go +++ b/internal/runtime/native_token_test.go @@ -6,14 +6,16 @@ import ( ) func TestEstimateTokens(t *testing.T) { + // Pinned values from cl100k_base tokenizer. 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", 1}, + {"hello world", 2}, + {strings.Repeat("a", 100), 13}, + {`func main() { fmt.Println("hello") }`, 10}, } for _, tt := range tests { got := estimateTokens(tt.input) @@ -25,13 +27,13 @@ func TestEstimateTokens(t *testing.T) { 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)}, // 50 tokens + 4 overhead + {Role: "user", Content: "hello"}, // 1 token + 4 overhead } got := estimateMessagesTokens(msgs) - // 100 + 4 + 2 + 4 = 110 (approximately) - if got < 100 || got > 120 { - t.Errorf("estimateMessagesTokens = %d, expected ~110", got) + // 50 + 4 + 1 + 4 = 59 + if got < 55 || got > 65 { + t.Errorf("estimateMessagesTokens = %d, expected ~59", got) } }