From f21dc0efde4c59513576a7a0bd54a5e7655b8bdf Mon Sep 17 00:00:00 2001 From: Luxferre Date: Fri, 11 Sep 2026 10:55:19 +0300 Subject: [PATCH] added correct reasoning token stats --- main.go | 46 ++++++++++++++++++++++++++------ main_test.go | 74 ++++++++++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 112 insertions(+), 8 deletions(-) diff --git a/main.go b/main.go index 7374672..caab5d7 100644 --- a/main.go +++ b/main.go @@ -857,7 +857,13 @@ type Usage struct { CachedTokens int `json:"cached_tokens"` } `json:"prompt_tokens_details"` CachedTokens int `json:"cached_tokens"` - Model string `json:"-"` + // ContextTokens is a local estimate of the total tokens occupying the context + // window for this sample, including any reasoning_content that is resent as + // part of the context. It acts as a floor beneath PromptTokens so reasoning + // tokens are always counted in the context window statistics even when the + // provider under-reports prompt_tokens. + ContextTokens int `json:"-"` + Model string `json:"-"` } func (u Usage) Cached() int { @@ -882,26 +888,39 @@ func estTokens(msgs []Message) int { return t } +// ctxTokens returns the number of tokens occupying the context window for a +// usage sample. We prefer the provider's reported prompt token count, but never +// go below a local estimate of the messages we actually sent (which includes any +// reasoning_content resent as context) so reasoning tokens are always counted in +// the context window statistics. +func ctxTokens(u Usage) int { + if u.ContextTokens > u.PromptTokens { + return u.ContextTokens + } + return u.PromptTokens +} + func contextPct(u Usage, cw int) float64 { if cw <= 0 { cw = 262144 } - return float64(u.PromptTokens) * 100.0 / float64(cw) + return float64(ctxTokens(u)) * 100.0 / float64(cw) } func formatUsage(u Usage, cw int) string { - pct := contextPct(u, cw) + base := ctxTokens(u) + pct := float64(base) * 100.0 / float64(cw) cached := u.Cached() if cached > 0 { uncached := u.PromptTokens - cached if uncached < 0 { uncached = 0 } if u.Model != "" { - return fmt.Sprintf("[%s: %d prompt (%d cached, %d uncached) + %d completion | context: %d/%d (%.1f%%)]", u.Model, u.PromptTokens, cached, uncached, u.CompletionTokens, u.PromptTokens, cw, pct) + return fmt.Sprintf("[%s: %d prompt (%d cached, %d uncached) + %d completion | context: %d/%d (%.1f%%)]", u.Model, u.PromptTokens, cached, uncached, u.CompletionTokens, base, cw, pct) } - return fmt.Sprintf("[%d prompt (%d cached, %d uncached) + %d completion | context: %d/%d (%.1f%%)]", u.PromptTokens, cached, uncached, u.CompletionTokens, u.PromptTokens, cw, pct) + return fmt.Sprintf("[%d prompt (%d cached, %d uncached) + %d completion | context: %d/%d (%.1f%%)]", u.PromptTokens, cached, uncached, u.CompletionTokens, base, cw, pct) } if u.Model != "" { - return fmt.Sprintf("[%s: %d prompt + %d completion | context: %d/%d (%.1f%%)]", u.Model, u.PromptTokens, u.CompletionTokens, u.PromptTokens, cw, pct) + return fmt.Sprintf("[%s: %d prompt + %d completion | context: %d/%d (%.1f%%)]", u.Model, u.PromptTokens, u.CompletionTokens, base, cw, pct) } - return fmt.Sprintf("[%d prompt + %d completion | context: %d/%d (%.1f%%)]", u.PromptTokens, u.CompletionTokens, u.PromptTokens, cw, pct) + return fmt.Sprintf("[%d prompt + %d completion | context: %d/%d (%.1f%%)]", u.PromptTokens, u.CompletionTokens, base, cw, pct) } type streamDelta struct { @@ -1078,6 +1097,9 @@ func llm(ctx context.Context, cfg *Cfg, msgs []Message, tools []map[string]any) if m.ReasoningContent == "" { m.ReasoningContent = cr.Choices[0].Message.Thought } u := cr.Usage if cr.Model != "" { u.Model = cr.Model } + // Local estimate of the full context (including any resent reasoning) so it + // is counted in the context window statistics as a floor under prompt_tokens. + u.ContextTokens = estTokens(msgs) if u.PromptTokens == 0 { u.PromptTokens = estTokens(msgs) u.CompletionTokens = estTokens([]Message{m}) @@ -1091,6 +1113,11 @@ func llm(ctx context.Context, cfg *Cfg, msgs []Message, tools []map[string]any) u.CompletionTokens = estTokens([]Message{m}) u.TotalTokens = u.PromptTokens + u.CompletionTokens } + if err == nil { + // Local estimate of the full context (including any resent reasoning) so it + // is counted in the context window statistics as a floor under prompt_tokens. + u.ContextTokens = estTokens(msgs) + } return m, u, err } @@ -1327,6 +1354,9 @@ func AL(ctx context.Context, cfg *Cfg, msgs []Message) ([]Message, Usage, error) turnUsage.PromptTokens = u.PromptTokens turnUsage.CompletionTokens += u.CompletionTokens turnUsage.TotalTokens += u.TotalTokens + // The last call sends the full conversation, so its context estimate is the + // authoritative one for the turn's context window statistics. + turnUsage.ContextTokens = u.ContextTokens if u.Cached() > 0 { turnUsage.CachedTokens = u.Cached() } if u.Model != "" { turnUsage.Model = u.Model } for j := range m.ToolCalls { @@ -2217,7 +2247,7 @@ func main() { fmt.Println(c(formatUsage(usg, cfg.ContextWindow), 2)) pct := contextPct(usg, cfg.ContextWindow) if pct >= 60.0 && len(msgs) > 1 { - fmt.Print(c(fmt.Sprintf("Context usage is at %.1f%% (%d / %d tokens). Compact conversation? [Y/n]: ", pct, usg.PromptTokens, cfg.ContextWindow), 33)) + fmt.Print(c(fmt.Sprintf("Context usage is at %.1f%% (%d / %d tokens). Compact conversation? [Y/n]: ", pct, ctxTokens(usg), cfg.ContextWindow), 33)) if ans, ok := readPlain(""); ok { ans = strings.TrimSpace(strings.ToLower(ans)) if ans == "" || ans == "y" || ans == "yes" { diff --git a/main_test.go b/main_test.go index 79a77cf..60008b7 100644 --- a/main_test.go +++ b/main_test.go @@ -2848,3 +2848,77 @@ func TestToolsDir(t *testing.T) { } } + +// TestContextStatsCountReasoning verifies that reasoning tokens which are resent +// as part of the context are counted in the context window statistics. The +// provider may under-report prompt_tokens (omitting the reasoning_content we +// echo back), so the statistics must never drop below a local estimate of the +// messages we actually send, which includes the reasoning content. +func TestContextStatsCountReasoning(t *testing.T) { + msgs := []Message{ + {Role: "system", Content: strp("sys")}, + {Role: "user", Content: strp("question")}, + {Role: "assistant", Content: strp("answer"), ReasoningContent: "long chain of thought that consumes many tokens"}, + } + // Local estimate of the context, including the reasoning content above. + est := estTokens(msgs) + if est <= estTokens([]Message{msgs[0], msgs[1], {Role: "assistant", Content: strp("answer")}}) { + t.Fatalf("estTokens should grow when reasoning content is present (got %d)", est) + } + + // Simulate a provider that reports prompt_tokens WITHOUT the reasoning tokens. + underReported := Usage{PromptTokens: est - 40, CompletionTokens: 10, ContextTokens: est} + if got := ctxTokens(underReported); got != est { + t.Errorf("ctxTokens = %d, want %d (reasoning must be counted)", got, est) + } + + // Simulate an accurate provider that already counts reasoning in prompt_tokens. + accurate := Usage{PromptTokens: est, CompletionTokens: 10, ContextTokens: est} + if got := ctxTokens(accurate); got != est { + t.Errorf("ctxTokens = %d, want %d", got, est) + } + + // The displayed context numerator must reflect the reasoning-inclusive count. + s := formatUsage(underReported, 200000) + want := fmt.Sprintf("[%d prompt + 10 completion | context: %d/200000 (%.1f%%)]", underReported.PromptTokens, est, float64(est)*100.0/200000) + if s != want { + t.Errorf("formatUsage under-reported =\n %q\nwant\n %q", s, want) + } +} + +// TestLLMContextTokensIncludesReasoning verifies that llm() populates +// ContextTokens with a local estimate that includes any reasoning_content we +// resend, even when the provider omits usage entirely (the fallback path). +func TestLLMContextTokensIncludesReasoning(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + // No usage block: forces the local-estimate fallback path. + w.Write([]byte(`{"choices":[{"message":{"role":"assistant","content":"hi","reasoning_content":"think hard about this for a while"}}]}`)) + })) + defer srv.Close() + + cfg := defCfg + cfg.Endpoint = srv.URL + cfg.Stream = false + msgs := []Message{ + {Role: "user", Content: strp("hi")}, + {Role: "assistant", Content: strp("prev"), ReasoningContent: "prior reasoning that is resent as context"}, + } + m, u, err := llm(context.Background(), &cfg, msgs, TOOLS) + if err != nil { + t.Fatalf("llm: %v", err) + } + if m.ReasoningContent != "think hard about this for a while" { + t.Errorf("reasoning not parsed: %q", m.ReasoningContent) + } + if u.ContextTokens <= 0 { + t.Fatalf("ContextTokens not populated, got %d", u.ContextTokens) + } + // ContextTokens must include the resent reasoning content. + withoutReasoning := estTokens([]Message{{Role: "user", Content: strp("hi")}, {Role: "assistant", Content: strp("prev")}}) + if u.ContextTokens <= withoutReasoning { + t.Errorf("ContextTokens=%d should exceed estimate without reasoning=%d", u.ContextTokens, withoutReasoning) + } + if ctxTokens(u) != u.ContextTokens { + t.Errorf("ctxTokens=%d should equal ContextTokens=%d when prompt_tokens is absent", ctxTokens(u), u.ContextTokens) + } +}