added correct reasoning token stats

main.go

@@ -857,6 +857,12 @@ type Usage struct {
     CachedTokens int `json:"cached_tokens"`
   } `json:"prompt_tokens_details"`
   CachedTokens int `json:"cached_tokens"`
+  // 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:"-"`
 }
 
@@ -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" {

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