diff --git a/internal/redisqueue/plugin.go b/internal/redisqueue/plugin.go index 029dd13f1..1ade177e9 100644 --- a/internal/redisqueue/plugin.go +++ b/internal/redisqueue/plugin.go @@ -56,10 +56,14 @@ func (p *usageQueuePlugin) HandleUsage(ctx context.Context, record coreusage.Rec if reasoningEffort == "" { reasoningEffort = coreusage.ReasoningEffortFromContext(ctx) } - serviceTier := strings.TrimSpace(record.ServiceTier) - if serviceTier == "" { - serviceTier = coreusage.ServiceTierFromContext(ctx) + requestServiceTier := strings.TrimSpace(record.RequestServiceTier) + if requestServiceTier == "" { + requestServiceTier = strings.TrimSpace(record.ServiceTier) } + if requestServiceTier == "" { + requestServiceTier = coreusage.ServiceTierFromContext(ctx) + } + responseServiceTier := strings.TrimSpace(record.ResponseServiceTier) tokens := tokenStats{ InputTokens: record.Detail.InputTokens, @@ -96,17 +100,19 @@ func (p *usageQueuePlugin) HandleUsage(ctx context.Context, record coreusage.Rec } payload, err := json.Marshal(queuedUsageDetail{ - requestDetail: detail, - Provider: provider, - ExecutorType: executorType, - Model: modelName, - Alias: aliasName, - Endpoint: resolveEndpoint(ctx), - AuthType: authType, - APIKey: apiKey, - RequestID: requestID, - ReasoningEffort: reasoningEffort, - ServiceTier: serviceTier, + requestDetail: detail, + Provider: provider, + ExecutorType: executorType, + Model: modelName, + Alias: aliasName, + Endpoint: resolveEndpoint(ctx), + AuthType: authType, + APIKey: apiKey, + RequestID: requestID, + ReasoningEffort: reasoningEffort, + ServiceTier: requestServiceTier, + RequestServiceTier: requestServiceTier, + ResponseServiceTier: responseServiceTier, }) if err != nil { return @@ -116,16 +122,18 @@ func (p *usageQueuePlugin) HandleUsage(ctx context.Context, record coreusage.Rec type queuedUsageDetail struct { requestDetail - Provider string `json:"provider"` - ExecutorType string `json:"executor_type"` - Model string `json:"model"` - Alias string `json:"alias"` - Endpoint string `json:"endpoint"` - AuthType string `json:"auth_type"` - APIKey string `json:"api_key"` - RequestID string `json:"request_id"` - ReasoningEffort string `json:"reasoning_effort"` - ServiceTier string `json:"service_tier"` + Provider string `json:"provider"` + ExecutorType string `json:"executor_type"` + Model string `json:"model"` + Alias string `json:"alias"` + Endpoint string `json:"endpoint"` + AuthType string `json:"auth_type"` + APIKey string `json:"api_key"` + RequestID string `json:"request_id"` + ReasoningEffort string `json:"reasoning_effort"` + ServiceTier string `json:"service_tier"` + RequestServiceTier string `json:"request_service_tier"` + ResponseServiceTier string `json:"response_service_tier,omitempty"` } type requestDetail struct { diff --git a/internal/redisqueue/plugin_test.go b/internal/redisqueue/plugin_test.go index 16c0a270a..8735c552a 100644 --- a/internal/redisqueue/plugin_test.go +++ b/internal/redisqueue/plugin_test.go @@ -25,18 +25,19 @@ func TestUsageQueuePluginPayloadIncludesStableFieldsAndSuccess(t *testing.T) { plugin := &usageQueuePlugin{} plugin.HandleUsage(ctx, coreusage.Record{ - Provider: "openai", - ExecutorType: "KimiExecutor", - Model: "gpt-5.4", - Alias: "client-gpt", - APIKey: "test-key", - AuthIndex: "0", - AuthType: "apikey", - Source: "user@example.com", - ReasoningEffort: "medium", - ServiceTier: "priority", - RequestedAt: time.Date(2026, 4, 25, 0, 0, 0, 0, time.UTC), - Latency: 1500 * time.Millisecond, + Provider: "openai", + ExecutorType: "KimiExecutor", + Model: "gpt-5.4", + Alias: "client-gpt", + APIKey: "test-key", + AuthIndex: "0", + AuthType: "apikey", + Source: "user@example.com", + ReasoningEffort: "medium", + ServiceTier: "priority", + ResponseServiceTier: "default", + RequestedAt: time.Date(2026, 4, 25, 0, 0, 0, 0, time.UTC), + Latency: 1500 * time.Millisecond, Detail: coreusage.Detail{ InputTokens: 10, OutputTokens: 20, @@ -57,6 +58,8 @@ func TestUsageQueuePluginPayloadIncludesStableFieldsAndSuccess(t *testing.T) { requireStringField(t, payload, "request_id", "ctx-request-id") requireStringField(t, payload, "reasoning_effort", "medium") requireStringField(t, payload, "service_tier", "priority") + requireStringField(t, payload, "request_service_tier", "priority") + requireStringField(t, payload, "response_service_tier", "default") requireHeaderField(t, payload, "response_headers", "X-Upstream-Request-Id", []string{"upstream-req-1"}) requireHeaderField(t, payload, "response_headers", "Retry-After", []string{"30"}) requireBoolField(t, payload, "failed", false) diff --git a/internal/runtime/executor/helps/usage_helpers.go b/internal/runtime/executor/helps/usage_helpers.go index afa323aa2..6234aa1ed 100644 --- a/internal/runtime/executor/helps/usage_helpers.go +++ b/internal/runtime/executor/helps/usage_helpers.go @@ -260,23 +260,25 @@ func (r *UsageReporter) buildRecordForModel(model string, detail usage.Detail, f return usage.Record{Model: model, Detail: detail, Failed: failed, Fail: fail} } return usage.Record{ - Provider: r.provider, - ExecutorType: r.executorType, - Model: model, - Alias: r.alias, - Source: r.source, - APIKey: r.apiKey, - AuthID: r.authID, - AuthIndex: r.authIndex, - AuthType: r.authType, - ReasoningEffort: r.reasoning, - ServiceTier: r.serviceTier, - RequestedAt: r.requestedAt, - Latency: r.latency(), - TTFT: r.ttftDuration(), - Failed: failed, - Fail: fail, - Detail: detail, + Provider: r.provider, + ExecutorType: r.executorType, + Model: model, + Alias: r.alias, + Source: r.source, + APIKey: r.apiKey, + AuthID: r.authID, + AuthIndex: r.authIndex, + AuthType: r.authType, + ReasoningEffort: r.reasoning, + ServiceTier: r.serviceTier, + RequestServiceTier: r.serviceTier, + ResponseServiceTier: strings.TrimSpace(detail.ResponseServiceTier), + RequestedAt: r.requestedAt, + Latency: r.latency(), + TTFT: r.ttftDuration(), + Failed: failed, + Fail: fail, + Detail: detail, } } @@ -484,7 +486,9 @@ func ParseCodexUsage(data []byte) (usage.Detail, bool) { if !hasOpenAIStyleUsageTokenFields(usageNode) { return usage.Detail{}, false } - return parseOpenAIStyleUsageNode(usageNode), true + detail := parseOpenAIStyleUsageNode(usageNode) + detail.ResponseServiceTier = extractResponseServiceTier(data) + return detail, true } func ParseCodexImageToolUsage(data []byte) (usage.Detail, bool) { @@ -500,7 +504,9 @@ func ParseOpenAIUsage(data []byte) usage.Detail { if !hasOpenAIStyleUsageTokenFields(usageNode) { return usage.Detail{} } - return parseOpenAIStyleUsageNode(usageNode) + detail := parseOpenAIStyleUsageNode(usageNode) + detail.ResponseServiceTier = extractResponseServiceTier(data) + return detail } func hasOpenAIStyleUsageTokenFields(usageNode gjson.Result) bool { @@ -562,7 +568,9 @@ func ParseOpenAIStreamUsage(line []byte) (usage.Detail, bool) { if !hasOpenAIStyleUsageTokenFields(usageNode) { return usage.Detail{}, false } - return parseOpenAIStyleUsageNode(usageNode), true + detail := parseOpenAIStyleUsageNode(usageNode) + detail.ResponseServiceTier = extractResponseServiceTier(payload) + return detail, true } func ParseClaudeUsage(data []byte) usage.Detail { @@ -643,9 +651,25 @@ func ParseInteractionsUsage(data []byte) usage.Detail { return usage.Detail{} } if node.Get("promptTokenCount").Exists() || node.Get("candidatesTokenCount").Exists() { - return parseGeminiFamilyUsageDetail(node) + detail := parseGeminiFamilyUsageDetail(node) + detail.ResponseServiceTier = extractResponseServiceTier(data) + return detail } - return parseInteractionsUsageDetail(node) + detail := parseInteractionsUsageDetail(node) + detail.ResponseServiceTier = extractResponseServiceTier(data) + return detail +} + +func extractResponseServiceTier(payload []byte) string { + if len(payload) == 0 || !gjson.ValidBytes(payload) { + return "" + } + for _, path := range []string{"response.service_tier", "service_tier", "interaction.service_tier"} { + if tier := strings.TrimSpace(gjson.GetBytes(payload, path).String()); tier != "" { + return tier + } + } + return "" } func ParseInteractionsStreamUsage(line []byte) (usage.Detail, bool) { diff --git a/internal/runtime/executor/helps/usage_helpers_test.go b/internal/runtime/executor/helps/usage_helpers_test.go index 1730cfd67..b7107752f 100644 --- a/internal/runtime/executor/helps/usage_helpers_test.go +++ b/internal/runtime/executor/helps/usage_helpers_test.go @@ -32,7 +32,7 @@ func TestParseOpenAIUsageChatCompletions(t *testing.T) { } func TestParseOpenAIUsageResponses(t *testing.T) { - data := []byte(`{"usage":{"input_tokens":10,"output_tokens":20,"total_tokens":30,"input_tokens_details":{"cached_tokens":7},"output_tokens_details":{"reasoning_tokens":9}}}`) + data := []byte(`{"service_tier":"default","usage":{"input_tokens":10,"output_tokens":20,"total_tokens":30,"input_tokens_details":{"cached_tokens":7},"output_tokens_details":{"reasoning_tokens":9}}}`) detail := ParseOpenAIUsage(data) if detail.InputTokens != 10 { t.Fatalf("input tokens = %d, want %d", detail.InputTokens, 10) @@ -49,10 +49,13 @@ func TestParseOpenAIUsageResponses(t *testing.T) { if detail.ReasoningTokens != 9 { t.Fatalf("reasoning tokens = %d, want %d", detail.ReasoningTokens, 9) } + if detail.ResponseServiceTier != "default" { + t.Fatalf("response service tier = %q, want default", detail.ResponseServiceTier) + } } func TestParseCodexUsageIncludesCacheWriteTokens(t *testing.T) { - data := []byte(`{"response":{"usage":{"input_tokens":100,"output_tokens":20,"total_tokens":120,"input_tokens_details":{"cached_tokens":30,"cache_write_tokens":40}}}}`) + data := []byte(`{"response":{"service_tier":"priority","usage":{"input_tokens":100,"output_tokens":20,"total_tokens":120,"input_tokens_details":{"cached_tokens":30,"cache_write_tokens":40}}}}`) detail, ok := ParseCodexUsage(data) if !ok { t.Fatal("ParseCodexUsage() ok = false, want true") @@ -72,6 +75,9 @@ func TestParseCodexUsageIncludesCacheWriteTokens(t *testing.T) { if detail.TotalTokens != 120 { t.Fatalf("total tokens = %d, want 120", detail.TotalTokens) } + if detail.ResponseServiceTier != "priority" { + t.Fatalf("response service tier = %q, want priority", detail.ResponseServiceTier) + } } func TestParseOpenAIUsageIgnoresNullUsage(t *testing.T) { @@ -90,7 +96,7 @@ func TestParseOpenAIStreamUsageIgnoresNullUsage(t *testing.T) { } func TestParseOpenAIStreamUsageResponsesFields(t *testing.T) { - line := []byte(`data: {"id":"chunk_1","object":"chat.completion.chunk","choices":[],"usage":{"input_tokens":8,"output_tokens":5,"total_tokens":13,"input_tokens_details":{"cached_tokens":3},"output_tokens_details":{"reasoning_tokens":2}}}`) + line := []byte(`data: {"id":"chunk_1","object":"chat.completion.chunk","service_tier":"flex","choices":[],"usage":{"input_tokens":8,"output_tokens":5,"total_tokens":13,"input_tokens_details":{"cached_tokens":3},"output_tokens_details":{"reasoning_tokens":2}}}`) detail, ok := ParseOpenAIStreamUsage(line) if !ok { t.Fatal("ParseOpenAIStreamUsage() ok = false, want true") @@ -110,6 +116,9 @@ func TestParseOpenAIStreamUsageResponsesFields(t *testing.T) { if detail.ReasoningTokens != 2 { t.Fatalf("reasoning tokens = %d, want %d", detail.ReasoningTokens, 2) } + if detail.ResponseServiceTier != "flex" { + t.Fatalf("response service tier = %q, want flex", detail.ResponseServiceTier) + } } func TestStreamUsageBufferKeepsLastUsage(t *testing.T) { @@ -324,10 +333,16 @@ func TestUsageReporterBuildRecordIncludesServiceTier(t *testing.T) { ctx := usage.WithServiceTier(context.Background(), "priority") reporter := NewUsageReporter(ctx, "openai", "gpt-5.4", nil) - record := reporter.buildRecord(usage.Detail{TotalTokens: 3}, false) + record := reporter.buildRecord(usage.Detail{TotalTokens: 3, ResponseServiceTier: "default"}, false) if record.ServiceTier != "priority" { t.Fatalf("service tier = %q, want %q", record.ServiceTier, "priority") } + if record.RequestServiceTier != "priority" { + t.Fatalf("request service tier = %q, want priority", record.RequestServiceTier) + } + if record.ResponseServiceTier != "default" { + t.Fatalf("response service tier = %q, want default", record.ResponseServiceTier) + } } func TestUsageReporterSetTranslatedReasoningEffortUpdatesServiceTier(t *testing.T) { diff --git a/sdk/cliproxy/usage/manager.go b/sdk/cliproxy/usage/manager.go index b7798dc29..5e84344df 100644 --- a/sdk/cliproxy/usage/manager.go +++ b/sdk/cliproxy/usage/manager.go @@ -29,12 +29,16 @@ type Record struct { ReasoningEffort string // ServiceTier stores the client-requested service tier for request event logs. ServiceTier string - RequestedAt time.Time - Latency time.Duration - TTFT time.Duration - Failed bool - Fail Failure - Detail Detail + // RequestServiceTier explicitly aliases the client-requested service tier. + RequestServiceTier string + // ResponseServiceTier stores the final tier reported by the upstream response. + ResponseServiceTier string + RequestedAt time.Time + Latency time.Duration + TTFT time.Duration + Failed bool + Fail Failure + Detail Detail // ResponseHeaders stores a snapshot of upstream response headers for usage sinks. ResponseHeaders http.Header } @@ -54,6 +58,7 @@ type Detail struct { CacheReadTokens int64 CacheCreationTokens int64 TotalTokens int64 + ResponseServiceTier string } type requestedModelAliasContextKey struct{}