fix(subscription): 收口观察与本地计数边界
This commit is contained in:
27 files changed
+688
-64
No files matched your search
@@ -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),
|
||||
|
||||
@@ -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(`{}`)
|
||||
}
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
|
||||
@@ -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),
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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")
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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))
|
||||
|
||||
@@ -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 + `"`
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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"`
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
|
||||
|
||||
@@ -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"
|
||||
|
||||
@@ -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)},
|
||||
|
||||
@@ -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) {
|
||||
|
||||
@@ -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.
|
||||
|
||||
+40
-11
@@ -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
|
||||
}
|
||||
+27
-2
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -276,8 +276,10 @@ const failureLabel = computed(() =>
|
||||
}`,
|
||||
)
|
||||
const observationErrorKeys: Readonly<Record<string, string>> = {
|
||||
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',
|
||||
}
|
||||
|
||||
@@ -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',
|
||||
},
|
||||
|
||||
@@ -471,7 +471,11 @@ export default {
|
||||
'認証情報の更新を開始できません。しばらくしてから再試行してください。',
|
||||
},
|
||||
observationError: {
|
||||
accessDenied:
|
||||
'上流がサブスクリプション情報へのアクセスを拒否しました。権限またはプランを確認してください',
|
||||
authorizationFailed: '上流のサブスクリプション認証が無効です。再ログインしてください',
|
||||
partial:
|
||||
'一部のサブスクリプション情報を取得できません。直近の有効な利用枠を保持しています',
|
||||
upstreamFailed: '上流のサブスクリプション利用枠情報を取得できません',
|
||||
payloadInvalid: '上流から認識できないサブスクリプション利用枠情報が返されました',
|
||||
},
|
||||
|
||||
@@ -449,7 +449,9 @@ export default {
|
||||
refreshStartFailed: '当前无法开始刷新凭据,请稍后重试。',
|
||||
},
|
||||
observationError: {
|
||||
accessDenied: '上游拒绝读取订阅权益,请检查账号权限或套餐',
|
||||
authorizationFailed: '上游订阅授权已失效,请重新登录账号',
|
||||
partial: '部分上游订阅信息暂不可用,正在保留最近有效额度',
|
||||
upstreamFailed: '无法读取上游订阅权益信息',
|
||||
payloadInvalid: '上游返回了无法识别的订阅权益信息',
|
||||
},
|
||||
|
||||
Reference in new issue
Block a user