package redisqueue import ( "context" "encoding/json" "net/http" "net/http/httptest" "testing" "time" "github.com/gin-gonic/gin" internallogging "github.com/router-for-me/CLIProxyAPI/v7/internal/logging" coresession "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/session" coreusage "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/usage") func TestUsageQueuePluginPayloadIncludesStableFieldsAndSuccess(t *testing.T) { withEnabledQueue(t, func() { ctx := internallogging.WithRequestID(context.Background(), "ctx-request-id") ctx = internallogging.WithEndpoint(ctx, "POST /v1/chat/completions") ctx = internallogging.WithClientRequestMetadata(ctx, internallogging.ClientRequestMetadata{ ClientIP: "192.0.2.10", XForwardedFor: "203.0.113.5, 198.51.100.8", UserAgent: "test-client/1.0", }) ctx = internallogging.WithResponseStatusHolder(ctx) internallogging.SetResponseStatus(ctx, http.StatusOK) responseHeaders := http.Header{} responseHeaders.Add("X-Upstream-Request-Id", "upstream-req-1") responseHeaders.Add("Retry-After", "30") plugin := &usageQueuePlugin{} plugin.HandleUsage(ctx, coreusage.Record{ Provider: "openai", ExecutorType: "KimiExecutor", Model: "gpt-5.4", Alias: "client-gpt", APIKey: "test-key", AuthIndex: "0", AccessTokenSHA256: "token-version-hash", AuthType: "apikey", Source: "user@example.com", ReasoningEffort: "medium", ServiceTier: "auto", ResponseServiceTier: "default", Generate: coreusage.GenerateFlag(true), RequestedAt: time.Date(2026, 4, 25, 0, 0, 0, 0, time.UTC), Latency: 1500 * time.Millisecond, Detail: coreusage.Detail{ InputTokens: 10, OutputTokens: 20, TotalTokens: 30, }, ResponseHeaders: responseHeaders.Clone(), }) responseHeaders.Set("Retry-After", "999") payload := popSinglePayload(t) requireStringField(t, payload, "provider", "openai") requireStringField(t, payload, "executor_type", "KimiExecutor") requireStringField(t, payload, "model", "gpt-5.4") requireStringField(t, payload, "alias", "client-gpt") requireStringField(t, payload, "endpoint", "POST /v1/chat/completions") requireStringField(t, payload, "auth_type", "apikey") requireStringField(t, payload, "access_token_sha256", "token-version-hash") requireMissingField(t, payload, "user_api_key") requireStringField(t, payload, "request_id", "ctx-request-id") requireStringField(t, payload, "client_ip", "192.0.2.10") requireStringField(t, payload, "x_forwarded_for", "203.0.113.5, 198.51.100.8") requireStringField(t, payload, "user_agent", "test-client/1.0") requireStringField(t, payload, "reasoning_effort", "medium") requireStringField(t, payload, "service_tier", "auto") requireMissingField(t, payload, "request_service_tier") requireStringField(t, payload, "response_service_tier", "default") requireIntField(t, payload, "accounting_version", coreusage.TokenAccountingSchemaVersion) requireTokenBreakdown(t, payload, coreusage.TokenAccountingQualityComplete, 30) requireTokensBoolField(t, payload, "cache_read_tokens_present", true) 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) requireBoolField(t, payload, "generate", true) requireBoolField(t, payload, "stream", false) requireFailField(t, payload, http.StatusOK, "") })} func TestUsageQueuePluginNormalizesDirectSDKUsageByProvider(t *testing.T) { tests := []struct { provider string wantTotal int }{ {provider: "openai", wantTotal: 130}, {provider: "gemini", wantTotal: 142}, } for _, tt := range tests { t.Run(tt.provider, func(t *testing.T) { withEnabledQueue(t, func() { ctx := internallogging.WithResponseStatusHolder(context.Background()) internallogging.SetResponseStatus(ctx, http.StatusOK) (&usageQueuePlugin{}).HandleUsage(ctx, coreusage.Record{ Provider: tt.provider, Model: "direct-sdk-model", Detail: coreusage.Detail{ InputTokens: 100, OutputTokens: 30, ReasoningTokens: 12, }, }) payload := popSinglePayload(t) requireIntField(t, requireTokensPayload(t, payload), "total_tokens", tt.wantTotal) requireTokenBreakdown(t, payload, coreusage.TokenAccountingQualityComplete, int64(tt.wantTotal)) }) }) }} func TestUsageQueuePluginPayloadIncludesGenerateFalse(t *testing.T) { withEnabledQueue(t, func() { ctx := internallogging.WithResponseStatusHolder(context.Background()) internallogging.SetResponseStatus(ctx, http.StatusOK) (&usageQueuePlugin{}).HandleUsage(ctx, coreusage.Record{ Provider: "openai", Model: "gpt-5.4", Generate: coreusage.GenerateFlag(false), Detail: coreusage.Detail{ InputTokens: 1, TotalTokens: 1, }, }) payload := popSinglePayload(t) requireBoolField(t, payload, "generate", false) })} func TestUsageQueuePluginPayloadDefaultsGenerateTrueWhenOmitted(t *testing.T) { withEnabledQueue(t, func() { ctx := internallogging.WithResponseStatusHolder(context.Background()) internallogging.SetResponseStatus(ctx, http.StatusOK) // Legacy callers construct usage.Record without Generate; omission must publish as true. (&usageQueuePlugin{}).HandleUsage(ctx, coreusage.Record{ Provider: "openai", Model: "gpt-5.4", Detail: coreusage.Detail{ InputTokens: 1, TotalTokens: 1, }, }) payload := popSinglePayload(t) requireBoolField(t, payload, "generate", true) })} func TestUsageQueuePluginPublishesStreamFlag(t *testing.T) { for _, stream := range []bool{true, false} { t.Run(map[bool]string{true: "stream_true", false: "stream_false"}[stream], func(t *testing.T) { withEnabledQueue(t, func() { ctx := internallogging.WithResponseStatusHolder(context.Background()) internallogging.SetResponseStatus(ctx, http.StatusOK) (&usageQueuePlugin{}).HandleUsage(ctx, coreusage.Record{ Provider: "openai", Model: "gpt-5.4", Stream: stream, Detail: coreusage.Detail{ InputTokens: 1, TotalTokens: 1, }, }) payload := popSinglePayload(t) requireBoolField(t, payload, "stream", stream) }) }) }} func TestUsageQueuePluginPreservesLegacyCachedOnlyUsage(t *testing.T) { withEnabledQueue(t, func() { ctx := internallogging.WithResponseStatusHolder(context.Background()) internallogging.SetResponseStatus(ctx, http.StatusOK) (&usageQueuePlugin{}).HandleUsage(ctx, coreusage.Record{ Provider: "openai", Model: "gpt-5.4", Detail: coreusage.Detail{ CachedTokens: 13, }, }) payload := popSinglePayload(t) requireTokensBoolField(t, payload, "cache_read_tokens_present", true) tokens := requireTokensPayload(t, payload) requireIntField(t, tokens, "cache_read_tokens", 13) requireIntField(t, tokens, "total_tokens", 13) requireTokenBreakdown(t, payload, coreusage.TokenAccountingQualityUnclassified, 13) })} func TestUsageQueuePluginEmitsSingleCanonicalAutoTier(t *testing.T) { withEnabledQueue(t, func() { ctx := coreusage.WithServiceTier(context.Background(), coreusage.AutoServiceTier) ctx = internallogging.WithResponseStatusHolder(ctx) internallogging.SetResponseStatus(ctx, http.StatusOK) (&usageQueuePlugin{}).HandleUsage(ctx, coreusage.Record{ Provider: "openai", Model: "gpt-5.4", Detail: coreusage.Detail{ InputTokens: 1, TotalTokens: 1, }, }) payload := popSinglePayload(t) requireStringField(t, payload, "service_tier", "auto") requireMissingField(t, payload, "request_service_tier") })} func TestUsageQueuePluginAcceptsDeprecatedRequestTierRecordField(t *testing.T) { withEnabledQueue(t, func() { ctx := internallogging.WithResponseStatusHolder(context.Background()) internallogging.SetResponseStatus(ctx, http.StatusOK) (&usageQueuePlugin{}).HandleUsage(ctx, coreusage.Record{ Provider: "openai", Model: "gpt-5.4", RequestServiceTier: "priority", Detail: coreusage.Detail{InputTokens: 1, TotalTokens: 1}, }) payload := popSinglePayload(t) requireStringField(t, payload, "service_tier", "priority") requireMissingField(t, payload, "request_service_tier") })} func TestUsageQueuePluginAsyncUsesRecordResponseHeaders(t *testing.T) { withEnabledQueue(t, func() { ctx := internallogging.WithRequestID(context.Background(), "ctx-request-id") ctx = internallogging.WithEndpoint(ctx, "POST /v1/chat/completions") ctx = internallogging.WithResponseStatusHolder(ctx) ctx = internallogging.WithResponseHeadersHolder(ctx) internallogging.SetResponseStatus(ctx, http.StatusOK) initialHeaders := http.Header{} initialHeaders.Set("X-Upstream-Request-Id", "upstream-req-1") internallogging.SetResponseHeaders(ctx, initialHeaders) mgr := coreusage.NewManager(16) defer mgr.Stop() mgr.Register(pluginFunc(func(ctx context.Context, _ coreusage.Record) { nextHeaders := http.Header{} nextHeaders.Set("X-Upstream-Request-Id", "upstream-req-2") internallogging.SetResponseHeaders(ctx, nextHeaders) })) mgr.Register(&usageQueuePlugin{}) mgr.Publish(ctx, coreusage.Record{ Provider: "openai", Model: "gpt-5.4", Alias: "client-gpt", APIKey: "test-key", AuthIndex: "0", AuthType: "apikey", Source: "user@example.com", RequestedAt: time.Date(2026, 4, 25, 0, 0, 0, 0, time.UTC), Latency: 1500 * time.Millisecond, Detail: coreusage.Detail{ InputTokens: 10, OutputTokens: 20, TotalTokens: 30, }, ResponseHeaders: internallogging.GetResponseHeaders(ctx), }) payload := waitForSinglePayload(t, 2*time.Second) requireHeaderField(t, payload, "response_headers", "X-Upstream-Request-Id", []string{"upstream-req-1"}) })} func TestUsageQueuePluginPayloadIncludesStableFieldsAndFailureAndGinRequestID(t *testing.T) { withEnabledQueue(t, func() { ctx := internallogging.WithRequestID(context.Background(), "gin-request-id") ctx = internallogging.WithEndpoint(ctx, "GET /v1/responses") ctx = internallogging.WithResponseStatusHolder(ctx) internallogging.SetResponseStatus(ctx, http.StatusInternalServerError) plugin := &usageQueuePlugin{} plugin.HandleUsage(ctx, coreusage.Record{ Provider: "openai", Model: "gpt-5.4-mini", Alias: "client-mini", APIKey: "test-key", AuthIndex: "0", AuthType: "apikey", Source: "user@example.com", RequestedAt: time.Date(2026, 4, 25, 0, 0, 0, 0, time.UTC), Latency: 2500 * time.Millisecond, Fail: coreusage.Failure{ StatusCode: http.StatusInternalServerError, Body: "upstream failed", }, Detail: coreusage.Detail{ InputTokens: 10, OutputTokens: 20, TotalTokens: 30, }, }) payload := popSinglePayload(t) requireStringField(t, payload, "provider", "openai") requireStringField(t, payload, "model", "gpt-5.4-mini") requireStringField(t, payload, "alias", "client-mini") requireStringField(t, payload, "endpoint", "GET /v1/responses") requireStringField(t, payload, "auth_type", "apikey") requireMissingField(t, payload, "user_api_key") requireStringField(t, payload, "request_id", "gin-request-id") requireBoolField(t, payload, "failed", true) requireFailField(t, payload, http.StatusInternalServerError, "upstream failed") })} func TestUsageQueuePluginAsyncIgnoresRecycledGinContext(t *testing.T) { withEnabledQueue(t, func() { ginCtx := newTestGinContext(t, http.MethodPost, "/v1/chat/completions", http.StatusOK) ctx := context.WithValue(context.Background(), "gin", ginCtx) ctx = internallogging.WithRequestID(ctx, "ctx-request-id") ctx = internallogging.WithEndpoint(ctx, "POST /v1/chat/completions") ctx = internallogging.WithResponseStatusHolder(ctx) internallogging.SetResponseStatus(ctx, http.StatusInternalServerError) mgr := coreusage.NewManager(16) defer mgr.Stop() mgr.Register(pluginFunc(func(_ context.Context, _ coreusage.Record) { ginCtx.Request = httptest.NewRequest(http.MethodGet, "http://example.com/v1/responses", nil) ginCtx.Status(http.StatusOK) })) mgr.Register(&usageQueuePlugin{}) mgr.Publish(ctx, coreusage.Record{ Provider: "openai", Model: "gpt-5.4", Alias: "client-gpt", APIKey: "test-key", AuthIndex: "0", AuthType: "apikey", Source: "user@example.com", RequestedAt: time.Date(2026, 4, 25, 0, 0, 0, 0, time.UTC), Latency: 1500 * time.Millisecond, Fail: coreusage.Failure{ StatusCode: http.StatusBadGateway, Body: "bad gateway", }, Detail: coreusage.Detail{ InputTokens: 10, OutputTokens: 20, TotalTokens: 30, }, }) payload := waitForSinglePayload(t, 2*time.Second) requireStringField(t, payload, "endpoint", "POST /v1/chat/completions") requireStringField(t, payload, "alias", "client-gpt") requireMissingField(t, payload, "user_api_key") requireStringField(t, payload, "request_id", "ctx-request-id") requireBoolField(t, payload, "failed", true) requireFailField(t, payload, http.StatusBadGateway, "bad gateway") })} func withEnabledQueue(t *testing.T, fn func()) { t.Helper() prevQueueEnabled := Enabled() prevUsageEnabled := UsageStatisticsEnabled() SetEnabled(false) SetEnabled(true) SetUsageStatisticsEnabled(true) defer func() { SetEnabled(false) SetEnabled(prevQueueEnabled) SetUsageStatisticsEnabled(prevUsageEnabled) }() fn()} func newTestGinContext(t *testing.T, method, path string, status int) *gin.Context { t.Helper() gin.SetMode(gin.TestMode) recorder := httptest.NewRecorder() ginCtx, _ := gin.CreateTestContext(recorder) ginCtx.Request = httptest.NewRequest(method, "http://example.com"+path, nil) if status != 0 { ginCtx.Status(status) } return ginCtx} func popSinglePayload(t *testing.T) map[string]json.RawMessage { t.Helper() items := PopOldest(10) if len(items) != 1 { t.Fatalf("PopOldest() items = %d, want 1", len(items)) } var payload map[string]json.RawMessage if err := json.Unmarshal(items[0], &payload); err != nil { t.Fatalf("unmarshal payload: %v", err) } return payload} func waitForSinglePayload(t *testing.T, timeout time.Duration) map[string]json.RawMessage { t.Helper() deadline := time.Now().Add(timeout) for time.Now().Before(deadline) { items := PopOldest(10) if len(items) == 0 { time.Sleep(10 * time.Millisecond) continue } if len(items) != 1 { t.Fatalf("PopOldest() items = %d, want 1", len(items)) } var payload map[string]json.RawMessage if err := json.Unmarshal(items[0], &payload); err != nil { t.Fatalf("unmarshal payload: %v", err) } return payload } t.Fatalf("timeout waiting for queued payload") return nil} func requireStringField(t *testing.T, payload map[string]json.RawMessage, key, want string) { t.Helper() raw, ok := payload[key] if !ok { t.Fatalf("payload missing %q", key) } var got string if err := json.Unmarshal(raw, &got); err != nil { t.Fatalf("unmarshal %q: %v", key, err) } if got != want { t.Fatalf("%s = %q, want %q", key, got, want) }} func requireIntField(t *testing.T, payload map[string]json.RawMessage, key string, want int) { t.Helper() raw, ok := payload[key] if !ok { t.Fatalf("payload missing %q", key) } var got int if err := json.Unmarshal(raw, &got); err != nil { t.Fatalf("unmarshal %q: %v", key, err) } if got != want { t.Fatalf("%s = %d, want %d", key, got, want) }} func requireTokenBreakdown(t *testing.T, payload map[string]json.RawMessage, quality coreusage.TokenAccountingQuality, total int64) { t.Helper() raw, ok := payload["token_breakdown"] if !ok { t.Fatal("payload missing token_breakdown") } var breakdown coreusage.TokenBreakdown if err := json.Unmarshal(raw, &breakdown); err != nil { t.Fatalf("unmarshal token_breakdown: %v", err) } if !breakdown.Valid() || breakdown.Quality != quality || breakdown.TotalTokens != total { t.Fatalf("token_breakdown = %+v, want quality=%s total=%d", breakdown, quality, total) }} func requireMissingField(t *testing.T, payload map[string]json.RawMessage, key string) { t.Helper() if _, ok := payload[key]; ok {Lines 1–500 of 710. Load more
package redisqueue import ( "context" "encoding/json" "net/http" "net/http/httptest" "testing" "time" "github.com/gin-gonic/gin" internallogging "github.com/router-for-me/CLIProxyAPI/v7/internal/logging" coresession "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/session" coreusage "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/usage" ) func TestUsageQueuePluginPayloadIncludesStableFieldsAndSuccess(t *testing.T) { withEnabledQueue(t, func() { ctx := internallogging.WithRequestID(context.Background(), "ctx-request-id") ctx = internallogging.WithEndpoint(ctx, "POST /v1/chat/completions") ctx = internallogging.WithClientRequestMetadata(ctx, internallogging.ClientRequestMetadata{ ClientIP: "192.0.2.10", XForwardedFor: "203.0.113.5, 198.51.100.8", UserAgent: "test-client/1.0", }) ctx = internallogging.WithResponseStatusHolder(ctx) internallogging.SetResponseStatus(ctx, http.StatusOK) responseHeaders := http.Header{} responseHeaders.Add("X-Upstream-Request-Id", "upstream-req-1") responseHeaders.Add("Retry-After", "30") plugin := &usageQueuePlugin{} plugin.HandleUsage(ctx, coreusage.Record{ Provider: "openai", ExecutorType: "KimiExecutor", Model: "gpt-5.4", Alias: "client-gpt", APIKey: "test-key", AuthIndex: "0", AccessTokenSHA256: "token-version-hash", AuthType: "apikey", Source: "user@example.com", ReasoningEffort: "medium", ServiceTier: "auto", ResponseServiceTier: "default", Generate: coreusage.GenerateFlag(true), RequestedAt: time.Date(2026, 4, 25, 0, 0, 0, 0, time.UTC), Latency: 1500 * time.Millisecond, Detail: coreusage.Detail{ InputTokens: 10, OutputTokens: 20, TotalTokens: 30, }, ResponseHeaders: responseHeaders.Clone(), }) responseHeaders.Set("Retry-After", "999") payload := popSinglePayload(t) requireStringField(t, payload, "provider", "openai") requireStringField(t, payload, "executor_type", "KimiExecutor") requireStringField(t, payload, "model", "gpt-5.4") requireStringField(t, payload, "alias", "client-gpt") requireStringField(t, payload, "endpoint", "POST /v1/chat/completions") requireStringField(t, payload, "auth_type", "apikey") requireStringField(t, payload, "access_token_sha256", "token-version-hash") requireMissingField(t, payload, "user_api_key") requireStringField(t, payload, "request_id", "ctx-request-id") requireStringField(t, payload, "client_ip", "192.0.2.10") requireStringField(t, payload, "x_forwarded_for", "203.0.113.5, 198.51.100.8") requireStringField(t, payload, "user_agent", "test-client/1.0") requireStringField(t, payload, "reasoning_effort", "medium") requireStringField(t, payload, "service_tier", "auto") requireMissingField(t, payload, "request_service_tier") requireStringField(t, payload, "response_service_tier", "default") requireIntField(t, payload, "accounting_version", coreusage.TokenAccountingSchemaVersion) requireTokenBreakdown(t, payload, coreusage.TokenAccountingQualityComplete, 30) requireTokensBoolField(t, payload, "cache_read_tokens_present", true) 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) requireBoolField(t, payload, "generate", true) requireBoolField(t, payload, "stream", false) requireFailField(t, payload, http.StatusOK, "") }) } func TestUsageQueuePluginNormalizesDirectSDKUsageByProvider(t *testing.T) { tests := []struct { provider string wantTotal int }{ {provider: "openai", wantTotal: 130}, {provider: "gemini", wantTotal: 142}, } for _, tt := range tests { t.Run(tt.provider, func(t *testing.T) { withEnabledQueue(t, func() { ctx := internallogging.WithResponseStatusHolder(context.Background()) internallogging.SetResponseStatus(ctx, http.StatusOK) (&usageQueuePlugin{}).HandleUsage(ctx, coreusage.Record{ Provider: tt.provider, Model: "direct-sdk-model", Detail: coreusage.Detail{ InputTokens: 100, OutputTokens: 30, ReasoningTokens: 12, }, }) payload := popSinglePayload(t) requireIntField(t, requireTokensPayload(t, payload), "total_tokens", tt.wantTotal) requireTokenBreakdown(t, payload, coreusage.TokenAccountingQualityComplete, int64(tt.wantTotal)) }) }) } } func TestUsageQueuePluginPayloadIncludesGenerateFalse(t *testing.T) { withEnabledQueue(t, func() { ctx := internallogging.WithResponseStatusHolder(context.Background()) internallogging.SetResponseStatus(ctx, http.StatusOK) (&usageQueuePlugin{}).HandleUsage(ctx, coreusage.Record{ Provider: "openai", Model: "gpt-5.4", Generate: coreusage.GenerateFlag(false), Detail: coreusage.Detail{ InputTokens: 1, TotalTokens: 1, }, }) payload := popSinglePayload(t) requireBoolField(t, payload, "generate", false) }) } func TestUsageQueuePluginPayloadDefaultsGenerateTrueWhenOmitted(t *testing.T) { withEnabledQueue(t, func() { ctx := internallogging.WithResponseStatusHolder(context.Background()) internallogging.SetResponseStatus(ctx, http.StatusOK) // Legacy callers construct usage.Record without Generate; omission must publish as true. (&usageQueuePlugin{}).HandleUsage(ctx, coreusage.Record{ Provider: "openai", Model: "gpt-5.4", Detail: coreusage.Detail{ InputTokens: 1, TotalTokens: 1, }, }) payload := popSinglePayload(t) requireBoolField(t, payload, "generate", true) }) } func TestUsageQueuePluginPublishesStreamFlag(t *testing.T) { for _, stream := range []bool{true, false} { t.Run(map[bool]string{true: "stream_true", false: "stream_false"}[stream], func(t *testing.T) { withEnabledQueue(t, func() { ctx := internallogging.WithResponseStatusHolder(context.Background()) internallogging.SetResponseStatus(ctx, http.StatusOK) (&usageQueuePlugin{}).HandleUsage(ctx, coreusage.Record{ Provider: "openai", Model: "gpt-5.4", Stream: stream, Detail: coreusage.Detail{ InputTokens: 1, TotalTokens: 1, }, }) payload := popSinglePayload(t) requireBoolField(t, payload, "stream", stream) }) }) } } func TestUsageQueuePluginPreservesLegacyCachedOnlyUsage(t *testing.T) { withEnabledQueue(t, func() { ctx := internallogging.WithResponseStatusHolder(context.Background()) internallogging.SetResponseStatus(ctx, http.StatusOK) (&usageQueuePlugin{}).HandleUsage(ctx, coreusage.Record{ Provider: "openai", Model: "gpt-5.4", Detail: coreusage.Detail{ CachedTokens: 13, }, }) payload := popSinglePayload(t) requireTokensBoolField(t, payload, "cache_read_tokens_present", true) tokens := requireTokensPayload(t, payload) requireIntField(t, tokens, "cache_read_tokens", 13) requireIntField(t, tokens, "total_tokens", 13) requireTokenBreakdown(t, payload, coreusage.TokenAccountingQualityUnclassified, 13) }) } func TestUsageQueuePluginEmitsSingleCanonicalAutoTier(t *testing.T) { withEnabledQueue(t, func() { ctx := coreusage.WithServiceTier(context.Background(), coreusage.AutoServiceTier) ctx = internallogging.WithResponseStatusHolder(ctx) internallogging.SetResponseStatus(ctx, http.StatusOK) (&usageQueuePlugin{}).HandleUsage(ctx, coreusage.Record{ Provider: "openai", Model: "gpt-5.4", Detail: coreusage.Detail{ InputTokens: 1, TotalTokens: 1, }, }) payload := popSinglePayload(t) requireStringField(t, payload, "service_tier", "auto") requireMissingField(t, payload, "request_service_tier") }) } func TestUsageQueuePluginAcceptsDeprecatedRequestTierRecordField(t *testing.T) { withEnabledQueue(t, func() { ctx := internallogging.WithResponseStatusHolder(context.Background()) internallogging.SetResponseStatus(ctx, http.StatusOK) (&usageQueuePlugin{}).HandleUsage(ctx, coreusage.Record{ Provider: "openai", Model: "gpt-5.4", RequestServiceTier: "priority", Detail: coreusage.Detail{InputTokens: 1, TotalTokens: 1}, }) payload := popSinglePayload(t) requireStringField(t, payload, "service_tier", "priority") requireMissingField(t, payload, "request_service_tier") }) } func TestUsageQueuePluginAsyncUsesRecordResponseHeaders(t *testing.T) { withEnabledQueue(t, func() { ctx := internallogging.WithRequestID(context.Background(), "ctx-request-id") ctx = internallogging.WithEndpoint(ctx, "POST /v1/chat/completions") ctx = internallogging.WithResponseStatusHolder(ctx) ctx = internallogging.WithResponseHeadersHolder(ctx) internallogging.SetResponseStatus(ctx, http.StatusOK) initialHeaders := http.Header{} initialHeaders.Set("X-Upstream-Request-Id", "upstream-req-1") internallogging.SetResponseHeaders(ctx, initialHeaders) mgr := coreusage.NewManager(16) defer mgr.Stop() mgr.Register(pluginFunc(func(ctx context.Context, _ coreusage.Record) { nextHeaders := http.Header{} nextHeaders.Set("X-Upstream-Request-Id", "upstream-req-2") internallogging.SetResponseHeaders(ctx, nextHeaders) })) mgr.Register(&usageQueuePlugin{}) mgr.Publish(ctx, coreusage.Record{ Provider: "openai", Model: "gpt-5.4", Alias: "client-gpt", APIKey: "test-key", AuthIndex: "0", AuthType: "apikey", Source: "user@example.com", RequestedAt: time.Date(2026, 4, 25, 0, 0, 0, 0, time.UTC), Latency: 1500 * time.Millisecond, Detail: coreusage.Detail{ InputTokens: 10, OutputTokens: 20, TotalTokens: 30, }, ResponseHeaders: internallogging.GetResponseHeaders(ctx), }) payload := waitForSinglePayload(t, 2*time.Second) requireHeaderField(t, payload, "response_headers", "X-Upstream-Request-Id", []string{"upstream-req-1"}) }) } func TestUsageQueuePluginPayloadIncludesStableFieldsAndFailureAndGinRequestID(t *testing.T) { withEnabledQueue(t, func() { ctx := internallogging.WithRequestID(context.Background(), "gin-request-id") ctx = internallogging.WithEndpoint(ctx, "GET /v1/responses") ctx = internallogging.WithResponseStatusHolder(ctx) internallogging.SetResponseStatus(ctx, http.StatusInternalServerError) plugin := &usageQueuePlugin{} plugin.HandleUsage(ctx, coreusage.Record{ Provider: "openai", Model: "gpt-5.4-mini", Alias: "client-mini", APIKey: "test-key", AuthIndex: "0", AuthType: "apikey", Source: "user@example.com", RequestedAt: time.Date(2026, 4, 25, 0, 0, 0, 0, time.UTC), Latency: 2500 * time.Millisecond, Fail: coreusage.Failure{ StatusCode: http.StatusInternalServerError, Body: "upstream failed", }, Detail: coreusage.Detail{ InputTokens: 10, OutputTokens: 20, TotalTokens: 30, }, }) payload := popSinglePayload(t) requireStringField(t, payload, "provider", "openai") requireStringField(t, payload, "model", "gpt-5.4-mini") requireStringField(t, payload, "alias", "client-mini") requireStringField(t, payload, "endpoint", "GET /v1/responses") requireStringField(t, payload, "auth_type", "apikey") requireMissingField(t, payload, "user_api_key") requireStringField(t, payload, "request_id", "gin-request-id") requireBoolField(t, payload, "failed", true) requireFailField(t, payload, http.StatusInternalServerError, "upstream failed") }) } func TestUsageQueuePluginAsyncIgnoresRecycledGinContext(t *testing.T) { withEnabledQueue(t, func() { ginCtx := newTestGinContext(t, http.MethodPost, "/v1/chat/completions", http.StatusOK) ctx := context.WithValue(context.Background(), "gin", ginCtx) ctx = internallogging.WithRequestID(ctx, "ctx-request-id") ctx = internallogging.WithEndpoint(ctx, "POST /v1/chat/completions") ctx = internallogging.WithResponseStatusHolder(ctx) internallogging.SetResponseStatus(ctx, http.StatusInternalServerError) mgr := coreusage.NewManager(16) defer mgr.Stop() mgr.Register(pluginFunc(func(_ context.Context, _ coreusage.Record) { ginCtx.Request = httptest.NewRequest(http.MethodGet, "http://example.com/v1/responses", nil) ginCtx.Status(http.StatusOK) })) mgr.Register(&usageQueuePlugin{}) mgr.Publish(ctx, coreusage.Record{ Provider: "openai", Model: "gpt-5.4", Alias: "client-gpt", APIKey: "test-key", AuthIndex: "0", AuthType: "apikey", Source: "user@example.com", RequestedAt: time.Date(2026, 4, 25, 0, 0, 0, 0, time.UTC), Latency: 1500 * time.Millisecond, Fail: coreusage.Failure{ StatusCode: http.StatusBadGateway, Body: "bad gateway", }, Detail: coreusage.Detail{ InputTokens: 10, OutputTokens: 20, TotalTokens: 30, }, }) payload := waitForSinglePayload(t, 2*time.Second) requireStringField(t, payload, "endpoint", "POST /v1/chat/completions") requireStringField(t, payload, "alias", "client-gpt") requireMissingField(t, payload, "user_api_key") requireStringField(t, payload, "request_id", "ctx-request-id") requireBoolField(t, payload, "failed", true) requireFailField(t, payload, http.StatusBadGateway, "bad gateway") }) } func withEnabledQueue(t *testing.T, fn func()) { t.Helper() prevQueueEnabled := Enabled() prevUsageEnabled := UsageStatisticsEnabled() SetEnabled(false) SetEnabled(true) SetUsageStatisticsEnabled(true) defer func() { SetEnabled(false) SetEnabled(prevQueueEnabled) SetUsageStatisticsEnabled(prevUsageEnabled) }() fn() } func newTestGinContext(t *testing.T, method, path string, status int) *gin.Context { t.Helper() gin.SetMode(gin.TestMode) recorder := httptest.NewRecorder() ginCtx, _ := gin.CreateTestContext(recorder) ginCtx.Request = httptest.NewRequest(method, "http://example.com"+path, nil) if status != 0 { ginCtx.Status(status) } return ginCtx } func popSinglePayload(t *testing.T) map[string]json.RawMessage { t.Helper() items := PopOldest(10) if len(items) != 1 { t.Fatalf("PopOldest() items = %d, want 1", len(items)) } var payload map[string]json.RawMessage if err := json.Unmarshal(items[0], &payload); err != nil { t.Fatalf("unmarshal payload: %v", err) } return payload } func waitForSinglePayload(t *testing.T, timeout time.Duration) map[string]json.RawMessage { t.Helper() deadline := time.Now().Add(timeout) for time.Now().Before(deadline) { items := PopOldest(10) if len(items) == 0 { time.Sleep(10 * time.Millisecond) continue } if len(items) != 1 { t.Fatalf("PopOldest() items = %d, want 1", len(items)) } var payload map[string]json.RawMessage if err := json.Unmarshal(items[0], &payload); err != nil { t.Fatalf("unmarshal payload: %v", err) } return payload } t.Fatalf("timeout waiting for queued payload") return nil } func requireStringField(t *testing.T, payload map[string]json.RawMessage, key, want string) { t.Helper() raw, ok := payload[key] if !ok { t.Fatalf("payload missing %q", key) } var got string if err := json.Unmarshal(raw, &got); err != nil { t.Fatalf("unmarshal %q: %v", key, err) } if got != want { t.Fatalf("%s = %q, want %q", key, got, want) } } func requireIntField(t *testing.T, payload map[string]json.RawMessage, key string, want int) { t.Helper() raw, ok := payload[key] if !ok { t.Fatalf("payload missing %q", key) } var got int if err := json.Unmarshal(raw, &got); err != nil { t.Fatalf("unmarshal %q: %v", key, err) } if got != want { t.Fatalf("%s = %d, want %d", key, got, want) } } func requireTokenBreakdown(t *testing.T, payload map[string]json.RawMessage, quality coreusage.TokenAccountingQuality, total int64) { t.Helper() raw, ok := payload["token_breakdown"] if !ok { t.Fatal("payload missing token_breakdown") } var breakdown coreusage.TokenBreakdown if err := json.Unmarshal(raw, &breakdown); err != nil { t.Fatalf("unmarshal token_breakdown: %v", err) } if !breakdown.Valid() || breakdown.Quality != quality || breakdown.TotalTokens != total { t.Fatalf("token_breakdown = %+v, want quality=%s total=%d", breakdown, quality, total) } } func requireMissingField(t *testing.T, payload map[string]json.RawMessage, key string) { t.Helper() if _, ok := payload[key]; ok {