Merge pull request #2014 from badgerbees/fix/context-pruning-guards

fix(agent): include SystemParts in token estimation and add reasoning guards

pkg/agent/context_budget.go

@@ -90,13 +90,28 @@ func findSafeBoundary(history []providers.Message, targetIndex int) int {
 // including Content, ReasoningContent, ToolCalls arguments, ToolCallID
 // metadata, and Media items. Uses a heuristic of 2.5 characters per token.
 func estimateMessageTokens(msg providers.Message) int {
-	chars := utf8.RuneCountInString(msg.Content)
+	contentChars := utf8.RuneCountInString(msg.Content)
 
-	// ReasoningContent (extended thinking / chain-of-thought) can be
-	// substantial and is stored in session history via AddFullMessage.
-	if msg.ReasoningContent != "" {
-		chars += utf8.RuneCountInString(msg.ReasoningContent)
+	// SystemParts are structured system blocks used for cache-aware adapters.
+	// They carry the same content as Content, but in multiple blocks.
+	// We estimate them as an alternative representation, not additive.
+	systemPartsChars := 0
+	if len(msg.SystemParts) > 0 {
+		for _, part := range msg.SystemParts {
+			systemPartsChars += utf8.RuneCountInString(part.Text)
 		}
+		// Per-part overhead for JSON structure (type, text, cache_control).
+		const perPartOverhead = 20
+		systemPartsChars += len(msg.SystemParts) * perPartOverhead
+	}
+
+	// Use the larger of the two representations to stay conservative.
+	chars := contentChars
+	if systemPartsChars > chars {
+		chars = systemPartsChars
+	}
+
+	chars += utf8.RuneCountInString(msg.ReasoningContent)
 
 	for _, tc := range msg.ToolCalls {
 		chars += len(tc.ID) + len(tc.Type)

pkg/agent/context_budget_test.go

@@ -529,6 +529,26 @@ func TestEstimateMessageTokens_MediaItems(t *testing.T) {
 	}
 }
 
+func TestEstimateMessageTokens_SystemParts(t *testing.T) {
+	plain := providers.Message{Role: "system", Content: "instructions"}
+	withParts := providers.Message{
+		Role:    "system",
+		Content: "instructions",
+		SystemParts: []providers.ContentBlock{
+			{Type: "text", Text: "some more system context"},
+			{Type: "text", Text: "even more cached blocks"},
+		},
+	}
+
+	plainTokens := estimateMessageTokens(plain)
+	partsTokens := estimateMessageTokens(withParts)
+
+	if partsTokens <= plainTokens {
+		t.Errorf("system message with SystemParts (%d) should exceed plain message (%d)",
+			partsTokens, plainTokens)
+	}
+}
+
 // --- estimateToolDefsTokens tests ---
 
 func TestEstimateToolDefsTokens(t *testing.T) {