feat(natsstore): lease Antigravity refreshes so one replica refreshes

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.

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

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

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

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

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

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

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

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