feat(usage): report request and response service tiers

This commit is contained in:
Supra4E8C
2026-07-11 00:49:30 +08:00
parent 3554b63721
commit dc4be1674a
5 changed files with 123 additions and 68 deletions

View File

@@ -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 {

View File

@@ -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)

View File

@@ -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) {

View File

@@ -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) {

View File

@@ -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{}