feat(auth): guard single-use refresh grants with the shared journal

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{

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())
+	}
+}

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()

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)

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)
+	}
+}