improved compacting algo and token visibility

main_test.go

@@ -59,7 +59,7 @@ func TestParseStreamReasoningNoDuplication(t *testing.T) {
     `data: [DONE]`,
   }, "\n")
 
-  msg, err := parseStream(context.Background(), bytes.NewBufferString(sseData))
+  msg, _, err := parseStream(context.Background(), bytes.NewBufferString(sseData))
   if err != nil {
     t.Fatalf("unexpected parseStream error: %v", err)
   }
@@ -91,7 +91,7 @@ func TestParseStreamReasoningAlias(t *testing.T) {
     `data: [DONE]`,
   }, "\n")
 
-  msg, err := parseStream(context.Background(), bytes.NewBufferString(sseData))
+  msg, _, err := parseStream(context.Background(), bytes.NewBufferString(sseData))
   if err != nil {
     t.Fatalf("unexpected parseStream error: %v", err)
   }
@@ -109,7 +109,7 @@ func TestParseStreamContentOnly(t *testing.T) {
     `data: [DONE]`,
   }, "\n")
 
-  msg, err := parseStream(context.Background(), bytes.NewBufferString(sseData))
+  msg, _, err := parseStream(context.Background(), bytes.NewBufferString(sseData))
   if err != nil {
     t.Fatalf("unexpected parseStream error: %v", err)
   }
@@ -130,7 +130,7 @@ func TestParseStreamReasoningOnly(t *testing.T) {
     `data: [DONE]`,
   }, "\n")
 
-  msg, err := parseStream(context.Background(), bytes.NewBufferString(sseData))
+  msg, _, err := parseStream(context.Background(), bytes.NewBufferString(sseData))
   if err != nil {
     t.Fatalf("unexpected parseStream error: %v", err)
   }
@@ -144,7 +144,7 @@ func TestParseStreamReasoningOnly(t *testing.T) {
 
 func TestParseStreamEmpty(t *testing.T) {
   for _, in := range []string{"", "\n\n", "event: message\n\n"} {
-    msg, err := parseStream(context.Background(), strings.NewReader(in))
+    msg, _, err := parseStream(context.Background(), strings.NewReader(in))
     if err != nil {
       t.Fatalf("unexpected parseStream error for input %q: %v", in, err)
     }
@@ -162,7 +162,7 @@ func TestParseStreamToolCallSplitAcrossChunks(t *testing.T) {
     `data: [DONE]`,
   }, "\n")
 
-  msg, err := parseStream(context.Background(), bytes.NewBufferString(sseData))
+  msg, _, err := parseStream(context.Background(), bytes.NewBufferString(sseData))
   if err != nil {
     t.Fatalf("unexpected parseStream error: %v", err)
   }
@@ -188,7 +188,7 @@ func TestParseStreamMultipleToolCallsKeepFirstAppearanceOrder(t *testing.T) {
     `data: [DONE]`,
   }, "\n")
 
-  msg, err := parseStream(context.Background(), bytes.NewBufferString(sseData))
+  msg, _, err := parseStream(context.Background(), bytes.NewBufferString(sseData))
   if err != nil {
     t.Fatalf("unexpected parseStream error: %v", err)
   }
@@ -210,7 +210,7 @@ func TestParseStreamJunkAndNoChoicesIgnored(t *testing.T) {
     `data: [DONE]`,
   }, "\n")
 
-  msg, err := parseStream(context.Background(), bytes.NewBufferString(sseData))
+  msg, _, err := parseStream(context.Background(), bytes.NewBufferString(sseData))
   if err != nil {
     t.Fatalf("unexpected parseStream error: %v", err)
   }
@@ -226,7 +226,7 @@ func TestParseStreamReasoningAfterContent(t *testing.T) {
     `data: [DONE]`,
   }, "\n")
 
-  msg, err := parseStream(context.Background(), bytes.NewBufferString(sseData))
+  msg, _, err := parseStream(context.Background(), bytes.NewBufferString(sseData))
   if err != nil {
     t.Fatalf("unexpected parseStream error: %v", err)
   }
