Add projects

backend/internal/redisqueue/plugin_test.go

@@ -0,0 +1,561 @@
+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"
+	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)
+		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 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 {
+		t.Fatalf("payload unexpectedly contains %q", key)
+	}
+}
+
+type pluginFunc func(context.Context, coreusage.Record)
+
+func (fn pluginFunc) HandleUsage(ctx context.Context, record coreusage.Record) {
+	fn(ctx, record)
+}
+
+func requireBoolField(t *testing.T, payload map[string]json.RawMessage, key string, want bool) {
+	t.Helper()
+
+	raw, ok := payload[key]
+	if !ok {
+		t.Fatalf("payload missing %q", key)
+	}
+	var got bool
+	if err := json.Unmarshal(raw, &got); err != nil {
+		t.Fatalf("unmarshal %q: %v", key, err)
+	}
+	if got != want {
+		t.Fatalf("%s = %t, want %t", key, got, want)
+	}
+}
+
+func requireTokensPayload(t *testing.T, payload map[string]json.RawMessage) map[string]json.RawMessage {
+	t.Helper()
+	raw, ok := payload["tokens"]
+	if !ok {
+		t.Fatal("payload missing tokens")
+	}
+	var tokens map[string]json.RawMessage
+	if errUnmarshal := json.Unmarshal(raw, &tokens); errUnmarshal != nil {
+		t.Fatalf("unmarshal tokens: %v", errUnmarshal)
+	}
+	return tokens
+}
+
+func requireTokensBoolField(t *testing.T, payload map[string]json.RawMessage, key string, want bool) {
+	t.Helper()
+	requireBoolField(t, requireTokensPayload(t, payload), key, want)
+}
+
+func requireFailField(t *testing.T, payload map[string]json.RawMessage, wantStatus int, wantBody string) {
+	t.Helper()
+
+	raw, ok := payload["fail"]
+	if !ok {
+		t.Fatalf("payload missing %q", "fail")
+	}
+	var got struct {
+		StatusCode int    `json:"status_code"`
+		Body       string `json:"body"`
+	}
+	if err := json.Unmarshal(raw, &got); err != nil {
+		t.Fatalf("unmarshal fail: %v", err)
+	}
+	if got.StatusCode != wantStatus || got.Body != wantBody {
+		t.Fatalf("fail = {status_code:%d body:%q}, want {status_code:%d body:%q}", got.StatusCode, got.Body, wantStatus, wantBody)
+	}
+}
+
+func requireHeaderField(t *testing.T, payload map[string]json.RawMessage, field, key string, want []string) {
+	t.Helper()
+
+	raw, ok := payload[field]
+	if !ok {
+		t.Fatalf("payload missing %q", field)
+	}
+	var headers map[string][]string
+	if err := json.Unmarshal(raw, &headers); err != nil {
+		t.Fatalf("unmarshal %q: %v", field, err)
+	}
+	got, ok := headers[key]
+	if !ok {
+		t.Fatalf("%s missing header %q", field, key)
+	}
+	if len(got) != len(want) {
+		t.Fatalf("%s[%q] = %v, want %v", field, key, got, want)
+	}
+	for i := range want {
+		if got[i] != want[i] {
+			t.Fatalf("%s[%q] = %v, want %v", field, key, got, want)
+		}
+	}
+}