feat: use tiktoken cl100k_base for token estimation

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
 )
 

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=

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,25 +515,29 @@ 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,
 	}
 
 	state := &State{
 		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"},
 		},

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
 }
 

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)
 	}
 }