@@ -707,11 +707,11 @@ func TestAutosave(t *testing.T) {
 // ---------- summarize / compact (error paths only, no network) ----------
 
 func TestSummarizeEmpty(t *testing.T) {
-  if _, err := summarize(context.Background(), &Cfg{}, nil); err == nil || !strings.Contains(err.Error(), "no conversation") {
-    t.Errorf("expected no-conversation error, got %v", err)
+  if _, err := summarize(context.Background(), &Cfg{}, nil); err == nil || !strings.Contains(err.Error(), "no system message") {
+    t.Errorf("expected no-system-message error, got %v", err)
   }
-  if _, err := summarize(context.Background(), &Cfg{}, []Message{{Role: "system", Content: strp("sys")}}); err == nil {
-    t.Errorf("expected error for system-only conversation")
+  if _, err := summarize(context.Background(), &Cfg{}, []Message{{Role: "system", Content: strp("sys")}}); err == nil || !strings.Contains(err.Error(), "nothing to compact") {
+    t.Errorf("expected nothing-to-compact error for system-only conversation, got %v", err)
   }
 }
 
@@ -724,12 +724,19 @@ func TestCompactNoSystem(t *testing.T) {
     t.Errorf("expected original messages on error")
   }
   msgs2, _, err2 := compact(context.Background(), &Cfg{}, []Message{{Role: "user", Content: strp("x")}})
-  if err2 == nil {
-    t.Errorf("expected error when first message is not system")
+  if err2 == nil || !strings.Contains(err2.Error(), "no system message") {
+    t.Errorf("expected error when first message is not system, got %v", err2)
   }
   if len(msgs2) != 1 {
     t.Errorf("expected original messages returned, got %d", len(msgs2))
   }
+  msgs3, _, err3 := compact(context.Background(), &Cfg{}, []Message{{Role: "system", Content: strp("sys")}})
+  if err3 == nil || !strings.Contains(err3.Error(), "nothing to compact") {
+    t.Errorf("expected nothing to compact error, got %v", err3)
+  }
+  if len(msgs3) != 1 {
+    t.Errorf("expected original messages returned, got %d", len(msgs3))
+  }
 }
 
 // ---------- history ----------
@@ -943,7 +950,7 @@ func TestLLMNonStreamingAndHeaders(t *testing.T) {
   cfg.Endpoint = srv.URL
   cfg.Stream = false
   cfg.APIKey = "secret"
-  m, err := llm(context.Background(), &cfg, []Message{{Role: "user", Content: strp("hi")}}, TOOLS)
+  m, _, err := llm(context.Background(), &cfg, []Message{{Role: "user", Content: strp("hi")}}, TOOLS)
   if err != nil {
     t.Fatalf("llm: %v", err)
   }
@@ -973,7 +980,7 @@ func TestLLMNoAuthHeaderWhenNoKey(t *testing.T) {
   cfg.Endpoint = srv.URL
   cfg.Stream = false
   cfg.APIKey = "-"
-  if _, err := llm(context.Background(), &cfg, []Message{{Role: "user", Content: strp("hi")}}, nil); err != nil {
+  if _, _, err := llm(context.Background(), &cfg, []Message{{Role: "user", Content: strp("hi")}}, nil); err != nil {
     t.Fatalf("llm: %v", err)
   }
   if gotAuth != "" {
@@ -994,7 +1001,7 @@ func TestLLMStreaming(t *testing.T) {
   cfg.Endpoint = srv.URL
   cfg.Stream = true
   cfg.APIKey = "-"
-  m, err := llm(context.Background(), &cfg, []Message{{Role: "user", Content: strp("hi")}}, nil)
+  m, _, err := llm(context.Background(), &cfg, []Message{{Role: "user", Content: strp("hi")}}, nil)
   if err != nil {
     t.Fatalf("llm: %v", err)
   }
@@ -1018,7 +1025,7 @@ func TestLLM4xxReturnsImmediately(t *testing.T) {
   cfg.Stream = false
   cfg.APIKey = "-"
   start := time.Now()
-  _, err := llm(context.Background(), &cfg, []Message{{Role: "user", Content: strp("hi")}}, nil)
+  _, _, err := llm(context.Background(), &cfg, []Message{{Role: "user", Content: strp("hi")}}, nil)
   if err == nil || !strings.Contains(err.Error(), "Invalid assistant message") {
     t.Fatalf("expected 400 error, got %v", err)
   }
@@ -1044,7 +1051,7 @@ func TestLLMRetriesOn5xxThenSucceeds(t *testing.T) {
   cfg.Endpoint = srv.URL
   cfg.Stream = false
   cfg.APIKey = "-"
-  m, err := llm(context.Background(), &cfg, []Message{{Role: "user", Content: strp("hi")}}, nil)
+  m, _, err := llm(context.Background(), &cfg, []Message{{Role: "user", Content: strp("hi")}}, nil)
   if err != nil {
     t.Fatalf("llm after retries: %v", err)
   }
@@ -1066,7 +1073,7 @@ func TestLLMEmptyChoices(t *testing.T) {
   cfg.Endpoint = srv.URL
   cfg.Stream = false
   cfg.APIKey = "-"
-  if _, err := llm(context.Background(), &cfg, []Message{{Role: "user", Content: strp("hi")}}, nil); err == nil || !strings.Contains(err.Error(), "empty choices") {
+  if _, _, err := llm(context.Background(), &cfg, []Message{{Role: "user", Content: strp("hi")}}, nil); err == nil || !strings.Contains(err.Error(), "empty choices") {
     t.Errorf("expected empty choices error, got %v", err)
   }
 }
@@ -1093,7 +1100,7 @@ func TestALToolLoop(t *testing.T) {
   cfg.Endpoint = srv.URL
   cfg.Stream = false
   cfg.APIKey = "-"
-  msgs, err := AL(context.Background(), &cfg, []Message{{Role: "system", Content: strp("sys")}, {Role: "user", Content: strp("run")}}, "sys", 0)
+  msgs, _, err := AL(context.Background(), &cfg, []Message{{Role: "system", Content: strp("sys")}, {Role: "user", Content: strp("run")}}, "sys", 0)
   if err != nil {
     t.Fatalf("AL: %v", err)
   }
@@ -1155,7 +1162,7 @@ func TestALRunSubagent(t *testing.T) {
   cfg.Endpoint = srv.URL
   cfg.Stream = false
   cfg.APIKey = "-"
-  msgs, err := AL(context.Background(), &cfg, []Message{{Role: "system", Content: strp("sys")}, {Role: "user", Content: strp("parent task")}}, "sys", 0)
+  msgs, _, err := AL(context.Background(), &cfg, []Message{{Role: "system", Content: strp("sys")}, {Role: "user", Content: strp("parent task")}}, "sys", 0)
   if err != nil {
     t.Fatalf("AL: %v", err)
   }
@@ -1198,7 +1205,7 @@ func TestALSubagentDepthLimit(t *testing.T) {
   cfg.Endpoint = srv.URL
   cfg.Stream = false
   cfg.APIKey = "-"
-  msgs, err := AL(context.Background(), &cfg, []Message{{Role: "user", Content: strp("go")}}, "sys", MAX_DEPTH)
+  msgs, _, err := AL(context.Background(), &cfg, []Message{{Role: "user", Content: strp("go")}}, "sys", MAX_DEPTH)
   if err != nil {
     t.Fatalf("AL: %v", err)
   }
@@ -1245,7 +1252,7 @@ func TestALStripsInvalidAssistantAndRetries(t *testing.T) {
   cfg.Endpoint = srv.URL
   cfg.Stream = false
   cfg.APIKey = "-"
-  msgs, err := AL(context.Background(), &cfg, []Message{{Role: "system", Content: strp("sys")}, {Role: "user", Content: strp("go")}}, "sys", 0)
+  msgs, _, err := AL(context.Background(), &cfg, []Message{{Role: "system", Content: strp("sys")}, {Role: "user", Content: strp("go")}}, "sys", 0)
   if err != nil {
     t.Fatalf("AL: %v", err)
   }
@@ -1269,7 +1276,8 @@ func TestSummarizeHappyPath(t *testing.T) {
   cfg.Endpoint = srv.URL
   cfg.Stream = false
   cfg.APIKey = "-"
-  s, err := summarize(context.Background(), &cfg, []Message{{Role: "user", Content: strp("hello world")}})
+  orig := []Message{{Role: "system", Content: strp("sys")}, {Role: "user", Content: strp("hello world")}}
+  s, err := summarize(context.Background(), &cfg, orig)
   if err != nil {
     t.Fatalf("summarize: %v", err)
   }
@@ -1279,7 +1287,13 @@ func TestSummarizeHappyPath(t *testing.T) {
 }
 
 func TestCompactHappyPath(t *testing.T) {
+  var receivedMessages []Message
   srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+    var req struct {
+      Messages []Message `json:"messages"`
+    }
+    json.NewDecoder(r.Body).Decode(&req)
+    receivedMessages = req.Messages
     w.Write([]byte(`{"choices":[{"message":{"role":"assistant","content":"the summary"}}]}`))
   }))
   defer srv.Close()
@@ -1302,6 +1316,19 @@ func TestCompactHappyPath(t *testing.T) {
   if msgs[1].Role != "user" || msgs[1].Content == nil || !strings.Contains(*msgs[1].Content, "the summary") {
     t.Errorf("continuation message = %+v", msgs[1])
   }
+  // Verify request sent to LLM contains the original conversation prefix plus compaction prompt
+  if len(receivedMessages) != 3 {
+    t.Fatalf("expected 3 messages sent to LLM, got %d", len(receivedMessages))
+  }
+  if receivedMessages[0].Role != "system" || *receivedMessages[0].Content != "sys" {
+    t.Errorf("message 0 mismatch: %+v", receivedMessages[0])
+  }
+  if receivedMessages[1].Role != "user" || *receivedMessages[1].Content != "hello world" {
+    t.Errorf("message 1 mismatch: %+v", receivedMessages[1])
+  }
+  if receivedMessages[2].Role != "user" || !strings.Contains(*receivedMessages[2].Content, "compaction engine") {
+    t.Errorf("message 2 mismatch (expected compaction prompt): %+v", receivedMessages[2])
+  }
 }
 
 // ---------- setCfg and LLM parameter forwarding ----------
@@ -1352,7 +1379,7 @@ func TestLLMForwardsRelevantParameters(t *testing.T) {
   }, "\n"))
 
   cfg := getCfg(cfgFile)
-  _, err := llm(context.Background(), &cfg, []Message{{Role: "user", Content: strp("hi")}}, nil)
+  _, _, err := llm(context.Background(), &cfg, []Message{{Role: "user", Content: strp("hi")}}, nil)
   if err != nil {
     t.Fatalf("llm: %v", err)
   }
@@ -1644,7 +1671,7 @@ func TestALContextCancellation(t *testing.T) {
 
   origMsgs := []Message{{Role: "system", Content: strp("sys")}, {Role: "user", Content: strp("hello")}}
   inputMsgs := append([]Message{}, origMsgs...)
-  msgs, err := AL(ctx, &cfg, inputMsgs, "sys", 0)
+  msgs, _, err := AL(ctx, &cfg, inputMsgs, "sys", 0)
   if err == nil {
     t.Fatalf("expected context cancellation error, got nil")
   }
@@ -1682,7 +1709,7 @@ func TestLLMContextCancellation(t *testing.T) {
     cancel()
   }()
 
-  _, err := llm(ctx, &cfg, []Message{{Role: "user", Content: strp("hi")}}, nil)
+  _, _, err := llm(ctx, &cfg, []Message{{Role: "user", Content: strp("hi")}}, nil)
   if err == nil {
     t.Fatalf("expected context cancellation error, got nil")
   }
@@ -1691,6 +1718,139 @@ func TestLLMContextCancellation(t *testing.T) {
   }
 }
 
+func TestTokenCounterAndUsageNonStreaming(t *testing.T) {
+  srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+    w.Write([]byte(`{
+      "choices":[{"message":{"role":"assistant","content":"hello world"}}],
+      "usage":{
+        "prompt_tokens": 120,
+        "completion_tokens": 30,
+        "total_tokens": 150,
+        "prompt_tokens_details": {"cached_tokens": 80}
+      }
+    }`))
+  }))
+  defer srv.Close()
+
+  cfg := defCfg
+  cfg.Endpoint = srv.URL
+  cfg.Stream = false
+  cfg.APIKey = "-"
+
+  m, u, err := llm(context.Background(), &cfg, []Message{{Role: "user", Content: strp("hi")}}, nil)
+  if err != nil {
+    t.Fatalf("llm: %v", err)
+  }
+  if m.Content == nil || *m.Content != "hello world" {
+    t.Errorf("unexpected content: %+v", m.Content)
+  }
+  if u.PromptTokens != 120 || u.CompletionTokens != 30 || u.TotalTokens != 150 {
+    t.Errorf("unexpected usage: %+v", u)
+  }
+  if u.Cached() != 80 {
+    t.Errorf("expected cached tokens 80, got %d", u.Cached())
+  }
+}
+
+func TestTokenCounterAndUsageStreaming(t *testing.T) {
+  sseData := strings.Join([]string{
+    `data: {"choices":[{"delta":{"content":"streaming "}}]}`,
+    `data: {"choices":[{"delta":{"content":"response"}}]}`,
+    `data: {"choices":[],"usage":{"prompt_tokens":250,"completion_tokens":45,"total_tokens":295,"cached_tokens":100}}`,
+    `data: [DONE]`,
+  }, "\n")
+
+  srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+    w.Header().Set("Content-Type", "text/event-stream")
+    w.Write([]byte(sseData))
+  }))
+  defer srv.Close()
+
+  cfg := defCfg
+  cfg.Endpoint = srv.URL
+  cfg.Stream = true
+  cfg.APIKey = "-"
+
+  m, u, err := llm(context.Background(), &cfg, []Message{{Role: "user", Content: strp("hi")}}, nil)
+  if err != nil {
+    t.Fatalf("llm streaming: %v", err)
+  }
+  if m.Content == nil || *m.Content != "streaming response" {
+    t.Errorf("unexpected content: %+v", m.Content)
+  }
+  if u.PromptTokens != 250 || u.CompletionTokens != 45 || u.TotalTokens != 295 {
+    t.Errorf("unexpected usage: %+v", u)
+  }
+  if u.Cached() != 100 {
+    t.Errorf("expected cached tokens 100, got %d", u.Cached())
+  }
+}
+
+func TestFormatUsage(t *testing.T) {
+  // With cached tokens
+  u1 := Usage{PromptTokens: 1000, CompletionTokens: 200, TotalTokens: 1200, CachedTokens: 800}
+  s1 := formatUsage(u1, 200000)
+  if s1 != "[tokens: 1000 prompt (800 cached, 200 uncached) + 200 completion | context: 1000/200000 (0.5%)]" {
+    t.Errorf("formatUsage u1 = %q", s1)
+  }
+
+  // Without cached tokens
+  u2 := Usage{PromptTokens: 120000, CompletionTokens: 500, TotalTokens: 120500}
+  s2 := formatUsage(u2, 200000)
+  if s2 != "[tokens: 120000 prompt + 500 completion | context: 120000/200000 (60.0%)]" {
+    t.Errorf("formatUsage u2 = %q", s2)
+  }
+}
+
+func TestFetchContextWindowFromModelsAPI(t *testing.T) {
+  srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+    if r.URL.Path != "/models" {
+      http.NotFound(w, r)
+      return
+    }
+    w.Write([]byte(`{
+      "data": [
+        {"id": "other-model", "context_window": 32000},
+        {"id": "my-target-model", "max_context_length": 131072}
+      ]
+    }`))
+  }))
+  defer srv.Close()
+
+  cfg := defCfg
+  cfg.Endpoint = srv.URL
+  cfg.Model = "my-target-model"
+  cfg.APIKey = "-"
+
+  cw := fetchContextWindow(&cfg)
+  if cw != 131072 {
+    t.Errorf("expected context window 131072 from /models API, got %d", cw)
+  }
+}
+
+func TestFetchContextWindowFallbackConfig(t *testing.T) {
+  srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+    http.Error(w, "server error", 500)
+  }))
+  defer srv.Close()
+
+  // Case 1: Config specifies context_window
+  cfg1 := defCfg
+  cfg1.Endpoint = srv.URL
+  cfg1.Raw = map[string]string{"context_window": "65536"}
+  if cw := fetchContextWindow(&cfg1); cw != 65536 {
+    t.Errorf("expected fallback to raw context_window 65536, got %d", cw)
+  }
+
+  // Case 2: Config does not specify context_window -> default 200000
+  cfg2 := defCfg
+  cfg2.Endpoint = srv.URL
+  cfg2.Raw = map[string]string{}
+  if cw := fetchContextWindow(&cfg2); cw != 200000 {
+    t.Errorf("expected default 200000, got %d", cw)
+  }
+}
+