diff --git a/internal/runtime/executor/refresh_policy.go b/internal/runtime/executor/refresh_policy.go index df2d2426a58..f7a7c3802d5 100644 --- a/internal/runtime/executor/refresh_policy.go +++ b/internal/runtime/executor/refresh_policy.go @@ -4,12 +4,12 @@ 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 +// Antigravity takes a short lease so one replica refreshes it. Every executor // that declares nothing is treated as single-use. -// RefreshPolicy declares that Antigravity grants are reusable. +// RefreshPolicy declares that Antigravity grants are reusable and refreshed under a lease. func (e *AntigravityExecutor) RefreshPolicy() cliproxyauth.RefreshPolicy { - return cliproxyauth.RefreshNone + return cliproxyauth.RefreshLease } // RefreshPolicy declares that Devin grants are reusable. diff --git a/internal/store/natsstore/kv.go b/internal/store/natsstore/kv.go index 3c98f9bf859..dd1c198068e 100644 --- a/internal/store/natsstore/kv.go +++ b/internal/store/natsstore/kv.go @@ -118,6 +118,26 @@ func (k *KV) Create(ctx context.Context, id string, value []byte) (uint64, error return rev, nil } +// CreateWithTTL writes a value only if the key has no live value, and the +// server removes it after ttl. The expiry is fixed at creation: an Update drops +// it, so a key that must expire again is deleted and created anew. It needs a +// bucket with a marker TTL. +func (k *KV) CreateWithTTL(ctx context.Context, id string, value []byte, ttl time.Duration) (uint64, error) { + key, err := EncodeKey(id) + if err != nil { + return 0, err + } + stored, err := k.wrap(value, 0) + if err != nil { + return 0, fmt.Errorf("failed to encode %s for %s: %w", id, k.bucket, err) + } + rev, err := k.kv.Create(ctx, key, stored, jetstream.KeyTTL(ttl)) + if err != nil { + return 0, k.writeError("create", id, err) + } + return rev, nil +} + // CreateOnce writes a value only if the key was never written. Unlike Create it // treats a delete marker as a write, so deleting a key does not reopen it. func (k *KV) CreateOnce(ctx context.Context, id string, value []byte) (uint64, error) { diff --git a/internal/store/natsstore/refresh_lease.go b/internal/store/natsstore/refresh_lease.go new file mode 100644 index 00000000000..6fba4c90ab6 --- /dev/null +++ b/internal/store/natsstore/refresh_lease.go @@ -0,0 +1,34 @@ +package natsstore + +import ( + "context" + "errors" + "time" + + cliproxyauth "github.com/router-for-me/CLIProxyAPI/v8/sdk/cliproxy/auth" +) + +var _ cliproxyauth.RefreshLeaser = (*Store)(nil) + +// AcquireRefreshLease takes the right to refresh a credential for at most ttl. +// The lease lives in its own bucket so the refresh journal, whose claims must +// never expire, keeps a bucket without any key TTL. A lease is never renewed or +// taken over: it runs out on the server and then the next Create wins, so a +// replica that dies mid-refresh frees the credential without anyone judging +// whether it is dead. +func (s *Store) AcquireRefreshLease(ctx context.Context, authID string, ttl time.Duration) (func(), error) { + rev, err := s.leases.CreateWithTTL(ctx, authID, []byte(s.replicaID), ttl) + if err != nil { + if errors.Is(err, ErrConflict) { + return nil, cliproxyauth.ErrRefreshLeaseHeld + } + return nil, err + } + return func() { + // Revision-checked, so a lease that ran out and was taken by another + // replica is not deleted from under it. + ctx, cancel := context.WithTimeout(context.WithoutCancel(ctx), 5*time.Second) + defer cancel() + _ = s.leases.Delete(ctx, authID, rev) + }, nil +} diff --git a/internal/store/natsstore/refresh_lease_test.go b/internal/store/natsstore/refresh_lease_test.go new file mode 100644 index 00000000000..49d90a67484 --- /dev/null +++ b/internal/store/natsstore/refresh_lease_test.go @@ -0,0 +1,82 @@ +package natsstore + +import ( + "context" + "errors" + "testing" + "time" + + cliproxyauth "github.com/router-for-me/CLIProxyAPI/v8/sdk/cliproxy/auth" +) + +func TestRefreshLeaseIsHeldByOneReplicaAtATime(t *testing.T) { + f := newStoreFixture(t) + a, b := f.open(), f.open() + ctx := context.Background() + + release, err := a.AcquireRefreshLease(ctx, "antigravity/one.json", time.Minute) + if err != nil { + t.Fatal(err) + } + if _, err := b.AcquireRefreshLease(ctx, "antigravity/one.json", time.Minute); !errors.Is(err, cliproxyauth.ErrRefreshLeaseHeld) { + t.Fatalf("second replica err = %v, want held", err) + } + if _, err := b.AcquireRefreshLease(ctx, "antigravity/other.json", time.Minute); err != nil { + t.Fatalf("lease for another credential: %v", err) + } + + release() + if _, err := b.AcquireRefreshLease(ctx, "antigravity/one.json", time.Minute); err != nil { + t.Fatalf("after release: %v", err) + } +} + +func TestRefreshLeaseRunsOutWhenTheHolderDies(t *testing.T) { + f := newStoreFixture(t) + a, b := f.open(), f.open() + ctx := context.Background() + + if _, err := a.AcquireRefreshLease(ctx, "antigravity/two.json", time.Second); err != nil { + t.Fatal(err) + } + waitFor(t, "lease to run out", func() bool { + release, err := b.AcquireRefreshLease(ctx, "antigravity/two.json", time.Minute) + if err == nil { + release() + } + return err == nil + }) +} + +func TestStaleReleaseDoesNotDeleteANewHoldersLease(t *testing.T) { + f := newStoreFixture(t) + a, b := f.open(), f.open() + ctx := context.Background() + + staleRelease, err := a.AcquireRefreshLease(ctx, "antigravity/three.json", time.Second) + if err != nil { + t.Fatal(err) + } + var fresh func() + waitFor(t, "lease to run out", func() bool { + fresh, err = b.AcquireRefreshLease(ctx, "antigravity/three.json", time.Minute) + return err == nil + }) + defer fresh() + + staleRelease() + if _, err := a.AcquireRefreshLease(ctx, "antigravity/three.json", time.Minute); !errors.Is(err, cliproxyauth.ErrRefreshLeaseHeld) { + t.Fatalf("err = %v, want the new holder's lease to survive a stale release", err) + } +} + +func TestRefreshJournalBucketNeverTakesAKeyTTL(t *testing.T) { + f := newStoreFixture(t) + s := f.open() + if _, err := s.leases.CreateWithTTL(context.Background(), "probe", []byte("x"), time.Minute); err != nil { + t.Fatalf("lease bucket must take key TTLs: %v", err) + } + if _, err := s.journal.kv.CreateWithTTL(context.Background(), "probe", []byte("x"), time.Minute); err == nil { + t.Fatal("journal bucket accepted a key TTL; its claims must never expire") + } +} diff --git a/internal/store/natsstore/store.go b/internal/store/natsstore/store.go index 48e290bf161..ed8d4f17c4d 100644 --- a/internal/store/natsstore/store.go +++ b/internal/store/natsstore/store.go @@ -31,6 +31,7 @@ type Store struct { config *ConfigStore journal *RefreshJournal refresh *RefreshQueue + leases *KV spoolRoot string configPath string @@ -77,6 +78,7 @@ func NewStore(nc *nats.Conn, boot *Bootstrapped, settings StoreSettings, spoolDi config: cfg, journal: NewRefreshJournal(NewKV(settings.BucketName(BucketRefresh), boot.Bucket(settings, BucketRefresh), origin)), refresh: boot.Refresh, + leases: NewKV(settings.BucketName(BucketLease), boot.Bucket(settings, BucketLease), 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 588fbb8429c..3e58746477a 100644 --- a/sdk/cliproxy/auth/conductor_refresh.go +++ b/sdk/cliproxy/auth/conductor_refresh.go @@ -1282,6 +1282,13 @@ func (m *Manager) refreshAuthForRequestAtEpoch(ctx context.Context, id, failedAc if registrationEpoch != 0 && auth.RegistrationEpoch != registrationEpoch { return nil, errors.New("auth registration changed before refresh") } + // A request that needs a usable token now refreshes even when another + // replica holds the lease; the grant is reusable, so it only costs a call. + release, leaseHeld := acquireRefreshLease(ctx, store, auth, exec, forceRefresh || failedAccessToken != "") + if leaseHeld { + return auth.Clone(), nil + } + defer release() if hasDisabledInvalidGrantFailure(auth) && !forceRefresh { return nil, errors.New("auth is disabled with invalid grant") } diff --git a/sdk/cliproxy/auth/refresh_lease.go b/sdk/cliproxy/auth/refresh_lease.go new file mode 100644 index 00000000000..11176a99c7d --- /dev/null +++ b/sdk/cliproxy/auth/refresh_lease.go @@ -0,0 +1,55 @@ +package auth + +import ( + "context" + "errors" + "time" + + log "github.com/sirupsen/logrus" +) + +// refreshLeaseTTL is how long a replica holds the right to refresh a +// RefreshLease credential. It covers a provider call, and a replica that dies +// holding it frees the credential by itself when it runs out. +const refreshLeaseTTL = 30 * time.Second + +// ErrRefreshLeaseHeld is returned by a RefreshLeaser when another replica holds the lease. +var ErrRefreshLeaseHeld = errors.New("refresh lease is held") + +// RefreshLeaser is implemented by a store that can hand out refresh leases +// shared by every replica using it. +type RefreshLeaser interface { + // AcquireRefreshLease takes the lease for a credential for at most ttl. + // It returns ErrRefreshLeaseHeld when another replica holds it. The + // returned function gives the lease back and is safe to call after it ran out. + AcquireRefreshLease(ctx context.Context, authID string, ttl time.Duration) (release func(), err error) +} + +func noRelease() {} + +// acquireRefreshLease takes the lease a RefreshLease provider needs before +// refreshing. held is true when another replica holds it and the caller should +// leave the refresh to that replica. Everything else proceeds: a provider +// without the policy, a store without leases, or a lease that cannot be taken +// for any other reason, since the grant is reusable and a second refresh is +// harmless. needed is for callers that cannot do without a fresh token: they +// refresh even while another replica holds the lease. +func acquireRefreshLease(ctx context.Context, store Store, auth *Auth, exec ProviderExecutor, needed bool) (release func(), held bool) { + if RefreshPolicyOf(exec) != RefreshLease { + return noRelease, false + } + leaser, ok := store.(RefreshLeaser) + if !ok { + return noRelease, false + } + release, err := leaser.AcquireRefreshLease(ctx, auth.ID, refreshLeaseTTL) + switch { + case err == nil: + return release, false + case errors.Is(err, ErrRefreshLeaseHeld): + return noRelease, !needed + default: + log.WithFields(log.Fields{"auth_id": auth.ID, "provider": auth.Provider}).Debugf("refresh lease unavailable, refreshing without it: %v", err) + return noRelease, false + } +} diff --git a/sdk/cliproxy/auth/refresh_lease_test.go b/sdk/cliproxy/auth/refresh_lease_test.go new file mode 100644 index 00000000000..7d3e1cd75f3 --- /dev/null +++ b/sdk/cliproxy/auth/refresh_lease_test.go @@ -0,0 +1,97 @@ +package auth + +import ( + "context" + "errors" + "sync/atomic" + "testing" + "time" + + "github.com/router-for-me/CLIProxyAPI/v8/internal/auth/codex" +) + +// leasingStore is a store that hands out refresh leases as a test directs. +type leasingStore struct { + queueHostStore + acquireErr error + acquired atomic.Int32 + released atomic.Int32 + ttl atomic.Int64 +} + +func (s *leasingStore) AcquireRefreshLease(_ context.Context, _ string, ttl time.Duration) (func(), error) { + s.acquired.Add(1) + s.ttl.Store(int64(ttl)) + if s.acquireErr != nil { + return nil, s.acquireErr + } + return func() { s.released.Add(1) }, nil +} + +// leasedGrantExecutor declares its grant is reusable but refreshed under a lease. +type leasedGrantExecutor struct{ countingRefreshExecutor } + +func (*leasedGrantExecutor) RefreshPolicy() RefreshPolicy { return RefreshLease } + +func refreshWithLease(t *testing.T, store *leasingStore, force bool) *leasedGrantExecutor { + t.Helper() + exec := &leasedGrantExecutor{countingRefreshExecutor{id: queueRefreshProvider}} + manager, _ := newQueueManager(t, store, exec) + registerQueueAuth(t, manager, "a", queueRefreshProvider, time.Now()) + manager.mu.RLock() + epoch := manager.auths["a"].RegistrationEpoch + manager.mu.RUnlock() + ctx := context.Background() + if force { + ctx = withForceRefresh(ctx) + } + if _, err := manager.refreshAuthForRequestAtEpoch(ctx, "a", "", epoch); err != nil { + t.Fatalf("refresh: %v", err) + } + return exec +} + +func TestLeasedRefreshHoldsTheLeaseForTheProviderCall(t *testing.T) { + store := &leasingStore{} + exec := refreshWithLease(t, store, false) + if got := exec.refreshCalls.Load(); got != 1 { + t.Fatalf("refreshed %d times, want 1", got) + } + if store.acquired.Load() != 1 || store.released.Load() != 1 { + t.Fatalf("acquired %d, released %d; want 1 and 1", store.acquired.Load(), store.released.Load()) + } + if got := time.Duration(store.ttl.Load()); got != refreshLeaseTTL { + t.Fatalf("lease ttl = %v, want %v", got, refreshLeaseTTL) + } +} + +func TestLeasedRefreshLeavesTheWorkToTheHolder(t *testing.T) { + store := &leasingStore{acquireErr: ErrRefreshLeaseHeld} + exec := refreshWithLease(t, store, false) + if got := exec.refreshCalls.Load(); got != 0 { + t.Fatalf("refreshed %d times while another replica held the lease", got) + } +} + +func TestLeasedRefreshProceedsWhenARequestNeedsAFreshToken(t *testing.T) { + store := &leasingStore{acquireErr: ErrRefreshLeaseHeld} + exec := refreshWithLease(t, store, true) + if got := exec.refreshCalls.Load(); got != 1 { + t.Fatalf("refreshed %d times, want 1: the grant is reusable and the request cannot wait", got) + } +} + +func TestLeasedRefreshProceedsWhenTheLeaseCannotBeTaken(t *testing.T) { + store := &leasingStore{acquireErr: errors.New("store unavailable")} + exec := refreshWithLease(t, store, false) + if got := exec.refreshCalls.Load(); got != 1 { + t.Fatalf("refreshed %d times, want 1", got) + } +} + +func TestLeasePolicyTakesNoJournalClaim(t *testing.T) { + journal, err := refreshJournalFor(sharedJournalStore{RefreshJournal: codex.NewMemoryRefreshJournal()}, &Auth{Provider: "antigravity"}, &leasedGrantExecutor{}) + if err != nil || journal != nil { + t.Fatalf("journal = %v, err = %v; a lease replaces the journal claim", journal, err) + } +} diff --git a/sdk/cliproxy/auth/refresh_policy.go b/sdk/cliproxy/auth/refresh_policy.go index 2331400c861..18d4bccce9c 100644 --- a/sdk/cliproxy/auth/refresh_policy.go +++ b/sdk/cliproxy/auth/refresh_policy.go @@ -17,6 +17,10 @@ const ( RefreshSingleUse RefreshPolicy = iota // RefreshNone is for grants that stay valid after use, so sending one twice is harmless. RefreshNone + // RefreshLease is for grants that stay valid after use but where a refresh is + // worth doing once: the replica that refreshes holds a short lease, and + // replicas that find it held leave the refresh to the holder. + RefreshLease ) // RefreshPolicyProvider is implemented by executors whose grants are not single-use. @@ -43,7 +47,7 @@ type SharedRefreshJournal interface { // 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 { + if RefreshPolicyOf(exec) != RefreshSingleUse { return nil, nil } if shared, ok := store.(SharedRefreshJournal); ok && shared.SharesRefreshJournal() {