From 2388b726122542730b008f0c948422f263bfa96c Mon Sep 17 00:00:00 2001 From: Yordis Prieto Date: Sat, 7 Nov 2026 21:26:14 -0500 Subject: [PATCH] feat(auth): guard single-use refresh grants with the shared journal Claude, Kimi and xAI rotate their refresh tokens but had no journal, so two replicas could spend the same grant and lose the credential. Each provider now declares whether its grant is single-use, and a store whose journal is shared guards every one that is. Stores without a shared journal behave as before. Signed-off-by: Yordis Prieto --- internal/auth/claude/anthropic_auth.go | 30 ++++- internal/auth/claude/refresh_journal_test.go | 122 +++++++++++++++++++ internal/auth/codex/refresh_journal.go | 36 ++++++ internal/auth/kimi/kimi.go | 55 ++++++++- internal/auth/kimi/kimi_refresh_test.go | 47 +++++++ internal/auth/xai/xai.go | 32 ++++- internal/auth/xai/xai_auth_test.go | 51 ++++++++ internal/runtime/executor/refresh_policy.go | 24 ++++ internal/store/natsstore/refresh_journal.go | 22 ++++ internal/store/natsstore/store.go | 2 + sdk/cliproxy/auth/conductor_refresh.go | 18 +-- sdk/cliproxy/auth/refresh_policy.go | 64 ++++++++++ sdk/cliproxy/auth/refresh_policy_test.go | 75 ++++++++++++ 13 files changed, 563 insertions(+), 15 deletions(-) create mode 100644 internal/auth/claude/refresh_journal_test.go create mode 100644 internal/runtime/executor/refresh_policy.go create mode 100644 sdk/cliproxy/auth/refresh_policy.go create mode 100644 sdk/cliproxy/auth/refresh_policy_test.go 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) + } +}