diff --git a/pkg/agent/context_budget.go b/pkg/agent/context_budget.go index c87695c7a..3398d7863 100644 --- a/pkg/agent/context_budget.go +++ b/pkg/agent/context_budget.go @@ -90,14 +90,29 @@ 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) if tc.Function != nil { 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) {