From 93f391a6bf7a845c3d59dda2da032ee460118e6c Mon Sep 17 00:00:00 2001
From: Badgerbees <rapabelias@gmail.com>
Date: Thu, 26 Mar 2026 00:31:22 +0700
Subject: [PATCH] fix(agent): include SystemParts in token estimation and add
 reasoning guards

---
 pkg/agent/context_budget.go      |  9 +++++----
 pkg/agent/context_budget_test.go | 20 ++++++++++++++++++++
 2 files changed, 25 insertions(+), 4 deletions(-)

diff --git a/pkg/agent/context_budget.go b/pkg/agent/context_budget.go
index c87695c7a..416ec4a5b 100644
--- a/pkg/agent/context_budget.go
+++ b/pkg/agent/context_budget.go
@@ -91,11 +91,12 @@ func findSafeBoundary(history []providers.Message, targetIndex int) int {
 // metadata, and Media items. Uses a heuristic of 2.5 characters per token.
 func estimateMessageTokens(msg providers.Message) int {
 	chars := utf8.RuneCountInString(msg.Content)
+	chars += utf8.RuneCountInString(msg.ReasoningContent)
 
-	// 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 that can be substantial
+	// when using instruction-heavy agents or KV-cache-aware adapters.
+	for _, part := range msg.SystemParts {
+		chars += utf8.RuneCountInString(part.Text)
 	}
 
 	for _, tc := range msg.ToolCalls {
diff --git a/pkg/agent/context_budget_test.go b/pkg/agent/context_budget_test.go
index 870f0fbe6..22cbdc0db 100644
--- a/pkg/agent/context_budget_test.go
+++ b/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) {