MUL-5652: fix(agent): accumulate CacheWriteTokens in ACP backend promptDone handlers

server/pkg/agent/grok.go

@@ -449,15 +449,10 @@ func (b *grokBackend) Execute(ctx context.Context, prompt string, opts ExecOptio
 				if effectiveModel == "" {
 					effectiveModel = pr.modelID
 				}
+				// The shared merge also preserves xAI's authoritative cost for
+				// long-context turns, which token counts alone cannot reconstruct.
 				c.usageMu.Lock()
-				c.usage.InputTokens += pr.usage.InputTokens
-				c.usage.OutputTokens += pr.usage.OutputTokens
-				c.usage.CacheReadTokens += pr.usage.CacheReadTokens
-				// xAI prices the turn itself and reports the result here.
-				// Carrying it through is the only way the ≥200K long-context
-				// surcharge reaches the bill — token counts alone cannot
-				// reconstruct which tier a request hit.
-				c.usage.CostUSDTicks += pr.usage.CostUSDTicks
+				mergeACPUsageSnapshot(&c.usage, &c.usageTotalTokens, pr.usage, pr.usageTotalTokens)
 				c.usageMu.Unlock()
 			default:
 			}

server/pkg/agent/grok_test.go

@@ -107,9 +107,11 @@ while IFS= read -r line; do
       printf '{"jsonrpc":"2.0","method":"session/update","params":{"sessionId":"ses_new","update":{"sessionUpdate":"tool_call_update","toolCallId":"tc-1","status":"completed","name":"Shell","output":"hi\\n"}}}\n'
       printf '{"jsonrpc":"2.0","method":"session/update","params":{"sessionId":"ses_new","update":{"sessionUpdate":"agent_message_chunk","content":{"type":"text","text":"pong"}}}}\n'
       if [ -n "$GROK_USAGE" ]; then
-        # Match live Grok Build ACP (0.2.x): metering lives under result._meta,
-        # not a top-level usage field or sessionUpdate=usage_update.
-        printf '{"jsonrpc":"2.0","id":%s,"result":{"stopReason":"end_turn","_meta":{"sessionId":"ses_new","modelId":"grok-4.5","inputTokens":120,"outputTokens":30,"cachedReadTokens":20,"usage":{"inputTokens":120,"outputTokens":30,"totalTokens":150,"cachedReadTokens":20,"modelCalls":1,"costUsdTicks":98765}}}}\n' "$id"
+        # Grok Build ACP (0.2.x) puts metering under result._meta rather than a
+        # top-level usage field or usage_update. The cache-write extension is
+        # synthetic, so keep its four token buckets mutually exclusive and
+        # totalTokens self-consistent: 100 + 30 + 20 + 8 = 158.
+        printf '{"jsonrpc":"2.0","id":%s,"result":{"stopReason":"end_turn","_meta":{"sessionId":"ses_new","modelId":"grok-4.5","inputTokens":100,"outputTokens":30,"cachedReadTokens":20,"cacheWriteTokens":8,"usage":{"inputTokens":100,"outputTokens":30,"totalTokens":158,"cachedReadTokens":20,"cacheWriteTokens":8,"modelCalls":1,"costUsdTicks":98765}}}}\n' "$id"
       else
         printf '{"jsonrpc":"2.0","id":%s,"result":{"stopReason":"end_turn"}}\n' "$id"
       fi
