From 2e9c65057b8ca5abc8aa54b21856b869c733d52c Mon Sep 17 00:00:00 2001
From: Greg Slepak <contact@taoeffect.com>
Date: Tue, 12 May 2026 15:56:50 -0700
Subject: [PATCH] fix(agent): correct fallback usage accounting

---
 internal/agent/agent.go               |  6 +--
 internal/agent/usage_fallback.go      |  8 +++-
 internal/agent/usage_fallback_test.go | 63 +++++++++++++++++++++++++++
 3 files changed, 71 insertions(+), 6 deletions(-)

diff --git a/internal/agent/agent.go b/internal/agent/agent.go
index c54331910b13a0e5a7ef747451be32895596ab65..fcedcc55615db8afd29f109231cae627b67ac98f 100644
--- a/internal/agent/agent.go
+++ b/internal/agent/agent.go
@@ -1166,11 +1166,9 @@ func (a *sessionAgent) updateSessionUsage(model Model, session *session.Session,
 	}
 
 	session.Cost += cost
-	if usage.OutputTokens != 0 {
+	if !usageIsZero(usage) {
 		session.CompletionTokens = usage.OutputTokens
-	}
-	if promptTokens := usage.InputTokens + usage.CacheReadTokens; promptTokens != 0 {
-		session.PromptTokens = promptTokens
+		session.PromptTokens = usage.InputTokens + usage.CacheReadTokens
 	}
 }
 
diff --git a/internal/agent/usage_fallback.go b/internal/agent/usage_fallback.go
index 78903c57a6c8c54a3598964c1c132579d0cad36d..3cf347c38f91a72968c3537a5d954bb6feb1d0b9 100644
--- a/internal/agent/usage_fallback.go
+++ b/internal/agent/usage_fallback.go
@@ -78,9 +78,13 @@ func estimateStepCompletionTokens(step fantasy.StepResult) int64 {
 		case *fantasy.ToolCallContent:
 			tokens += estimateToolCallTokens(c.ToolName, c.Input)
 		case fantasy.ToolResultContent:
-			tokens += estimateToolResultContentTokens(c.ToolCallID, c.ToolName, c.ClientMetadata, c.Result)
+			if c.ProviderExecuted {
+				tokens += estimateToolResultContentTokens(c.ToolCallID, c.ToolName, c.ClientMetadata, c.Result)
+			}
 		case *fantasy.ToolResultContent:
-			tokens += estimateToolResultContentTokens(c.ToolCallID, c.ToolName, c.ClientMetadata, c.Result)
+			if c.ProviderExecuted {
+				tokens += estimateToolResultContentTokens(c.ToolCallID, c.ToolName, c.ClientMetadata, c.Result)
+			}
 		}
 	}
 	return tokens
diff --git a/internal/agent/usage_fallback_test.go b/internal/agent/usage_fallback_test.go
index 96263f214a1be035dc5978827c286b81dde2e0ba..ec925c4148da922edec666a5d8bfec350c9295bf 100644
--- a/internal/agent/usage_fallback_test.go
+++ b/internal/agent/usage_fallback_test.go
@@ -145,6 +145,51 @@ func TestFallbackStepUsageEstimatesToolResults(t *testing.T) {
 	require.Equal(t, usage.InputTokens, usage.TotalTokens)
 }
 
+func TestFallbackStepUsageSkipsClientToolResultsAsOutput(t *testing.T) {
+	t.Parallel()
+
+	step := fantasy.StepResult{
+		Response: fantasy.Response{
+			Content: fantasy.ResponseContent{
+				fantasy.ToolResultContent{
+					ToolCallID: "tool-call-1",
+					ToolName:   "bash",
+					Result: fantasy.ToolResultOutputContentText{
+						Text: "large client-executed payload that should not count as model output tokens",
+					},
+				},
+			},
+		},
+	}
+
+	usage, estimated := fallbackStepUsage(nil, step)
+	require.False(t, estimated)
+	require.Zero(t, usage.OutputTokens)
+}
+
+func TestFallbackStepUsageCountsProviderToolResultsAsOutput(t *testing.T) {
+	t.Parallel()
+
+	step := fantasy.StepResult{
+		Response: fantasy.Response{
+			Content: fantasy.ResponseContent{
+				fantasy.ToolResultContent{
+					ToolCallID:       "tool-call-1",
+					ToolName:         "web_search",
+					ProviderExecuted: true,
+					ClientMetadata:   "provider metadata",
+					Result:           fantasy.ToolResultOutputContentText{Text: "provider-executed result"},
+				},
+			},
+		},
+	}
+
+	usage, estimated := fallbackStepUsage(nil, step)
+	require.True(t, estimated)
+	require.Positive(t, usage.OutputTokens)
+	require.Equal(t, usage.OutputTokens, usage.TotalTokens)
+}
+
 func TestFallbackStepUsageReturnsZeroWithoutContent(t *testing.T) {
 	t.Parallel()
 
@@ -187,6 +232,24 @@ func TestUpdateSessionUsageKeepsCountersForZeroUsage(t *testing.T) {
 	require.Equal(t, int64(456), currentSession.CompletionTokens)
 }
 
+func TestUpdateSessionUsageReplacesCountersForPartialUsage(t *testing.T) {
+	t.Parallel()
+
+	agent := &sessionAgent{}
+	currentSession := &session.Session{
+		ID:               "session-id",
+		PromptTokens:     123,
+		CompletionTokens: 456,
+	}
+	model := Model{CatwalkCfg: catwalk.Model{CostPer1MIn: 10, CostPer1MOut: 20}}
+	usage := fantasy.Usage{InputTokens: 789}
+
+	agent.updateSessionUsage(model, currentSession, usage, nil, false)
+
+	require.Equal(t, int64(789), currentSession.PromptTokens)
+	require.Zero(t, currentSession.CompletionTokens)
+}
+
 func TestUpdateSessionUsageAddsProviderCost(t *testing.T) {
 	t.Parallel()