[management, proxy] Add prompt-cache token and cost accounting to agent network usage (#6900)

Co-authored-by: braginini <bangvalo@gmail.com>

proxy/internal/llm/bedrock.go

@@ -56,10 +56,12 @@ type bedrockResponse struct {
 		OutputTokens             int64 `json:"output_tokens"`
 		CacheReadInputTokens     int64 `json:"cache_read_input_tokens"`
 		CacheCreationInputTokens int64 `json:"cache_creation_input_tokens"`
-		// Converse — camelCase.
-		InputTokensCamel  int64 `json:"inputTokens"`
-		OutputTokensCamel int64 `json:"outputTokens"`
-		TotalTokensCamel  int64 `json:"totalTokens"`
+		// Converse — camelCase; cache buckets are additive to inputTokens (AWS names the write bucket cacheWriteInputTokens).
+		InputTokensCamel      int64 `json:"inputTokens"`
+		OutputTokensCamel     int64 `json:"outputTokens"`
+		TotalTokensCamel      int64 `json:"totalTokens"`
+		CacheReadTokensCamel  int64 `json:"cacheReadInputTokens"`
+		CacheWriteTokensCamel int64 `json:"cacheWriteInputTokens"`
 	} `json:"usage"`
 }
 
@@ -83,16 +85,18 @@ func (BedrockParser) ParseResponse(status int, contentType string, body []byte)
 	}
 	inTok := firstNonZero(resp.Usage.InputTokens, resp.Usage.InputTokensCamel)
 	outTok := firstNonZero(resp.Usage.OutputTokens, resp.Usage.OutputTokensCamel)
+	cacheRead := firstNonZero(resp.Usage.CacheReadInputTokens, resp.Usage.CacheReadTokensCamel)
+	cacheWrite := firstNonZero(resp.Usage.CacheCreationInputTokens, resp.Usage.CacheWriteTokensCamel)
 	total := resp.Usage.TotalTokensCamel
 	if total == 0 {
-		total = inTok + outTok + resp.Usage.CacheReadInputTokens + resp.Usage.CacheCreationInputTokens
+		total = inTok + outTok + cacheRead + cacheWrite
 	}
 	return Usage{
 		InputTokens:         inTok,
 		OutputTokens:        outTok,
 		TotalTokens:         total,
-		CachedInputTokens:   resp.Usage.CacheReadInputTokens,
-		CacheCreationTokens: resp.Usage.CacheCreationInputTokens,
+		CachedInputTokens:   cacheRead,
+		CacheCreationTokens: cacheWrite,
 	}, nil
 }
 

proxy/internal/llm/bedrock_test.go

@@ -26,6 +26,18 @@ func TestBedrockParser_ParseResponse_Converse(t *testing.T) {
 	require.Equal(t, int64(14), u.TotalTokens, "converse uses provider total")
 }
 
