diff --git a/internal/claude/account.go b/internal/claude/account.go index 115f8c02..fe33f67f 100644 --- a/internal/claude/account.go +++ b/internal/claude/account.go @@ -37,7 +37,7 @@ type UsageWindow struct { } type ExtraUsage struct { - Enabled bool + Enabled *bool MonthlyLimit *float64 UsedCredits *float64 Utilization *float64 @@ -69,6 +69,8 @@ type AccountObservation struct { Profile AccountProfile Usage Usage Header http.Header + AccountObserved bool + UsageObserved bool IncompleteSources []string } @@ -109,6 +111,8 @@ func accountObservationFromBridge(value cpaembedded.ClaudeAccountObservation) Ac ExtraUsageEnabled: cloneBool(profile.ExtraUsageEnabled), }, Header: value.Header.Clone(), + AccountObserved: value.AccountObserved, + UsageObserved: value.UsageObserved, IncompleteSources: append([]string(nil), value.IncompleteSources...), } result.Usage = usageFromBridge(value.Usage) @@ -126,7 +130,7 @@ func usageFromBridge(value cpaembedded.ClaudeUsage) Usage { } if value.ExtraUsage != nil { result.ExtraUsage = &ExtraUsage{ - Enabled: value.ExtraUsage.Enabled, + Enabled: cloneBool(value.ExtraUsage.Enabled), MonthlyLimit: cloneFloat(value.ExtraUsage.MonthlyLimit), UsedCredits: cloneFloat(value.ExtraUsage.UsedCredits), Utilization: cloneFloat(value.ExtraUsage.Utilization), diff --git a/internal/control/credential_observations.go b/internal/control/credential_observations.go index 1118651d..e1d5c92f 100644 --- a/internal/control/credential_observations.go +++ b/internal/control/credential_observations.go @@ -21,8 +21,9 @@ import ( ) const ( - observationRefreshFloor = 5 * time.Minute - observationFreshTTL = time.Hour + observationRefreshFloor = 5 * time.Minute + observationFreshTTL = time.Hour + observationAuthRetryBackoff = time.Hour ) type ObservationPlanSummary struct { @@ -204,7 +205,10 @@ func (s *Service) refreshCredentialObservationOnce( defer cancelObserve() channelID := channel.ID(group.ChannelID) observation, observeErr := s.observeSubscriptionAccount(observeContext, channelID, preparedCredential) + authRefreshVersion := previous.LastAuthRefreshSecretVersion if subscriptionObservationHTTPStatus(observeErr) == http.StatusUnauthorized && + (previous.LastAuthRefreshSecretVersion == nil || + *previous.LastAuthRefreshSecretVersion != credential.SecretVersion) && s.prepareSubscriptionCredential != nil { preparedCredential, err = s.prepareStoredSubscriptionCredentialWithForce( observeContext, @@ -215,38 +219,45 @@ func (s *Service) refreshCredentialObservationOnce( if err != nil { return CredentialObservationResponse{}, err } + var refreshed models.Credential + if err := s.db.WithContext(observeContext).Select("secret_version").First(&refreshed, credential.ID).Error; err != nil { + return CredentialObservationResponse{}, app_errors.ParseDBError(err) + } + version := refreshed.SecretVersion + authRefreshVersion = &version observation, observeErr = s.observeSubscriptionAccount(observeContext, channelID, preparedCredential) } if observeErr != nil { if errors.Is(observeErr, subscriptionruntime.ErrObservationPayloadInvalid) { return s.recordCredentialObservationFailure( ctx, credential, previous, attemptMS, nextAllowedMS, - "observation_payload_invalid", "normalize subscription information", + "observation_payload_invalid", "normalize subscription information", nil, ) } if status := subscriptionObservationHTTPStatus(observeErr); status == http.StatusUnauthorized || status == http.StatusForbidden { + nextAllowedMS = now.Add(observationAuthRetryBackoff).UnixMilli() + code, summary := "observation_authorization_failed", "authorize subscription information" + if status == http.StatusForbidden { + code, summary = "observation_access_denied", "access subscription information" + } return s.recordCredentialObservationFailure( ctx, credential, previous, attemptMS, nextAllowedMS, - "observation_authorization_failed", "authorize subscription information", + code, summary, authRefreshVersion, ) } return s.recordCredentialObservationFailure( ctx, credential, previous, attemptMS, nextAllowedMS, - "observation_upstream_failed", "refresh subscription information", + "observation_upstream_failed", "refresh subscription information", nil, ) } var snapshot CredentialObservationSnapshot if err := json.Unmarshal(observation.Payload, &snapshot); err != nil || snapshot.QuotaWindows == nil { return s.recordCredentialObservationFailure( ctx, credential, previous, attemptMS, nextAllowedMS, - "observation_payload_invalid", "normalize subscription information", + "observation_payload_invalid", "normalize subscription information", nil, ) } - encoded, err := json.Marshal(snapshot) - if err != nil { - return CredentialObservationResponse{}, app_errors.ErrInternalServer - } version := previous.ObservationVersion + 1 if previous.ObservationVersion == 0 { version = 1 @@ -255,12 +266,35 @@ func (s *Service) refreshCredentialObservationOnce( var freshUntilMS *int64 lastErrorCode := "" if observation.Partial { - state = models.CredentialObservationStale lastErrorCode = "observation_partial" + var previousSnapshot CredentialObservationSnapshot + previousSnapshotOK := len(previous.SnapshotJSON) > 0 && + json.Unmarshal(previous.SnapshotJSON, &previousSnapshot) == nil + previousQuotaFresh := previousSnapshotOK && + previous.State == models.CredentialObservationFresh && + previous.FreshUntilMS != nil && *previous.FreshUntilMS > attemptMS + switch { + case !observation.QuotaObserved && previousQuotaFresh: + snapshot = previousSnapshot + freshUntilMS = previous.FreshUntilMS + case observation.QuotaObserved: + if !observation.AccountObserved && previousSnapshotOK { + snapshot.Plan = previousSnapshot.Plan + snapshot.Account = previousSnapshot.Account + } + value := now.Add(observationFreshTTL).UnixMilli() + freshUntilMS = &value + default: + state = models.CredentialObservationStale + } } else { value := now.Add(observationFreshTTL).UnixMilli() freshUntilMS = &value } + encoded, err := json.Marshal(snapshot) + if err != nil { + return CredentialObservationResponse{}, app_errors.ErrInternalServer + } row := models.CredentialObservation{ CredentialID: credential.ID, IdentityFingerprint: credential.IdentityFingerprint, SchemaVersion: 1, ObservationVersion: version, SnapshotJSON: models.JSON(encoded), @@ -293,6 +327,7 @@ func (s *Service) recordCredentialObservationFailure( nextAllowedMS int64, code string, summary string, + authRefreshSecretVersion *uint64, ) (CredentialObservationResponse, error) { failed := previous failed.CredentialID = credential.ID @@ -301,11 +336,20 @@ func (s *Service) recordCredentialObservationFailure( if failed.ObservationVersion == 0 { failed.ObservationVersion = 1 } - failed.State = models.CredentialObservationError + if failed.State == models.CredentialObservationFresh && failed.FreshUntilMS != nil && + *failed.FreshUntilMS > attemptMS { + failed.State = models.CredentialObservationFresh + } else { + failed.State = models.CredentialObservationError + } failed.LastAttemptAtMS = &attemptMS failed.NextAllowedAtMS = &nextAllowedMS failed.LastErrorCode = code failed.UpdatedAtMS = attemptMS + if authRefreshSecretVersion != nil { + value := *authRefreshSecretVersion + failed.LastAuthRefreshSecretVersion = &value + } if len(failed.SnapshotJSON) == 0 { failed.SnapshotJSON = models.JSON(`{}`) } diff --git a/internal/control/credential_observations_test.go b/internal/control/credential_observations_test.go index e418e5f2..d0d650e5 100644 --- a/internal/control/credential_observations_test.go +++ b/internal/control/credential_observations_test.go @@ -171,7 +171,8 @@ func TestRefreshCredentialObservationPersistsLKGAndAllowsImmediateManualRefresh( return codex.AccountObservation{}, errors.New("upstream unavailable") }) failed, err := fixture.service.RefreshCredentialObservation(t.Context(), groupID, credentialID) - if err == nil || failed.Snapshot == nil || failed.State != string(models.CredentialObservationError) { + if err == nil || failed.Snapshot == nil || failed.State != string(models.CredentialObservationFresh) || + failed.LastErrorCode != "observation_upstream_failed" { t.Fatalf("failed = %#v, %v", failed, err) } if !reflect.DeepEqual(failed.Snapshot, first.Snapshot) { @@ -179,6 +180,51 @@ func TestRefreshCredentialObservationPersistsLKGAndAllowsImmediateManualRefresh( } } +func TestRefreshCredentialObservationKeepsFreshQuotaWhenPartialUsageIsMissing(t *testing.T) { + fixture, groupID, credentialID := newSubscriptionCredentialFixture(t) + now := time.Date(2026, time.August, 17, 10, 0, 0, 0, time.UTC) + fixture.service.now = func() time.Time { return now } + resetAt := now.Add(2 * time.Hour).Unix() + setCodexAccountObservation(fixture.service, func(context.Context, codex.Credential) (codex.AccountObservation, error) { + return codex.AccountObservation{Payload: []byte(fmt.Sprintf(`{ + "rate_limit":{"primary_window":{"limit_window_seconds":18000,"used_percent":100,"reset_at":%d}} + }`, resetAt))}, nil + }) + first, err := fixture.service.RefreshCredentialObservation(t.Context(), groupID, credentialID) + if err != nil { + t.Fatal(err) + } + if first.State != string(models.CredentialObservationFresh) || first.Snapshot == nil || len(first.Snapshot.QuotaWindows) != 1 || first.FreshUntilMS == nil { + t.Fatalf("initial observation = %#v", first) + } + if candidates := fixture.registry.CollectCredentialCandidates([]uint{groupID}, nil, now); len(candidates) != 0 { + t.Fatalf("exhausted candidates = %#v", candidates) + } + + now = now.Add(time.Minute) + fixture.service.observeSubscriptionAccount = func( + context.Context, + channel.ID, + subscriptionruntime.Credential, + ) (subscriptionruntime.Observation, error) { + return subscriptionruntime.Observation{ + Payload: []byte(`{"account_summary":{"display_name":"Updated owner"},"quota_windows":[]}`), + Partial: true, + AccountObserved: true, + }, nil + } + + result, err := fixture.service.RefreshCredentialObservation(t.Context(), groupID, credentialID) + if err != nil || result.State != string(models.CredentialObservationFresh) || + result.LastErrorCode != "observation_partial" || result.Snapshot == nil || + len(result.Snapshot.QuotaWindows) != 1 { + t.Fatalf("partial result = %#v / %#v / %v", result, result.Snapshot, err) + } + if candidates := fixture.registry.CollectCredentialCandidates([]uint{groupID}, nil, now); len(candidates) != 0 { + t.Fatalf("partial observation cleared quota protection: %#v", candidates) + } +} + func TestRefreshCredentialObservationRetriesOnceAfterUnauthorized(t *testing.T) { fixture, groupID, credentialID := newSubscriptionCredentialFixture(t) originalPrepare := fixture.service.prepareSubscriptionCredential @@ -245,6 +291,15 @@ func TestRefreshCredentialObservationClassifiesRepeatedAuthorizationFailure(t *t forcedRefreshes != 1 || observationCalls != 2 { t.Fatalf("result/error/refreshes/calls = %#v / %v / %d / %d", result, err, forcedRefreshes, observationCalls) } + var stored models.CredentialObservation + if err := fixture.db.Where("credential_id = ?", credentialID).Take(&stored).Error; err != nil || + stored.LastAuthRefreshSecretVersion == nil || *stored.LastAuthRefreshSecretVersion != 1 { + t.Fatalf("stored observation = %#v / %v", stored, err) + } + if _, err := fixture.service.RefreshCredentialObservation(t.Context(), groupID, credentialID); err == nil || + forcedRefreshes != 1 || observationCalls != 3 { + t.Fatalf("second result/refreshes/calls = %v / %d / %d", err, forcedRefreshes, observationCalls) + } } func TestRefreshCredentialObservationClassifiesNonRefreshableHTTPFailure(t *testing.T) { @@ -253,7 +308,7 @@ func TestRefreshCredentialObservationClassifiesNonRefreshableHTTPFailure(t *test statusCode int errorCode string }{ - {name: "forbidden", statusCode: 403, errorCode: "observation_authorization_failed"}, + {name: "forbidden", statusCode: 403, errorCode: "observation_access_denied"}, {name: "server error", statusCode: 500, errorCode: "observation_upstream_failed"}, } for _, test := range tests { @@ -581,8 +636,8 @@ func TestRefreshCredentialObservationPublishesAndClearsQuotaRoutingState(t *test if _, err := fixture.service.RefreshCredentialObservation(t.Context(), groupID, credentialID); err == nil { t.Fatal("forced failed refresh error = nil") } - if candidates := fixture.registry.CollectCredentialCandidates([]uint{groupID}, nil, now); len(candidates) != 1 { - t.Fatalf("failed observation fallback candidates = %#v", candidates) + if candidates := fixture.registry.CollectCredentialCandidates([]uint{groupID}, nil, now); len(candidates) != 0 { + t.Fatalf("failed observation lost quota protection = %#v", candidates) } } diff --git a/internal/execution/contracts.go b/internal/execution/contracts.go index f2df1bf5..5eae6b28 100644 --- a/internal/execution/contracts.go +++ b/internal/execution/contracts.go @@ -218,17 +218,19 @@ func (s AttemptSpec) Clone() AttemptSpec { return clone } -// DispatchState describes whether an attempt may have reached the upstream. +// DispatchState describes whether an attempt may have reached an upstream or +// was completed entirely within GPT-Load. type DispatchState string const ( DispatchNotSent DispatchState = "not_sent" DispatchMaybeSent DispatchState = "maybe_sent" + DispatchLocal DispatchState = "local" ) // Valid reports whether the dispatch state is recognized. func (s DispatchState) Valid() bool { - return s == DispatchNotSent || s == DispatchMaybeSent + return s == DispatchNotSent || s == DispatchMaybeSent || s == DispatchLocal } // MaxErrorSummaryLength bounds sanitized error messages retained by the contract. diff --git a/internal/execution/contracts_test.go b/internal/execution/contracts_test.go index f83bb0e0..4a8e2328 100644 --- a/internal/execution/contracts_test.go +++ b/internal/execution/contracts_test.go @@ -58,7 +58,7 @@ func TestOperationAndDispatchEnums(t *testing.T) { RouteRequirementNative.Allows(RouteConverted) { t.Fatal("route requirement mode policy is invalid") } - for _, state := range []DispatchState{DispatchNotSent, DispatchMaybeSent} { + for _, state := range []DispatchState{DispatchNotSent, DispatchMaybeSent, DispatchLocal} { if !state.Valid() { t.Fatalf("expected dispatch state %q to be valid", state) } @@ -448,6 +448,9 @@ func TestValidationAcceptsValidContractsAndRejectsInvalidFields(t *testing.T) { if err := (AttemptResult{DispatchState: DispatchMaybeSent, ResponseStarted: true, StatusCode: http.StatusOK}).Validate(); err != nil { t.Fatalf("valid AttemptResult.Validate() error = %v", err) } + if err := (AttemptResult{DispatchState: DispatchLocal, ResponseStarted: true, StatusCode: http.StatusOK, Body: []byte("local")}).Validate(); err != nil { + t.Fatalf("valid local AttemptResult.Validate() error = %v", err) + } if err := (StreamResult{DispatchState: DispatchNotSent, Error: &ErrorEvidence{Kind: ErrorKindTransport, Type: "transport", Summary: "dial failed"}}).Validate(); err != nil { t.Fatalf("valid StreamResult.Validate() error = %v", err) } diff --git a/internal/execution/cpa/adapter.go b/internal/execution/cpa/adapter.go index e4fcea8a..7b303a4e 100644 --- a/internal/execution/cpa/adapter.go +++ b/internal/execution/cpa/adapter.go @@ -93,6 +93,7 @@ func (a *Adapter) Execute(ctx context.Context, spec execution.AttemptSpec) execu execCtx, cancel := withRequestTimeout(ctx, spec.Timeouts.Request) defer cancel() var response providerResponse + localTokenCount := false if countTokensOperation(spec.Operation) { counter, ok := provider.(providerTokenCounter) if !ok { @@ -103,6 +104,17 @@ func (a *Adapter) Execute(ctx context.Context, spec execution.AttemptSpec) execu nil, ) } + if local, ok := counter.(providerLocalTokenCounter); ok { + localTokenCount = true + if err := local.ValidateLocalTokenCount(bridgeRequest(spec)); err != nil { + return unaryNotSent( + execution.ErrorKindInvalidRequest, + "local token count supports only stateless text input", + "local_token_count_unsupported_input", + err, + ) + } + } response, err = counter.CountTokens( execCtx, strconv.FormatUint(uint64(spec.Credential.ID), 10), @@ -118,6 +130,14 @@ func (a *Adapter) Execute(ctx context.Context, spec execution.AttemptSpec) execu ) } if err != nil { + if localTokenCount { + return unaryNotSent( + execution.ErrorKindInternal, + "local token count failed", + "local_token_count_failed", + err, + ) + } result := unaryExecutionError(execCtx, provider, err, credential) if result.Error != nil && execution.UpstreamCountTokensUnsupported( spec.Operation, @@ -133,8 +153,12 @@ func (a *Adapter) Execute(ctx context.Context, spec execution.AttemptSpec) execu } body := append([]byte(nil), response.Payload...) headers := convertedResponseHeaders(response.Headers, "application/json") + dispatchState := execution.DispatchMaybeSent + if response.Local { + dispatchState = execution.DispatchLocal + } return execution.AttemptResult{ - DispatchState: execution.DispatchMaybeSent, ResponseStarted: true, + DispatchState: dispatchState, ResponseStarted: true, UpstreamProtocol: provider.UpstreamProtocol(), AppliedReasoning: appliedReasoning(response.AppliedReasoningEffort), StatusCode: http.StatusOK, Header: headers, Body: body, Model: responseModel(body, spec.UpstreamModel), UpstreamRequestID: upstreamRequestID(headers), Usage: responseUsage(spec, body), diff --git a/internal/execution/cpa/adapter_test.go b/internal/execution/cpa/adapter_test.go index 2962a04e..0b13f78d 100644 --- a/internal/execution/cpa/adapter_test.go +++ b/internal/execution/cpa/adapter_test.go @@ -4,6 +4,7 @@ import ( "bytes" "context" "encoding/json" + "errors" "io" "net/http" "strings" @@ -210,22 +211,26 @@ func TestAdapterCountsCodexTokensForEverySupportedProtocol(t *testing.T) { operation execution.Operation routeMode execution.RouteMode wantFormat string + body string response string }{ { name: "OpenAI Responses", protocol: protocol.OpenAIResponses, operation: execution.OperationResponsesInputTokens, routeMode: execution.RouteNative, - wantFormat: "openai-response", response: `{"object":"response.input_tokens","input_tokens":7}`, + wantFormat: "openai-response", body: `{"model":"gpt-5","input":"hello"}`, + response: `{"object":"response.input_tokens","input_tokens":7}`, }, { name: "Anthropic", protocol: protocol.Anthropic, operation: execution.OperationCountTokens, routeMode: execution.RouteConverted, - wantFormat: "claude", response: `{"input_tokens":7}`, + wantFormat: "claude", body: `{"model":"gpt-5","messages":[{"role":"user","content":"hello"}]}`, + response: `{"input_tokens":7}`, }, { name: "Gemini", protocol: protocol.Gemini, operation: execution.OperationCountTokens, routeMode: execution.RouteConverted, - wantFormat: "gemini", response: `{"totalTokens":7}`, + wantFormat: "gemini", body: `{"contents":[{"role":"user","parts":[{"text":"hello"}]}]}`, + response: `{"totalTokens":7}`, }, } for _, test := range tests { @@ -237,9 +242,12 @@ func TestAdapterCountsCodexTokensForEverySupportedProtocol(t *testing.T) { spec.ClientProtocol = test.protocol spec.Operation = test.operation spec.RouteMode = test.routeMode + spec.Body = []byte(test.body) result := adapter.Execute(t.Context(), spec) - if result.Error != nil || fake.countCalls != 1 || fake.calls != 0 || + if result.Error != nil || result.DispatchState != execution.DispatchLocal || !result.ResponseStarted || + result.Header.Get(localTokenCountHeader) != "local-estimate" || + fake.countCalls != 1 || fake.calls != 0 || fake.request.Format != test.wantFormat || string(result.Body) != test.response { t.Fatalf("result=%#v calls=%d countCalls=%d request=%#v", result, fake.calls, fake.countCalls, fake.request) } @@ -247,6 +255,81 @@ func TestAdapterCountsCodexTokensForEverySupportedProtocol(t *testing.T) { } } +func TestAdapterRejectsUnsupportedLocalCodexTokenCountInput(t *testing.T) { + tests := []struct { + name string + protocol protocol.Protocol + operation execution.Operation + routeMode execution.RouteMode + model string + body string + }{ + { + name: "responses previous response", protocol: protocol.OpenAIResponses, + operation: execution.OperationResponsesInputTokens, routeMode: execution.RouteNative, + body: `{"model":"gpt-5","previous_response_id":"resp_123","input":"continue"}`, + }, + { + name: "responses image", protocol: protocol.OpenAIResponses, + operation: execution.OperationResponsesInputTokens, routeMode: execution.RouteNative, + body: `{"model":"gpt-5","input":[{"type":"message","role":"user","content":[{"type":"input_image","image_url":"https://example.test/image.png"}]}]}`, + }, + { + name: "anthropic document", protocol: protocol.Anthropic, + operation: execution.OperationCountTokens, routeMode: execution.RouteConverted, + body: `{"model":"gpt-5","messages":[{"role":"user","content":[{"type":"document","source":{"type":"base64","data":"AA=="}}]}]}`, + }, + { + name: "gemini inline data", protocol: protocol.Gemini, + operation: execution.OperationCountTokens, routeMode: execution.RouteConverted, + body: `{"contents":[{"role":"user","parts":[{"inlineData":{"mimeType":"image/png","data":"AA=="}}]}]}`, + }, + { + name: "unknown tokenizer model", protocol: protocol.OpenAIResponses, + operation: execution.OperationResponsesInputTokens, routeMode: execution.RouteNative, + model: "codex-unknown", body: `{"model":"codex-unknown","input":"hello"}`, + }, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + adapter, _, _, keyService, row := newAdapterFixture(t, credentialJSON("access", "refresh", time.Now().Add(time.Hour))) + fake := &fakeExecutor{countResult: codex.ExecuteResponse{Payload: []byte(`{"input_tokens":7}`)}} + setCodexExecutor(t, adapter, fake) + spec := validSpec(t, row, keyService) + spec.ClientProtocol = test.protocol + spec.Operation = test.operation + spec.RouteMode = test.routeMode + spec.Body = []byte(test.body) + if test.model != "" { + spec.ClientModel, spec.UpstreamModel = test.model, test.model + } + + result := adapter.Execute(t.Context(), spec) + if result.DispatchState != execution.DispatchNotSent || result.ResponseStarted || result.Error == nil || + result.Error.Kind != execution.ErrorKindInvalidRequest || result.Error.Code != "local_token_count_unsupported_input" || + fake.countCalls != 0 { + t.Fatalf("result=%#v countCalls=%d", result, fake.countCalls) + } + }) + } +} + +func TestAdapterMarksLocalCodexTokenCountFailureNotSent(t *testing.T) { + adapter, _, _, keyService, row := newAdapterFixture(t, credentialJSON("access", "refresh", time.Now().Add(time.Hour))) + fake := &fakeExecutor{err: errors.New("tokenizer failed")} + setCodexExecutor(t, adapter, fake) + spec := validSpec(t, row, keyService) + spec.Operation = execution.OperationResponsesInputTokens + spec.Path = "/v1/responses/input_tokens" + + result := adapter.Execute(t.Context(), spec) + if result.DispatchState != execution.DispatchNotSent || result.ResponseStarted || result.Error == nil || + result.Error.Kind != execution.ErrorKindInternal || result.Error.Code != "local_token_count_failed" || + fake.countCalls != 1 { + t.Fatalf("result=%#v countCalls=%d", result, fake.countCalls) + } +} + func TestAdapterStreamsEverySupportedClientProtocolThroughCPA(t *testing.T) { tests := []struct { name string diff --git a/internal/execution/cpa/codex_provider.go b/internal/execution/cpa/codex_provider.go index b9fcab7d..9a2137d5 100644 --- a/internal/execution/cpa/codex_provider.go +++ b/internal/execution/cpa/codex_provider.go @@ -2,10 +2,12 @@ package cpa import ( "context" + "encoding/json" "errors" "fmt" "net" "net/http" + "strings" "time" "gpt-load/internal/channel" @@ -27,6 +29,8 @@ type codexProviderBridge struct { executor codex.Executor } +const localTokenCountHeader = "X-GPT-Load-Token-Count" + func newCodexProviderBridge() *codexProviderBridge { return &codexProviderBridge{executor: codex.NewExecutor()} } @@ -74,12 +78,185 @@ func (bridge *codexProviderBridge) CountTokens( Model: request.Model, Payload: append([]byte(nil), request.Payload...), Format: request.Format, Headers: request.Headers.Clone(), OriginalRequest: append([]byte(nil), request.OriginalRequest...), }) + headers := response.Headers.Clone() + if headers == nil { + headers = make(http.Header) + } + headers.Set(localTokenCountHeader, "local-estimate") return providerResponse{ - Payload: append([]byte(nil), response.Payload...), Headers: response.Headers.Clone(), + Payload: append([]byte(nil), response.Payload...), Headers: headers, AppliedReasoningEffort: response.AppliedReasoningEffort, + Local: true, }, err } +func (*codexProviderBridge) ValidateLocalTokenCount(request providerRequest) error { + if !localTokenCountModelSupported(request.Model) { + return errors.New("local token count tokenizer is unavailable for model") + } + var root map[string]any + if err := json.Unmarshal(request.Payload, &root); err != nil || root == nil { + return errors.New("local token count request is invalid") + } + switch request.Format { + case "openai-response": + return validateLocalResponsesTokenCount(root) + case "claude": + return validateLocalClaudeTokenCount(root) + case "gemini": + return validateLocalGeminiTokenCount(root) + default: + return errors.New("local token count request format is unsupported") + } +} + +func localTokenCountModelSupported(model string) bool { + model = strings.ToLower(strings.TrimSpace(model)) + return strings.HasPrefix(model, "gpt-5") || strings.HasPrefix(model, "gpt-4.1") || + strings.HasPrefix(model, "gpt-4o") || strings.HasPrefix(model, "gpt-4") || + strings.HasPrefix(model, "gpt-3.5") || strings.HasPrefix(model, "gpt-3") +} + +func validateLocalResponsesTokenCount(root map[string]any) error { + if !onlyLocalTokenCountFields(root, "model", "input", "instructions") { + return errors.New("local token count responses fields are unsupported") + } + if instructions, ok := root["instructions"]; ok && !isString(instructions) { + return errors.New("local token count instructions must be text") + } + return validateLocalResponseInput(root["input"]) +} + +func validateLocalResponseInput(value any) error { + if value == nil || isString(value) { + return nil + } + items, ok := value.([]any) + if !ok { + return errors.New("local token count input must be text or messages") + } + for _, item := range items { + message, ok := item.(map[string]any) + if !ok || !onlyLocalTokenCountFields(message, "type", "role", "content") { + return errors.New("local token count input item is unsupported") + } + if kind, exists := message["type"]; exists && (!isString(kind) || kind.(string) != "message") { + return errors.New("local token count input item is unsupported") + } + if role, exists := message["role"]; exists && !isString(role) { + return errors.New("local token count message role is invalid") + } + if err := validateLocalTextContent(message["content"], "input_text", "text"); err != nil { + return err + } + } + return nil +} + +func validateLocalClaudeTokenCount(root map[string]any) error { + if !onlyLocalTokenCountFields(root, "model", "system", "messages") { + return errors.New("local token count Claude fields are unsupported") + } + if system, exists := root["system"]; exists { + if err := validateLocalTextContent(system, "text"); err != nil { + return err + } + } + messages, exists := root["messages"] + if !exists { + return nil + } + items, ok := messages.([]any) + if !ok { + return errors.New("local token count messages must be an array") + } + for _, item := range items { + message, ok := item.(map[string]any) + if !ok || !onlyLocalTokenCountFields(message, "role", "content") { + return errors.New("local token count message is unsupported") + } + if err := validateLocalTextContent(message["content"], "text"); err != nil { + return err + } + } + return nil +} + +func validateLocalGeminiTokenCount(root map[string]any) error { + if !onlyLocalTokenCountFields(root, "contents") { + return errors.New("local token count Gemini fields are unsupported") + } + contents, exists := root["contents"] + if !exists { + return nil + } + items, ok := contents.([]any) + if !ok { + return errors.New("local token count contents must be an array") + } + for _, item := range items { + content, ok := item.(map[string]any) + if !ok || !onlyLocalTokenCountFields(content, "role", "parts") { + return errors.New("local token count content is unsupported") + } + parts, ok := content["parts"].([]any) + if !ok { + return errors.New("local token count parts must be an array") + } + for _, part := range parts { + partObject, ok := part.(map[string]any) + if !ok || !onlyLocalTokenCountFields(partObject, "text") || !isString(partObject["text"]) { + return errors.New("local token count part is unsupported") + } + } + } + return nil +} + +func validateLocalTextContent(value any, allowedTypes ...string) error { + if isString(value) { + return nil + } + parts, ok := value.([]any) + if !ok { + return errors.New("local token count content must be text") + } + for _, part := range parts { + partObject, ok := part.(map[string]any) + if !ok || !onlyLocalTokenCountFields(partObject, "type", "text") || !isString(partObject["text"]) { + return errors.New("local token count content is unsupported") + } + kind, ok := partObject["type"].(string) + if !ok || !containsLocalTokenCountType(allowedTypes, kind) { + return errors.New("local token count content is unsupported") + } + } + return nil +} + +func onlyLocalTokenCountFields(value map[string]any, allowed ...string) bool { + for field := range value { + if !containsLocalTokenCountType(allowed, field) { + return false + } + } + return true +} + +func containsLocalTokenCountType(values []string, candidate string) bool { + for _, value := range values { + if value == candidate { + return true + } + } + return false +} + +func isString(value any) bool { + _, ok := value.(string) + return ok +} + func (*codexProviderBridge) ParseCredential(raw []byte) (providerCredential, error) { credential, err := codex.ParseCredentialJSON(raw) if err != nil { diff --git a/internal/execution/cpa/provider.go b/internal/execution/cpa/provider.go index e4a382d6..173efee5 100644 --- a/internal/execution/cpa/provider.go +++ b/internal/execution/cpa/provider.go @@ -57,6 +57,7 @@ type providerResponse struct { Payload []byte Headers http.Header AppliedReasoningEffort string + Local bool } type providerStreamResponse struct { @@ -89,6 +90,13 @@ type providerTokenCounter interface { CountTokens(context.Context, string, providerCredential, providerRequest) (providerResponse, error) } +// providerLocalTokenCounter is a deliberately narrow contract for providers +// whose CountTokens implementation never contacts an upstream service. +type providerLocalTokenCounter interface { + providerTokenCounter + ValidateLocalTokenCount(providerRequest) error +} + func indexProviderBridges(bridges ...providerBridge) map[channel.ProviderKind]providerBridge { indexed := make(map[channel.ProviderKind]providerBridge, len(bridges)) for _, bridge := range bridges { diff --git a/internal/execution/validation.go b/internal/execution/validation.go index 5078b63f..d0dee0e4 100644 --- a/internal/execution/validation.go +++ b/internal/execution/validation.go @@ -266,8 +266,8 @@ func validateResultMetadata( return validationError("dispatch_state", "not_sent cannot contain upstream response metadata") } } - if responseStarted && dispatchState != DispatchMaybeSent { - return validationError("dispatch_state", "response_started requires maybe_sent") + if responseStarted && dispatchState != DispatchMaybeSent && dispatchState != DispatchLocal { + return validationError("dispatch_state", "response_started requires maybe_sent or local") } if err := validateHeader(header); err != nil { return validationError("header", "contains an invalid name or value") diff --git a/internal/storage/db_test.go b/internal/storage/db_test.go index 4836470a..e9a1c644 100644 --- a/internal/storage/db_test.go +++ b/internal/storage/db_test.go @@ -623,6 +623,7 @@ func TestAutoMigrateCreatesUsageJournalAndMigrationLedger(t *testing.T) { "0005_credential_attempt_stats", "0006_upstream_protocol", "0007_model_price_mode_schedules", + "0008_observation_auth_refresh", } if !reflect.DeepEqual(migrationIDs, wantMigrationIDs) { t.Fatalf("schema_migrations IDs = %v, want %v", migrationIDs, wantMigrationIDs) diff --git a/internal/storage/migration.go b/internal/storage/migration.go index 0d448f94..0cf227e7 100644 --- a/internal/storage/migration.go +++ b/internal/storage/migration.go @@ -76,6 +76,12 @@ var migrations = []migration{ Validate: migrationfiles.Validate0007, ValidateRecoverable: migrationfiles.ValidateRecoverable0007, }, + { + ID: migrationfiles.ID0008, + Up: migrationfiles.Up0008, + Validate: migrationfiles.Validate0008, + ValidateRecoverable: migrationfiles.ValidateRecoverable0008, + }, } func applyMigrations(db *gorm.DB) error { diff --git a/internal/storage/migration_test.go b/internal/storage/migration_test.go index ce51f51c..eba786ce 100644 --- a/internal/storage/migration_test.go +++ b/internal/storage/migration_test.go @@ -19,6 +19,7 @@ func TestMigrationRegistryContainsOrderedMigrations(t *testing.T) { migrationfiles.ID0005, migrationfiles.ID0006, migrationfiles.ID0007, + migrationfiles.ID0008, } if len(migrations) != len(wantIDs) { t.Fatalf("migration registry length = %d, want %d", len(migrations), len(wantIDs)) diff --git a/internal/storage/migrations/0008_observation_auth_refresh.go b/internal/storage/migrations/0008_observation_auth_refresh.go new file mode 100644 index 00000000..06d88265 --- /dev/null +++ b/internal/storage/migrations/0008_observation_auth_refresh.go @@ -0,0 +1,89 @@ +package migrations + +import ( + "fmt" + "strings" + + "gorm.io/gorm" +) + +const ID0008 = "0008_observation_auth_refresh" + +const ( + credentialObservationsTable = "credential_observations" + lastAuthRefreshSecretVersionColumn0008 = "last_auth_refresh_secret_version" +) + +// Up0008 records the credential secret version already force-refreshed after +// an observation authorization failure, preventing repeated token rotation. +func Up0008(db *gorm.DB) error { + if err := validateObservationTable0008(db); err != nil { + return err + } + if !db.Migrator().HasColumn(credentialObservationsTable, lastAuthRefreshSecretVersionColumn0008) { + if err := db.Exec(fmt.Sprintf( + "ALTER TABLE %s ADD COLUMN %s BIGINT NULL", + quoteIdentifier0008(db, credentialObservationsTable), + quoteIdentifier0008(db, lastAuthRefreshSecretVersionColumn0008), + )).Error; err != nil { + return fmt.Errorf("add observation auth refresh version: %w", err) + } + } + return Validate0008(db) +} + +func ValidateRecoverable0008(db *gorm.DB) error { + if err := validateObservationTable0008(db); err != nil { + return err + } + if !db.Migrator().HasColumn(credentialObservationsTable, lastAuthRefreshSecretVersionColumn0008) { + return nil + } + return validateAuthRefreshSecretVersionColumn0008(db) +} + +func Validate0008(db *gorm.DB) error { + if err := validateObservationTable0008(db); err != nil { + return err + } + if !db.Migrator().HasColumn(credentialObservationsTable, lastAuthRefreshSecretVersionColumn0008) { + return fmt.Errorf("validate observation auth refresh schema: column %q is missing", lastAuthRefreshSecretVersionColumn0008) + } + return validateAuthRefreshSecretVersionColumn0008(db) +} + +func validateObservationTable0008(db *gorm.DB) error { + if db == nil || db.Dialector == nil || !db.Migrator().HasTable(credentialObservationsTable) { + return fmt.Errorf("validate observation auth refresh schema: table %q is missing", credentialObservationsTable) + } + return nil +} + +func validateAuthRefreshSecretVersionColumn0008(db *gorm.DB) error { + columns, err := db.Migrator().ColumnTypes(credentialObservationsTable) + if err != nil { + return fmt.Errorf("inspect observation auth refresh version column: %w", err) + } + for _, column := range columns { + if !strings.EqualFold(column.Name(), lastAuthRefreshSecretVersionColumn0008) { + continue + } + databaseType := strings.ToLower(strings.TrimSpace(column.DatabaseTypeName())) + if !strings.Contains(databaseType, "int") { + return fmt.Errorf("validate observation auth refresh version column: type %q is not integer", column.DatabaseTypeName()) + } + nullable, known := column.Nullable() + if !known || !nullable { + return fmt.Errorf("validate observation auth refresh version column: nullable is %t (known=%t), want true", nullable, known) + } + return nil + } + return fmt.Errorf("validate observation auth refresh version column: column %q is missing", lastAuthRefreshSecretVersionColumn0008) +} + +func quoteIdentifier0008(db *gorm.DB, value string) string { + if db != nil && db.Dialector != nil && strings.EqualFold(db.Dialector.Name(), "mysql") { + return "`" + value + "`" + } + return `"` + value + `"` +} diff --git a/internal/storage/migrations/0008_observation_auth_refresh_test.go b/internal/storage/migrations/0008_observation_auth_refresh_test.go new file mode 100644 index 00000000..cf27bd29 --- /dev/null +++ b/internal/storage/migrations/0008_observation_auth_refresh_test.go @@ -0,0 +1,30 @@ +package migrations_test + +import ( + "testing" + + "gorm.io/gorm" + + "gpt-load/internal/storage/migrations" +) + +func TestObservationAuthRefreshMigrationAddsNullableSecretVersion(t *testing.T) { + db := openInitialTestDatabase(t) + for _, up := range []func(*gorm.DB) error{ + migrations.Up0001, + migrations.Up0002, + } { + if err := up(db); err != nil { + t.Fatal(err) + } + } + if err := migrations.Up0008(db); err != nil { + t.Fatal(err) + } + if err := migrations.Validate0008(db); err != nil { + t.Fatal(err) + } + if err := migrations.Up0008(db); err != nil { + t.Fatal(err) + } +} diff --git a/internal/storage/models/subscription.go b/internal/storage/models/subscription.go index dcc3e0ca..565f08b4 100644 --- a/internal/storage/models/subscription.go +++ b/internal/storage/models/subscription.go @@ -43,17 +43,18 @@ const ( ) type CredentialObservation struct { - CredentialID uint `gorm:"primaryKey;not null"` - Credential *Credential `gorm:"foreignKey:CredentialID;references:ID;constraint:OnUpdate:CASCADE,OnDelete:CASCADE"` - IdentityFingerprint string `gorm:"type:varchar(128);not null"` - SchemaVersion uint `gorm:"not null;default:1"` - ObservationVersion uint64 `gorm:"not null;default:1"` - SnapshotJSON JSON `gorm:"column:snapshot_json;type:json;not null"` - State CredentialObservationState `gorm:"type:varchar(32);not null"` - ObservedAtMS *int64 `gorm:"column:observed_at_ms"` - FreshUntilMS *int64 `gorm:"column:fresh_until_ms"` - LastAttemptAtMS *int64 `gorm:"column:last_attempt_at_ms"` - NextAllowedAtMS *int64 `gorm:"column:next_allowed_at_ms"` - LastErrorCode string `gorm:"type:varchar(64);not null;default:''"` - UpdatedAtMS int64 `gorm:"column:updated_at_ms;not null"` + CredentialID uint `gorm:"primaryKey;not null"` + Credential *Credential `gorm:"foreignKey:CredentialID;references:ID;constraint:OnUpdate:CASCADE,OnDelete:CASCADE"` + IdentityFingerprint string `gorm:"type:varchar(128);not null"` + SchemaVersion uint `gorm:"not null;default:1"` + ObservationVersion uint64 `gorm:"not null;default:1"` + SnapshotJSON JSON `gorm:"column:snapshot_json;type:json;not null"` + State CredentialObservationState `gorm:"type:varchar(32);not null"` + ObservedAtMS *int64 `gorm:"column:observed_at_ms"` + FreshUntilMS *int64 `gorm:"column:fresh_until_ms"` + LastAttemptAtMS *int64 `gorm:"column:last_attempt_at_ms"` + NextAllowedAtMS *int64 `gorm:"column:next_allowed_at_ms"` + LastAuthRefreshSecretVersion *uint64 `gorm:"column:last_auth_refresh_secret_version"` + LastErrorCode string `gorm:"type:varchar(64);not null;default:''"` + UpdatedAtMS int64 `gorm:"column:updated_at_ms;not null"` } diff --git a/internal/subscription/runtime/claude_driver.go b/internal/subscription/runtime/claude_driver.go index 8f74764d..24e001ae 100644 --- a/internal/subscription/runtime/claude_driver.go +++ b/internal/subscription/runtime/claude_driver.go @@ -160,7 +160,9 @@ func (*claudeDriver) Observe(ctx context.Context, credential Credential) (Observ } return Observation{ Payload: normalized, Header: observed.Header.Clone(), - Partial: len(observed.IncompleteSources) > 0, + Partial: len(observed.IncompleteSources) > 0, + AccountObserved: observed.AccountObserved, + QuotaObserved: observed.UsageObserved, }, nil } diff --git a/internal/subscription/runtime/claude_quota.go b/internal/subscription/runtime/claude_quota.go index d7ff75d9..013f17f7 100644 --- a/internal/subscription/runtime/claude_quota.go +++ b/internal/subscription/runtime/claude_quota.go @@ -113,7 +113,9 @@ func normalizeClaudeAccount(profile claude.AccountProfile, extraUsage *claude.Ex ExtraUsageEnabled: cloneBoolPointer(profile.ExtraUsageEnabled), } if extraUsage != nil { - result.ExtraUsageEnabled = boolPointerValue(extraUsage.Enabled) + if extraUsage.Enabled != nil { + result.ExtraUsageEnabled = cloneBoolPointer(extraUsage.Enabled) + } if extraUsage.DisabledReason != nil { result.ExtraUsageDisabledReason = strings.TrimSpace(*extraUsage.DisabledReason) } @@ -223,7 +225,7 @@ func normalizeClaudeExtraUsage(extra *claude.ExtraUsage) (quotaWindow, error) { utilization := math.Min(1, *result.Used / *result.Limit) result.Utilization = &utilization } - if extra.Enabled && (result.Limit != nil || result.Utilization != nil) { + if extra.Enabled != nil && *extra.Enabled && (result.Limit != nil || result.Utilization != nil) { result.State = "available" if result.Remaining != nil && *result.Remaining <= 0 || result.Utilization != nil && *result.Utilization >= 1 { result.State = "exhausted" diff --git a/internal/subscription/runtime/claude_quota_test.go b/internal/subscription/runtime/claude_quota_test.go index 289e1a8b..741402e1 100644 --- a/internal/subscription/runtime/claude_quota_test.go +++ b/internal/subscription/runtime/claude_quota_test.go @@ -29,7 +29,7 @@ func TestNormalizeClaudeObservationIncludesAccountAndQuotaWindows(t *testing.T) Utilization: floatPointer(40), ResetsAt: stringPointer("2026-08-23T08:00:00Z"), }, ExtraUsage: &claude.ExtraUsage{ - Enabled: true, MonthlyLimit: floatPointer(100), UsedCredits: floatPointer(12.5), + Enabled: boolPointer(true), MonthlyLimit: floatPointer(100), UsedCredits: floatPointer(12.5), Utilization: floatPointer(12.5), Currency: stringPointer("USD"), }, Limits: []claude.ScopedLimit{{ @@ -89,7 +89,7 @@ func TestNormalizeClaudeObservationUsesUsageExtraStatus(t *testing.T) { raw, err := NormalizeClaudeObservation(claude.AccountObservation{ Profile: claude.AccountProfile{ExtraUsageEnabled: &enabled}, Usage: claude.Usage{ExtraUsage: &claude.ExtraUsage{ - Enabled: false, DisabledReason: &reason, + Enabled: boolPointer(false), DisabledReason: &reason, }}, }) if err != nil { @@ -110,6 +110,30 @@ func TestNormalizeClaudeObservationUsesUsageExtraStatus(t *testing.T) { } } +func TestNormalizeClaudeObservationKeepsProfileExtraStatusWhenUsageOmitsEnabled(t *testing.T) { + enabled := true + raw, err := NormalizeClaudeObservation(claude.AccountObservation{ + Profile: claude.AccountProfile{ExtraUsageEnabled: &enabled}, + Usage: claude.Usage{ExtraUsage: &claude.ExtraUsage{ + MonthlyLimit: floatPointer(100), + }}, + }) + if err != nil { + t.Fatal(err) + } + var snapshot struct { + Account struct { + ExtraUsageEnabled *bool `json:"extra_usage_enabled"` + } `json:"account_summary"` + } + if err := json.Unmarshal(raw, &snapshot); err != nil { + t.Fatal(err) + } + if snapshot.Account.ExtraUsageEnabled == nil || !*snapshot.Account.ExtraUsageEnabled { + t.Fatalf("account summary = %s", raw) + } +} + func TestNormalizeClaudeObservationKeepsQuotaWindowIDsUnique(t *testing.T) { raw, err := NormalizeClaudeObservation(claude.AccountObservation{Usage: claude.Usage{ FiveHour: &claude.UsageWindow{Utilization: floatPointer(10)}, diff --git a/internal/subscription/runtime/codex_driver.go b/internal/subscription/runtime/codex_driver.go index 00c2b431..b52a03a2 100644 --- a/internal/subscription/runtime/codex_driver.go +++ b/internal/subscription/runtime/codex_driver.go @@ -163,7 +163,7 @@ func (*codexDriver) Observe(ctx context.Context, credential Credential) (Observa if err != nil { return Observation{}, fmt.Errorf("%w: %v", ErrObservationPayloadInvalid, err) } - return Observation{Payload: normalized, Header: observed.Header.Clone()}, nil + return Observation{Payload: normalized, Header: observed.Header.Clone(), QuotaObserved: true}, nil } func (*codexDriver) Consume(ctx context.Context, credential Credential, requestID string) (ResetCreditResult, error) { diff --git a/internal/subscription/runtime/runtime.go b/internal/subscription/runtime/runtime.go index 5984e387..f7f4605c 100644 --- a/internal/subscription/runtime/runtime.go +++ b/internal/subscription/runtime/runtime.go @@ -125,9 +125,11 @@ var ErrObservationPayloadInvalid = errors.New("subscription observation payload // Observation contains the canonical, provider-neutral observation JSON. // Headers are retained only for bounded metadata such as retry timing. type Observation struct { - Payload []byte - Header http.Header - Partial bool + Payload []byte + Header http.Header + Partial bool + AccountObserved bool + QuotaObserved bool } // QuotaObservation is a narrow optional account observation capability. diff --git a/third_party/cpaembedded/embedded/claude_account.go b/third_party/cpaembedded/embedded/claude_account.go index c06c56fd..04b5bc5e 100644 --- a/third_party/cpaembedded/embedded/claude_account.go +++ b/third_party/cpaembedded/embedded/claude_account.go @@ -73,7 +73,7 @@ type ClaudeUsageWindow struct { // ClaudeExtraUsage describes optional paid usage credits. type ClaudeExtraUsage struct { - Enabled bool + Enabled *bool MonthlyLimit *float64 UsedCredits *float64 Utilization *float64 @@ -108,6 +108,8 @@ type ClaudeAccountObservation struct { Profile ClaudeAccountProfile Usage ClaudeUsage Header http.Header + AccountObserved bool + UsageObserved bool IncompleteSources []string } @@ -164,10 +166,10 @@ type claudeUsagePayload struct { CinderCove *claudeUsageWindowPayload `json:"cinder_cove"` ExtraUsage *claudeExtraUsagePayload `json:"extra_usage"` Limits []struct { - Kind string `json:"kind"` - Group string `json:"group"` - Percent float64 `json:"percent"` - ResetsAt *string `json:"resets_at"` + Kind string `json:"kind"` + Group string `json:"group"` + Percent *float64 `json:"percent"` + ResetsAt *string `json:"resets_at"` Scope *struct { Model *struct { DisplayName string `json:"display_name"` @@ -310,8 +312,11 @@ func ObserveClaudeAccount( } incomplete := make([]string, 0, 4) successfulSources := 0 + accountObserved, usageObserved := false, false failedSources, commonHTTPStatus := 0, 0 onlyHTTPFailures := true + authorizationFailures := 0 + sawUnauthorized := false recordFailure := func(sourceErr error) { failedSources++ var upstream *ClaudeUpstreamHTTPError @@ -319,6 +324,10 @@ func ObserveClaudeAccount( onlyHTTPFailures = false return } + if upstream.StatusCode == http.StatusUnauthorized || upstream.StatusCode == http.StatusForbidden { + authorizationFailures++ + sawUnauthorized = sawUnauthorized || upstream.StatusCode == http.StatusUnauthorized + } if commonHTTPStatus == 0 { commonHTTPStatus = upstream.StatusCode } else if commonHTTPStatus != upstream.StatusCode { @@ -336,6 +345,7 @@ func ObserveClaudeAccount( } else { profile = mergeClaudeAccountProfile(profile, observedProfile) successfulSources++ + accountObserved = true } } else if isClaudeObservationContextError(profileErr) { return ClaudeAccountObservation{}, profileErr @@ -370,6 +380,7 @@ func ObserveClaudeAccount( } if rolesUsable { successfulSources++ + accountObserved = true } if !rolesUsable || !rolesComplete { incomplete = append(incomplete, "roles") @@ -390,6 +401,7 @@ func ObserveClaudeAccount( bootstrapUsable, bootstrapComplete := applyClaudeBootstrapAccount(&profile, bootstrap.OAuthAccount) if bootstrapUsable { successfulSources++ + accountObserved = true } if !bootstrapUsable || !bootstrapComplete { incomplete = append(incomplete, "bootstrap") @@ -413,11 +425,19 @@ func ObserveClaudeAccount( recordFailure(err) } else if claudeUsageUsable(usage) { successfulSources++ + usageObserved = true } else { incomplete = append(incomplete, "usage") recordFailure(nil) } if successfulSources == 0 { + if failedSources == 4 && authorizationFailures == failedSources { + status := http.StatusForbidden + if sawUnauthorized { + status = http.StatusUnauthorized + } + return ClaudeAccountObservation{}, &ClaudeUpstreamHTTPError{StatusCode: status} + } if failedSources == 4 && onlyHTTPFailures && commonHTTPStatus != 0 { return ClaudeAccountObservation{}, &ClaudeUpstreamHTTPError{StatusCode: commonHTTPStatus} } @@ -425,6 +445,7 @@ func ObserveClaudeAccount( } return ClaudeAccountObservation{ Profile: profile, Usage: usage, Header: header, + AccountObserved: accountObserved, UsageObserved: usageObserved, IncompleteSources: incomplete, }, nil } @@ -711,7 +732,7 @@ func mapClaudeUsage(payload claudeUsagePayload) (ClaudeUsage, error) { return ClaudeUsage{}, err } result.ExtraUsage = &ClaudeExtraUsage{ - Enabled: payload.ExtraUsage.Enabled != nil && *payload.ExtraUsage.Enabled, + Enabled: cloneBool(payload.ExtraUsage.Enabled), MonthlyLimit: cloneFloat64(payload.ExtraUsage.MonthlyLimit), UsedCredits: cloneFloat64(payload.ExtraUsage.UsedCredits), Utilization: cloneFloat64(payload.ExtraUsage.Utilization), @@ -723,8 +744,8 @@ func mapClaudeUsage(payload claudeUsagePayload) (ClaudeUsage, error) { } result.Limits = make([]ClaudeScopedLimit, 0, len(payload.Limits)) for _, limit := range payload.Limits { - if math.IsNaN(limit.Percent) || math.IsInf(limit.Percent, 0) || - limit.Percent < 0 || limit.Percent > 100 { + if limit.Percent == nil || math.IsNaN(*limit.Percent) || math.IsInf(*limit.Percent, 0) || + *limit.Percent < 0 || *limit.Percent > 100 { return ClaudeUsage{}, fmt.Errorf("Claude scoped limit percent is invalid") } kind, err := claudeSafeScalar("limit kind", limit.Kind, true) @@ -739,7 +760,7 @@ func mapClaudeUsage(payload claudeUsagePayload) (ClaudeUsage, error) { if err != nil { return ClaudeUsage{}, err } - item := ClaudeScopedLimit{Kind: kind, Group: group, Percent: limit.Percent, ResetsAt: reset} + item := ClaudeScopedLimit{Kind: kind, Group: group, Percent: *limit.Percent, ResetsAt: reset} if limit.Scope != nil { if limit.Scope.Model != nil { item.ModelDisplayName, err = claudeSafeScalar( @@ -765,8 +786,8 @@ func mapClaudeUsage(payload claudeUsagePayload) (ClaudeUsage, error) { func claudeExtraUsagePayloadUsable(payload *claudeExtraUsagePayload) bool { return payload != nil && (payload.Enabled != nil || payload.MonthlyLimit != nil || - payload.UsedCredits != nil || payload.Utilization != nil || payload.Currency != nil || - payload.DisabledReason != nil) + payload.UsedCredits != nil || payload.Utilization != nil || + payload.DisabledReason != nil && strings.TrimSpace(*payload.DisabledReason) != "") } func mapClaudeUsageWindow(name string, payload *claudeUsageWindowPayload) (*ClaudeUsageWindow, error) { @@ -872,3 +893,11 @@ func cloneFloat64(value *float64) *float64 { cloned := *value return &cloned } + +func cloneBool(value *bool) *bool { + if value == nil { + return nil + } + cloned := *value + return &cloned +} diff --git a/third_party/cpaembedded/embedded/claude_test.go b/third_party/cpaembedded/embedded/claude_test.go index 038c9260..5f5fe31b 100644 --- a/third_party/cpaembedded/embedded/claude_test.go +++ b/third_party/cpaembedded/embedded/claude_test.go @@ -900,7 +900,8 @@ func TestObserveClaudeAccountCombinesProfileRolesBootstrapAndUsage(t *testing.T) if observation.Usage.FiveHour == nil || observation.Usage.FiveHour.Utilization == nil || *observation.Usage.FiveHour.Utilization != 25 || observation.Usage.ExtraUsage == nil || observation.Usage.ExtraUsage.MonthlyLimit == nil || - len(observation.Usage.Limits) != 1 || observation.Header.Get("Request-Id") != "usage-request-one" { + len(observation.Usage.Limits) != 1 || !observation.AccountObserved || !observation.UsageObserved || + observation.Header.Get("Request-Id") != "usage-request-one" { t.Fatalf("usage/header = %#v / %v", observation.Usage, observation.Header) } } @@ -940,6 +941,28 @@ func TestObserveClaudeAccountFailsWhenEverySourceIsUnavailable(t *testing.T) { } } +func TestObserveClaudeAccountClassifiesMixedAuthorizationFailures(t *testing.T) { + credential := testClaudeExecutionCredential() + server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { + status := http.StatusForbidden + if request.URL.Path == "/profile" || request.URL.Path == "/bootstrap" { + status = http.StatusUnauthorized + } + writer.WriteHeader(status) + })) + defer server.Close() + + _, err := ObserveClaudeAccount(t.Context(), credential, ClaudeOptions{ + ProfileURL: server.URL + "/profile", RolesURL: server.URL + "/roles", + BootstrapURL: server.URL + "/bootstrap", UsageURL: server.URL + "/usage", + HTTPClient: server.Client(), + }) + var upstream *ClaudeUpstreamHTTPError + if !errors.As(err, &upstream) || upstream.StatusCode != http.StatusUnauthorized { + t.Fatalf("observation error = %v", err) + } +} + func TestObserveClaudeAccountRejectsSourceWithoutUsableData(t *testing.T) { credential := testClaudeExecutionCredential() tests := []struct { @@ -952,6 +975,8 @@ func TestObserveClaudeAccountRejectsSourceWithoutUsableData(t *testing.T) { {name: "empty usage", path: "/usage", payload: `{}`}, {name: "empty usage window", path: "/usage", payload: `{"five_hour":{}}`}, {name: "empty extra usage", path: "/usage", payload: `{"extra_usage":{}}`}, + {name: "currency-only extra usage", path: "/usage", payload: `{"extra_usage":{"currency":""}}`}, + {name: "limit without percent", path: "/usage", payload: `{"limits":[{"kind":"weekly","group":"opus"}]}`}, {name: "unknown usage schema", path: "/usage", payload: `{"new_unknown_field":"value"}`}, } for _, test := range tests { @@ -1104,7 +1129,7 @@ func TestObserveClaudeAccountKeepsUsageWhenProfileIsUnavailable(t *testing.T) { if observation.Profile.Email != credential.Email || observation.Profile.OrganizationName != "Bootstrap Org" || observation.Profile.SeatTier != "team_standard" || - observation.Usage.FiveHour == nil { + observation.Usage.FiveHour == nil || !observation.AccountObserved || !observation.UsageObserved { t.Fatalf("observation = %#v", observation) } } diff --git a/web/src/features/groups/credentials/SubscriptionAccountCard.vue b/web/src/features/groups/credentials/SubscriptionAccountCard.vue index b0a5c4db..60966ee3 100644 --- a/web/src/features/groups/credentials/SubscriptionAccountCard.vue +++ b/web/src/features/groups/credentials/SubscriptionAccountCard.vue @@ -276,8 +276,10 @@ const failureLabel = computed(() => }`, ) const observationErrorKeys: Readonly> = { + observation_access_denied: 'group.credentials.subscription.observationError.accessDenied', observation_authorization_failed: 'group.credentials.subscription.observationError.authorizationFailed', + observation_partial: 'group.credentials.subscription.observationError.partial', observation_upstream_failed: 'group.credentials.subscription.observationError.upstreamFailed', observation_payload_invalid: 'group.credentials.subscription.observationError.payloadInvalid', } diff --git a/web/src/i18n/locales/en-US/group.ts b/web/src/i18n/locales/en-US/group.ts index e87905cb..f3a2c8df 100644 --- a/web/src/i18n/locales/en-US/group.ts +++ b/web/src/i18n/locales/en-US/group.ts @@ -466,8 +466,12 @@ export default { refreshStartFailed: 'The credential refresh could not start. Try again shortly.', }, observationError: { + accessDenied: + 'The upstream denied access to subscription information. Check account permissions or plan.', authorizationFailed: 'The upstream subscription authorization is no longer valid. Sign in again.', + partial: + 'Some subscription information is unavailable; the last valid quota is retained.', upstreamFailed: 'Unable to read upstream subscription entitlement information', payloadInvalid: 'The upstream returned unrecognized subscription entitlement information', }, diff --git a/web/src/i18n/locales/ja-JP/group.ts b/web/src/i18n/locales/ja-JP/group.ts index bff878bb..0280ea12 100644 --- a/web/src/i18n/locales/ja-JP/group.ts +++ b/web/src/i18n/locales/ja-JP/group.ts @@ -471,7 +471,11 @@ export default { '認証情報の更新を開始できません。しばらくしてから再試行してください。', }, observationError: { + accessDenied: + '上流がサブスクリプション情報へのアクセスを拒否しました。権限またはプランを確認してください', authorizationFailed: '上流のサブスクリプション認証が無効です。再ログインしてください', + partial: + '一部のサブスクリプション情報を取得できません。直近の有効な利用枠を保持しています', upstreamFailed: '上流のサブスクリプション利用枠情報を取得できません', payloadInvalid: '上流から認識できないサブスクリプション利用枠情報が返されました', }, diff --git a/web/src/i18n/locales/zh-CN/group.ts b/web/src/i18n/locales/zh-CN/group.ts index a63f207c..601d5b32 100644 --- a/web/src/i18n/locales/zh-CN/group.ts +++ b/web/src/i18n/locales/zh-CN/group.ts @@ -449,7 +449,9 @@ export default { refreshStartFailed: '当前无法开始刷新凭据，请稍后重试。', }, observationError: { + accessDenied: '上游拒绝读取订阅权益，请检查账号权限或套餐', authorizationFailed: '上游订阅授权已失效，请重新登录账号', + partial: '部分上游订阅信息暂不可用，正在保留最近有效额度', upstreamFailed: '无法读取上游订阅权益信息', payloadInvalid: '上游返回了无法识别的订阅权益信息', },