feat(auth): introduce mutation locking and persistence delta merging

- Add `lockAuthMutation` and `lockAuthMutationContext` with context cancellation support for mutation gates and reload barriers
- Introduce `persistLocked` to unlock manager state during storage I/O and prevent blocking runtime operations
- Implement `mergeAuthSaveDelta` to merge store save changes while preserving concurrent edits

Closes: #6329
This commit is contained in:
Luis Pater
2026-10-04 15:14:16 +08:00
parent 1672be5d25
commit 8a01bc022b
11 changed files with 1155 additions and 37 deletions

View File

@@ -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,

View File

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

View File

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

View File

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

View File

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

View File

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

View File

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

View File

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

View File

@@ -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

View File

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

View File

@@ -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