+// Converse camelCase cache fields must land in the billed Usage buckets, same as the InvokeModel snake_case fields.
+func TestBedrockParser_ParseResponse_ConverseCacheBuckets(t *testing.T) {
+	body := []byte(`{"usage":{"inputTokens":11,"outputTokens":3,"cacheReadInputTokens":7,"cacheWriteInputTokens":9}}`)
+	u, err := BedrockParser{}.ParseResponse(200, "application/json", body)
+	require.NoError(t, err)
+	require.Equal(t, int64(11), u.InputTokens, "converse input tokens")
+	require.Equal(t, int64(3), u.OutputTokens, "converse output tokens")
+	require.Equal(t, int64(7), u.CachedInputTokens, "converse cache-read tokens")
+	require.Equal(t, int64(9), u.CacheCreationTokens, "converse cache-write tokens")
+	require.Equal(t, int64(11+3+7+9), u.TotalTokens, "total backfill is additive when the provider omits totalTokens")
+}
+
 func TestBedrockParser_ParseResponse_StreamingUnsupported(t *testing.T) {
 	_, err := BedrockParser{}.ParseResponse(200, "application/vnd.amazon.eventstream", []byte("binary"))
 	require.ErrorIs(t, err, ErrStreamingUnsupported, "event-stream must route to the streaming accumulator")

proxy/internal/llm/pricing/pricing.go

@@ -128,6 +128,46 @@ type Table struct {
 //   - Other providers: cached and cacheCreation are ignored; cost is
 //     inTokens*InputPer1K + outTokens*OutputPer1K.
 func (t *Table) Cost(provider, model string, inTokens, outTokens, cachedInput, cacheCreation int64) (float64, bool) {
+	c, ok := t.Costs(provider, model, inTokens, outTokens, cachedInput, cacheCreation)
+	return c.TotalUSD, ok
+}
+
+// Costs is a per-request cost split. The four per-bucket fields are the base
+// of the breakdown — one per token bucket the provider bills separately — and
+// the two aggregates are derived from them:
+//
+//	TotalUSD = InputUSD + CachedInputUSD + CacheCreationUSD + OutputUSD
+//	CacheUSD = CachedInputUSD + CacheCreationUSD
+//
+// InputUSD is always the cost of the *non-cached* input bucket, for both
+// provider shapes: on OpenAI the cached subset is carved out of inTokens and
+// billed as CachedInputUSD, so the two never double-count. Buckets a provider
+// doesn't bill are zero, which keeps the identities above true everywhere.
+type Costs struct {
+	InputUSD         float64
+	CachedInputUSD   float64
+	CacheCreationUSD float64
+	OutputUSD        float64
+	TotalUSD         float64
+	CacheUSD         float64
+}
+
+// newCosts assembles a split from its per-bucket parts, deriving the two
+// aggregates so TotalUSD and CacheUSD can never drift from the breakdown.
+func newCosts(input, cachedInput, cacheCreation, output float64) Costs {
+	return Costs{
+		InputUSD:         input,
+		CachedInputUSD:   cachedInput,
+		CacheCreationUSD: cacheCreation,
+		OutputUSD:        output,
+		TotalUSD:         input + cachedInput + cacheCreation + output,
+		CacheUSD:         cachedInput + cacheCreation,
+	}
+}
+
+// Costs returns the estimated USD cost split for the given token counts, with
+// the same semantics as Cost.
+func (t *Table) Costs(provider, model string, inTokens, outTokens, cachedInput, cacheCreation int64) (Costs, bool) {
 	// Clamp negatives to zero before any pricing math so a malformed
 	// upstream count can never produce a negative cost.
 	if inTokens < 0 {
@@ -143,15 +183,15 @@ func (t *Table) Cost(provider, model string, inTokens, outTokens, cachedInput, c
 		cacheCreation = 0
 	}
 	if t == nil {
-		return 0, false
+		return Costs{}, false
 	}
 	byModel, ok := t.entries[provider]
 	if !ok {
-		return 0, false
+		return Costs{}, false
 	}
 	entry, ok := byModel[model]
 	if !ok {
-		return 0, false
+		return Costs{}, false
 	}
 	output := (float64(outTokens) / 1000.0) * entry.OutputPer1K
 	switch provider {
@@ -168,7 +208,7 @@ func (t *Table) Cost(provider, model string, inTokens, outTokens, cachedInput, c
 		}
 		nonCached := float64(inTokens-clamped) / 1000.0 * entry.InputPer1K
 		cached := float64(clamped) / 1000.0 * cachedRate
-		return nonCached + cached + output, true
+		return newCosts(nonCached, cached, 0, output), true
 	case "anthropic", "bedrock":
 		// Bedrock-Anthropic returns the same additive cache buckets as
 		// first-party Anthropic; non-Anthropic Bedrock models simply report
@@ -184,10 +224,10 @@ func (t *Table) Cost(provider, model string, inTokens, outTokens, cachedInput, c
 		input := float64(inTokens) / 1000.0 * entry.InputPer1K
 		read := float64(cachedInput) / 1000.0 * readRate
 		create := float64(cacheCreation) / 1000.0 * createRate
-		return input + read + create + output, true
+		return newCosts(input, read, create, output), true
 	default:
 		input := float64(inTokens) / 1000.0 * entry.InputPer1K
-		return input + output, true
+		return newCosts(input, 0, 0, output), true
 	}
 }