[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.
This commit is contained in:
Maycon Santos
2026-07-25 19:25:24 +00:00
parent 0c26e364eb
commit 6faffe5716
5 changed files with 146 additions and 99 deletions
@@ -18,21 +18,24 @@ import (
// contract between the proxy and management; management flattens them into // contract between the proxy and management; management flattens them into
// queryable columns. Keep in sync with the proxy side. // queryable columns. Keep in sync with the proxy side.
const ( const (
metaKeyProvider = "llm.provider" metaKeyProvider = "llm.provider"
metaKeyModel = "llm.model" metaKeyModel = "llm.model"
metaKeyResolvedProviderID = "llm.resolved_provider_id" metaKeyResolvedProviderID = "llm.resolved_provider_id"
metaKeySelectedPolicyID = "llm.selected_policy_id" metaKeySelectedPolicyID = "llm.selected_policy_id"
metaKeyPolicyDecision = "llm_policy.decision" metaKeyPolicyDecision = "llm_policy.decision"
metaKeyPolicyReason = "llm_policy.reason" metaKeyPolicyReason = "llm_policy.reason"
metaKeyInputTokens = "llm.input_tokens" //nolint:gosec // metadata key name, not a credential metaKeyInputTokens = "llm.input_tokens" //nolint:gosec // metadata key name, not a credential
metaKeyOutputTokens = "llm.output_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 metaKeyTotalTokens = "llm.total_tokens" //nolint:gosec // metadata key name, not a credential
metaKeyCostUSDTotal = "cost.usd_total" metaKeyCachedInputTokens = "llm.cached_input_tokens" //nolint:gosec // metadata key name, not a credential
metaKeyStream = "llm.stream" metaKeyCacheCreationTokens = "llm.cache_creation_tokens" //nolint:gosec // metadata key name, not a credential
metaKeySessionID = "llm.session_id" metaKeyCostUSDTotal = "cost.usd_total"
metaKeyAuthorisingGroups = "llm.authorising_groups" metaKeyCostUSDCache = "cost.usd_cache"
metaKeyRequestPrompt = "llm.request_prompt" metaKeyStream = "llm.stream"
metaKeyResponseCompletion = "llm.response_completion" 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 // IngestAccessLog flattens the metadata-bearing reverse-proxy access-log entry
@@ -108,20 +111,23 @@ func flattenAccessLog(e *accesslogs.AccessLogEntry) (*types.AgentNetworkAccessLo
BytesUpload: e.BytesUpload, BytesUpload: e.BytesUpload,
BytesDownload: e.BytesDownload, BytesDownload: e.BytesDownload,
Provider: meta[metaKeyProvider], Provider: meta[metaKeyProvider],
Model: meta[metaKeyModel], Model: meta[metaKeyModel],
SessionID: meta[metaKeySessionID], SessionID: meta[metaKeySessionID],
ResolvedProviderID: meta[metaKeyResolvedProviderID], ResolvedProviderID: meta[metaKeyResolvedProviderID],
SelectedPolicyID: meta[metaKeySelectedPolicyID], SelectedPolicyID: meta[metaKeySelectedPolicyID],
Decision: meta[metaKeyPolicyDecision], Decision: meta[metaKeyPolicyDecision],
DenyReason: meta[metaKeyPolicyReason], DenyReason: meta[metaKeyPolicyReason],
InputTokens: parseMetaInt(meta, metaKeyInputTokens), InputTokens: parseMetaInt(meta, metaKeyInputTokens),
OutputTokens: parseMetaInt(meta, metaKeyOutputTokens), OutputTokens: parseMetaInt(meta, metaKeyOutputTokens),
TotalTokens: parseMetaInt(meta, metaKeyTotalTokens), TotalTokens: parseMetaInt(meta, metaKeyTotalTokens),
CostUSD: parseMetaFloat(meta, metaKeyCostUSDTotal), CachedInputTokens: parseMetaInt(meta, metaKeyCachedInputTokens),
Stream: parseMetaBool(meta, metaKeyStream), CacheCreationTokens: parseMetaInt(meta, metaKeyCacheCreationTokens),
RequestPrompt: meta[metaKeyRequestPrompt], CostUSD: parseMetaFloat(meta, metaKeyCostUSDTotal),
ResponseCompletion: meta[metaKeyResponseCompletion], CacheCostUSD: parseMetaFloat(meta, metaKeyCostUSDCache),
Stream: parseMetaBool(meta, metaKeyStream),
RequestPrompt: meta[metaKeyRequestPrompt],
ResponseCompletion: meta[metaKeyResponseCompletion],
} }
var groups []types.AgentNetworkAccessLogGroup var groups []types.AgentNetworkAccessLogGroup
@@ -140,18 +146,21 @@ func flattenAccessLog(e *accesslogs.AccessLogEntry) (*types.AgentNetworkAccessLo
// log's ID so the two correlate. // log's ID so the two correlate.
func usageFromFlattenedLog(e *types.AgentNetworkAccessLog, groups []types.AgentNetworkAccessLogGroup) (*types.AgentNetworkUsage, []types.AgentNetworkUsageGroup) { func usageFromFlattenedLog(e *types.AgentNetworkAccessLog, groups []types.AgentNetworkAccessLogGroup) (*types.AgentNetworkUsage, []types.AgentNetworkUsageGroup) {
usage := &types.AgentNetworkUsage{ usage := &types.AgentNetworkUsage{
ID: e.ID, ID: e.ID,
AccountID: e.AccountID, AccountID: e.AccountID,
Timestamp: e.Timestamp, Timestamp: e.Timestamp,
UserID: e.UserID, UserID: e.UserID,
ResolvedProviderID: e.ResolvedProviderID, ResolvedProviderID: e.ResolvedProviderID,
Provider: e.Provider, Provider: e.Provider,
Model: e.Model, Model: e.Model,
SessionID: e.SessionID, SessionID: e.SessionID,
InputTokens: e.InputTokens, InputTokens: e.InputTokens,
OutputTokens: e.OutputTokens, OutputTokens: e.OutputTokens,
TotalTokens: e.TotalTokens, TotalTokens: e.TotalTokens,
CostUSD: e.CostUSD, CachedInputTokens: e.CachedInputTokens,
CacheCreationTokens: e.CacheCreationTokens,
CostUSD: e.CostUSD,
CacheCostUSD: e.CacheCostUSD,
} }
usageGroups := make([]types.AgentNetworkUsageGroup, 0, len(groups)) usageGroups := make([]types.AgentNetworkUsageGroup, 0, len(groups))
@@ -28,17 +28,20 @@ func newIngestTestEntry() *accesslogs.AccessLogEntry {
UserId: "user-1", UserId: "user-1",
AgentNetwork: true, AgentNetwork: true,
Metadata: map[string]string{ Metadata: map[string]string{
metaKeyProvider: "openai", metaKeyProvider: "openai",
metaKeyModel: "gpt-5.4", metaKeyModel: "gpt-5.4",
metaKeyResolvedProviderID: "prov-1", metaKeyResolvedProviderID: "prov-1",
metaKeySessionID: "sess-1", metaKeySessionID: "sess-1",
metaKeyInputTokens: "100", metaKeyInputTokens: "100",
metaKeyOutputTokens: "50", metaKeyOutputTokens: "50",
metaKeyTotalTokens: "150", metaKeyTotalTokens: "1174",
metaKeyCostUSDTotal: "0.0123", metaKeyCachedInputTokens: "256",
metaKeyStream: "true", metaKeyCacheCreationTokens: "768",
metaKeyRequestPrompt: "hello", metaKeyCostUSDTotal: "0.0123",
metaKeyResponseCompletion: "world", metaKeyCostUSDCache: "0.0029",
metaKeyStream: "true",
metaKeyRequestPrompt: "hello",
metaKeyResponseCompletion: "world",
// repeated id must be de-duplicated before the group rows insert. // repeated id must be de-duplicated before the group rows insert.
metaKeyAuthorisingGroups: "grp-eng,grp-eng,grp-ops", 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") 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(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(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.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{}) logs, _, err := s.GetAgentNetworkAccessLogs(ctx, store.LockingStrengthNone, testAccountID, types.AgentNetworkAccessLogFilter{})
require.NoError(t, err) 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.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") 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, "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, "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.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") assert.True(t, logs[0].Stream, "stream flag must flatten from metadata")
@@ -41,8 +41,12 @@ type AgentNetworkAccessLog struct {
InputTokens int64 InputTokens int64
OutputTokens int64 OutputTokens int64
TotalTokens int64 TotalTokens int64
CostUSD float64 // Prompt-cache buckets: read + write token counts and the portion of CostUSD they account for.
Stream bool CachedInputTokens int64
CacheCreationTokens int64
CostUSD float64
CacheCostUSD float64
Stream bool
// Prompt capture. Only populated when prompt collection is enabled // Prompt capture. Only populated when prompt collection is enabled
// (account master switch AND policy guardrail). Heavy free text. // (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. // ToAPIResponse renders the flattened entry as the API representation.
func (a *AgentNetworkAccessLog) ToAPIResponse() api.AgentNetworkAccessLog { func (a *AgentNetworkAccessLog) ToAPIResponse() api.AgentNetworkAccessLog {
out := api.AgentNetworkAccessLog{ out := api.AgentNetworkAccessLog{
Id: a.ID, Id: a.ID,
ServiceId: a.ServiceID, ServiceId: a.ServiceID,
Timestamp: a.Timestamp, Timestamp: a.Timestamp,
StatusCode: a.StatusCode, StatusCode: a.StatusCode,
DurationMs: int(a.Duration.Milliseconds()), DurationMs: int(a.Duration.Milliseconds()),
InputTokens: a.InputTokens, InputTokens: a.InputTokens,
OutputTokens: a.OutputTokens, OutputTokens: a.OutputTokens,
TotalTokens: a.TotalTokens, TotalTokens: a.TotalTokens,
CostUsd: a.CostUSD, CachedInputTokens: a.CachedInputTokens,
Stream: &a.Stream, CacheCreationTokens: a.CacheCreationTokens,
CostUsd: a.CostUSD,
CacheCostUsd: a.CacheCostUSD,
Stream: &a.Stream,
} }
out.UserId = strPtr(a.UserID) 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 // summary plus its ordered entries. Assembled in Go from a page of entries — it
// is not a stored table. // is not a stored table.
type AgentNetworkAccessLogSession struct { type AgentNetworkAccessLogSession struct {
SessionID string // empty for a session-less (singleton) request SessionID string // empty for a session-less (singleton) request
UserID string UserID string
GroupIDs []string // union of the entries' authorising groups GroupIDs []string // union of the entries' authorising groups
StartedAt time.Time StartedAt time.Time
EndedAt time.Time EndedAt time.Time
RequestCount int RequestCount int
InputTokens int64 InputTokens int64
OutputTokens int64 OutputTokens int64
TotalTokens int64 TotalTokens int64
CostUSD float64 CachedInputTokens int64
Providers []string // distinct vendors seen in the session CacheCreationTokens int64
Models []string // distinct models seen in the session CostUSD float64
Decision string // "deny" if any entry was denied, else "allow" CacheCostUSD float64
Entries []*AgentNetworkAccessLog 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 // 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.InputTokens += e.InputTokens
sess.OutputTokens += e.OutputTokens sess.OutputTokens += e.OutputTokens
sess.TotalTokens += e.TotalTokens sess.TotalTokens += e.TotalTokens
sess.CachedInputTokens += e.CachedInputTokens
sess.CacheCreationTokens += e.CacheCreationTokens
sess.CostUSD += e.CostUSD sess.CostUSD += e.CostUSD
sess.CacheCostUSD += e.CacheCostUSD
if e.Timestamp.Before(sess.StartedAt) { if e.Timestamp.Before(sess.StartedAt) {
sess.StartedAt = e.Timestamp sess.StartedAt = e.Timestamp
} }
@@ -248,15 +261,18 @@ func (sess *AgentNetworkAccessLogSession) ToAPIResponse() api.AgentNetworkAccess
} }
out := api.AgentNetworkAccessLogSession{ out := api.AgentNetworkAccessLogSession{
StartedAt: sess.StartedAt, StartedAt: sess.StartedAt,
EndedAt: sess.EndedAt, EndedAt: sess.EndedAt,
RequestCount: sess.RequestCount, RequestCount: sess.RequestCount,
InputTokens: sess.InputTokens, InputTokens: sess.InputTokens,
OutputTokens: sess.OutputTokens, OutputTokens: sess.OutputTokens,
TotalTokens: sess.TotalTokens, TotalTokens: sess.TotalTokens,
CostUsd: sess.CostUSD, CachedInputTokens: sess.CachedInputTokens,
Decision: sess.Decision, CacheCreationTokens: sess.CacheCreationTokens,
Entries: entries, CostUsd: sess.CostUSD,
CacheCostUsd: sess.CacheCostUSD,
Decision: sess.Decision,
Entries: entries,
} }
out.SessionId = strPtr(sess.SessionID) out.SessionId = strPtr(sess.SessionID)
out.UserId = strPtr(sess.UserID) out.UserId = strPtr(sess.UserID)
@@ -25,8 +25,12 @@ type AgentNetworkUsage struct {
InputTokens int64 InputTokens int64
OutputTokens int64 OutputTokens int64
TotalTokens int64 TotalTokens int64
CostUSD float64 // Prompt-cache buckets: read + write token counts and the portion of CostUSD they account for.
CreatedAt time.Time CachedInputTokens int64
CacheCreationTokens int64
CostUSD float64
CacheCostUSD float64
CreatedAt time.Time
} }
// TableName keeps usage records in their own stripped table. Named // TableName keeps usage records in their own stripped table. Named
@@ -33,21 +33,27 @@ func ParseUsageGranularity(s string) UsageGranularity {
// AgentNetworkUsageBucket is one aggregated usage time bucket. PeriodStart is // AgentNetworkUsageBucket is one aggregated usage time bucket. PeriodStart is
// the UTC start of the bucket as YYYY-MM-DD. // the UTC start of the bucket as YYYY-MM-DD.
type AgentNetworkUsageBucket struct { type AgentNetworkUsageBucket struct {
PeriodStart string PeriodStart string
InputTokens int64 InputTokens int64
OutputTokens int64 OutputTokens int64
TotalTokens int64 TotalTokens int64
CostUSD float64 CachedInputTokens int64
CacheCreationTokens int64
CostUSD float64
CacheCostUSD float64
} }
// ToAPIResponse renders the bucket as the API representation. // ToAPIResponse renders the bucket as the API representation.
func (b *AgentNetworkUsageBucket) ToAPIResponse() api.AgentNetworkUsageBucket { func (b *AgentNetworkUsageBucket) ToAPIResponse() api.AgentNetworkUsageBucket {
return api.AgentNetworkUsageBucket{ return api.AgentNetworkUsageBucket{
PeriodStart: b.PeriodStart, PeriodStart: b.PeriodStart,
InputTokens: b.InputTokens, InputTokens: b.InputTokens,
OutputTokens: b.OutputTokens, OutputTokens: b.OutputTokens,
TotalTokens: b.TotalTokens, TotalTokens: b.TotalTokens,
CostUsd: b.CostUSD, 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.InputTokens += r.InputTokens
b.OutputTokens += r.OutputTokens b.OutputTokens += r.OutputTokens
b.TotalTokens += r.TotalTokens b.TotalTokens += r.TotalTokens
b.CachedInputTokens += r.CachedInputTokens
b.CacheCreationTokens += r.CacheCreationTokens
b.CostUSD += r.CostUSD b.CostUSD += r.CostUSD
b.CacheCostUSD += r.CacheCostUSD
} }
out := make([]*AgentNetworkUsageBucket, 0, len(byPeriod)) out := make([]*AgentNetworkUsageBucket, 0, len(byPeriod))