From cdfa3b9f02ff405bfa1ea065bed6d1ecda268856 Mon Sep 17 00:00:00 2001 From: Yordis Prieto Date: Sat, 7 Nov 2026 22:03:06 -0500 Subject: [PATCH] feat(natsstore): lease Antigravity refreshes so one replica refreshes Antigravity grants stay valid after use, so the journal claim is wrong for them, yet every replica refreshing the same expiring token wastes provider calls. A short lease in its own bucket lets one replica refresh while the others adopt the result. Per-key TTL is fixed at creation and an update drops it, so a lease is never renewed: it runs out on the server and the next create wins. The journal bucket stays free of TTL because its claims must never expire. Signed-off-by: Yordis Prieto --- internal/runtime/executor/refresh_policy.go | 6 +- internal/store/natsstore/kv.go | 20 ++++ internal/store/natsstore/refresh_lease.go | 34 +++++++ .../store/natsstore/refresh_lease_test.go | 82 ++++++++++++++++ internal/store/natsstore/store.go | 2 + sdk/cliproxy/auth/conductor_refresh.go | 7 ++ sdk/cliproxy/auth/refresh_lease.go | 55 +++++++++++ sdk/cliproxy/auth/refresh_lease_test.go | 97 +++++++++++++++++++ sdk/cliproxy/auth/refresh_policy.go | 6 +- 9 files changed, 305 insertions(+), 4 deletions(-) create mode 100644 internal/store/natsstore/refresh_lease.go create mode 100644 internal/store/natsstore/refresh_lease_test.go create mode 100644 sdk/cliproxy/auth/refresh_lease.go create mode 100644 sdk/cliproxy/auth/refresh_lease_test.go 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() {