@@ -610,12 +612,14 @@ func TestGrokPropagatesMCPAndUsage(t *testing.T) {
 	if !ok {
 		t.Fatalf("usage missing grok-4.5 key: %+v", result.Usage)
 	}
-	// The fixture's totalTokens (150) equals input + output, so its 20 cached
-	// reads sit inside inputTokens and are billed once: input is stored as the
-	// uncached remainder 120 - 20 = 100.
+	// The fixture uses mutually exclusive counters whose sum matches
+	// totalTokens, so no bucket is charged twice.
 	if usage.InputTokens != 100 || usage.OutputTokens != 30 || usage.CacheReadTokens != 20 {
 		t.Fatalf("unexpected usage: %+v", usage)
 	}
+	if usage.CacheWriteTokens != 8 {
+		t.Fatalf("CacheWriteTokens = %d, want 8", usage.CacheWriteTokens)
+	}
 	// xAI's own price for the turn has to survive the whole backend, not just
 	// the parser: it is the only figure carrying the ≥200K prompt surcharge,
 	// and everything downstream falls back to a rate-table guess without it.
@@ -670,6 +674,9 @@ func TestGrokAttributesUsageOnResumeWithoutConfiguredModel(t *testing.T) {
 	if usage.InputTokens != 100 || usage.OutputTokens != 30 || usage.CacheReadTokens != 20 {
 		t.Fatalf("unexpected usage: %+v", usage)
 	}
+	if usage.CacheWriteTokens != 8 {
+		t.Fatalf("CacheWriteTokens = %d, want 8", usage.CacheWriteTokens)
+	}
 }
 
 func TestGrokTimeoutAndCancellation(t *testing.T) {

server/pkg/agent/hermes.go

@@ -557,11 +557,11 @@ func (b *hermesBackend) Execute(ctx context.Context, prompt string, opts ExecOpt
 					finalStatus = "aborted"
 					finalError = "hermes cancelled the prompt"
 				}
-				// Merge usage from the PromptResponse.
+				// Prompt responses and usage_update notifications can carry the
+				// same cumulative counters. Reconcile whole normalized snapshots
+				// so a runtime that emits both is charged exactly once.
 				c.usageMu.Lock()
-				c.usage.InputTokens += pr.usage.InputTokens
-				c.usage.OutputTokens += pr.usage.OutputTokens
-				c.usage.CacheReadTokens += pr.usage.CacheReadTokens
+				mergeACPUsageSnapshot(&c.usage, &c.usageTotalTokens, pr.usage, pr.usageTotalTokens)
 				c.usageMu.Unlock()
 			default:
 			}
@@ -741,8 +741,9 @@ func waitForHermesPipeDrain(readerDone, stderrDone <-chan struct{}, timeout time
 // ── hermesClient: ACP JSON-RPC 2.0 transport ──
 
 type hermesPromptResult struct {
-	stopReason string
-	usage      TokenUsage
+	stopReason       string
+	usage            TokenUsage
+	usageTotalTokens int64
 	// modelID is the model the agent actually billed this turn against, as
 	// reported on `result._meta.modelId`. Empty for agents that don't report
 	// it. Backends use it to attribute usage when the session handshake
@@ -777,8 +778,9 @@ type hermesClient struct {
 	toolMu       sync.Mutex
 	pendingTools map[string]*pendingToolCall
 
-	usageMu sync.Mutex
-	usage   TokenUsage
+	usageMu          sync.Mutex
+	usage            TokenUsage
+	usageTotalTokens int64
 }
 
 // pendingToolCall buffers state for a tool call while its arguments
@@ -1213,16 +1215,19 @@ func (c *hermesClient) extractPromptResult(data json.RawMessage) {
 		modelID:    parseACPModelIDFromMeta(resp.Meta),
 	}
 	if len(resp.Usage) > 0 && string(resp.Usage) != "null" {
-		pr.usage = parseACPTokenUsage(resp.Usage)
+		snapshot := parseACPTokenUsageSnapshot(resp.Usage)
+		pr.usage = snapshot.usage
+		pr.usageTotalTokens = snapshot.totalTokens
 	}
 	// Prefer the standard top-level ACP `usage` field when present. Some
 	// agents (notably xAI Grok Build) put per-turn metering only under
 	// result._meta — either as `_meta.usage` or as flat token counters on
 	// `_meta` itself. Without this fallback, tasks complete with an empty
 	// usage map and Multica's Usage/cost dashboards stay at zero.
 	if !acpTokenUsagePresent(pr.usage) {
-		if metaUsage := parseACPTokenUsageFromMeta(resp.Meta); acpTokenUsagePresent(metaUsage) {
-			pr.usage = metaUsage
+		if metaUsage := parseACPTokenUsageSnapshotFromMeta(resp.Meta); acpTokenUsagePresent(metaUsage.usage) {
+			pr.usage = metaUsage.usage
+			pr.usageTotalTokens = metaUsage.totalTokens
 		}
 	}
 
@@ -1236,28 +1241,28 @@ func acpTokenUsagePresent(u TokenUsage) bool {
 	return u.InputTokens > 0 || u.OutputTokens > 0 || u.CacheReadTokens > 0 || u.CacheWriteTokens > 0
 }
 
-// parseACPTokenUsageFromMeta extracts token usage from an ACP result `_meta`
-// object. Grok Build returns shapes like:
+// parseACPTokenUsageSnapshotFromMeta extracts token usage from an ACP result
+// `_meta` object. Grok Build returns shapes like:
 //
 //	{"inputTokens":…,"outputTokens":…,"cachedReadTokens":…,"usage":{…}}
 //
 // Prefer the nested `usage` object when it carries counters; otherwise parse
 // the flat `_meta` fields with the same alias rules as top-level usage.
-func parseACPTokenUsageFromMeta(meta json.RawMessage) TokenUsage {
+func parseACPTokenUsageSnapshotFromMeta(meta json.RawMessage) acpTokenUsageSnapshot {
 	if len(meta) == 0 || string(meta) == "null" {
-		return TokenUsage{}
+		return acpTokenUsageSnapshot{}
 	}
 	var envelope struct {
 		Usage json.RawMessage `json:"usage"`
 	}
 	if err := json.Unmarshal(meta, &envelope); err == nil {
 		if len(envelope.Usage) > 0 && string(envelope.Usage) != "null" {
-			if u := parseACPTokenUsage(envelope.Usage); acpTokenUsagePresent(u) {
-				return u
+			if snapshot := parseACPTokenUsageSnapshot(envelope.Usage); acpTokenUsagePresent(snapshot.usage) {
+				return snapshot
 			}
 		}
 	}
-	return parseACPTokenUsage(meta)
+	return parseACPTokenUsageSnapshot(meta)
 }
 
 // parseACPModelIDFromMeta pulls the model id off an ACP result `_meta`
@@ -1740,35 +1745,65 @@ func (c *hermesClient) handleUsageUpdate(data json.RawMessage) {
 	if err := json.Unmarshal(data, &msg); err != nil {
 		return
 	}
-	usage := parseACPTokenUsage(msg.Usage)
+	usage := parseACPTokenUsageSnapshot(msg.Usage)
 
 	c.usageMu.Lock()
-	// Usage updates from ACP are cumulative snapshots, so take the latest.
-	if usage.InputTokens > c.usage.InputTokens {
-		c.usage.InputTokens = usage.InputTokens
-	}
-	if usage.OutputTokens > c.usage.OutputTokens {
-		c.usage.OutputTokens = usage.OutputTokens
-	}
-	if usage.CacheReadTokens > c.usage.CacheReadTokens {
-		c.usage.CacheReadTokens = usage.CacheReadTokens
-	}
-	if usage.CacheWriteTokens > c.usage.CacheWriteTokens {
-		c.usage.CacheWriteTokens = usage.CacheWriteTokens
+	mergeACPUsageSnapshot(&c.usage, &c.usageTotalTokens, usage.usage, usage.totalTokens)
+	c.usageMu.Unlock()
+}
+
+// mergeACPUsageSnapshot reconciles cumulative usage reported through ACP's
+// usage_update and PromptResponse paths. A snapshot with totalTokens is
+// self-describing and has already been normalized into mutually exclusive
+// token buckets, so keep its counters together instead of mixing them with an
+// ambiguous snapshot that may still include cached tokens in inputTokens.
+// Among self-describing snapshots the largest cumulative total wins; without
+// one, per-field maxima retain the most complete available counters. Provider
+// cost is cumulative too and always uses the largest reported value.
+func mergeACPUsageSnapshot(dst *TokenUsage, dstTotalTokens *int64, next TokenUsage, nextTotalTokens int64) {
+	if nextTotalTokens > 0 {
+		if *dstTotalTokens == 0 || nextTotalTokens >= *dstTotalTokens {
+			dst.InputTokens = next.InputTokens
+			dst.OutputTokens = next.OutputTokens
+			dst.CacheReadTokens = next.CacheReadTokens
+			dst.CacheWriteTokens = next.CacheWriteTokens
+			*dstTotalTokens = nextTotalTokens
+		}
+	} else if *dstTotalTokens == 0 {
+		if next.InputTokens > dst.InputTokens {
+			dst.InputTokens = next.InputTokens
+		}
+		if next.OutputTokens > dst.OutputTokens {
+			dst.OutputTokens = next.OutputTokens
+		}
+		if next.CacheReadTokens > dst.CacheReadTokens {
+			dst.CacheReadTokens = next.CacheReadTokens
+		}
+		if next.CacheWriteTokens > dst.CacheWriteTokens {
+			dst.CacheWriteTokens = next.CacheWriteTokens
+		}
 	}
-	if usage.CostUSDTicks > c.usage.CostUSDTicks {
-		c.usage.CostUSDTicks = usage.CostUSDTicks
+	if next.CostUSDTicks > dst.CostUSDTicks {
+		dst.CostUSDTicks = next.CostUSDTicks
 	}
-	c.usageMu.Unlock()
 }
 
 func parseACPTokenUsage(data json.RawMessage) TokenUsage {
+	return parseACPTokenUsageSnapshot(data).usage
+}
+
+type acpTokenUsageSnapshot struct {
+	usage       TokenUsage
+	totalTokens int64
+}
+
+func parseACPTokenUsageSnapshot(data json.RawMessage) acpTokenUsageSnapshot {
 	if len(data) == 0 || string(data) == "null" {
-		return TokenUsage{}
+		return acpTokenUsageSnapshot{}
 	}
 	var fields map[string]json.RawMessage
 	if err := json.Unmarshal(data, &fields); err != nil {
-		return TokenUsage{}
+		return acpTokenUsageSnapshot{}
 	}
 	usage := TokenUsage{
 		InputTokens:  acpUsageInt64(fields, "inputTokens", "input_tokens"),
@@ -1791,7 +1826,11 @@ func parseACPTokenUsage(data json.RawMessage) TokenUsage {
 		// counts (see TokenUsage.CostUSDTicks).
 		CostUSDTicks: acpUsageInt64(fields, "costUsdTicks", "cost_usd_ticks"),
 	}
-	return excludeACPCachedInput(usage, acpUsageInt64(fields, "totalTokens", "total_tokens"))
+	totalTokens := acpUsageInt64(fields, "totalTokens", "total_tokens")
+	return acpTokenUsageSnapshot{
+		usage:       excludeACPCachedInput(usage, totalTokens),
+		totalTokens: totalTokens,
+	}
 }
 
 // excludeACPCachedInput re-buckets a usage record whose `inputTokens` already