[management] persist prompt-cache tokens and cache cost per request

Flatten llm.cached_input_tokens / llm.cache_creation_tokens and cost.usd_cache from the proxy metadata into the agent-network access-log and request-usage tables, sum them in session summaries and usage buckets, and render them in the API responses.
2026-07-28 00:42:37 -04:00 · 2026-07-25 19:25:24 +00:00
parent 0c26e364eb
commit 6faffe5716
5 changed files with 146 additions and 99 deletions
--- a/management/internals/modules/agentnetwork/accesslog_ingest.go
+++ b/management/internals/modules/agentnetwork/accesslog_ingest.go
@@ -18,21 +18,24 @@ import (
 // contract between the proxy and management; management flattens them into
 // queryable columns. Keep in sync with the proxy side.
 const (
-	metaKeyProvider           = "llm.provider"
-	metaKeyModel              = "llm.model"
-	metaKeyResolvedProviderID = "llm.resolved_provider_id"
-	metaKeySelectedPolicyID   = "llm.selected_policy_id"
-	metaKeyPolicyDecision     = "llm_policy.decision"
-	metaKeyPolicyReason       = "llm_policy.reason"
-	metaKeyInputTokens        = "llm.input_tokens"  //nolint:gosec // metadata key name, not a credential
-	metaKeyOutputTokens       = "llm.output_tokens" //nolint:gosec // metadata key name, not a credential
-	metaKeyTotalTokens        = "llm.total_tokens"  //nolint:gosec // metadata key name, not a credential
-	metaKeyCostUSDTotal       = "cost.usd_total"
-	metaKeyStream             = "llm.stream"
-	metaKeySessionID          = "llm.session_id"
-	metaKeyAuthorisingGroups  = "llm.authorising_groups"
-	metaKeyRequestPrompt      = "llm.request_prompt"
-	metaKeyResponseCompletion = "llm.response_completion"
+	metaKeyProvider            = "llm.provider"
+	metaKeyModel               = "llm.model"
+	metaKeyResolvedProviderID  = "llm.resolved_provider_id"
+	metaKeySelectedPolicyID    = "llm.selected_policy_id"
+	metaKeyPolicyDecision      = "llm_policy.decision"
+	metaKeyPolicyReason        = "llm_policy.reason"
+	metaKeyInputTokens         = "llm.input_tokens"          //nolint:gosec // metadata key name, not a credential
+	metaKeyOutputTokens        = "llm.output_tokens"         //nolint:gosec // metadata key name, not a credential
+	metaKeyTotalTokens         = "llm.total_tokens"          //nolint:gosec // metadata key name, not a credential
+	metaKeyCachedInputTokens   = "llm.cached_input_tokens"   //nolint:gosec // metadata key name, not a credential
+	metaKeyCacheCreationTokens = "llm.cache_creation_tokens" //nolint:gosec // metadata key name, not a credential
+	metaKeyCostUSDTotal        = "cost.usd_total"
+	metaKeyCostUSDCache        = "cost.usd_cache"
+	metaKeyStream              = "llm.stream"
+	metaKeySessionID           = "llm.session_id"
+	metaKeyAuthorisingGroups   = "llm.authorising_groups"
+	metaKeyRequestPrompt       = "llm.request_prompt"
+	metaKeyResponseCompletion  = "llm.response_completion"
 )

 // IngestAccessLog flattens the metadata-bearing reverse-proxy access-log entry
@@ -108,20 +111,23 @@ func flattenAccessLog(e *accesslogs.AccessLogEntry) (*types.AgentNetworkAccessLo
 		BytesUpload:   e.BytesUpload,
 		BytesDownload: e.BytesDownload,

-		Provider:           meta[metaKeyProvider],
-		Model:              meta[metaKeyModel],
-		SessionID:          meta[metaKeySessionID],
-		ResolvedProviderID: meta[metaKeyResolvedProviderID],
-		SelectedPolicyID:   meta[metaKeySelectedPolicyID],
-		Decision:           meta[metaKeyPolicyDecision],
-		DenyReason:         meta[metaKeyPolicyReason],
-		InputTokens:        parseMetaInt(meta, metaKeyInputTokens),
-		OutputTokens:       parseMetaInt(meta, metaKeyOutputTokens),
-		TotalTokens:        parseMetaInt(meta, metaKeyTotalTokens),
-		CostUSD:            parseMetaFloat(meta, metaKeyCostUSDTotal),
-		Stream:             parseMetaBool(meta, metaKeyStream),
-		RequestPrompt:      meta[metaKeyRequestPrompt],
-		ResponseCompletion: meta[metaKeyResponseCompletion],
+		Provider:            meta[metaKeyProvider],
+		Model:               meta[metaKeyModel],
+		SessionID:           meta[metaKeySessionID],
+		ResolvedProviderID:  meta[metaKeyResolvedProviderID],
+		SelectedPolicyID:    meta[metaKeySelectedPolicyID],
+		Decision:            meta[metaKeyPolicyDecision],
+		DenyReason:          meta[metaKeyPolicyReason],
+		InputTokens:         parseMetaInt(meta, metaKeyInputTokens),
+		OutputTokens:        parseMetaInt(meta, metaKeyOutputTokens),
+		TotalTokens:         parseMetaInt(meta, metaKeyTotalTokens),
+		CachedInputTokens:   parseMetaInt(meta, metaKeyCachedInputTokens),
+		CacheCreationTokens: parseMetaInt(meta, metaKeyCacheCreationTokens),
+		CostUSD:             parseMetaFloat(meta, metaKeyCostUSDTotal),
+		CacheCostUSD:        parseMetaFloat(meta, metaKeyCostUSDCache),
+		Stream:              parseMetaBool(meta, metaKeyStream),
+		RequestPrompt:       meta[metaKeyRequestPrompt],
+		ResponseCompletion:  meta[metaKeyResponseCompletion],
 	}

 	var groups []types.AgentNetworkAccessLogGroup
@@ -140,18 +146,21 @@ func flattenAccessLog(e *accesslogs.AccessLogEntry) (*types.AgentNetworkAccessLo
 // log's ID so the two correlate.
 func usageFromFlattenedLog(e *types.AgentNetworkAccessLog, groups []types.AgentNetworkAccessLogGroup) (*types.AgentNetworkUsage, []types.AgentNetworkUsageGroup) {
 	usage := &types.AgentNetworkUsage{
-		ID:                 e.ID,
-		AccountID:          e.AccountID,
-		Timestamp:          e.Timestamp,
-		UserID:             e.UserID,
-		ResolvedProviderID: e.ResolvedProviderID,
-		Provider:           e.Provider,
-		Model:              e.Model,
-		SessionID:          e.SessionID,
-		InputTokens:        e.InputTokens,
-		OutputTokens:       e.OutputTokens,
-		TotalTokens:        e.TotalTokens,
-		CostUSD:            e.CostUSD,
+		ID:                  e.ID,
+		AccountID:           e.AccountID,
+		Timestamp:           e.Timestamp,
+		UserID:              e.UserID,
+		ResolvedProviderID:  e.ResolvedProviderID,
+		Provider:            e.Provider,
+		Model:               e.Model,
+		SessionID:           e.SessionID,
+		InputTokens:         e.InputTokens,
+		OutputTokens:        e.OutputTokens,
+		TotalTokens:         e.TotalTokens,
+		CachedInputTokens:   e.CachedInputTokens,
+		CacheCreationTokens: e.CacheCreationTokens,
+		CostUSD:             e.CostUSD,
+		CacheCostUSD:        e.CacheCostUSD,
 	}

 	usageGroups := make([]types.AgentNetworkUsageGroup, 0, len(groups))
--- a/management/internals/modules/agentnetwork/accesslog_ingest_realstore_test.go
+++ b/management/internals/modules/agentnetwork/accesslog_ingest_realstore_test.go
@@ -28,17 +28,20 @@ func newIngestTestEntry() *accesslogs.AccessLogEntry {
 		UserId:       "user-1",
 		AgentNetwork: true,
 		Metadata: map[string]string{
-			metaKeyProvider:           "openai",
-			metaKeyModel:              "gpt-5.4",
-			metaKeyResolvedProviderID: "prov-1",
-			metaKeySessionID:          "sess-1",
-			metaKeyInputTokens:        "100",
-			metaKeyOutputTokens:       "50",
-			metaKeyTotalTokens:        "150",
-			metaKeyCostUSDTotal:       "0.0123",
-			metaKeyStream:             "true",
-			metaKeyRequestPrompt:      "hello",
-			metaKeyResponseCompletion: "world",
+			metaKeyProvider:            "openai",
+			metaKeyModel:               "gpt-5.4",
+			metaKeyResolvedProviderID:  "prov-1",
+			metaKeySessionID:           "sess-1",
+			metaKeyInputTokens:         "100",
+			metaKeyOutputTokens:        "50",
+			metaKeyTotalTokens:         "1174",
+			metaKeyCachedInputTokens:   "256",
+			metaKeyCacheCreationTokens: "768",
+			metaKeyCostUSDTotal:        "0.0123",
+			metaKeyCostUSDCache:        "0.0029",
+			metaKeyStream:              "true",
+			metaKeyRequestPrompt:       "hello",
+			metaKeyResponseCompletion:  "world",
 			// repeated id must be de-duplicated before the group rows insert.
 			metaKeyAuthorisingGroups: "grp-eng,grp-eng,grp-ops",
 		},
@@ -65,7 +68,10 @@ func TestIngestAccessLog_RealStore_LogCollectionOff(t *testing.T) {
 	require.Len(t, usage, 1, "usage row must be written even with log collection off")
 	assert.Equal(t, int64(100), usage[0].InputTokens, "input tokens must round-trip from metadata")
 	assert.Equal(t, int64(50), usage[0].OutputTokens, "output tokens must round-trip from metadata")
+	assert.Equal(t, int64(256), usage[0].CachedInputTokens, "cache-read tokens must round-trip from metadata")
+	assert.Equal(t, int64(768), usage[0].CacheCreationTokens, "cache-write tokens must round-trip from metadata")
 	assert.InDelta(t, 0.0123, usage[0].CostUSD, 1e-9, "cost must round-trip from metadata")
+	assert.InDelta(t, 0.0029, usage[0].CacheCostUSD, 1e-9, "cache cost must round-trip from metadata")

 	logs, _, err := s.GetAgentNetworkAccessLogs(ctx, store.LockingStrengthNone, testAccountID, types.AgentNetworkAccessLogFilter{})
 	require.NoError(t, err)
@@ -96,6 +102,9 @@ func TestIngestAccessLog_RealStore_LogCollectionOn(t *testing.T) {
 	require.Equal(t, int64(1), total, "exactly one access-log row expected")
 	require.Len(t, logs, 1, "full access-log row must be written when log collection is on")
 	assert.Equal(t, "gpt-5.4", logs[0].Model, "model must flatten from metadata")
+	assert.Equal(t, int64(256), logs[0].CachedInputTokens, "cache-read tokens must flatten from metadata")
+	assert.Equal(t, int64(768), logs[0].CacheCreationTokens, "cache-write tokens must flatten from metadata")
+	assert.InDelta(t, 0.0029, logs[0].CacheCostUSD, 1e-9, "cache cost must flatten from metadata")
 	assert.Equal(t, "hello", logs[0].RequestPrompt, "prompt must be retained when log collection is on")
 	assert.Equal(t, "world", logs[0].ResponseCompletion, "completion must be retained when log collection is on")
 	assert.True(t, logs[0].Stream, "stream flag must flatten from metadata")
--- a/management/internals/modules/agentnetwork/types/accesslog.go
+++ b/management/internals/modules/agentnetwork/types/accesslog.go
@@ -41,8 +41,12 @@ type AgentNetworkAccessLog struct {
 	InputTokens        int64
 	OutputTokens       int64
 	TotalTokens        int64
-	CostUSD            float64
-	Stream             bool
+	// Prompt-cache buckets: read + write token counts and the portion of CostUSD they account for.
+	CachedInputTokens   int64
+	CacheCreationTokens int64
+	CostUSD             float64
+	CacheCostUSD        float64
+	Stream              bool

 	// Prompt capture. Only populated when prompt collection is enabled
 	// (account master switch AND policy guardrail). Heavy free text.
@@ -63,16 +67,19 @@ func (AgentNetworkAccessLog) TableName() string { return "agent_network_access_l
 // ToAPIResponse renders the flattened entry as the API representation.
 func (a *AgentNetworkAccessLog) ToAPIResponse() api.AgentNetworkAccessLog {
 	out := api.AgentNetworkAccessLog{
-		Id:           a.ID,
-		ServiceId:    a.ServiceID,
-		Timestamp:    a.Timestamp,
-		StatusCode:   a.StatusCode,
-		DurationMs:   int(a.Duration.Milliseconds()),
-		InputTokens:  a.InputTokens,
-		OutputTokens: a.OutputTokens,
-		TotalTokens:  a.TotalTokens,
-		CostUsd:      a.CostUSD,
-		Stream:       &a.Stream,
+		Id:                  a.ID,
+		ServiceId:           a.ServiceID,
+		Timestamp:           a.Timestamp,
+		StatusCode:          a.StatusCode,
+		DurationMs:          int(a.Duration.Milliseconds()),
+		InputTokens:         a.InputTokens,
+		OutputTokens:        a.OutputTokens,
+		TotalTokens:         a.TotalTokens,
+		CachedInputTokens:   a.CachedInputTokens,
+		CacheCreationTokens: a.CacheCreationTokens,
+		CostUsd:             a.CostUSD,
+		CacheCostUsd:        a.CacheCostUSD,
+		Stream:              &a.Stream,
 	}

 	out.UserId = strPtr(a.UserID)
@@ -112,20 +119,23 @@ func strPtr(s string) *string {
 // summary plus its ordered entries. Assembled in Go from a page of entries — it
 // is not a stored table.
 type AgentNetworkAccessLogSession struct {
-	SessionID    string // empty for a session-less (singleton) request
-	UserID       string
-	GroupIDs     []string // union of the entries' authorising groups
-	StartedAt    time.Time
-	EndedAt      time.Time
-	RequestCount int
-	InputTokens  int64
-	OutputTokens int64
-	TotalTokens  int64
-	CostUSD      float64
-	Providers    []string // distinct vendors seen in the session
-	Models       []string // distinct models seen in the session
-	Decision     string   // "deny" if any entry was denied, else "allow"
-	Entries      []*AgentNetworkAccessLog
+	SessionID           string // empty for a session-less (singleton) request
+	UserID              string
+	GroupIDs            []string // union of the entries' authorising groups
+	StartedAt           time.Time
+	EndedAt             time.Time
+	RequestCount        int
+	InputTokens         int64
+	OutputTokens        int64
+	TotalTokens         int64
+	CachedInputTokens   int64
+	CacheCreationTokens int64
+	CostUSD             float64
+	CacheCostUSD        float64
+	Providers           []string // distinct vendors seen in the session
+	Models              []string // distinct models seen in the session
+	Decision            string   // "deny" if any entry was denied, else "allow"
+	Entries             []*AgentNetworkAccessLog
 }

 // sessionKey is the grouping key for an entry: its session id, or — when the
@@ -205,7 +215,10 @@ func (sess *AgentNetworkAccessLogSession) foldEntry(sk *sessionSeen, e *AgentNet
 	sess.InputTokens += e.InputTokens
 	sess.OutputTokens += e.OutputTokens
 	sess.TotalTokens += e.TotalTokens
+	sess.CachedInputTokens += e.CachedInputTokens
+	sess.CacheCreationTokens += e.CacheCreationTokens
 	sess.CostUSD += e.CostUSD
+	sess.CacheCostUSD += e.CacheCostUSD
 	if e.Timestamp.Before(sess.StartedAt) {
 		sess.StartedAt = e.Timestamp
 	}
@@ -248,15 +261,18 @@ func (sess *AgentNetworkAccessLogSession) ToAPIResponse() api.AgentNetworkAccess
 	}

 	out := api.AgentNetworkAccessLogSession{
-		StartedAt:    sess.StartedAt,
-		EndedAt:      sess.EndedAt,
-		RequestCount: sess.RequestCount,
-		InputTokens:  sess.InputTokens,
-		OutputTokens: sess.OutputTokens,
-		TotalTokens:  sess.TotalTokens,
-		CostUsd:      sess.CostUSD,
-		Decision:     sess.Decision,
-		Entries:      entries,
+		StartedAt:           sess.StartedAt,
+		EndedAt:             sess.EndedAt,
+		RequestCount:        sess.RequestCount,
+		InputTokens:         sess.InputTokens,
+		OutputTokens:        sess.OutputTokens,
+		TotalTokens:         sess.TotalTokens,
+		CachedInputTokens:   sess.CachedInputTokens,
+		CacheCreationTokens: sess.CacheCreationTokens,
+		CostUsd:             sess.CostUSD,
+		CacheCostUsd:        sess.CacheCostUSD,
+		Decision:            sess.Decision,
+		Entries:             entries,
 	}
 	out.SessionId = strPtr(sess.SessionID)
 	out.UserId = strPtr(sess.UserID)
--- a/management/internals/modules/agentnetwork/types/usage.go
+++ b/management/internals/modules/agentnetwork/types/usage.go
@@ -25,8 +25,12 @@ type AgentNetworkUsage struct {
 	InputTokens        int64
 	OutputTokens       int64
 	TotalTokens        int64
-	CostUSD            float64
-	CreatedAt          time.Time
+	// Prompt-cache buckets: read + write token counts and the portion of CostUSD they account for.
+	CachedInputTokens   int64
+	CacheCreationTokens int64
+	CostUSD             float64
+	CacheCostUSD        float64
+	CreatedAt           time.Time
 }

 // TableName keeps usage records in their own stripped table. Named
--- a/management/internals/modules/agentnetwork/types/usageoverview.go
+++ b/management/internals/modules/agentnetwork/types/usageoverview.go
@@ -33,21 +33,27 @@ func ParseUsageGranularity(s string) UsageGranularity {
 // AgentNetworkUsageBucket is one aggregated usage time bucket. PeriodStart is
 // the UTC start of the bucket as YYYY-MM-DD.
 type AgentNetworkUsageBucket struct {
-	PeriodStart  string
-	InputTokens  int64
-	OutputTokens int64
-	TotalTokens  int64
-	CostUSD      float64
+	PeriodStart         string
+	InputTokens         int64
+	OutputTokens        int64
+	TotalTokens         int64
+	CachedInputTokens   int64
+	CacheCreationTokens int64
+	CostUSD             float64
+	CacheCostUSD        float64
 }

 // ToAPIResponse renders the bucket as the API representation.
 func (b *AgentNetworkUsageBucket) ToAPIResponse() api.AgentNetworkUsageBucket {
 	return api.AgentNetworkUsageBucket{
-		PeriodStart:  b.PeriodStart,
-		InputTokens:  b.InputTokens,
-		OutputTokens: b.OutputTokens,
-		TotalTokens:  b.TotalTokens,
-		CostUsd:      b.CostUSD,
+		PeriodStart:         b.PeriodStart,
+		InputTokens:         b.InputTokens,
+		OutputTokens:        b.OutputTokens,
+		TotalTokens:         b.TotalTokens,
+		CachedInputTokens:   b.CachedInputTokens,
+		CacheCreationTokens: b.CacheCreationTokens,
+		CostUsd:             b.CostUSD,
+		CacheCostUsd:        b.CacheCostUSD,
 	}
 }

@@ -84,7 +90,10 @@ func AggregateUsageByGranularity(rows []*AgentNetworkUsage, g UsageGranularity)
 		b.InputTokens += r.InputTokens
 		b.OutputTokens += r.OutputTokens
 		b.TotalTokens += r.TotalTokens
+		b.CachedInputTokens += r.CachedInputTokens
+		b.CacheCreationTokens += r.CacheCreationTokens
 		b.CostUSD += r.CostUSD
+		b.CacheCostUSD += r.CacheCostUSD
 	}

 	out := make([]*AgentNetworkUsageBucket, 0, len(byPeriod))