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.
This commit is contained in:
Luis Pater
2026-05-28 22:15:54 +08:00
parent 65e760aa1a
commit 71c185f614
10 changed files with 180 additions and 7 deletions
+6
View File
@@ -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 {
+2
View File
@@ -33,6 +33,7 @@ func TestUsageQueuePluginPayloadIncludesStableFieldsAndSuccess(t *testing.T) {
AuthType: "apikey",
Source: "[email protected]",
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)
@@ -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 {
@@ -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",
+20
View File
@@ -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
@@ -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")
}
}
+24 -1
View File
@@ -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 ""
@@ -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")
}
}
+3
View File
@@ -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"
+48 -6
View File
@@ -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)