package service import ( "context" "net/http" "net/http/httptest" "strconv" "strings" "testing" "time" "github.com/Wei-Shaw/sub2api/internal/pkg/kirocooldown" "github.com/stretchr/testify/require" ) type kiroUsageCooldownStore struct { state *kirocooldown.State err error } func (s *kiroUsageCooldownStore) ReserveRequest(context.Context, string) (time.Duration, error) { return 0, nil } func (s *kiroUsageCooldownStore) MarkSuccess(context.Context, string) error { return nil } func (s *kiroUsageCooldownStore) Mark429(context.Context, string) (time.Duration, error) { return 0, nil } func (s *kiroUsageCooldownStore) MarkSuspended(context.Context, string) (time.Duration, error) { return 0, nil } func (s *kiroUsageCooldownStore) GetState(context.Context, string) (*kirocooldown.State, error) { return s.state, s.err } func (s *kiroUsageCooldownStore) ClearEarliestTransientCooldown(context.Context, []string) (bool, error) { return false, nil } func kiroFloatPtr(v float64) *float64 { return &v } func TestChannel_IsWebSearchEmulationEnabled_Kiro(t *testing.T) { c := &Channel{ FeaturesConfig: map[string]any{ featureKeyWebSearchEmulation: map[string]any{"kiro": true}, }, } require.True(t, c.IsWebSearchEmulationEnabled("kiro")) } func TestOpenAIGatewayServiceRecordUsage_NormalizesKiroBillingModel(t *testing.T) { usageRepo := &openAIRecordUsageLogRepoStub{inserted: true} userRepo := &openAIRecordUsageUserRepoStub{} subRepo := &openAIRecordUsageSubRepoStub{} svc := newOpenAIRecordUsageServiceForTest(usageRepo, userRepo, subRepo, nil) svc.billingService = NewBillingService(svc.cfg, &PricingService{ pricingData: map[string]*LiteLLMModelPricing{ "claude-sonnet-4-6": { InputCostPerToken: 2.5e-6, OutputCostPerToken: 10e-6, }, }, }) expectedCost, err := svc.billingService.CalculateCost("claude-sonnet-4-6", UsageTokens{ InputTokens: 20, OutputTokens: 10, }, 1.1) require.NoError(t, err) err = svc.RecordUsage(context.Background(), &OpenAIRecordUsageInput{ Result: &OpenAIForwardResult{ RequestID: "resp_kiro_billing_normalized", Model: "claude-sonnet-4-6", UpstreamModel: "claude-sonnet-4.6", Usage: OpenAIUsage{ InputTokens: 20, OutputTokens: 10, }, Duration: time.Second, }, APIKey: &APIKey{ID: 10}, User: &User{ID: 20}, Account: &Account{ID: 30, Platform: PlatformKiro}, }) require.NoError(t, err) require.NotNil(t, usageRepo.lastLog) require.Equal(t, "claude-sonnet-4-6", usageRepo.lastLog.Model) require.Equal(t, "claude-sonnet-4-6", usageRepo.lastLog.RequestedModel) require.NotNil(t, usageRepo.lastLog.UpstreamModel) require.Equal(t, "claude-sonnet-4.6", *usageRepo.lastLog.UpstreamModel) require.InDelta(t, expectedCost.ActualCost, usageRepo.lastLog.ActualCost, 1e-12) require.InDelta(t, expectedCost.TotalCost, usageRepo.lastLog.TotalCost, 1e-12) } func TestAccountUsageService_GetUsage_KiroMapsCredits(t *testing.T) { account := Account{ ID: 701, Platform: PlatformKiro, Type: AccountTypeOAuth, Credentials: map[string]any{ "access_token": "kiro-access-token", "provider": "Github", "auth_method": "social", "profile_arn": "arn:aws:codewhisperer:us-east-1:123456789012:profile/SOCIAL", }, } repo := &stubOpenAIAccountRepo{accounts: []Account{account}} svc := NewAccountUsageService(repo, nil, nil, nil, nil, NewUsageCache(), nil, nil) resetAt := time.Now().Add(10 * 24 * time.Hour).Unix() bonusExpiry := time.Now().Add(7 * 24 * time.Hour).Unix() server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { require.Equal(t, "/getUsageLimits", r.URL.Path) require.Equal(t, "arn:aws:codewhisperer:us-east-1:123456789012:profile/SOCIAL", r.URL.Query().Get("profileArn")) require.Equal(t, kiroUsageOrigin, r.URL.Query().Get("origin")) require.Equal(t, kiroUsageResourceType, r.URL.Query().Get("resourceType")) require.Equal(t, "Bearer kiro-access-token", r.Header.Get("Authorization")) require.Equal(t, "*/*", r.Header.Get("Accept")) require.True(t, strings.Contains(r.Header.Get("User-Agent"), "KiroIDE-")) require.True(t, strings.Contains(r.Header.Get("X-Amz-User-Agent"), "KiroIDE-")) require.Equal(t, "vibe", r.Header.Get("x-amzn-kiro-agent-mode")) require.Equal(t, "true", r.Header.Get("x-amzn-codewhisperer-optout")) require.NotEmpty(t, r.Header.Get("Amz-Sdk-Invocation-Id")) w.Header().Set("Content-Type", "application/json") _, _ = w.Write([]byte(`{ "nextDateReset": ` + strconv.FormatInt(resetAt, 10) + `, "overageConfiguration": {"overageStatus":"ENABLED"}, "subscriptionInfo": {"subscriptionTitle":"KIRO PRO+","type":"Q_DEVELOPER_STANDALONE_PRO_PLUS"}, "usageBreakdownList": [{ "currency":"USD", "currentOveragesWithPrecision":2, "currentUsageWithPrecision":125, "freeTrialInfo":{ "currentUsageWithPrecision":25, "freeTrialExpiry":` + strconv.FormatInt(bonusExpiry, 10) + `, "freeTrialStatus":"ACTIVE", "usageLimitWithPrecision":500 }, "nextDateReset": ` + strconv.FormatInt(resetAt, 10) + `, "overageCharges":0.08, "resourceType":"CREDIT", "usageLimitWithPrecision":2000 }] }`)) })) defer server.Close() prevResolver := resolveKiroRuntimeEndpoint resolveKiroRuntimeEndpoint = func(_ string) string { return server.URL } defer func() { resolveKiroRuntimeEndpoint = prevResolver }() usage, err := svc.GetUsage(context.Background(), account.ID) require.NoError(t, err) require.NotNil(t, usage) require.Equal(t, "active", usage.Source) require.Equal(t, "KIRO PRO+", usage.KiroSubscriptionName) require.Equal(t, "Q_DEVELOPER_STANDALONE_PRO_PLUS", usage.KiroSubscriptionType) require.True(t, usage.KiroOveragesEnabled) require.NotNil(t, usage.KiroCredit) require.Equal(t, 125.0, usage.KiroCredit.CurrentUsage) require.Equal(t, 2000.0, usage.KiroCredit.UsageLimit) require.InDelta(t, 6.25, usage.KiroCredit.PercentageUsed, 0.001) require.NotNil(t, usage.KiroBonus) require.Equal(t, 25.0, usage.KiroBonus.CurrentUsage) require.Equal(t, 500.0, usage.KiroBonus.UsageLimit) require.NotNil(t, usage.KiroOverage) require.Equal(t, "$", usage.KiroOverage.CurrencySymbol) require.Equal(t, 2.0, usage.KiroOverage.CurrentOverages) require.Equal(t, 0.08, usage.KiroOverage.OverageCharges) require.NotNil(t, usage.KiroResetAt) require.Equal(t, kiroQuotaStateOverageActive, usage.KiroQuotaState) require.Equal(t, "overages_enabled", usage.KiroQuotaReason) require.NotNil(t, usage.KiroQuotaResetAt) } func TestAccountUsageService_GetUsage_KiroActiveUsesCachedSnapshotWithinTTL(t *testing.T) { account := Account{ ID: 702, Platform: PlatformKiro, Type: AccountTypeOAuth, Credentials: map[string]any{ "access_token": "kiro-access-token", "provider": "Github", "auth_method": "social", }, } repo := &stubOpenAIAccountRepo{accounts: []Account{account}} svc := NewAccountUsageService(repo, nil, nil, nil, nil, NewUsageCache(), nil, nil) successServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { w.Header().Set("Content-Type", "application/json") _, _ = w.Write([]byte(`{ "subscriptionInfo": {"subscriptionTitle":"KIRO PRO+"}, "usageBreakdownList": [{ "currency":"USD", "currentUsageWithPrecision":300, "usageLimitWithPrecision":2000, "resourceType":"CREDIT" }] }`)) })) defer successServer.Close() prevResolver := resolveKiroRuntimeEndpoint resolveKiroRuntimeEndpoint = func(_ string) string { return successServer.URL } defer func() { resolveKiroRuntimeEndpoint = prevResolver }() firstUsage, err := svc.GetUsage(context.Background(), account.ID) require.NoError(t, err) require.NotNil(t, firstUsage) require.NotNil(t, firstUsage.KiroCredit) require.Equal(t, 300.0, firstUsage.KiroCredit.CurrentUsage) failingServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { http.Error(w, `{"message":"temporary failure"}`, http.StatusInternalServerError) })) defer failingServer.Close() resolveKiroRuntimeEndpoint = func(_ string) string { return failingServer.URL } usage, err := svc.GetUsage(context.Background(), account.ID) require.NoError(t, err) require.NotNil(t, usage) require.NotNil(t, usage.KiroCredit) require.Equal(t, 300.0, usage.KiroCredit.CurrentUsage) require.Empty(t, usage.Error) require.Empty(t, usage.ErrorCode) } func TestAccountUsageService_GetUsage_KiroBuilderIDWithoutProfileArnOmitsProfileArn(t *testing.T) { account := Account{ ID: 703, Platform: PlatformKiro, Type: AccountTypeOAuth, Credentials: map[string]any{ "access_token": "kiro-access-token", "provider": "BuilderId", "auth_method": "idc", "region": "us-east-1", }, } repo := &stubOpenAIAccountRepo{accounts: []Account{account}} svc := NewAccountUsageService(repo, nil, nil, nil, nil, NewUsageCache(), nil, nil) server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { require.Equal(t, "/getUsageLimits", r.URL.Path) require.Empty(t, r.URL.Query().Get("profileArn")) w.Header().Set("Content-Type", "application/json") _, _ = w.Write([]byte(`{ "subscriptionInfo": {"subscriptionTitle":"KIRO PRO+"}, "usageBreakdownList": [{ "currency":"USD", "currentUsageWithPrecision":42, "usageLimitWithPrecision":2000, "resourceType":"CREDIT" }] }`)) })) defer server.Close() prevResolver := resolveKiroRuntimeEndpoint resolveKiroRuntimeEndpoint = func(_ string) string { return server.URL } defer func() { resolveKiroRuntimeEndpoint = prevResolver }() usage, err := svc.GetUsage(context.Background(), account.ID) require.NoError(t, err) require.NotNil(t, usage) require.NotNil(t, usage.KiroCredit) require.Equal(t, 42.0, usage.KiroCredit.CurrentUsage) } func TestAccountUsageService_GetUsage_KiroEnterpriseUsesCredentialProfileArn(t *testing.T) { account := Account{ ID: 707, Platform: PlatformKiro, Type: AccountTypeOAuth, Credentials: map[string]any{ "access_token": "kiro-access-token", "provider": "AWS", "auth_method": "idc", "region": "us-east-1", "start_url": "https://d-example.awsapps.com/start", "profile_arn": "arn:aws:codewhisperer:us-east-1:123456789012:profile/REALENTERPRISE", }, } repo := &stubOpenAIAccountRepo{accounts: []Account{account}} svc := NewAccountUsageService(repo, nil, nil, nil, nil, NewUsageCache(), nil, nil) const resolvedProfileArn = "arn:aws:codewhisperer:us-east-1:123456789012:profile/REALENTERPRISE" server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { require.Equal(t, "/getUsageLimits", r.URL.Path) require.Equal(t, resolvedProfileArn, r.URL.Query().Get("profileArn")) w.Header().Set("Content-Type", "application/json") _, _ = w.Write([]byte(`{ "subscriptionInfo": {"subscriptionTitle":"KIRO PRO+"}, "usageBreakdownList": [{ "currency":"USD", "currentUsageWithPrecision":64, "usageLimitWithPrecision":2000, "resourceType":"CREDIT" }] }`)) })) defer server.Close() prevResolver := resolveKiroRuntimeEndpoint resolveKiroRuntimeEndpoint = func(_ string) string { return server.URL } defer func() { resolveKiroRuntimeEndpoint = prevResolver }() usage, err := svc.GetUsage(context.Background(), account.ID) require.NoError(t, err) require.NotNil(t, usage) require.NotNil(t, usage.KiroCredit) require.Equal(t, 64.0, usage.KiroCredit.CurrentUsage) } func TestAccountUsageService_GetUsage_KiroUsesAPIRegionForUsageRequest(t *testing.T) { account := Account{ ID: 709, Platform: PlatformKiro, Type: AccountTypeOAuth, Credentials: map[string]any{ "access_token": "kiro-access-token", "provider": "AWS", "auth_method": "idc", "api_region": "eu-west-1", "region": "ap-northeast-2", "start_url": "https://d-example.awsapps.com/start", "profile_arn": "arn:aws:codewhisperer:eu-west-1:123456789012:profile/REALAPIREGION", }, } repo := &stubOpenAIAccountRepo{accounts: []Account{account}} svc := NewAccountUsageService(repo, nil, nil, nil, nil, NewUsageCache(), nil, nil) const resolvedProfileArn = "arn:aws:codewhisperer:eu-west-1:123456789012:profile/REALAPIREGION" gotRegions := make([]string, 0, 2) server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { require.Equal(t, "/getUsageLimits", r.URL.Path) require.Equal(t, resolvedProfileArn, r.URL.Query().Get("profileArn")) w.Header().Set("Content-Type", "application/json") _, _ = w.Write([]byte(`{ "subscriptionInfo": {"subscriptionTitle":"KIRO PRO+"}, "usageBreakdownList": [{ "currency":"USD", "currentUsageWithPrecision":11, "usageLimitWithPrecision":2000, "resourceType":"CREDIT" }] }`)) })) defer server.Close() prevResolver := resolveKiroRuntimeEndpoint resolveKiroRuntimeEndpoint = func(region string) string { gotRegions = append(gotRegions, region) return server.URL } defer func() { resolveKiroRuntimeEndpoint = prevResolver }() usage, err := svc.GetUsage(context.Background(), account.ID) require.NoError(t, err) require.NotNil(t, usage) require.Equal(t, []string{"eu-west-1"}, gotRegions) } func TestAccountUsageService_GetUsage_KiroOmitsProfileArnAndUsesDefaultRegionWithoutAPIRegionOrProfileArn(t *testing.T) { account := Account{ ID: 710, Platform: PlatformKiro, Type: AccountTypeOAuth, Credentials: map[string]any{ "access_token": "kiro-access-token", "provider": "AWS", "auth_method": "idc", "region": "ap-northeast-2", "start_url": "https://d-example.awsapps.com/start", }, } repo := &stubOpenAIAccountRepo{accounts: []Account{account}} svc := NewAccountUsageService(repo, nil, nil, nil, nil, NewUsageCache(), nil, nil) gotRegions := make([]string, 0, 2) server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { require.Equal(t, "/getUsageLimits", r.URL.Path) require.Empty(t, r.URL.Query().Get("profileArn")) w.Header().Set("Content-Type", "application/json") _, _ = w.Write([]byte(`{ "subscriptionInfo": {"subscriptionTitle":"KIRO PRO+"}, "usageBreakdownList": [{ "currency":"USD", "currentUsageWithPrecision":7, "usageLimitWithPrecision":2000, "resourceType":"CREDIT" }] }`)) })) defer server.Close() prevResolver := resolveKiroRuntimeEndpoint resolveKiroRuntimeEndpoint = func(region string) string { gotRegions = append(gotRegions, region) return server.URL } defer func() { resolveKiroRuntimeEndpoint = prevResolver }() usage, err := svc.GetUsage(context.Background(), account.ID) require.NoError(t, err) require.NotNil(t, usage) require.Equal(t, []string{kiroDefaultRegion}, gotRegions) } func TestAccountUsageService_GetUsage_KiroIncludesRuntimeCooldownState(t *testing.T) { account := Account{ ID: 704, Platform: PlatformKiro, Type: AccountTypeOAuth, Credentials: map[string]any{ "access_token": "kiro-access-token", "provider": "Github", "auth_method": "social", }, } repo := &stubOpenAIAccountRepo{accounts: []Account{account}} svc := NewAccountUsageService(repo, nil, nil, nil, nil, NewUsageCache(), nil, nil). SetKiroCooldownStore(&kiroUsageCooldownStore{ state: &kirocooldown.State{ Active: true, Reason: kirocooldown.CooldownReason429, CooldownUntil: time.Now().Add(90 * time.Second), Remaining: 90 * time.Second, }, }) server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { w.Header().Set("Content-Type", "application/json") _, _ = w.Write([]byte(`{ "subscriptionInfo": {"subscriptionTitle":"KIRO PRO+"}, "usageBreakdownList": [{ "currency":"USD", "currentUsageWithPrecision":42, "usageLimitWithPrecision":2000, "resourceType":"CREDIT" }] }`)) })) defer server.Close() prevResolver := resolveKiroRuntimeEndpoint resolveKiroRuntimeEndpoint = func(_ string) string { return server.URL } defer func() { resolveKiroRuntimeEndpoint = prevResolver }() usage, err := svc.GetUsage(context.Background(), account.ID) require.NoError(t, err) require.Equal(t, "cooldown", usage.KiroRuntimeState) require.Equal(t, kirocooldown.CooldownReason429, usage.KiroRuntimeReason) require.NotNil(t, usage.KiroRuntimeResetAt) } func TestBuildKiroDegradedUsage_ClassifiesProfileError(t *testing.T) { info := buildKiroDegradedUsage(&kiroUsageHTTPError{ StatusCode: http.StatusBadRequest, Body: `{"message":"profileArn is required for this request."}`, }) require.Equal(t, errorCodeForbidden, info.ErrorCode) require.False(t, info.NeedsReauth) } func TestBuildKiroDegradedUsage_ClassifiesOverageExhausted(t *testing.T) { info := buildKiroDegradedUsage(&kiroUsageHTTPError{ StatusCode: http.StatusTooManyRequests, Body: `{"message":"overage exhausted for this billing window"}`, }) require.Equal(t, errorCodeNetworkError, info.ErrorCode) require.Equal(t, kiroQuotaStateOverageExhausted, info.KiroQuotaState) require.Contains(t, info.KiroQuotaReason, "overage exhausted") } func TestAccountUsageService_GetUsage_KiroCachesErrorSnapshotWhenRefreshFailsWithoutPriorSuccess(t *testing.T) { account := Account{ ID: 708, Platform: PlatformKiro, Type: AccountTypeOAuth, Credentials: map[string]any{ "access_token": "kiro-access-token", "provider": "Github", "auth_method": "social", }, } repo := &stubOpenAIAccountRepo{accounts: []Account{account}} svc := NewAccountUsageService(repo, nil, nil, nil, nil, NewUsageCache(), nil, nil) requestCount := 0 server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { requestCount++ http.Error(w, `{"message":"FEATURE_NOT_SUPPORTED","reason":"FEATURE_NOT_SUPPORTED"}`, http.StatusForbidden) })) defer server.Close() prevResolver := resolveKiroRuntimeEndpoint resolveKiroRuntimeEndpoint = func(_ string) string { return server.URL } defer func() { resolveKiroRuntimeEndpoint = prevResolver }() firstUsage, err := svc.GetUsage(context.Background(), account.ID) require.NoError(t, err) require.NotNil(t, firstUsage) require.Equal(t, errorCodeForbidden, firstUsage.ErrorCode) secondUsage, err := svc.GetUsage(context.Background(), account.ID) require.NoError(t, err) require.NotNil(t, secondUsage) require.Equal(t, errorCodeForbidden, secondUsage.ErrorCode) require.Equal(t, 1, requestCount) } func TestMapKiroUsageToInfo_CreditsExhaustedWithoutOverages(t *testing.T) { info := mapKiroUsageToInfo(&kiroUsageLimitsResponse{ NextDateReset: "2099-03-13T12:00:00Z", OverageConfiguration: kiroOverageConfiguration{ OverageStatus: "DISABLED", }, UsageBreakdownList: []kiroUsageBreakdown{ { ResourceType: "CREDIT", CurrentUsageWithPrecision: kiroFloatPtr(2000), UsageLimitWithPrecision: kiroFloatPtr(2000), CurrentOveragesWithPrecision: kiroFloatPtr(0), }, }, }) require.Equal(t, kiroQuotaStateCreditsExhausted, info.KiroQuotaState) require.Equal(t, "credits_exhausted", info.KiroQuotaReason) require.NotNil(t, info.KiroQuotaResetAt) } func TestAccountUsageService_EnrichAccountWithKiroRuntimeState(t *testing.T) { svc := NewAccountUsageService(nil, nil, nil, nil, nil, NewUsageCache(), nil, nil). SetKiroCooldownStore(&kiroUsageCooldownStore{ state: &kirocooldown.State{ Active: true, Reason: kirocooldown.CooldownReason429, CooldownUntil: time.Now().Add(2 * time.Minute), Remaining: 2 * time.Minute, }, }) account := &Account{ ID: 705, Platform: PlatformKiro, Type: AccountTypeOAuth, Credentials: map[string]any{"access_token": "kiro-access-token"}, } svc.EnrichAccountWithKiroRuntimeState(context.Background(), account) require.Equal(t, "cooldown", account.KiroRuntimeState) require.Equal(t, kirocooldown.CooldownReason429, account.KiroRuntimeReason) require.NotNil(t, account.KiroRuntimeResetAt) } func TestAccountUsageService_EnrichAccountWithKiroRuntimeStateIncludesCachedQuotaState(t *testing.T) { account := Account{ ID: 706, Platform: PlatformKiro, Type: AccountTypeOAuth, Credentials: map[string]any{ "access_token": "kiro-access-token", "provider": "Github", "auth_method": "social", }, } repo := &stubOpenAIAccountRepo{accounts: []Account{account}} svc := NewAccountUsageService(repo, nil, nil, nil, nil, NewUsageCache(), nil, nil) server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { w.Header().Set("Content-Type", "application/json") _, _ = w.Write([]byte(`{ "nextDateReset":"2099-03-13T12:00:00Z", "overageConfiguration":{"overageStatus":"ENABLED"}, "subscriptionInfo": {"subscriptionTitle":"KIRO PRO+"}, "usageBreakdownList": [{ "currency":"USD", "currentUsageWithPrecision":2000, "currentOveragesWithPrecision":4, "overageCharges":0.2, "usageLimitWithPrecision":2000, "resourceType":"CREDIT" }] }`)) })) defer server.Close() prevResolver := resolveKiroRuntimeEndpoint resolveKiroRuntimeEndpoint = func(_ string) string { return server.URL } defer func() { resolveKiroRuntimeEndpoint = prevResolver }() _, err := svc.GetUsage(context.Background(), account.ID) require.NoError(t, err) target := &Account{ ID: account.ID, Platform: PlatformKiro, Type: AccountTypeOAuth, Credentials: map[string]any{"access_token": "kiro-access-token"}, } svc.EnrichAccountWithKiroRuntimeState(context.Background(), target) require.Equal(t, kiroQuotaStateOverageActive, target.KiroQuotaState) require.Equal(t, "overages_enabled", target.KiroQuotaReason) require.NotNil(t, target.KiroQuotaResetAt) }