diff --git a/sdk/cliproxy/auth/conductor.go b/sdk/cliproxy/auth/conductor.go index 050da3295..785efcd70 100644 --- a/sdk/cliproxy/auth/conductor.go +++ b/sdk/cliproxy/auth/conductor.go @@ -2,6 +2,7 @@ package auth import ( "context" + "math" "net/http" "sync" "sync/atomic" @@ -10,6 +11,7 @@ import ( internalconfig "github.com/router-for-me/CLIProxyAPI/v8/internal/config" cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v8/sdk/cliproxy/executor" "github.com/router-for-me/CLIProxyAPI/v8/sdk/pluginapi" + "golang.org/x/sync/semaphore" ) // ProviderExecutor defines the contract required by Manager to execute provider calls. @@ -208,6 +210,11 @@ type Manager struct { refreshLocks sync.Map // persistLocks serializes disk persistence per auth ID and guards against out-of-order writes. persistLocks sync.Map + // authMutationLocks coordinate credential mutations without blocking unrelated readers. + authMutationLocks sync.Map + // authLoadGate allows concurrent credential transactions, but excludes whole-store reloads. + // Mutations acquire one permit; Load acquires all permits before reading the store. + authLoadGate *semaphore.Weighted } // NewManager constructs a manager with optional custom selector and hook. @@ -220,6 +227,7 @@ func NewManager(store Store, selector Selector, hook Hook) *Manager { } manager := &Manager{ store: store, + authLoadGate: semaphore.NewWeighted(math.MaxInt64), executors: make(map[string]ProviderExecutor), selector: selector, hook: hook, diff --git a/sdk/cliproxy/auth/conductor_cancellation_test.go b/sdk/cliproxy/auth/conductor_cancellation_test.go new file mode 100644 index 000000000..c783f6e3c --- /dev/null +++ b/sdk/cliproxy/auth/conductor_cancellation_test.go @@ -0,0 +1,180 @@ +package auth + +import ( + "context" + "errors" + "math" + "sync/atomic" + "testing" + "testing/synctest" +) + +type cancellationStore struct { + saves atomic.Int32 + lists atomic.Int32 + entered chan struct{} + release chan struct{} +} + +func (s *cancellationStore) Save(_ context.Context, auth *Auth) (string, error) { + s.saves.Add(1) + if auth.Label == "blocked" { + close(s.entered) + <-s.release + } + return "", nil +} +func (s *cancellationStore) List(context.Context) ([]*Auth, error) { + s.lists.Add(1) + return nil, nil +} +func (*cancellationStore) Delete(context.Context, string) error { return nil } + +type cancellationHook struct{ calls atomic.Int32 } + +func (h *cancellationHook) OnAuthRegistered(context.Context, *Auth) { h.calls.Add(1) } +func (h *cancellationHook) OnAuthUpdated(context.Context, *Auth) { h.calls.Add(1) } +func (*cancellationHook) OnResult(context.Context, Result) {} + +func assertCanceledMutation(t *testing.T, done <-chan error) { + t.Helper() + select { + case err := <-done: + if !errors.Is(err, context.Canceled) { + t.Errorf("expected context cancellation, got %v", err) + } + default: + t.Error("canceled operation still waiting for lock") + } +} + +func TestManagerCanceledLoadUnblocksUnrelatedMutation(t *testing.T) { + synctest.Test(t, func(t *testing.T) { + store := &cancellationStore{entered: make(chan struct{}), release: make(chan struct{})} + manager := NewManager(store, nil, nil) + saveDone := make(chan error, 1) + go func() { + _, err := manager.Register(context.Background(), &Auth{ID: "A", Provider: "audit", Label: "blocked", Metadata: map[string]any{"type": "audit"}}) + saveDone <- err + }() + <-store.entered + ctx, cancel := context.WithCancel(context.Background()) + loadDone := make(chan error, 1) + go func() { loadDone <- manager.Load(ctx) }() + synctest.Wait() + bDone := make(chan error, 1) + go func() { + _, err := manager.Register(context.Background(), &Auth{ID: "B", Provider: "audit"}) + bDone <- err + }() + synctest.Wait() + select { + case <-bDone: + t.Error("B overtook queued reload") + default: + } + cancel() + synctest.Wait() + assertCanceledMutation(t, loadDone) + select { + case err := <-bDone: + if err != nil { + t.Error(err) + } + default: + t.Error("canceled reload still blocks unrelated B") + } + if store.lists.Load() != 0 { + t.Error("canceled reload called List") + } + close(store.release) + synctest.Wait() + if err := <-saveDone; err != nil { + t.Error(err) + } + if !manager.authLoadGate.TryAcquire(math.MaxInt64) { + t.Error("reload barrier permit leaked") + } else { + manager.authLoadGate.Release(math.MaxInt64) + } + }) +} + +func TestManagerCanceledMutationWait(t *testing.T) { + for _, operation := range []string{"Register", "Update", "Prepare", "Refresh"} { + for _, waitOn := range []string{"barrier", "gate"} { + t.Run(operation+"/"+waitOn, func(t *testing.T) { + synctest.Test(t, func(t *testing.T) { + store := &cancellationStore{} + hook := &cancellationHook{} + manager := NewManager(store, nil, hook) + base, err := manager.Register(WithSkipPersist(context.Background()), &Auth{ID: "A", Provider: "audit", Metadata: map[string]any{"type": "audit"}}) + if err != nil { + t.Fatal(err) + } + hook.calls.Store(0) + var release func() + if waitOn == "barrier" { + if errAcquire := manager.authLoadGate.Acquire(context.Background(), math.MaxInt64); errAcquire != nil { + t.Fatal(errAcquire) + } + release = func() { manager.authLoadGate.Release(math.MaxInt64) } + } else { + release = manager.lockAuthMutation(base.ID) + } + ctx, cancel := context.WithCancel(context.Background()) + done := make(chan error, 1) + go func() { + var errMutation error + switch operation { + case "Register": + _, errMutation = manager.Register(ctx, base.Clone()) + case "Update": + _, errMutation = manager.Update(ctx, base.Clone()) + case "Prepare": + _, errMutation = manager.UpdatePreparedAuth(ctx, base, base.Clone()) + case "Refresh": + _, errMutation = manager.UpdateRefreshedAuth(ctx, base, base.Clone()) + } + done <- errMutation + }() + synctest.Wait() + cancel() + synctest.Wait() + assertCanceledMutation(t, done) + release() + synctest.Wait() + if store.saves.Load() != 0 || hook.calls.Load() != 0 { + t.Error("canceled mutation invoked Save or hook") + } + if got, _ := manager.GetByID(base.ID); got.Generation != base.Generation || got.RegistrationEpoch != base.RegistrationEpoch { + t.Error("canceled mutation published auth") + } + if !manager.authLoadGate.TryAcquire(math.MaxInt64) { + t.Error("mutation leaked reload barrier permit") + } else { + manager.authLoadGate.Release(math.MaxInt64) + } + value, _ := manager.authMutationLocks.Load(base.ID) + if len(value.(chan struct{})) != 0 { + t.Error("mutation leaked credential gate") + } + }) + }) + } + } +} + +func TestManagerMutationNilContext(t *testing.T) { + manager := NewManager(&cancellationStore{}, nil, nil) + auth, err := manager.Register(nil, &Auth{ID: "A", Provider: "audit", Metadata: map[string]any{"type": "audit"}}) + if err != nil { + t.Fatal(err) + } + if _, errUpdate := manager.Update(nil, auth); errUpdate != nil { + t.Fatal(errUpdate) + } + if errLoad := manager.Load(nil); errLoad != nil { + t.Fatal(errLoad) + } +} diff --git a/sdk/cliproxy/auth/conductor_cooldown.go b/sdk/cliproxy/auth/conductor_cooldown.go index 33f8b5d8d..8e6f08461 100644 --- a/sdk/cliproxy/auth/conductor_cooldown.go +++ b/sdk/cliproxy/auth/conductor_cooldown.go @@ -475,6 +475,8 @@ func (m *Manager) ResetQuota(ctx context.Context, authID string) (*Auth, []strin registeredModels := modelsForRegisteredAuth(authID) cooldownStateChanged := false + releaseMutation := m.lockAuthMutation(authID) + defer releaseMutation() m.mu.Lock() auth, ok := m.auths[authID] if !ok || auth == nil { @@ -524,8 +526,9 @@ func (m *Manager) ResetQuota(ctx context.Context, authID string) (*Auth, []strin cooldownRecordsAfter := m.cooldownStateRecordsForAuthLocked(auth, now) cooldownStateChanged = !cooldownStateRecordsEqual(cooldownRecordsBefore, cooldownRecordsAfter) } - errPersist := m.persist(ctx, auth) + errPersist := m.persistLocked(ctx, auth) m.mu.Unlock() + releaseMutation() defer func() { if cooldownStateChanged { @@ -766,6 +769,8 @@ func (m *Manager) MarkResult(ctx context.Context, result Result) { cooldownStateChanged := false now := time.Now() + releaseMutation := m.lockAuthMutation(result.AuthID) + defer releaseMutation() m.mu.Lock() if auth, ok := m.auths[result.AuthID]; ok && auth != nil { if modelKey == "" && strings.TrimSpace(result.RouteModel) != "" { @@ -1016,7 +1021,7 @@ func (m *Manager) MarkResult(ctx context.Context, result Result) { } } - _ = m.persist(ctx, auth) + _ = m.persistLocked(ctx, auth) authSnapshot = auth.Clone() if trackCooldownState { cooldownRecordsAfter := m.cooldownStateRecordsForAuthLocked(auth, now) @@ -1024,6 +1029,7 @@ func (m *Manager) MarkResult(ctx context.Context, result Result) { } } m.mu.Unlock() + releaseMutation() if m.scheduler != nil && authSnapshot != nil { var targetModels []string if !result.CredentialScope && modelKey != "" { @@ -1096,6 +1102,8 @@ func (m *Manager) recordAvailabilityNeutralResult(ctx context.Context, result Re } var authSnapshot *Auth + releaseMutation := m.lockAuthMutation(result.AuthID) + defer releaseMutation() m.mu.Lock() if auth, ok := m.auths[result.AuthID]; ok && auth != nil { now := time.Now() @@ -1107,10 +1115,11 @@ func (m *Manager) recordAvailabilityNeutralResult(ctx context.Context, result Re } auth.Generation++ auth.UpdatedAt = now - _ = m.persist(ctx, auth) + _ = m.persistLocked(ctx, auth) authSnapshot = auth.Clone() } m.mu.Unlock() + releaseMutation() m.hook.OnResult(ctx, result) m.publishErrorEvent(result, authSnapshot) diff --git a/sdk/cliproxy/auth/conductor_lifecycle.go b/sdk/cliproxy/auth/conductor_lifecycle.go index 3fa801b26..597a7d7c3 100644 --- a/sdk/cliproxy/auth/conductor_lifecycle.go +++ b/sdk/cliproxy/auth/conductor_lifecycle.go @@ -3,6 +3,7 @@ package auth import ( "context" "fmt" + "math" "strings" "sync" "time" @@ -88,6 +89,11 @@ func (m *Manager) Register(ctx context.Context, auth *Auth) (*Auth, error) { if auth.ID == "" { auth.ID = uuid.NewString() } + releaseMutation, errMutation := m.lockAuthMutationContext(ctx, auth.ID) + if errMutation != nil { + return nil, errMutation + } + defer releaseMutation() now := time.Now() if auth.Generation == 0 { auth.Generation = 1 @@ -114,6 +120,16 @@ func (m *Manager) Register(ctx context.Context, auth *Auth) (*Auth, error) { m.authEpochs[auth.ID]++ auth.RegistrationEpoch = m.authEpochs[auth.ID] auth.Generation = 1 + // Serialize this credential, but release the manager lock during store I/O. + // Persist failures stay non-fatal, but must not be silent: a restart would + // lose the credential that only exists in memory. + if errPersist := m.persistLocked(ctx, auth); errPersist != nil { + log.WithFields(log.Fields{"auth_id": auth.ID, "credential": auth.ID, "provider": auth.Provider}).Warnf("failed to persist registered auth %s (%s): %v", auth.Provider, auth.ID, errPersist) + } + if m.authEpochs[auth.ID] != auth.RegistrationEpoch { + m.mu.Unlock() + return nil, fmt.Errorf("auth %s changed registration during persistence", auth.ID) + } authClone := auth.Clone() m.auths[auth.ID] = authClone m.notifyAuthChangeLocked(auth.ID) @@ -123,6 +139,7 @@ func (m *Manager) Register(ctx context.Context, auth *Auth) (*Auth, error) { schedulerSnapshot = authClone.Clone() } m.mu.Unlock() + releaseMutation() if !shouldDeferAPIKeyModelAliasRebuild(ctx) { m.rebuildAPIKeyModelAliasFromRuntimeConfig() } @@ -131,11 +148,6 @@ func (m *Manager) Register(ctx context.Context, auth *Auth) (*Auth, error) { } m.structuralEpoch.Add(1) m.queueRefreshReschedule(auth.ID) - // Persist failures stay non-fatal, but must not be silent: a restart would - // lose the credential that only exists in memory. - if errPersist := m.persist(ctx, auth); errPersist != nil { - log.WithFields(log.Fields{"auth_id": auth.ID, "credential": auth.ID, "provider": auth.Provider}).Warnf("failed to persist registered auth %s (%s): %v", auth.Provider, auth.ID, errPersist) - } m.hook.OnAuthRegistered(ctx, auth.Clone()) if cooldownStateChanged { m.persistCooldownStates(context.Background()) @@ -176,6 +188,11 @@ func (m *Manager) updateInternal(ctx context.Context, base, auth *Auth, mode upd if errWeight := ValidateAuthWeight(auth); errWeight != nil { return nil, fmt.Errorf("update auth: %w", errWeight) } + releaseMutation, errMutation := m.lockAuthMutationContext(ctx, auth.ID) + if errMutation != nil { + return nil, errMutation + } + defer releaseMutation() m.mu.Lock() existing, ok := m.auths[auth.ID] if !ok || existing == nil { @@ -260,15 +277,26 @@ func (m *Manager) updateInternal(ctx context.Context, base, auth *Auth, mode upd cooldownStateChanged = clearCooldownStateForAuth(auth, now) || cooldownStateChanged } auth.EnsureIndex() - // A minted Meta key must reach the configured store before requests can use it. - // Keep the epoch check, save and installation together so a concurrent reload - // or removal cannot let an obsolete mint overwrite the credential on disk. + // Save before publication, including the transactional Meta mint path. + // Runtime-only changes made during I/O are merged below, never overwritten. + existingBeforeSave := existing.Clone() persistMetaMint := (mode == updateModePrepare || mode == updateModeRefresh) && strings.EqualFold(strings.TrimSpace(auth.Provider), "meta") - if persistMetaMint { - if errPersist := m.persist(ctx, auth); errPersist != nil { + if errPersist := m.persistLocked(ctx, auth); errPersist != nil { + // A minted Meta key must reach the store before requests can use it. + if persistMetaMint { m.mu.Unlock() return nil, fmt.Errorf("persist meta auth: %w", errPersist) } + // Ordinary persistence failures remain non-fatal, but never silent. + log.WithFields(log.Fields{"auth_id": auth.ID, "credential": auth.ID, "provider": auth.Provider}).Warnf("failed to persist updated auth %s (%s): %v", auth.Provider, auth.ID, errPersist) + } + if m.authEpochs[auth.ID] != auth.RegistrationEpoch { + m.mu.Unlock() + return nil, fmt.Errorf("auth %s changed registration during persistence", auth.ID) + } + mergeAuthSaveDelta(auth, existingBeforeSave, existing, false) + if existing.Generation >= auth.Generation { + auth.Generation = existing.Generation + 1 } authClone := auth.Clone() m.auths[auth.ID] = authClone @@ -279,6 +307,7 @@ func (m *Manager) updateInternal(ctx context.Context, base, auth *Auth, mode upd schedulerSnapshot = authClone.Clone() } m.mu.Unlock() + releaseMutation() if !shouldDeferAPIKeyModelAliasRebuild(ctx) { m.rebuildAPIKeyModelAliasFromRuntimeConfig() } @@ -287,13 +316,6 @@ func (m *Manager) updateInternal(ctx context.Context, base, auth *Auth, mode upd } m.structuralEpoch.Add(1) m.queueRefreshReschedule(auth.ID) - if !persistMetaMint { - // Persist failures stay non-fatal, but must not be silent: after a token - // refresh the rotated credentials only exist in memory until persisted. - if errPersist := m.persist(ctx, auth); errPersist != nil { - log.WithFields(log.Fields{"auth_id": auth.ID, "credential": auth.ID, "provider": auth.Provider}).Warnf("failed to persist updated auth %s (%s): %v", auth.Provider, auth.ID, errPersist) - } - } m.hook.OnAuthUpdated(ctx, auth.Clone()) if cooldownStateChanged { m.persistCooldownStates(context.Background()) @@ -312,6 +334,8 @@ func (m *Manager) Remove(ctx context.Context, id string) { return } _ = ctx + releaseMutation := m.lockAuthMutation(id) + defer releaseMutation() m.mu.Lock() existing := m.auths[id] @@ -343,6 +367,7 @@ func (m *Manager) Remove(ctx context.Context, id string) { m.authEpochs[id]++ tombstoneEpoch := m.authEpochs[id] m.mu.Unlock() + releaseMutation() if !shouldDeferAPIKeyModelAliasRebuild(ctx) { m.rebuildAPIKeyModelAliasFromRuntimeConfig() @@ -376,16 +401,27 @@ func (m *Manager) invalidateSessionAffinity(authID string) { // Load resets manager state from the backing store. func (m *Manager) Load(ctx context.Context) error { - m.mu.Lock() - if m.store == nil { - m.mu.Unlock() + // Exclude complete Save/publication transactions, including new registrations + // not yet present in auths. Never wait for this barrier while holding m.mu. + lockCtx := ctx + if lockCtx == nil { + lockCtx = context.Background() + } + if errAcquire := m.authLoadGate.Acquire(lockCtx, math.MaxInt64); errAcquire != nil { + return errAcquire + } + defer m.authLoadGate.Release(math.MaxInt64) + m.mu.RLock() + store := m.store + m.mu.RUnlock() + if store == nil { return nil } - items, err := m.store.List(ctx) + items, err := store.List(ctx) if err != nil { - m.mu.Unlock() return err } + m.mu.Lock() previousAuths := m.auths m.auths = make(map[string]*Auth, len(items)) if m.authEpochs == nil { @@ -448,7 +484,10 @@ type authPersistLock struct { } func (m *Manager) persist(ctx context.Context, auth *Auth) error { - if m.store == nil || auth == nil { + m.mu.RLock() + store := m.store + m.mu.RUnlock() + if store == nil || auth == nil { return nil } if errWeight := ValidateAuthWeight(auth); errWeight != nil { @@ -483,13 +522,13 @@ func (m *Manager) persist(ctx context.Context, auth *Auth) error { if shouldSkipPersist(ctx) { return nil } - _, err := m.store.Save(ctx, auth) + _, err := store.Save(ctx, auth) return err } if shouldSkipPersist(ctx) { return nil } - _, err := m.store.Save(ctx, auth) + _, err := store.Save(ctx, auth) return err } diff --git a/sdk/cliproxy/auth/conductor_load_persistence_test.go b/sdk/cliproxy/auth/conductor_load_persistence_test.go new file mode 100644 index 000000000..6507b550b --- /dev/null +++ b/sdk/cliproxy/auth/conductor_load_persistence_test.go @@ -0,0 +1,196 @@ +package auth + +import ( + "context" + "errors" + "reflect" + "sync" + "testing" + "testing/synctest" + + cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v8/sdk/cliproxy/executor" +) + +// reloadAuthStore snapshots records independently of blocked storage I/O. +type reloadAuthStore struct { + mu sync.Mutex + records map[string]*Auth + saveEntered chan struct{} + saveRelease chan struct{} + listEntered chan struct{} + listRelease chan struct{} +} + +func (*reloadAuthStore) Delete(context.Context, string) error { return nil } +func (s *reloadAuthStore) Save(_ context.Context, auth *Auth) (string, error) { + if s.saveEntered != nil { + close(s.saveEntered) + <-s.saveRelease + } + auth.FileName = auth.ID + ".json" + s.mu.Lock() + s.records[auth.ID] = auth.Clone() + s.mu.Unlock() + return auth.FileName, nil +} +func (s *reloadAuthStore) List(context.Context) ([]*Auth, error) { + s.mu.Lock() + items := make([]*Auth, 0, len(s.records)) + for _, auth := range s.records { + items = append(items, auth.Clone()) + } + s.mu.Unlock() + close(s.listEntered) + if s.listRelease != nil { + <-s.listRelease + } + return items, nil +} + +func assertUnrelatedAuthReadable(t *testing.T, manager *Manager) { + t.Helper() + done := make(chan error, 1) + go func() { + if _, ok := manager.GetByID("B"); !ok { + done <- errors.New("B disappeared") + return + } + selected, err := manager.SelectAuth(context.Background(), "audit-b", "", cliproxyexecutor.Options{}) + if err == nil && (selected == nil || selected.ID != "B") { + err = errors.New("did not select B") + } + done <- err + }() + synctest.Wait() + select { + case err := <-done: + if err != nil { + t.Error(err) + } + default: + t.Error("storage I/O blocked reading/selecting unrelated B") + } +} + +func TestManagerBlockedSaveSerializesLoad(t *testing.T) { + for _, operation := range []string{"Register", "ReRegister", "Update", "Prepare", "Refresh", "MarkResult"} { + t.Run(operation, func(t *testing.T) { + synctest.Test(t, func(t *testing.T) { + ctx := context.Background() + store := &reloadAuthStore{records: make(map[string]*Auth), saveEntered: make(chan struct{}), saveRelease: make(chan struct{}), listEntered: make(chan struct{})} + manager := NewManager(store, nil, nil) + manager.RegisterExecutor(&aliasRoutingExecutor{id: "audit-b"}) + b, err := manager.Register(WithSkipPersist(ctx), &Auth{ID: "B", Provider: "audit-b", Status: StatusActive}) + if err != nil { + t.Fatal(err) + } + store.records[b.ID] = b.Clone() + base := &Auth{ID: "A", Provider: "meta", Status: StatusActive, Metadata: map[string]any{"access_token": "old"}} + if operation != "Register" { + base, err = manager.Register(WithSkipPersist(ctx), base) + if err != nil { + t.Fatal(err) + } + store.records[base.ID] = base.Clone() + } + updated := base.Clone() + updated.Metadata["access_token"] = "minted" + saveDone := make(chan error, 1) + go func() { + var errSave error + switch operation { + case "Register", "ReRegister": + _, errSave = manager.Register(ctx, updated) + case "Update": + _, errSave = manager.Update(ctx, updated) + case "Prepare": + _, errSave = manager.UpdatePreparedAuth(ctx, base, updated) + case "Refresh": + _, errSave = manager.UpdateRefreshedAuth(ctx, base, updated) + case "MarkResult": + manager.MarkResult(ctx, Result{AuthID: base.ID, Provider: base.Provider, Success: true}) + } + saveDone <- errSave + }() + <-store.saveEntered + loadDone := make(chan error, 1) + go func() { loadDone <- manager.Load(ctx) }() + synctest.Wait() + select { + case <-store.listEntered: + t.Error("Load read stale records before Save/publication completed") + default: + } + assertUnrelatedAuthReadable(t, manager) + close(store.saveRelease) + if errSave := <-saveDone; errSave != nil { + t.Errorf("save/publication failed: %v", errSave) + } + if errLoad := <-loadDone; errLoad != nil { + t.Fatal(errLoad) + } + got, ok := manager.GetByID(base.ID) + if !ok { + t.Fatal("Load lost the saved credential") + } + store.mu.Lock() + saved := store.records[base.ID].Clone() + store.mu.Unlock() + if got.RegistrationEpoch <= saved.RegistrationEpoch { + t.Error("Load did not run after save/publication") + } + saved.RegistrationEpoch, saved.Generation = got.RegistrationEpoch, got.Generation + if !reflect.DeepEqual(saved, got) { + t.Errorf("store and manager diverged: saved=%+v got=%+v", saved, got) + } + manager.scheduler.mu.Lock() + scheduled := manager.scheduler.providers["meta"].auths[base.ID].auth.Clone() + manager.scheduler.mu.Unlock() + if !reflect.DeepEqual(got, scheduled) { + t.Error("scheduler differs from reloaded credential") + } + }) + }) + } +} + +func TestManagerBlockedLoadSerializesNewRegistration(t *testing.T) { + synctest.Test(t, func(t *testing.T) { + ctx := context.Background() + store := &reloadAuthStore{records: make(map[string]*Auth), saveEntered: make(chan struct{}), saveRelease: make(chan struct{}), listEntered: make(chan struct{}), listRelease: make(chan struct{})} + manager := NewManager(store, nil, nil) + manager.RegisterExecutor(&aliasRoutingExecutor{id: "audit-b"}) + b, err := manager.Register(WithSkipPersist(ctx), &Auth{ID: "B", Provider: "audit-b", Status: StatusActive}) + if err != nil { + t.Fatal(err) + } + store.records[b.ID] = b.Clone() + loadDone := make(chan error, 1) + go func() { loadDone <- manager.Load(ctx) }() + <-store.listEntered + saveDone := make(chan error, 1) + go func() { + _, errRegister := manager.Register(ctx, &Auth{ID: "new", Provider: "audit", Metadata: map[string]any{"type": "audit"}}) + saveDone <- errRegister + }() + synctest.Wait() + select { + case <-store.saveEntered: + t.Error("new registration overtook Load's whole-map replacement") + default: + } + assertUnrelatedAuthReadable(t, manager) + close(store.listRelease) + if errLoad := <-loadDone; errLoad != nil { + t.Fatal(errLoad) + } + <-store.saveEntered + close(store.saveRelease) + if errSave := <-saveDone; errSave != nil { + t.Fatal(errSave) + } + if got, ok := manager.GetByID("new"); !ok || got.FileName != "new.json" { + t.Fatal("Load discarded a new registration") + } + }) +} diff --git a/sdk/cliproxy/auth/conductor_persistence.go b/sdk/cliproxy/auth/conductor_persistence.go new file mode 100644 index 000000000..72739c054 --- /dev/null +++ b/sdk/cliproxy/auth/conductor_persistence.go @@ -0,0 +1,113 @@ +package auth + +import ( + "context" + "reflect" + "sync" +) + +// lockAuthMutation acquires the reload barrier, then the credential gate, before +// m.mu. Keep both until publication completes. Load never takes credential or +// refresh locks. Store callbacks may read the manager, but must not mutate it. +func (m *Manager) lockAuthMutation(id string) func() { + release, _ := m.lockAuthMutationContext(context.Background(), id) + return release +} + +// lockAuthMutationContext lets fallible mutations cancel either lock wait. +func (m *Manager) lockAuthMutationContext(ctx context.Context, id string) (func(), error) { + if ctx == nil { + ctx = context.Background() + } + if errAcquire := m.authLoadGate.Acquire(ctx, 1); errAcquire != nil { + return nil, errAcquire + } + value, _ := m.authMutationLocks.LoadOrStore(id, make(chan struct{}, 1)) + gate := value.(chan struct{}) + select { + case gate <- struct{}{}: + // Cancellation may race with the gate becoming available. + if errContext := ctx.Err(); errContext != nil { + <-gate + m.authLoadGate.Release(1) + return nil, errContext + } + case <-ctx.Done(): + m.authLoadGate.Release(1) + return nil, ctx.Err() + } + var once sync.Once + return func() { + once.Do(func() { + <-gate + m.authLoadGate.Release(1) + }) + }, nil +} + +// persistLocked requires m.mu and the credential mutation gate. Never pass a +// published object to Store.Save: readers and runtime-only writers remain live +// during I/O. Only the store's delta is merged back, preserving intervening edits. +func (m *Manager) persistLocked(ctx context.Context, auth *Auth) error { + before := auth.Clone() + saved := auth.Clone() + m.mu.Unlock() + err := m.persist(ctx, saved) + m.mu.Lock() + mergeAuthSaveDelta(auth, before, saved, true) + return err +} + +// mergeAuthSaveDelta handles all exported value fields, not a store-specific +// allowlist. Map entries are independent fields; nested values are atomic. +// Identity/version fields belong exclusively to the manager. Store-owned shared +// objects must be replaced rather than mutated in place (see Store.Save). +func mergeAuthSaveDelta(target, before, after *Auth, preserveConcurrent bool) { + dst := reflect.ValueOf(target).Elem() + base := reflect.ValueOf(before).Elem() + next := reflect.ValueOf(after).Elem() + for i := 0; i < dst.NumField(); i++ { + field := dst.Type().Field(i) + if field.PkgPath != "" || field.Name == "ID" || field.Name == "RegistrationEpoch" || field.Name == "Generation" { + continue + } + d, b, n := dst.Field(i), base.Field(i), next.Field(i) + if reflect.DeepEqual(b.Interface(), n.Interface()) { + continue + } + if reflect.DeepEqual(d.Interface(), b.Interface()) { + d.Set(n) + continue + } + if d.Kind() == reflect.Map { + keys := make(map[any]reflect.Value) + for _, key := range b.MapKeys() { + keys[key.Interface()] = key + } + for _, key := range n.MapKeys() { + keys[key.Interface()] = key + } + for _, key := range keys { + oldValue, newValue := b.MapIndex(key), n.MapIndex(key) + if equalSaveValue(oldValue, newValue) || (preserveConcurrent && !equalSaveValue(d.MapIndex(key), oldValue)) { + continue + } + if d.IsNil() { + d.Set(reflect.MakeMap(d.Type())) + } + d.SetMapIndex(key, newValue) + } + continue + } + if !preserveConcurrent { + d.Set(n) + } + } +} + +func equalSaveValue(a, b reflect.Value) bool { + if !a.IsValid() || !b.IsValid() { + return a.IsValid() == b.IsValid() + } + return reflect.DeepEqual(a.Interface(), b.Interface()) +} diff --git a/sdk/cliproxy/auth/conductor_refresh.go b/sdk/cliproxy/auth/conductor_refresh.go index a92e6b5c2..285f83e43 100644 --- a/sdk/cliproxy/auth/conductor_refresh.go +++ b/sdk/cliproxy/auth/conductor_refresh.go @@ -786,7 +786,9 @@ func (m *Manager) refreshAuthForRequestAtEpoch(ctx context.Context, id, failedAc if m.shouldRefresh(updated, now) { updated.NextRefreshAfter = now.Add(refreshIneffectiveBackoff) } - saved, errUpdate := m.UpdateRefreshedAuth(ctx, base, updated) + // A successful refresh may have rotated upstream tokens after cancellation. + // Commit its result with context values intact; keep epoch checks and save ordering. + saved, errUpdate := m.UpdateRefreshedAuth(context.WithoutCancel(ctx), base, updated) if errUpdate != nil { // Warn, not debug: a restart after a lost persist fails with invalid_grant. log.WithFields(log.Fields{"auth_id": auth.ID, "credential": auth.ID, "provider": auth.Provider}).Warnf("persist refreshed auth %s (%s) failed: %v", auth.Provider, auth.ID, errUpdate) diff --git a/sdk/cliproxy/auth/conductor_saved_fields_test.go b/sdk/cliproxy/auth/conductor_saved_fields_test.go new file mode 100644 index 000000000..ec88b3325 --- /dev/null +++ b/sdk/cliproxy/auth/conductor_saved_fields_test.go @@ -0,0 +1,521 @@ +package auth + +import ( + "context" + "errors" + "reflect" + "sync" + "testing" + "testing/synctest" + "time" + + cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v8/sdk/cliproxy/executor" +) + +type enrichingAuthStore struct{ err error } + +func (*enrichingAuthStore) List(context.Context) ([]*Auth, error) { return nil, nil } +func (*enrichingAuthStore) Delete(context.Context, string) error { return nil } +func (s *enrichingAuthStore) Save(_ context.Context, auth *Auth) (string, error) { + if auth.Attributes == nil { + auth.Attributes = make(map[string]string) + } + auth.Attributes[AttributePath] = "/fixture/" + auth.ID + ".json" + auth.Attributes[AttributeSource] = auth.Attributes[AttributePath] + auth.Attributes[AttributeSourceBackend] = AuthSourceFile + auth.FileName = auth.ID + ".json" + auth.Metadata["disabled"] = auth.Disabled + return auth.Attributes[AttributePath], s.err +} + +type savedFieldsHook struct{ auth *Auth } + +func (h *savedFieldsHook) OnAuthRegistered(_ context.Context, auth *Auth) { h.auth = auth.Clone() } +func (h *savedFieldsHook) OnAuthUpdated(_ context.Context, auth *Auth) { h.auth = auth.Clone() } +func (*savedFieldsHook) OnResult(context.Context, Result) {} + +func TestManagerSavedFieldsPublished(t *testing.T) { + for _, saveErr := range []error{nil, errors.New("save failed after enrichment")} { + name := "success" + if saveErr != nil { + name = "failure" + } + t.Run(name, func(t *testing.T) { + testManagerSavedFieldsPublished(t, saveErr) + }) + } +} + +func testManagerSavedFieldsPublished(t *testing.T, saveErr error) { + for _, operation := range []string{"Register", "Update", "Prepare", "Refresh"} { + t.Run(operation, func(t *testing.T) { + ctx := context.Background() + hook := &savedFieldsHook{} + manager := NewManager(&enrichingAuthStore{err: saveErr}, nil, hook) + input := &Auth{ID: "saved-fields", Provider: "audit", Status: StatusActive, Metadata: map[string]any{"type": "audit"}} + var returned *Auth + var err error + if operation == "Register" { + returned, err = manager.Register(ctx, input) + } else { + base, errRegister := manager.Register(ctx, input) + if errRegister != nil { + t.Fatal(errRegister) + } + input = base.Clone() + input.FileName = "" + input.Attributes = nil + delete(input.Metadata, "disabled") + switch operation { + case "Update": + returned, err = manager.Update(ctx, input) + case "Prepare": + returned, err = manager.UpdatePreparedAuth(ctx, base, input) + case "Refresh": + returned, err = manager.UpdateRefreshedAuth(ctx, base, input) + } + } + if err != nil { + t.Fatal(err) + } + published, ok := manager.GetByID(input.ID) + if !ok { + t.Fatal("missing published auth") + } + for name, snapshot := range map[string]*Auth{"manager": published, "return": returned, "hook": hook.auth} { + if snapshot == nil || snapshot.FileName != "saved-fields.json" || snapshot.Attributes[AttributePath] != "/fixture/saved-fields.json" || snapshot.Attributes[AttributeSource] != "/fixture/saved-fields.json" || snapshot.Attributes[AttributeSourceBackend] != AuthSourceFile || snapshot.Metadata["disabled"] != false { + t.Errorf("%s snapshot missing Store.Save enrichment: %+v", name, snapshot) + } + if !reflect.DeepEqual(snapshot, published) { + t.Errorf("%s snapshot differs from manager", name) + } + } + manager.scheduler.mu.Lock() + scheduled := manager.scheduler.providers["audit"].auths[input.ID].auth.Clone() + manager.scheduler.mu.Unlock() + if !reflect.DeepEqual(scheduled, published) { + t.Error("scheduler snapshot differs from manager") + } + }) + } +} + +// blockingEnrichingAuthStore pauses one save before enriching the auth. +type blockingEnrichingAuthStore struct { + enrichingAuthStore + entered chan struct{} + release chan struct{} + once sync.Once + mu sync.Mutex + saved []*Auth +} + +func (s *blockingEnrichingAuthStore) Save(ctx context.Context, auth *Auth) (string, error) { + s.once.Do(func() { + close(s.entered) + <-s.release + }) + path, err := s.enrichingAuthStore.Save(ctx, auth) + s.mu.Lock() + s.saved = append(s.saved, auth.Clone()) + s.mu.Unlock() + return path, err +} + +func TestManagerSavedFieldsConcurrentLifecycle(t *testing.T) { + for _, saveErr := range []error{nil, errors.New("save failed after enrichment")} { + name := "success" + if saveErr != nil { + name = "failure" + } + t.Run(name, func(t *testing.T) { testManagerSavedFieldsConcurrentLifecycle(t, saveErr) }) + } +} + +func testManagerSavedFieldsConcurrentLifecycle(t *testing.T, saveErr error) { + for _, operation := range []string{"Register", "Update"} { + for _, concurrent := range []string{"MarkResult", "Update", "Remove", "ReRegister"} { + t.Run(operation+"/"+concurrent, func(t *testing.T) { + synctest.Test(t, func(t *testing.T) { + ctx := context.Background() + store := &blockingEnrichingAuthStore{enrichingAuthStore: enrichingAuthStore{err: saveErr}, entered: make(chan struct{}), release: make(chan struct{})} + var releaseOnce sync.Once + release := func() { releaseOnce.Do(func() { close(store.release) }) } + defer release() + manager := NewManager(store, nil, nil) + input := &Auth{ID: "saved-fields", Provider: "audit", Status: StatusActive, Label: "old", Metadata: map[string]any{"type": "audit"}} + if operation == "Update" { + var errRegister error + input, errRegister = manager.Register(WithSkipPersist(ctx), input) + if errRegister != nil { + t.Fatal(errRegister) + } + } + next := input.Clone() + next.Label = "new" + done := make(chan error, 1) + go func() { + var err error + if operation == "Register" { + _, err = manager.Register(ctx, input) + } else { + _, err = manager.Update(ctx, input) + } + done <- err + }() + <-store.entered + started := make(chan struct{}) + concurrentDone := make(chan error, 1) + go func() { + close(started) + var err error + switch concurrent { + case "MarkResult": + manager.MarkResult(ctx, Result{AuthID: next.ID, Provider: next.Provider, Model: "model", Error: &Error{HTTPStatus: 429, Message: "quota exceeded"}}) + case "Update": + _, err = manager.Update(ctx, next) + case "Remove": + manager.Remove(ctx, next.ID) + case "ReRegister": + manager.Remove(ctx, next.ID) + _, err = manager.Register(ctx, next) + } + concurrentDone <- err + }() + <-started + synctest.Wait() + select { + case <-concurrentDone: + t.Fatal("same-credential mutation completed before the active save") + default: + } + release() + if errDone := <-done; errDone != nil { + t.Fatal(errDone) + } + if errConcurrent := <-concurrentDone; errConcurrent != nil { + t.Fatal(errConcurrent) + } + published, ok := manager.GetByID(input.ID) + manager.scheduler.mu.Lock() + var scheduled *Auth + if provider := manager.scheduler.providers["audit"]; provider != nil { + if meta := provider.auths[input.ID]; meta != nil { + scheduled = meta.auth.Clone() + } + } + manager.scheduler.mu.Unlock() + if concurrent == "Remove" { + if ok || scheduled != nil { + t.Fatal("completed save resurrected removed auth") + } + return + } + if !ok || published.FileName != "saved-fields.json" || published.Attributes[AttributePath] != "/fixture/saved-fields.json" { + t.Fatalf("missing saved fields: %+v", published) + } + if concurrent == "MarkResult" { + state := published.ModelStates["model"] + if published.Failed != 1 || published.LastError == nil || state == nil || !state.Quota.Exceeded || !state.Unavailable { + t.Fatalf("save overwrote concurrent result: %+v", published) + } + } else if published.Label != "new" { + t.Fatalf("save overwrote newer lifecycle state: %+v", published) + } + if !reflect.DeepEqual(scheduled, published) { + t.Fatal("scheduler retained an obsolete save snapshot") + } + store.mu.Lock() + saved := store.saved[len(store.saved)-1].Clone() + store.mu.Unlock() + if !reflect.DeepEqual(saved, published) { + t.Fatal("store retained an obsolete auth version") + } + if concurrent == "ReRegister" { + if published.RegistrationEpoch <= input.RegistrationEpoch { + t.Fatal("re-registration did not advance epoch") + } + if _, errUpdate := manager.UpdateRefreshedAuth(ctx, input, input.Clone()); errUpdate == nil { + t.Fatal("accepted stale registration after re-register") + } + } + }) + }) + } + } +} + +func TestManagerBlockedSaveAllowsUnrelatedCredential(t *testing.T) { + for _, operation := range []string{"Register", "Update"} { + t.Run(operation, func(t *testing.T) { + ctx := context.Background() + store := &blockingEnrichingAuthStore{entered: make(chan struct{}), release: make(chan struct{})} + manager := NewManager(store, nil, nil) + manager.RegisterExecutor(&aliasRoutingExecutor{id: "audit"}) + if _, err := manager.Register(WithSkipPersist(ctx), &Auth{ID: "B", Provider: "audit", Status: StatusActive}); err != nil { + t.Fatal(err) + } + input := &Auth{ID: "A", Provider: "audit-a", Status: StatusActive, Metadata: map[string]any{"type": "audit-a"}} + if operation == "Update" { + if _, err := manager.Register(WithSkipPersist(ctx), input); err != nil { + t.Fatal(err) + } + } + done := make(chan struct{}) + go func() { + defer close(done) + if operation == "Register" { + _, _ = manager.Register(ctx, input) + } else { + _, _ = manager.Update(ctx, input) + } + }() + <-store.entered + readDone := make(chan error, 1) + go func() { + if _, ok := manager.GetByID("B"); !ok { + readDone <- errors.New("missing B") + return + } + selected, err := manager.SelectAuth(ctx, "audit", "", cliproxyexecutor.Options{}) + if err == nil && (selected == nil || selected.ID != "B") { + err = errors.New("did not select B") + } + readDone <- err + }() + select { + case err := <-readDone: + if err != nil { + t.Error(err) + } + case <-time.After(2 * time.Second): + t.Error("blocked Save(A) prevented reading/selecting unrelated B") + } + close(store.release) + <-done + }) + } +} + +type callbackAuthStore struct{ save func(*Auth) error } + +func (*callbackAuthStore) List(context.Context) ([]*Auth, error) { return nil, nil } +func (*callbackAuthStore) Delete(context.Context, string) error { return nil } +func (s *callbackAuthStore) Save(_ context.Context, auth *Auth) (string, error) { + return "", s.save(auth) +} + +func TestManagerSaveReadCallbackAndCustomFields(t *testing.T) { + synctest.Test(t, func(t *testing.T) { + ctx := context.Background() + store := &callbackAuthStore{} + manager := NewManager(store, nil, nil) + store.save = func(auth *Auth) error { + manager.GetByID(auth.ID) + manager.List() + auth.Prefix = "custom-prefix" + auth.ProxyURL = "http://custom.invalid" + auth.Metadata["custom"] = map[string]any{"value": "saved"} + delete(auth.Metadata, "remove") + return nil + } + for _, operation := range []string{"Register", "Update"} { + input := &Auth{ID: "callback", Provider: "audit", Status: StatusActive, Metadata: map[string]any{"remove": true}} + var result *Auth + var err error + if operation == "Register" { + result, err = manager.Register(ctx, input) + } else { + result, err = manager.Update(ctx, input) + } + if err != nil { + t.Fatal(err) + } + got, _ := manager.GetByID(input.ID) + if got.Prefix != "custom-prefix" || got.ProxyURL != "http://custom.invalid" || got.Metadata["custom"] == nil { + t.Fatalf("custom store fields missing: %+v", got) + } + if _, ok := got.Metadata["remove"]; ok { + t.Fatal("store deletion lost") + } + if !reflect.DeepEqual(result, got) { + t.Fatal("return differs from manager") + } + } + }) +} + +func TestManagerSavePreservesConcurrentRefreshState(t *testing.T) { + for _, operation := range []string{"Update", "MarkResult"} { + t.Run(operation, func(t *testing.T) { + synctest.Test(t, func(t *testing.T) { + ctx := context.Background() + store := &blockingEnrichingAuthStore{entered: make(chan struct{}), release: make(chan struct{})} + manager := NewManager(store, nil, nil) + base, err := manager.Register(WithSkipPersist(ctx), &Auth{ID: "runtime", Provider: "audit", Status: StatusActive, Metadata: map[string]any{"type": "audit"}}) + if err != nil { + t.Fatal(err) + } + done := make(chan struct{}) + go func() { + defer close(done) + if operation == "Update" { + _, _ = manager.Update(ctx, base.Clone()) + } else { + manager.MarkResult(ctx, Result{AuthID: base.ID, Provider: base.Provider, Success: true}) + } + }() + <-store.entered + job := manager.markRefreshPending(nil, base.ID, base.RegistrationEpoch, time.Now()) + if job == nil { + t.Fatal("failed to schedule concurrent refresh") + } + close(store.release) + <-done + got, _ := manager.GetByID(base.ID) + if !got.NextRefreshAfter.Equal(job.pendingUntil) { + t.Fatal("save lost concurrent refresh scheduling") + } + if got.Attributes[AttributePath] == "" { + t.Fatal("save enrichment lost") + } + if operation == "MarkResult" && got.Success != 1 { + t.Fatal("save lost result counter") + } + }) + }) + } +} + +func TestManagerSavedFieldsSkipPersistence(t *testing.T) { + ctx := context.Background() + store := &callbackAuthStore{save: func(*Auth) error { t.Fatal("unexpected Save"); return nil }} + for _, mode := range []string{"context", "nil-metadata", "runtime-only", "config-api-key", "plugin-virtual"} { + t.Run(mode, func(t *testing.T) { + manager := NewManager(store, nil, nil) + input := &Auth{ID: mode, Provider: "audit", Status: StatusActive, Metadata: map[string]any{"type": "audit"}} + saveCtx := ctx + switch mode { + case "context": + saveCtx = WithSkipPersist(ctx) + case "nil-metadata": + input.Metadata = nil + case "runtime-only": + input.Attributes = map[string]string{"runtime_only": "true"} + case "config-api-key": + input.Attributes = map[string]string{"source": "config", "api_key": "test"} + case "plugin-virtual": + MarkPluginVirtualAuth(input, "fixture", 0) + } + for _, operation := range []string{"Register", "Update"} { + var result *Auth + var err error + if operation == "Register" { + result, err = manager.Register(saveCtx, input) + } else { + result, err = manager.Update(saveCtx, input) + } + if err != nil { + t.Fatal(err) + } + got, _ := manager.GetByID(input.ID) + if !reflect.DeepEqual(got, result) { + t.Fatal("skip return differs from manager") + } + } + }) + } +} + +func TestMergeAuthSaveDeltaMapChanges(t *testing.T) { + t.Run("empty-map", func(t *testing.T) { + before := &Auth{} + after := &Auth{Attributes: map[string]string{}, Metadata: map[string]any{}} + target := before.Clone() + mergeAuthSaveDelta(target, before, after, true) + if !reflect.DeepEqual(target, after) { + t.Fatal("lost explicit empty map enrichment") + } + }) + t.Run("delete-map-preserves-concurrent-key", func(t *testing.T) { + before := &Auth{Metadata: map[string]any{"old": true}} + after := before.Clone() + after.Metadata = nil + target := before.Clone() + target.Metadata["concurrent"] = true + mergeAuthSaveDelta(target, before, after, true) + if len(target.Metadata) != 1 || target.Metadata["concurrent"] != true { + t.Fatalf("map deletion lost: %v", target.Metadata) + } + }) + t.Run("conflicts", func(t *testing.T) { + before := &Auth{Label: "old", Metadata: map[string]any{"same": "old", "other": "old"}} + after := before.Clone() + after.Label = "store" + after.Metadata["same"] = "store" + after.Metadata["other"] = "store" + target := before.Clone() + target.Label = "runtime" + target.Metadata["same"] = "runtime" + mergeAuthSaveDelta(target, before, after, true) + if target.Label != "runtime" || target.Metadata["same"] != "runtime" || target.Metadata["other"] != "store" { + t.Fatalf("incorrect merge: %+v", target) + } + }) +} + +func TestManagerSavedFieldsMetaMintPublication(t *testing.T) { + for _, operation := range []string{"Prepare", "Refresh"} { + for _, fail := range []bool{false, true} { + name := operation + "/success" + if fail { + name = operation + "/failure" + } + t.Run(name, func(t *testing.T) { + synctest.Test(t, func(t *testing.T) { + ctx := context.Background() + store := &blockingEnrichingAuthStore{entered: make(chan struct{}), release: make(chan struct{})} + if fail { + store.err = errors.New("mint save failed") + } + hook := &savedFieldsHook{} + manager := NewManager(store, nil, hook) + base, err := manager.Register(WithSkipPersist(ctx), &Auth{ID: "meta-mint", Provider: "meta", Status: StatusActive, Metadata: map[string]any{"access_token": "old"}}) + if err != nil { + t.Fatal(err) + } + updated := base.Clone() + updated.Metadata["access_token"] = "minted" + done := make(chan struct{}) + var returned *Auth + var errUpdate error + go func() { + defer close(done) + if operation == "Prepare" { + returned, errUpdate = manager.UpdatePreparedAuth(ctx, base, updated) + } else { + returned, errUpdate = manager.UpdateRefreshedAuth(ctx, base, updated) + } + }() + <-store.entered + pending, _ := manager.GetByID(base.ID) + if !reflect.DeepEqual(base, pending) { + t.Fatal("mint published before persistence completed") + } + close(store.release) + <-done + got, _ := manager.GetByID(base.ID) + if fail { + if errUpdate == nil || returned != nil || !reflect.DeepEqual(base, got) || !reflect.DeepEqual(base, hook.auth) { + t.Fatal("failed Meta mint changed published state") + } + } else { + if errUpdate != nil || got.Metadata["access_token"] != "minted" || got.FileName == "" || !reflect.DeepEqual(got, returned) || !reflect.DeepEqual(got, hook.auth) { + t.Fatalf("successful mint snapshots differ: %v", errUpdate) + } + } + }) + }) + } + } +} diff --git a/sdk/cliproxy/auth/conductor_selection.go b/sdk/cliproxy/auth/conductor_selection.go index 4464cac36..755e81c00 100644 --- a/sdk/cliproxy/auth/conductor_selection.go +++ b/sdk/cliproxy/auth/conductor_selection.go @@ -213,6 +213,8 @@ func (m *Manager) ReconcileRegistryModelStates(ctx context.Context, authID strin cooldownStateChanged bool ) + releaseMutation := m.lockAuthMutation(authID) + defer releaseMutation() m.mu.Lock() auth, ok := m.auths[authID] if ok && auth != nil { @@ -347,7 +349,7 @@ func (m *Manager) ReconcileRegistryModelStates(ctx context.Context, authID strin } auth.Generation++ auth.UpdatedAt = now - if errPersist := m.persist(context.Background(), auth); errPersist != nil { + if errPersist := m.persistLocked(context.Background(), auth); errPersist != nil { logEntryWithRequestID(ctx).WithField("auth_id", auth.ID).Warnf("failed to persist auth changes during model state reconciliation: %v", errPersist) } } @@ -361,6 +363,7 @@ func (m *Manager) ReconcileRegistryModelStates(ctx context.Context, authID strin } } m.mu.Unlock() + releaseMutation() if snapshot == nil { return diff --git a/sdk/cliproxy/auth/force_refresh_test.go b/sdk/cliproxy/auth/force_refresh_test.go index e970787a6..6b170717f 100644 --- a/sdk/cliproxy/auth/force_refresh_test.go +++ b/sdk/cliproxy/auth/force_refresh_test.go @@ -175,7 +175,8 @@ func (e *concurrencyTrackingRefreshExecutor) Refresh(_ context.Context, auth *Au if auth.Metadata == nil { auth.Metadata = make(map[string]any) } - auth.Metadata["access_token"] = "refreshed-token" + auth.Metadata["access_token"] = "rotated-access-" + auth.ID + auth.Metadata["refresh_token"] = "rotated-refresh-" + auth.ID return auth, nil } @@ -269,10 +270,34 @@ func TestManager_ForceRefreshAll_CanceledContextSkipsExecutors(t *testing.T) { } } -func TestManager_ForceRefreshAll_DynamicCancellationSkipsRemaining(t *testing.T) { - ctx, cancel := context.WithCancel(context.Background()) +// refreshCommitStore observes the commit context without writing credential files. +type refreshCommitStore struct { + *memoryAuthTestStore + checkContext func(context.Context) +} - manager := NewManager(nil, &RoundRobinSelector{}, nil) +func (s *refreshCommitStore) Save(ctx context.Context, auth *Auth) (string, error) { + s.checkContext(ctx) + return s.memoryAuthTestStore.Save(ctx, auth) +} + +func TestManager_ForceRefreshAll_DynamicCancellationSkipsRemaining(t *testing.T) { + type contextKey struct{} + ctx, cancel := context.WithCancel(context.WithValue(context.Background(), contextKey{}, "refresh-value")) + defer cancel() + + store := &refreshCommitStore{ + memoryAuthTestStore: newMemoryAuthTestStore(), + checkContext: func(ctx context.Context) { + if ctx.Err() != nil || ctx.Done() != nil { + t.Error("successful refresh commit must not inherit cancellation") + } + if ctx.Value(contextKey{}) != "refresh-value" { + t.Error("refresh commit lost context value") + } + }, + } + manager := NewManager(store, &RoundRobinSelector{}, nil) manager.runtimeConfig.Store(&internalconfig.Config{AuthAutoRefreshWorkers: 2}) executor := &concurrencyTrackingRefreshExecutor{ @@ -286,9 +311,9 @@ func TestManager_ForceRefreshAll_DynamicCancellationSkipsRemaining(t *testing.T) auth := &Auth{ ID: fmt.Sprintf("ag-%d", i), Provider: "antigravity", - Metadata: map[string]any{"refresh_token": fmt.Sprintf("ref-%d", i)}, + Metadata: map[string]any{"access_token": "old-access", "refresh_token": "old-refresh"}, } - if _, err := manager.Register(context.Background(), auth); err != nil { + if _, err := manager.Register(WithSkipPersist(context.Background()), auth); err != nil { t.Fatalf("register auth: %v", err) } } @@ -315,7 +340,8 @@ func TestManager_ForceRefreshAll_DynamicCancellationSkipsRemaining(t *testing.T) t.Fatalf("expected 6 results, got %d", len(results)) } - // Exactly 2 workers entered Refresh; remaining 4 were canceled in queue + // Exactly 2 workers entered Refresh; remaining 4 were canceled in queue. + // Successful rotations must be saved and published despite cancellation. if entered := executor.totalEntered.Load(); entered != 2 { t.Fatalf("expected exactly 2 entered calls, got %d", entered) } @@ -323,10 +349,25 @@ func TestManager_ForceRefreshAll_DynamicCancellationSkipsRemaining(t *testing.T) successCount := 0 cancelCount := 0 for _, res := range results { + current, exists := manager.GetByID(res.ID) + if !exists { + t.Fatalf("auth %s disappeared", res.ID) + } + store.mu.Lock() + persisted := store.auths[res.ID].Clone() + store.mu.Unlock() if res.Success { successCount++ + for name, auth := range map[string]*Auth{"manager": current, "store": persisted} { + if auth == nil || auth.Metadata["access_token"] != "rotated-access-"+res.ID || auth.Metadata["refresh_token"] != "rotated-refresh-"+res.ID { + t.Errorf("%s missing rotated tokens for %s", name, res.ID) + } + } } else if res.Error == context.Canceled.Error() { cancelCount++ + if persisted != nil || current.Metadata["access_token"] != "old-access" || current.Metadata["refresh_token"] != "old-refresh" { + t.Errorf("unstarted auth %s was modified or persisted", res.ID) + } } } if successCount != 2 { diff --git a/sdk/cliproxy/auth/store.go b/sdk/cliproxy/auth/store.go index 0594a77dd..83e401dfa 100644 --- a/sdk/cliproxy/auth/store.go +++ b/sdk/cliproxy/auth/store.go @@ -7,6 +7,12 @@ type Store interface { // List returns all auth records stored in the backend. List(ctx context.Context) ([]*Auth, error) // Save persists the provided auth record, replacing any existing one with same ID. + // It may enrich exported fields, but must not change ID, RegistrationEpoch or + // Generation. The manager merges these changes even on non-fatal save errors. + // Do not retain auth or mutate shared nested values in place: replace metadata + // values, pointers, Storage and Runtime instead. Auth.Clone shares these objects. + // Save may call manager read APIs, but must not synchronously mutate the + // manager or call Load (the reload barrier remains held until publication). Save(ctx context.Context, auth *Auth) (string, error) // Delete removes the auth record identified by id. Delete(ctx context.Context, id string) error