diff --git a/internal/auth/claude/anthropic_auth.go b/internal/auth/claude/anthropic_auth.go index c6dd901e0f2..369468825d5 100644 --- a/internal/auth/claude/anthropic_auth.go +++ b/internal/auth/claude/anthropic_auth.go @@ -14,6 +14,8 @@ import ( "sync" "time" + "github.com/MoeForge/moego/ojson" + "github.com/router-for-me/CLIProxyAPI/v8/internal/auth/codex" "github.com/router-for-me/CLIProxyAPI/v8/internal/config" log "github.com/sirupsen/logrus" "golang.org/x/sync/singleflight" @@ -510,7 +512,16 @@ func (o *ClaudeAuth) RefreshTokens(ctx context.Context, refreshToken string) (*C refreshCtx, cancelRefresh := context.WithTimeout(context.WithoutCancel(ctx), claudeRefreshTimeout) defer cancelRefresh() refreshCtx = context.WithValue(refreshCtx, claudeRefreshHandshakeTimeoutContextKey{}, claudeRefreshHandshakeTimeout) - return o.refreshTokensSingleFlight(refreshCtx, refreshToken) + // Without a journal bound by the store, refresh as before. With one, + // the grant is claimed first and never replayed after a lost response. + journal, journaled := codex.RefreshJournalFromContext(ctx) + if !journaled { + return o.refreshTokensSingleFlight(refreshCtx, refreshToken) + } + return codex.JournalRefresh(refreshCtx, journal, ClientID, refreshToken, func() (*ClaudeTokenData, error) { + tokenData, errRefresh := o.refreshTokensSingleFlight(refreshCtx, refreshToken) + return tokenData, codex.QualifyRefreshOutcome("response", errRefresh) + }, decodeRefreshReceipt) }) if err != nil { return nil, err @@ -522,6 +533,23 @@ func (o *ClaudeAuth) RefreshTokens(ctx context.Context, refreshToken string) (*C return tokenData, nil } +// decodeRefreshReceipt rebuilds the tokens a successful refresh saved, for +// replicas that lost the claim. +func decodeRefreshReceipt(value ojson.Result) (*ClaudeTokenData, error) { + if !value.IsObject() || value.Get("access_token").String() == "" || value.Get("expired").String() == "" { + return nil, fmt.Errorf("invalid refresh receipt") + } + return &ClaudeTokenData{ + AccessToken: value.Get("access_token").String(), + RefreshToken: value.Get("refresh_token").String(), + Email: value.Get("email").String(), + AccountUUID: value.Get("account_uuid").String(), + OrganizationUUID: value.Get("organization_uuid").String(), + OrganizationName: value.Get("organization_name").String(), + Expire: value.Get("expired").String(), + }, nil +} + func (o *ClaudeAuth) refreshTokensSingleFlight(ctx context.Context, refreshToken string) (*ClaudeTokenData, error) { if blockedUntil := claudeRefreshBlockedUntil(refreshToken); blockedUntil.After(time.Now()) { return nil, &refreshHTTPError{ diff --git a/internal/auth/claude/refresh_journal_test.go b/internal/auth/claude/refresh_journal_test.go new file mode 100644 index 00000000000..2b35ea47f3c --- /dev/null +++ b/internal/auth/claude/refresh_journal_test.go @@ -0,0 +1,122 @@ +package claude + +import ( + "context" + "errors" + "io" + "net/http" + "strings" + "sync/atomic" + "testing" + + "github.com/router-for-me/CLIProxyAPI/v8/internal/auth/codex" +) + +// claimSpy records whether the grant was claimed before any request was sent. +type claimSpy struct { + codex.RefreshJournal + claimed atomic.Bool +} + +func (s *claimSpy) RefreshClaim(ctx context.Context, key string) (bool, error) { + claimed, err := s.RefreshJournal.RefreshClaim(ctx, key) + if claimed { + s.claimed.Store(true) + } + return claimed, err +} + +func TestRefreshTokensJournalClaimsBeforeSendingAndServesLateCallers(t *testing.T) { + resetClaudeRefreshState() + defer resetClaudeRefreshState() + + spy := &claimSpy{RefreshJournal: codex.NewMemoryRefreshJournal()} + var calls atomic.Int32 + var claimedBeforeSend atomic.Bool + transport := roundTripFunc(func(req *http.Request) (*http.Response, error) { + calls.Add(1) + claimedBeforeSend.Store(spy.claimed.Load()) + if strings.Contains(req.URL.Host, "api.anthropic.com") { + // The profile lookup after the rotation is not part of the grant. + return &http.Response{StatusCode: http.StatusNotFound, Body: io.NopCloser(strings.NewReader("{}")), Header: make(http.Header)}, nil + } + return &http.Response{ + StatusCode: http.StatusOK, + Body: io.NopCloser(strings.NewReader(`{"access_token":"new-access","refresh_token":"new-refresh","expires_in":3600}`)), + Header: make(http.Header), + }, nil + }) + ctx := codex.WithRefreshJournal(context.Background(), spy) + first := &ClaudeAuth{httpClient: &http.Client{Transport: transport}} + got, err := first.RefreshTokens(ctx, "rotating-grant") + if err != nil || got.AccessToken != "new-access" || got.RefreshToken != "new-refresh" { + t.Fatalf("first refresh = %+v, %v", got, err) + } + if !claimedBeforeSend.Load() { + t.Fatal("the grant was sent before it was claimed") + } + sent := calls.Load() + + // Another replica sharing the journal adopts the saved receipt. + second := &ClaudeAuth{httpClient: &http.Client{Transport: roundTripFunc(func(*http.Request) (*http.Response, error) { + t.Error("a late caller sent the grant again") + return nil, errors.New("unexpected request") + })}} + late, err := second.RefreshTokens(ctx, "rotating-grant") + if err != nil || late.AccessToken != "new-access" || late.RefreshToken != "new-refresh" || late.Expire != got.Expire { + t.Fatalf("late refresh = %+v, %v", late, err) + } + if calls.Load() != sent { + t.Fatalf("requests after the receipt = %d, want %d", calls.Load(), sent) + } +} + +func TestRefreshTokensJournalNeverReplaysALostResponse(t *testing.T) { + resetClaudeRefreshState() + defer resetClaudeRefreshState() + + var calls atomic.Int32 + lost := &ClaudeAuth{httpClient: &http.Client{Transport: roundTripFunc(func(*http.Request) (*http.Response, error) { + calls.Add(1) + return &http.Response{StatusCode: http.StatusOK, Body: &errorReader{err: io.ErrUnexpectedEOF}, Header: make(http.Header)}, nil + })}} + ctx := codex.WithRefreshJournal(context.Background(), codex.NewMemoryRefreshJournal()) + if _, err := lost.RefreshTokensWithRetry(ctx, "lost-grant", 3); !codex.IsRefreshOutcomeUnknown(err) { + t.Fatalf("lost response = %v, want refresh_outcome_unknown", err) + } + for _, status := range []int{http.StatusInternalServerError, http.StatusTooManyRequests} { + other := &ClaudeAuth{httpClient: &http.Client{Transport: roundTripFunc(func(*http.Request) (*http.Response, error) { + calls.Add(1) + return &http.Response{StatusCode: status, Body: io.NopCloser(strings.NewReader("{}")), Header: make(http.Header)}, nil + })}} + if _, err := other.RefreshTokensWithRetry(ctx, "lost-grant", 3); !codex.IsRefreshOutcomeUnknown(err) { + t.Fatalf("retry after a lost response = %v, want refresh_outcome_unknown", err) + } + } + if calls.Load() != 1 { + t.Fatalf("requests = %d, want exactly the first", calls.Load()) + } +} + +func TestRefreshTokensJournalRecordsOnlyQualifiedRejections(t *testing.T) { + resetClaudeRefreshState() + defer resetClaudeRefreshState() + + journal := codex.NewMemoryRefreshJournal() + ctx := codex.WithRefreshJournal(context.Background(), journal) + status := http.StatusBadRequest + var calls atomic.Int32 + auth := &ClaudeAuth{httpClient: &http.Client{Transport: roundTripFunc(func(*http.Request) (*http.Response, error) { + calls.Add(1) + return &http.Response{StatusCode: status, Body: io.NopCloser(strings.NewReader(`{"error":"invalid_grant"}`)), Header: make(http.Header)}, nil + })}} + _, err := auth.RefreshTokens(ctx, "rejected-grant") + if err == nil || codex.IsRefreshOutcomeUnknown(err) { + t.Fatalf("rejected grant = %v, want a plain rejection", err) + } + // The rejection is final for every replica: nothing is sent again. + _, err = auth.RefreshTokens(ctx, "rejected-grant") + if err == nil || !strings.Contains(err.Error(), "invalid_grant") || calls.Load() != 1 { + t.Fatalf("second attempt = %v after %d requests, want the recorded rejection", err, calls.Load()) + } +} diff --git a/internal/auth/codex/refresh_journal.go b/internal/auth/codex/refresh_journal.go index e41cfa9881d..59200393b46 100644 --- a/internal/auth/codex/refresh_journal.go +++ b/internal/auth/codex/refresh_journal.go @@ -52,6 +52,14 @@ func WithRefreshJournal(ctx context.Context, journal RefreshJournal) context.Con return context.WithValue(ctx, refreshJournalContextKey{}, journal) } +// RefreshJournalFromContext returns the journal a caller bound with +// WithRefreshJournal. Providers without a journal of their own refresh with no +// claim unless one is bound, so a store opts them in rather than changing them. +func RefreshJournalFromContext(ctx context.Context) (RefreshJournal, bool) { + journal, ok := ctx.Value(refreshJournalContextKey{}).(RefreshJournal) + return journal, ok && journal != nil +} + // SetRefreshJournal configures the journal for callers outside the auth manager. func (o *CodexAuth) SetRefreshJournal(journal RefreshJournal) { o.refreshJournal = journal } @@ -143,6 +151,29 @@ func journalRefresh[T any](ctx context.Context, journal RefreshJournal, clientID return value, errRun } +// JournalRefresh runs one refresh of grant under the journal for any provider. +// The claim is recorded before run sends anything and the receipt is saved after +// it, so a response lost in between leaves a pending claim that no replica +// replays; late callers are served from the receipt chain. clientID keeps the +// grant digests of different providers apart. Its receipts need the same shape +// codex writes: value.expired as RFC 3339 and value.refresh_token for rotation, +// which decode turns back into T. +func JournalRefresh[T any](ctx context.Context, journal RefreshJournal, clientID, grant string, run func() (T, error), decode func(ojson.Result) (T, error)) (T, error) { + return journalRefresh(ctx, journal, clientID, grant, run, decode) +} + +// QualifyRefreshOutcome keeps a failure as a rejection only when the provider +// answered with a response known to leave the grant unconsumed, using the same +// rules as codex. Anything else, a transport error, an unreadable body or a +// bare 429 or 5xx, may have consumed the grant and becomes an unknown outcome +// at stage, which the journal records instead of letting the grant be replayed. +func QualifyRefreshOutcome(stage string, err error) error { + if err == nil || IsRefreshOutcomeUnknown(err) || isNonRetryableRefreshErr(err) { + return err + } + return &RefreshOutcomeUnknownError{Stage: stage, Cause: err} +} + func usableRefreshReceipt[T any](value T, expires string) (T, error) { expiry, errParse := time.Parse(time.RFC3339, expires) if errParse != nil || !expiry.After(time.Now()) { @@ -205,6 +236,11 @@ type memoryRefreshJournal struct { var processRefreshJournal memoryRefreshJournal +// NewMemoryRefreshJournal returns a journal that lives only in this process, for +// direct SDK callers and tests; it protects against replay within the process +// and never across replicas. +func NewMemoryRefreshJournal() RefreshJournal { return &memoryRefreshJournal{} } + func (j *memoryRefreshJournal) RefreshClaim(_ context.Context, key string) (bool, error) { j.mu.Lock() defer j.mu.Unlock() diff --git a/internal/auth/kimi/kimi.go b/internal/auth/kimi/kimi.go index 973c1b5240b..51384f159db 100644 --- a/internal/auth/kimi/kimi.go +++ b/internal/auth/kimi/kimi.go @@ -14,7 +14,9 @@ import ( "strings" "time" + "github.com/MoeForge/moego/ojson" "github.com/google/uuid" + "github.com/router-for-me/CLIProxyAPI/v8/internal/auth/codex" "github.com/router-for-me/CLIProxyAPI/v8/internal/buildinfo" "github.com/router-for-me/CLIProxyAPI/v8/internal/config" "github.com/router-for-me/CLIProxyAPI/v8/internal/util" @@ -586,7 +588,24 @@ func (c *DeviceFlowClient) RefreshToken(ctx context.Context, refreshToken string flightKey := c.tokenURL() + ":" + refreshToken result, err, _ := kimiRefreshGroup.Do(flightKey, func() (interface{}, error) { - return c.refreshTokenSingleFlight(context.WithoutCancel(ctx), refreshToken) + refreshCtx := context.WithoutCancel(ctx) + // Without a journal bound by the store, refresh as before. With one, + // the grant is claimed first and never replayed after a lost response. + journal, journaled := codex.RefreshJournalFromContext(ctx) + if !journaled { + return c.refreshTokenSingleFlight(refreshCtx, refreshToken) + } + receipt, errRefresh := codex.JournalRefresh(refreshCtx, journal, kimiClientID, refreshToken, func() (*kimiRefreshReceipt, error) { + tokenData, errSend := c.refreshTokenSingleFlight(refreshCtx, refreshToken) + if errSend = codex.QualifyRefreshOutcome("response", errSend); errSend != nil { + return nil, errSend + } + return newKimiRefreshReceipt(tokenData), nil + }, decodeKimiRefreshReceipt) + if errRefresh != nil { + return nil, errRefresh + } + return &receipt.KimiTokenData, nil }) if err != nil { return nil, err @@ -598,6 +617,40 @@ func (c *DeviceFlowClient) RefreshToken(ctx context.Context, refreshToken string return tokenData, nil } +// kimiRefreshReceipt is what a successful refresh saves. Kimi reports expiry as +// a Unix time, and the journal reads value.expired as RFC 3339, so the receipt +// carries both. +type kimiRefreshReceipt struct { + KimiTokenData + Expired string `json:"expired"` +} + +func newKimiRefreshReceipt(tokenData *KimiTokenData) *kimiRefreshReceipt { + receipt := &kimiRefreshReceipt{KimiTokenData: *tokenData} + if tokenData.ExpiresAt > 0 { + receipt.Expired = time.Unix(tokenData.ExpiresAt, 0).UTC().Format(time.RFC3339) + } + return receipt +} + +// decodeKimiRefreshReceipt rebuilds the tokens a successful refresh saved, for +// replicas that lost the claim. +func decodeKimiRefreshReceipt(value ojson.Result) (*kimiRefreshReceipt, error) { + if !value.IsObject() || value.Get("access_token").String() == "" || value.Get("expired").String() == "" { + return nil, fmt.Errorf("invalid refresh receipt") + } + return &kimiRefreshReceipt{ + KimiTokenData: KimiTokenData{ + AccessToken: value.Get("access_token").String(), + RefreshToken: value.Get("refresh_token").String(), + TokenType: value.Get("token_type").String(), + ExpiresAt: value.Get("expires_at").Int(), + Scope: value.Get("scope").String(), + }, + Expired: value.Get("expired").String(), + }, nil +} + func (c *DeviceFlowClient) refreshTokenSingleFlight(ctx context.Context, refreshToken string) (*KimiTokenData, error) { data := url.Values{} data.Set("client_id", kimiClientID) diff --git a/internal/auth/kimi/kimi_refresh_test.go b/internal/auth/kimi/kimi_refresh_test.go index d71fc4bc200..e85bf7c30ab 100644 --- a/internal/auth/kimi/kimi_refresh_test.go +++ b/internal/auth/kimi/kimi_refresh_test.go @@ -10,6 +10,7 @@ import ( "testing" "time" + "github.com/router-for-me/CLIProxyAPI/v8/internal/auth/codex" "golang.org/x/sync/singleflight" ) @@ -87,3 +88,49 @@ func TestRefreshToken_DeduplicatesConcurrentRefreshAcrossInstances(t *testing.T) t.Fatalf("expected both refresh callers to share a single upstream call, got %d", got) } } + +func TestRefreshToken_JournalSendsOnceAndNeverReplaysALostResponse(t *testing.T) { + resetKimiRefreshGroupForTest() + t.Cleanup(resetKimiRefreshGroupForTest) + + var calls atomic.Int32 + var lose atomic.Bool + transport := kimiRoundTripFunc(func(req *http.Request) (*http.Response, error) { + calls.Add(1) + if lose.Load() { + return &http.Response{StatusCode: http.StatusBadGateway, Body: io.NopCloser(strings.NewReader("{}")), Header: make(http.Header), Request: req}, nil + } + return &http.Response{ + StatusCode: http.StatusOK, + Body: io.NopCloser(strings.NewReader(`{"access_token":"new-access","refresh_token":"new-refresh","token_type":"Bearer","expires_in":3600,"scope":"kimi"}`)), + Header: make(http.Header), + Request: req, + }, nil + }) + ctx := codex.WithRefreshJournal(context.Background(), codex.NewMemoryRefreshJournal()) + client := &DeviceFlowClient{httpClient: &http.Client{Transport: transport}} + + first, err := client.RefreshToken(ctx, "old-refresh") + if err != nil || first.AccessToken != "new-access" || first.RefreshToken != "new-refresh" { + t.Fatalf("first refresh = %+v, %v", first, err) + } + // The receipt is rebuilt for a replica that lost the claim, expiry and scope included. + late, err := (&DeviceFlowClient{httpClient: &http.Client{Transport: transport}}).RefreshToken(ctx, "old-refresh") + if err != nil || late.AccessToken != "new-access" || late.RefreshToken != "new-refresh" || late.ExpiresAt != first.ExpiresAt || late.Scope != "kimi" { + t.Fatalf("late refresh = %+v, %v", late, err) + } + if calls.Load() != 1 { + t.Fatalf("requests = %d, want 1", calls.Load()) + } + + lose.Store(true) + before := calls.Load() + for range 2 { + if _, err := client.RefreshToken(ctx, "lost-refresh"); !codex.IsRefreshOutcomeUnknown(err) { + t.Fatalf("lost grant = %v, want refresh_outcome_unknown", err) + } + } + if got := calls.Load() - before; got != 1 { + t.Fatalf("requests for the lost grant = %d, want 1", got) + } +} diff --git a/internal/auth/xai/xai.go b/internal/auth/xai/xai.go index 7fe735fc4f0..e59016a4d64 100644 --- a/internal/auth/xai/xai.go +++ b/internal/auth/xai/xai.go @@ -11,6 +11,8 @@ import ( "time" "github.com/MoeForge/moego/base64" + "github.com/MoeForge/moego/ojson" + "github.com/router-for-me/CLIProxyAPI/v8/internal/auth/codex" "github.com/router-for-me/CLIProxyAPI/v8/internal/config" "github.com/router-for-me/CLIProxyAPI/v8/internal/util" log "github.com/sirupsen/logrus" @@ -358,7 +360,17 @@ func (a *XAIAuth) RefreshTokens(ctx context.Context, refreshToken, tokenEndpoint tokenEndpoint = strings.TrimSpace(tokenEndpoint) result, err, _ := xaiRefreshGroup.Do(refreshToken, func() (interface{}, error) { - return a.refreshTokensSingleFlight(context.WithoutCancel(ctx), refreshToken, tokenEndpoint) + refreshCtx := context.WithoutCancel(ctx) + // Without a journal bound by the store, refresh as before. With one, + // the grant is claimed first and never replayed after a lost response. + journal, journaled := codex.RefreshJournalFromContext(ctx) + if !journaled { + return a.refreshTokensSingleFlight(refreshCtx, refreshToken, tokenEndpoint) + } + return codex.JournalRefresh(refreshCtx, journal, ClientID, refreshToken, func() (*TokenData, error) { + tokenData, errRefresh := a.refreshTokensSingleFlight(refreshCtx, refreshToken, tokenEndpoint) + return tokenData, codex.QualifyRefreshOutcome("response", errRefresh) + }, decodeRefreshReceipt) }) if err != nil { return nil, err @@ -370,6 +382,24 @@ func (a *XAIAuth) RefreshTokens(ctx context.Context, refreshToken, tokenEndpoint return tokenData, nil } +// decodeRefreshReceipt rebuilds the tokens a successful refresh saved, for +// replicas that lost the claim. +func decodeRefreshReceipt(value ojson.Result) (*TokenData, error) { + if !value.IsObject() || value.Get("access_token").String() == "" || value.Get("expired").String() == "" { + return nil, fmt.Errorf("invalid refresh receipt") + } + return &TokenData{ + AccessToken: value.Get("access_token").String(), + RefreshToken: value.Get("refresh_token").String(), + IDToken: value.Get("id_token").String(), + TokenType: value.Get("token_type").String(), + ExpiresIn: int(value.Get("expires_in").Int()), + Expire: value.Get("expired").String(), + Email: value.Get("email").String(), + Subject: value.Get("sub").String(), + }, nil +} + func (a *XAIAuth) refreshTokensSingleFlight(ctx context.Context, refreshToken, tokenEndpoint string) (*TokenData, error) { form := url.Values{ "grant_type": {"refresh_token"}, diff --git a/internal/auth/xai/xai_auth_test.go b/internal/auth/xai/xai_auth_test.go index 92e9e29dce3..385f12786f3 100644 --- a/internal/auth/xai/xai_auth_test.go +++ b/internal/auth/xai/xai_auth_test.go @@ -13,6 +13,7 @@ import ( "time" "github.com/MoeForge/moego/base64" + "github.com/router-for-me/CLIProxyAPI/v8/internal/auth/codex" "golang.org/x/sync/singleflight" ) @@ -337,3 +338,53 @@ func fakeJWTWithEmail(email, subject string) string { payload := base64.RawURLEncoding.EncodeToString([]byte(`{"email":"` + email + `","sub":"` + subject + `"}`)) return header + "." + payload + ".sig" } + +func TestRefreshTokensJournalSendsOnceAndNeverReplaysALostResponse(t *testing.T) { + resetXAIRefreshGroupForTest() + t.Cleanup(resetXAIRefreshGroupForTest) + + var requests atomic.Int32 + lose := atomic.Bool{} + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + requests.Add(1) + if lose.Load() { + w.WriteHeader(http.StatusBadGateway) + return + } + w.Header().Set("Content-Type", "application/json") + _ = json.NewEncoder(w).Encode(map[string]any{ + "access_token": "new-access", + "refresh_token": "new-refresh", + "token_type": "Bearer", + "expires_in": 3600, + }) + })) + defer server.Close() + + journal := codex.NewMemoryRefreshJournal() + ctx := codex.WithRefreshJournal(context.Background(), journal) + auth := NewXAIAuth(nil) + first, err := auth.RefreshTokens(ctx, "old-refresh", server.URL) + if err != nil || first.AccessToken != "new-access" || first.RefreshToken != "new-refresh" { + t.Fatalf("first refresh = %+v, %v", first, err) + } + late, err := auth.RefreshTokens(ctx, "old-refresh", server.URL) + if err != nil || late.AccessToken != "new-access" || late.Expire != first.Expire { + t.Fatalf("late refresh = %+v, %v", late, err) + } + if requests.Load() != 1 { + t.Fatalf("requests = %d, want 1", requests.Load()) + } + + // A grant whose answer was lost is never sent again. + lose.Store(true) + before := requests.Load() + for range 2 { + if _, err := auth.RefreshTokens(ctx, "lost-refresh", server.URL); !codex.IsRefreshOutcomeUnknown(err) { + t.Fatalf("lost grant = %v, want refresh_outcome_unknown", err) + } + } + if got := requests.Load() - before; got != 1 { + t.Fatalf("requests for the lost grant = %d, want 1", got) + } +} diff --git a/internal/runtime/executor/refresh_policy.go b/internal/runtime/executor/refresh_policy.go new file mode 100644 index 00000000000..df2d2426a58 --- /dev/null +++ b/internal/runtime/executor/refresh_policy.go @@ -0,0 +1,24 @@ +package executor + +import cliproxyauth "github.com/router-for-me/CLIProxyAPI/v8/sdk/cliproxy/auth" + +// Providers whose refresh grants stay valid after use declare so here; a second +// replica sending the same grant is harmless, so they take no journal claim. +// Antigravity will take a short lease with the refresh queue. Every executor +// that declares nothing is treated as single-use. + +// RefreshPolicy declares that Antigravity grants are reusable. +func (e *AntigravityExecutor) RefreshPolicy() cliproxyauth.RefreshPolicy { + return cliproxyauth.RefreshNone +} + +// RefreshPolicy declares that Devin grants are reusable. +func (e *DevinExecutor) RefreshPolicy() cliproxyauth.RefreshPolicy { return cliproxyauth.RefreshNone } + +// RefreshPolicy declares that Meta grants are reusable. +func (e *MetaExecutor) RefreshPolicy() cliproxyauth.RefreshPolicy { return cliproxyauth.RefreshNone } + +// RefreshPolicy declares that Vertex credentials are reusable. +func (e *GeminiVertexExecutor) RefreshPolicy() cliproxyauth.RefreshPolicy { + return cliproxyauth.RefreshNone +} diff --git a/internal/store/natsstore/refresh_journal.go b/internal/store/natsstore/refresh_journal.go index ba8d7c8293d..03df262a193 100644 --- a/internal/store/natsstore/refresh_journal.go +++ b/internal/store/natsstore/refresh_journal.go @@ -96,3 +96,25 @@ func (j *RefreshJournal) RefreshComplete(ctx context.Context, key string, receip } return nil } + +// The store is the refresh journal of the credentials it holds, so the codex +// refresh and the conductor find it without knowing about NATS. + +// RefreshClaim records the grant as consumed, see RefreshJournal. +func (s *Store) RefreshClaim(ctx context.Context, key string) (bool, error) { + return s.journal.RefreshClaim(ctx, key) +} + +// RefreshRead returns the receipt of a claimed grant, see RefreshJournal. +func (s *Store) RefreshRead(ctx context.Context, key string) ([]byte, error) { + return s.journal.RefreshRead(ctx, key) +} + +// RefreshComplete saves the receipt of a claimed grant, see RefreshJournal. +func (s *Store) RefreshComplete(ctx context.Context, key string, receipt []byte) error { + return s.journal.RefreshComplete(ctx, key, receipt) +} + +// SharesRefreshJournal reports that every replica of the bucket claims in one +// journal, so the conductor guards every single-use provider with it. +func (s *Store) SharesRefreshJournal() bool { return true } diff --git a/internal/store/natsstore/store.go b/internal/store/natsstore/store.go index 1f7ed4911aa..70ca79db73e 100644 --- a/internal/store/natsstore/store.go +++ b/internal/store/natsstore/store.go @@ -29,6 +29,7 @@ type Store struct { replicaID string auth *AuthStore config *ConfigStore + journal *RefreshJournal spoolRoot string configPath string @@ -73,6 +74,7 @@ func NewStore(nc *nats.Conn, boot *Bootstrapped, settings StoreSettings, spoolDi replicaID: replicaID, auth: auth, config: cfg, + journal: NewRefreshJournal(NewKV(settings.BucketName(BucketRefresh), boot.Bucket(settings, BucketRefresh), origin)), spoolRoot: root, configPath: filepath.Join(configDir, "config.yaml"), }, nil diff --git a/sdk/cliproxy/auth/conductor_refresh.go b/sdk/cliproxy/auth/conductor_refresh.go index 9fedd9420b3..aa31a893a4a 100644 --- a/sdk/cliproxy/auth/conductor_refresh.go +++ b/sdk/cliproxy/auth/conductor_refresh.go @@ -1270,18 +1270,12 @@ func (m *Manager) refreshAuthForRequestAtEpoch(ctx context.Context, id, failedAc if auth == nil || exec == nil { return nil, errors.New("auth or executor not found") } - if strings.EqualFold(strings.TrimSpace(auth.Provider), "codex") { - if resolver, okResolver := store.(interface { - RefreshJournalFor(*Auth) (codex.RefreshJournal, error) - }); okResolver { - journal, errJournal := resolver.RefreshJournalFor(auth) - if errJournal != nil { - return nil, fmt.Errorf("resolve credential refresh journal: %w", errJournal) - } - ctx = codex.WithRefreshJournal(ctx, journal) - } else if journal, okJournal := store.(codex.RefreshJournal); okJournal { - ctx = codex.WithRefreshJournal(ctx, journal) - } + journal, errJournal := refreshJournalFor(store, auth, exec) + if errJournal != nil { + return nil, fmt.Errorf("resolve credential refresh journal: %w", errJournal) + } + if journal != nil { + ctx = codex.WithRefreshJournal(ctx, journal) } if registrationEpoch != 0 && auth.RegistrationEpoch != registrationEpoch { return nil, errors.New("auth registration changed before refresh") diff --git a/sdk/cliproxy/auth/refresh_policy.go b/sdk/cliproxy/auth/refresh_policy.go new file mode 100644 index 00000000000..2331400c861 --- /dev/null +++ b/sdk/cliproxy/auth/refresh_policy.go @@ -0,0 +1,64 @@ +package auth + +import ( + "strings" + + "github.com/router-for-me/CLIProxyAPI/v8/internal/auth/codex" +) + +// RefreshPolicy says how a provider's refresh grants must be guarded against two +// replicas sending the same grant. +type RefreshPolicy int + +const ( + // RefreshSingleUse is for grants that rotate or whose behavior is unknown: the + // grant is claimed in the journal before it is sent and the claim never + // expires. It is the zero value, so a provider that declares nothing fails safe. + RefreshSingleUse RefreshPolicy = iota + // RefreshNone is for grants that stay valid after use, so sending one twice is harmless. + RefreshNone +) + +// RefreshPolicyProvider is implemented by executors whose grants are not single-use. +type RefreshPolicyProvider interface { + RefreshPolicy() RefreshPolicy +} + +// RefreshPolicyOf returns the policy the executor declares, RefreshSingleUse when it declares none. +func RefreshPolicyOf(exec ProviderExecutor) RefreshPolicy { + if declared, ok := exec.(RefreshPolicyProvider); ok { + return declared.RefreshPolicy() + } + return RefreshSingleUse +} + +// SharedRefreshJournal is implemented by a store whose refresh journal is shared +// by every replica using it. Such a journal guards every single-use provider. +// A store without it keeps guarding only Codex, the one provider that had a +// journal before the stores shared one. +type SharedRefreshJournal interface { + codex.RefreshJournal + SharesRefreshJournal() bool +} + +// refreshJournalFor picks the journal bound to a refresh of auth, nil for none. +func refreshJournalFor(store Store, auth *Auth, exec ProviderExecutor) (codex.RefreshJournal, error) { + if RefreshPolicyOf(exec) == RefreshNone { + return nil, nil + } + if shared, ok := store.(SharedRefreshJournal); ok && shared.SharesRefreshJournal() { + return shared, nil + } + if !strings.EqualFold(strings.TrimSpace(auth.Provider), "codex") { + return nil, nil + } + if resolver, ok := store.(interface { + RefreshJournalFor(*Auth) (codex.RefreshJournal, error) + }); ok { + return resolver.RefreshJournalFor(auth) + } + if journal, ok := store.(codex.RefreshJournal); ok { + return journal, nil + } + return nil, nil +} diff --git a/sdk/cliproxy/auth/refresh_policy_test.go b/sdk/cliproxy/auth/refresh_policy_test.go new file mode 100644 index 00000000000..6b8c2c9e687 --- /dev/null +++ b/sdk/cliproxy/auth/refresh_policy_test.go @@ -0,0 +1,75 @@ +package auth + +import ( + "testing" + + "github.com/router-for-me/CLIProxyAPI/v8/internal/auth/codex" +) + +// policyExecutor is only inspected for its policy, so no method is ever called. +type policyExecutor struct { + ProviderExecutor +} + +type noneExecutor struct{ policyExecutor } + +func (noneExecutor) RefreshPolicy() RefreshPolicy { return RefreshNone } + +// plainStore is a store with no journal at all. +type plainStore struct{ Store } + +// codexJournalStore is a store that only ever journaled Codex. +type codexJournalStore struct { + plainStore + codex.RefreshJournal +} + +// sharedJournalStore is a store whose journal every replica shares. +type sharedJournalStore struct { + plainStore + codex.RefreshJournal +} + +func (sharedJournalStore) SharesRefreshJournal() bool { return true } + +func TestRefreshJournalForFollowsPolicyAndStore(t *testing.T) { + journal := codex.NewMemoryRefreshJournal() + cases := []struct { + name string + store Store + provider string + exec ProviderExecutor + want bool + }{ + {"shared codex", sharedJournalStore{RefreshJournal: journal}, "codex", policyExecutor{}, true}, + {"shared claude", sharedJournalStore{RefreshJournal: journal}, "claude", policyExecutor{}, true}, + {"shared kimi", sharedJournalStore{RefreshJournal: journal}, "kimi", policyExecutor{}, true}, + {"shared xai", sharedJournalStore{RefreshJournal: journal}, "xai", policyExecutor{}, true}, + {"shared unknown provider fails safe", sharedJournalStore{RefreshJournal: journal}, "newcomer", nil, true}, + {"shared reusable grant", sharedJournalStore{RefreshJournal: journal}, "antigravity", noneExecutor{}, false}, + {"shared reusable codex", sharedJournalStore{RefreshJournal: journal}, "codex", noneExecutor{}, false}, + {"unshared codex", codexJournalStore{RefreshJournal: journal}, "codex", policyExecutor{}, true}, + {"unshared claude keeps no journal", codexJournalStore{RefreshJournal: journal}, "claude", policyExecutor{}, false}, + {"no journal in the store", plainStore{}, "codex", policyExecutor{}, false}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + got, err := refreshJournalFor(tc.store, &Auth{Provider: tc.provider}, tc.exec) + if err != nil { + t.Fatalf("refreshJournalFor() error = %v", err) + } + if (got != nil) != tc.want { + t.Fatalf("journal bound = %v, want %v", got != nil, tc.want) + } + }) + } +} + +func TestRefreshPolicyOfDefaultsToSingleUse(t *testing.T) { + if got := RefreshPolicyOf(policyExecutor{}); got != RefreshSingleUse { + t.Fatalf("policy = %v, want single use", got) + } + if got := RefreshPolicyOf(noneExecutor{}); got != RefreshNone { + t.Fatalf("policy = %v, want none", got) + } +}