From 71c185f6144ca185aff18f1486d36d1d3504bc1f Mon Sep 17 00:00:00 2001 From: Luis Pater Date: Thu, 28 May 2026 22:15:54 +0800 Subject: [PATCH] feat(usage): add service tier tracking and defaults in usage reporting - Introduced `service_tier` metadata key to capture client-requested service tiers. - Updated usage records, context propagation, and plugins to include service tier data. - Added default handling logic for cases where `service_tier` is absent. - Implemented tests for `service_tier` extraction, defaults, and updates across components. --- internal/redisqueue/plugin.go | 6 +++ internal/redisqueue/plugin_test.go | 2 + .../runtime/executor/helps/usage_helpers.go | 17 ++++++ .../executor/helps/usage_helpers_test.go | 33 ++++++++++++ sdk/api/handlers/handlers.go | 20 +++++++ sdk/api/handlers/handlers_metadata_test.go | 22 ++++++++ sdk/cliproxy/auth/conductor.go | 25 ++++++++- sdk/cliproxy/auth/conductor_usage_test.go | 5 ++ sdk/cliproxy/executor/types.go | 3 ++ sdk/cliproxy/usage/manager.go | 54 ++++++++++++++++--- 10 files changed, 180 insertions(+), 7 deletions(-) diff --git a/internal/redisqueue/plugin.go b/internal/redisqueue/plugin.go index ac48d0c1..f6c8e52c 100644 --- a/internal/redisqueue/plugin.go +++ b/internal/redisqueue/plugin.go @@ -52,6 +52,10 @@ 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) + } tokens := tokenStats{ InputTokens: record.Detail.InputTokens, @@ -97,6 +101,7 @@ func (p *usageQueuePlugin) HandleUsage(ctx context.Context, record coreusage.Rec APIKey: apiKey, RequestID: requestID, ReasoningEffort: reasoningEffort, + ServiceTier: serviceTier, }) if err != nil { return @@ -114,6 +119,7 @@ type queuedUsageDetail struct { APIKey string `json:"api_key"` RequestID string `json:"request_id"` ReasoningEffort string `json:"reasoning_effort"` + ServiceTier string `json:"service_tier"` } type requestDetail struct { diff --git a/internal/redisqueue/plugin_test.go b/internal/redisqueue/plugin_test.go index 4917955c..09ee681a 100644 --- a/internal/redisqueue/plugin_test.go +++ b/internal/redisqueue/plugin_test.go @@ -33,6 +33,7 @@ func TestUsageQueuePluginPayloadIncludesStableFieldsAndSuccess(t *testing.T) { 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, Detail: coreusage.Detail{ @@ -53,6 +54,7 @@ func TestUsageQueuePluginPayloadIncludesStableFieldsAndSuccess(t *testing.T) { requireMissingField(t, payload, "user_api_key") requireStringField(t, payload, "request_id", "ctx-request-id") requireStringField(t, payload, "reasoning_effort", "medium") + requireStringField(t, payload, "service_tier", "priority") 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 295b797d..10c4108c 100644 --- a/internal/runtime/executor/helps/usage_helpers.go +++ b/internal/runtime/executor/helps/usage_helpers.go @@ -30,6 +30,7 @@ type UsageReporter struct { apiKey string source string reasoning string + serviceTier string requestedAt time.Time ttftMu sync.RWMutex ttft time.Duration @@ -53,6 +54,7 @@ func NewUsageReporter(ctx context.Context, provider, model string, auth *cliprox source: resolveUsageSource(auth, apiKey), authType: resolveUsageAuthType(auth), reasoning: usage.ReasoningEffortFromContext(ctx), + serviceTier: usage.ServiceTierFromContext(ctx), } if auth != nil { reporter.authID = auth.ID @@ -78,6 +80,7 @@ func (r *UsageReporter) SetTranslatedReasoningEffort(payload []byte, format stri return } r.reasoning = thinking.ExtractTranslatedReasoningEffort(payload, format) + r.serviceTier = extractServiceTierFromPayload(payload) } func (r *UsageReporter) TrackHTTPClient(client *http.Client) *http.Client { @@ -239,6 +242,7 @@ func (r *UsageReporter) buildRecordForModel(model string, detail usage.Detail, f AuthIndex: r.authIndex, AuthType: r.authType, ReasoningEffort: r.reasoning, + ServiceTier: r.serviceTier, RequestedAt: r.requestedAt, Latency: r.latency(), TTFT: r.ttftDuration(), @@ -248,6 +252,19 @@ func (r *UsageReporter) buildRecordForModel(model string, detail usage.Detail, f } } +func extractServiceTierFromPayload(payload []byte) string { + if len(payload) == 0 { + return usage.DefaultServiceTier + } + for _, path := range []string{"service_tier", "request.service_tier", "response.service_tier"} { + serviceTier := strings.TrimSpace(gjson.GetBytes(payload, path).String()) + if serviceTier != "" { + return serviceTier + } + } + return usage.DefaultServiceTier +} + func failFromErrors(errs ...error) usage.Failure { for _, err := range errs { if err == nil { diff --git a/internal/runtime/executor/helps/usage_helpers_test.go b/internal/runtime/executor/helps/usage_helpers_test.go index b14a389a..483d8ef5 100644 --- a/internal/runtime/executor/helps/usage_helpers_test.go +++ b/internal/runtime/executor/helps/usage_helpers_test.go @@ -241,6 +241,39 @@ func TestUsageReporterBuildRecordIncludesReasoningEffort(t *testing.T) { } } +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) + if record.ServiceTier != "priority" { + t.Fatalf("service tier = %q, want %q", record.ServiceTier, "priority") + } +} + +func TestUsageReporterSetTranslatedReasoningEffortUpdatesServiceTier(t *testing.T) { + reporter := NewUsageReporter(context.Background(), "openai", "gpt-5.4", nil) + + reporter.SetTranslatedReasoningEffort([]byte(`{"service_tier":"priority"}`), "openai") + + record := reporter.buildRecord(usage.Detail{TotalTokens: 3}, false) + if record.ServiceTier != "priority" { + t.Fatalf("service tier = %q, want %q", record.ServiceTier, "priority") + } +} + +func TestUsageReporterSetTranslatedReasoningEffortDefaultsServiceTierWhenRemoved(t *testing.T) { + ctx := usage.WithServiceTier(context.Background(), "priority") + reporter := NewUsageReporter(ctx, "openai", "gpt-5.4", nil) + + reporter.SetTranslatedReasoningEffort([]byte(`{"model":"gpt-5.4"}`), "openai") + + record := reporter.buildRecord(usage.Detail{TotalTokens: 3}, false) + if record.ServiceTier != usage.DefaultServiceTier { + t.Fatalf("service tier = %q, want %q", record.ServiceTier, usage.DefaultServiceTier) + } +} + func TestUsageReporterBuildAdditionalModelRecordSkipsZeroTokens(t *testing.T) { reporter := &UsageReporter{ provider: "codex", diff --git a/sdk/api/handlers/handlers.go b/sdk/api/handlers/handlers.go index 5a25681d..55b4d6ab 100644 --- a/sdk/api/handlers/handlers.go +++ b/sdk/api/handlers/handlers.go @@ -20,8 +20,10 @@ import ( "github.com/router-for-me/CLIProxyAPI/v7/internal/util" coreauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth" coreexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor" + coreusage "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/usage" "github.com/router-for-me/CLIProxyAPI/v7/sdk/config" sdktranslator "github.com/router-for-me/CLIProxyAPI/v7/sdk/translator" + "github.com/tidwall/gjson" "golang.org/x/net/context" ) @@ -242,6 +244,21 @@ func setReasoningEffortMetadata(meta map[string]any, handlerType, model string, meta[coreexecutor.ReasoningEffortMetadataKey] = effort } +func setServiceTierMetadata(meta map[string]any, rawJSON []byte) { + if meta == nil { + return + } + serviceTier := coreusage.DefaultServiceTier + node := gjson.GetBytes(rawJSON, "service_tier") + if node.Exists() { + value := strings.TrimSpace(node.String()) + if value != "" { + serviceTier = value + } + } + meta[coreexecutor.ServiceTierMetadataKey] = serviceTier +} + // headersFromContext extracts the original HTTP request headers from the gin context // embedded in the provided context. This allows session affinity selectors to read // client headers like X-Amp-Thread-Id. @@ -562,6 +579,7 @@ func (h *BaseAPIHandler) executeWithAuthManager(ctx context.Context, handlerType reqMeta := requestExecutionMetadata(ctx) reqMeta[coreexecutor.RequestedModelMetadataKey] = modelName setReasoningEffortMetadata(reqMeta, handlerType, normalizedModel, rawJSON) + setServiceTierMetadata(reqMeta, rawJSON) payload := rawJSON if len(payload) == 0 { payload = nil @@ -611,6 +629,7 @@ func (h *BaseAPIHandler) ExecuteCountWithAuthManager(ctx context.Context, handle reqMeta := requestExecutionMetadata(ctx) reqMeta[coreexecutor.RequestedModelMetadataKey] = modelName setReasoningEffortMetadata(reqMeta, handlerType, normalizedModel, rawJSON) + setServiceTierMetadata(reqMeta, rawJSON) payload := rawJSON if len(payload) == 0 { payload = nil @@ -673,6 +692,7 @@ func (h *BaseAPIHandler) executeStreamWithAuthManager(ctx context.Context, handl reqMeta := requestExecutionMetadata(ctx) reqMeta[coreexecutor.RequestedModelMetadataKey] = modelName setReasoningEffortMetadata(reqMeta, handlerType, normalizedModel, rawJSON) + setServiceTierMetadata(reqMeta, rawJSON) payload := rawJSON if len(payload) == 0 { payload = nil diff --git a/sdk/api/handlers/handlers_metadata_test.go b/sdk/api/handlers/handlers_metadata_test.go index d2bdab68..24a9130f 100644 --- a/sdk/api/handlers/handlers_metadata_test.go +++ b/sdk/api/handlers/handlers_metadata_test.go @@ -38,3 +38,25 @@ func TestSetReasoningEffortMetadataSupportsOpenAIResponses(t *testing.T) { t.Fatalf("ReasoningEffortMetadataKey = %v, want %q", got, "medium") } } + +func TestSetServiceTierMetadataExtractsValue(t *testing.T) { + meta := make(map[string]any) + + setServiceTierMetadata(meta, []byte(`{"service_tier":"priority"}`)) + + gotServiceTier := meta[coreexecutor.ServiceTierMetadataKey] + if gotServiceTier != "priority" { + t.Fatalf("ServiceTierMetadataKey = %v, want %q", gotServiceTier, "priority") + } +} + +func TestSetServiceTierMetadataDefaultsWhenMissing(t *testing.T) { + meta := make(map[string]any) + + setServiceTierMetadata(meta, []byte(`{"model":"gpt-5.4"}`)) + + gotServiceTier := meta[coreexecutor.ServiceTierMetadataKey] + if gotServiceTier != "default" { + t.Fatalf("ServiceTierMetadataKey = %v, want %q", gotServiceTier, "default") + } +} diff --git a/sdk/cliproxy/auth/conductor.go b/sdk/cliproxy/auth/conductor.go index ac1a9298..5413dcf4 100644 --- a/sdk/cliproxy/auth/conductor.go +++ b/sdk/cliproxy/auth/conductor.go @@ -1731,9 +1731,14 @@ func (m *Manager) prepareRequestAuth(ctx context.Context, executor ProviderExecu func contextWithRequestedModelAlias(ctx context.Context, opts cliproxyexecutor.Options, fallback string) context.Context { alias := requestedModelAliasFromOptions(opts, fallback) ctx = coreusage.WithRequestedModelAlias(ctx, alias) - if effort := reasoningEffortFromOptions(opts); effort != "" { + effort := reasoningEffortFromOptions(opts) + if effort != "" { ctx = coreusage.WithReasoningEffort(ctx, effort) } + serviceTier := serviceTierFromOptions(opts) + if serviceTier != "" { + ctx = coreusage.WithServiceTier(ctx, serviceTier) + } return ctx } @@ -1780,6 +1785,24 @@ func reasoningEffortFromOptions(opts cliproxyexecutor.Options) string { } } +func serviceTierFromOptions(opts cliproxyexecutor.Options) string { + if len(opts.Metadata) == 0 { + return "" + } + raw, ok := opts.Metadata[cliproxyexecutor.ServiceTierMetadataKey] + if !ok || raw == nil { + return "" + } + switch value := raw.(type) { + case string: + return strings.TrimSpace(value) + case []byte: + return strings.TrimSpace(string(value)) + default: + return "" + } +} + func pinnedAuthIDFromMetadata(meta map[string]any) string { if len(meta) == 0 { return "" diff --git a/sdk/cliproxy/auth/conductor_usage_test.go b/sdk/cliproxy/auth/conductor_usage_test.go index 23a70ea2..af6c1ee2 100644 --- a/sdk/cliproxy/auth/conductor_usage_test.go +++ b/sdk/cliproxy/auth/conductor_usage_test.go @@ -13,6 +13,7 @@ func TestContextWithRequestedModelAliasIncludesReasoningEffort(t *testing.T) { Metadata: map[string]any{ cliproxyexecutor.RequestedModelMetadataKey: "client-model", cliproxyexecutor.ReasoningEffortMetadataKey: "medium", + cliproxyexecutor.ServiceTierMetadataKey: "priority", }, }, "fallback-model") @@ -22,4 +23,8 @@ func TestContextWithRequestedModelAliasIncludesReasoningEffort(t *testing.T) { if got := coreusage.ReasoningEffortFromContext(ctx); got != "medium" { t.Fatalf("reasoning effort = %q, want %q", got, "medium") } + gotServiceTier := coreusage.ServiceTierFromContext(ctx) + if gotServiceTier != "priority" { + t.Fatalf("service tier = %q, want %q", gotServiceTier, "priority") + } } diff --git a/sdk/cliproxy/executor/types.go b/sdk/cliproxy/executor/types.go index fc003540..8f0fc567 100644 --- a/sdk/cliproxy/executor/types.go +++ b/sdk/cliproxy/executor/types.go @@ -20,6 +20,9 @@ const DisallowFreeAuthMetadataKey = "disallow_free_auth" // ReasoningEffortMetadataKey stores the client-requested reasoning effort for usage logs. const ReasoningEffortMetadataKey = "reasoning_effort" +// ServiceTierMetadataKey stores the client-requested service tier for usage logs. +const ServiceTierMetadataKey = "service_tier" + const ( // PinnedAuthMetadataKey locks execution to a specific auth ID. PinnedAuthMetadataKey = "pinned_auth_id" diff --git a/sdk/cliproxy/usage/manager.go b/sdk/cliproxy/usage/manager.go index 6113ca1e..6c113b12 100644 --- a/sdk/cliproxy/usage/manager.go +++ b/sdk/cliproxy/usage/manager.go @@ -10,6 +10,9 @@ import ( log "github.com/sirupsen/logrus" ) +// DefaultServiceTier is used when a request does not specify service_tier. +const DefaultServiceTier = "default" + // Record contains the usage statistics captured for a single provider request. type Record struct { Provider string @@ -22,12 +25,14 @@ type Record struct { Source string // ReasoningEffort stores the translated upstream thinking level for request event logs. ReasoningEffort string - RequestedAt time.Time - Latency time.Duration - TTFT time.Duration - Failed bool - Fail Failure - Detail Detail + // 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 // ResponseHeaders stores a snapshot of upstream response headers for usage sinks. ResponseHeaders http.Header } @@ -51,6 +56,7 @@ type Detail struct { type requestedModelAliasContextKey struct{} type reasoningEffortContextKey struct{} +type serviceTierContextKey struct{} // WithRequestedModelAlias stores the client-requested model name for usage sinks. func WithRequestedModelAlias(ctx context.Context, alias string) context.Context { @@ -108,6 +114,42 @@ func ReasoningEffortFromContext(ctx context.Context) string { } } +// WithServiceTier stores the client-requested service tier for usage sinks. +func WithServiceTier(ctx context.Context, tier string) context.Context { + if ctx == nil { + ctx = context.Background() + } + tier = strings.TrimSpace(tier) + if tier == "" { + tier = DefaultServiceTier + } + return context.WithValue(ctx, serviceTierContextKey{}, tier) +} + +// ServiceTierFromContext returns the client-requested service tier stored in ctx. +func ServiceTierFromContext(ctx context.Context) string { + if ctx == nil { + return DefaultServiceTier + } + raw := ctx.Value(serviceTierContextKey{}) + switch value := raw.(type) { + case string: + tier := strings.TrimSpace(value) + if tier == "" { + return DefaultServiceTier + } + return tier + case []byte: + tier := strings.TrimSpace(string(value)) + if tier == "" { + return DefaultServiceTier + } + return tier + default: + return DefaultServiceTier + } +} + // Plugin consumes usage records emitted by the proxy runtime. type Plugin interface { HandleUsage(ctx context.Context, record Record) -- 2.51.2