diff --git a/internal/api/handlers/management/auth_files_fields.go b/internal/api/handlers/management/auth_files_fields.go index f124fab04..53de7d68f 100644 --- a/internal/api/handlers/management/auth_files_fields.go +++ b/internal/api/handlers/management/auth_files_fields.go @@ -20,6 +20,7 @@ import ( "github.com/router-for-me/CLIProxyAPI/v7/internal/watcher/synthesizer" sdkAuth "github.com/router-for-me/CLIProxyAPI/v7/sdk/auth" coreauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth" + log "github.com/sirupsen/logrus" ) // PatchAuthFileStatus toggles the disabled state of an auth file @@ -50,6 +51,9 @@ func (h *Handler) PatchAuthFileStatus(c *gin.Context) { return } + h.authStatusMu.Lock() + defer h.authStatusMu.Unlock() + ctx := c.Request.Context() targetAuth, _ := h.lookupAuthFile(name, authIndex) @@ -108,10 +112,22 @@ func (h *Handler) PatchAuthFileStatus(c *gin.Context) { } applyAuthDisabledState(targetAuth, *req.Disabled) - if _, err := h.authManager.Update(ctx, targetAuth); err != nil { + updatedAuth, err := h.authManager.Update(ctx, targetAuth) + if err != nil { c.JSON(http.StatusInternalServerError, gin.H{"error": fmt.Sprintf("failed to update auth: %v", err)}) return } + if h.postAuthPersistHook != nil { + hookAuth := updatedAuth + if hookAuth == nil { + hookAuth = targetAuth + } + if errHook := h.postAuthPersistHook(ctx, hookAuth); errHook != nil { + log.Errorf("post-auth persist hook failed for status update on %s: %v", targetAuth.ID, errHook) + c.JSON(http.StatusInternalServerError, gin.H{"error": fmt.Sprintf("failed to synchronize auth runtime: %v", errHook)}) + return + } + } c.JSON(http.StatusOK, gin.H{"status": "ok", "disabled": *req.Disabled}) } @@ -146,9 +162,20 @@ func (h *Handler) patchPluginVirtualSourceStatus(ctx context.Context, targetAuth } applyAuthDisabledState(auth, disabled) auth.UpdatedAt = now - if _, errUpdate := h.authManager.Update(ctx, auth); errUpdate != nil { + updated, errUpdate := h.authManager.Update(ctx, auth) + if errUpdate != nil { return fmt.Errorf("failed to update auth %s: %w", auth.ID, errUpdate) } + if h.postAuthPersistHook != nil { + hookAuth := updated + if hookAuth == nil { + hookAuth = auth + } + if errHook := h.postAuthPersistHook(ctx, hookAuth); errHook != nil { + log.Errorf("post-auth persist hook failed for plugin virtual auth %s: %v", auth.ID, errHook) + return fmt.Errorf("failed to synchronize plugin virtual auth %s: %w", auth.ID, errHook) + } + } } return nil } diff --git a/internal/api/handlers/management/auth_files_status_sync_test.go b/internal/api/handlers/management/auth_files_status_sync_test.go new file mode 100644 index 000000000..7e1c40722 --- /dev/null +++ b/internal/api/handlers/management/auth_files_status_sync_test.go @@ -0,0 +1,351 @@ +package management + +import ( + "context" + "encoding/json" + "errors" + "net/http" + "net/http/httptest" + "os" + "path/filepath" + "strings" + "testing" + + "github.com/gin-gonic/gin" + "github.com/router-for-me/CLIProxyAPI/v7/internal/config" + "github.com/router-for-me/CLIProxyAPI/v7/internal/registry" + coreauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth" +) + +func TestPatchAuthFileStatusInvokesPostAuthPersistHook(t *testing.T) { + t.Setenv("MANAGEMENT_PASSWORD", "") + + authDir := t.TempDir() + fileName := "codex-test.json" + filePath := filepath.Join(authDir, fileName) + if errWrite := os.WriteFile(filePath, []byte(`{"type":"codex","disabled":false}`), 0o600); errWrite != nil { + t.Fatalf("write auth file: %v", errWrite) + } + + manager := coreauth.NewManager(nil, nil, nil) + auth := &coreauth.Auth{ + ID: fileName, + FileName: fileName, + Provider: "codex", + Status: coreauth.StatusActive, + Attributes: map[string]string{ + "path": filePath, + }, + } + if _, errRegister := manager.Register(context.Background(), auth); errRegister != nil { + t.Fatalf("register auth: %v", errRegister) + } + + h := NewHandlerWithoutConfigFilePath(&config.Config{AuthDir: authDir}, manager) + + var hookCalls []*coreauth.Auth + h.SetPostAuthPersistHook(func(_ context.Context, updated *coreauth.Auth) error { + if updated != nil { + hookCalls = append(hookCalls, updated.Clone()) + } + return nil + }) + + // 1. Disable the credential + rec := httptest.NewRecorder() + ctx, _ := gin.CreateTestContext(rec) + req := httptest.NewRequest(http.MethodPatch, "/v0/management/auth-files/status", strings.NewReader(`{"name":"codex-test.json","disabled":true}`)) + req.Header.Set("Content-Type", "application/json") + ctx.Request = req + + h.PatchAuthFileStatus(ctx) + if rec.Code != http.StatusOK { + t.Fatalf("status = %d, want %d body=%s", rec.Code, http.StatusOK, rec.Body.String()) + } + + if len(hookCalls) != 1 { + t.Fatalf("expected 1 hook call after disable, got %d", len(hookCalls)) + } + if !hookCalls[0].Disabled { + t.Fatalf("expected hook call auth to be disabled, got %+v", hookCalls[0]) + } + + // 2. Re-enable the credential + rec = httptest.NewRecorder() + ctx, _ = gin.CreateTestContext(rec) + req = httptest.NewRequest(http.MethodPatch, "/v0/management/auth-files/status", strings.NewReader(`{"name":"codex-test.json","disabled":false}`)) + req.Header.Set("Content-Type", "application/json") + ctx.Request = req + + h.PatchAuthFileStatus(ctx) + if rec.Code != http.StatusOK { + t.Fatalf("status = %d, want %d body=%s", rec.Code, http.StatusOK, rec.Body.String()) + } + + if len(hookCalls) != 2 { + t.Fatalf("expected 2 hook calls after re-enable, got %d", len(hookCalls)) + } + if hookCalls[1].Disabled { + t.Fatalf("expected second hook call auth to be enabled, got %+v", hookCalls[1]) + } +} + +func TestPatchAuthFileStatusRestoresModelsViaSyncHook(t *testing.T) { + t.Setenv("MANAGEMENT_PASSWORD", "") + + authDir := t.TempDir() + fileName := "codex-models.json" + filePath := filepath.Join(authDir, fileName) + if errWrite := os.WriteFile(filePath, []byte(`{"type":"codex","disabled":false}`), 0o600); errWrite != nil { + t.Fatalf("write auth file: %v", errWrite) + } + + reg := registry.GetGlobalRegistry() + authID := fileName + t.Cleanup(func() { + reg.UnregisterClient(authID) + }) + + manager := coreauth.NewManager(nil, nil, nil) + auth := &coreauth.Auth{ + ID: authID, + FileName: fileName, + Provider: "codex", + Status: coreauth.StatusActive, + Attributes: map[string]string{ + "path": filePath, + }, + } + if _, errRegister := manager.Register(context.Background(), auth); errRegister != nil { + t.Fatalf("register auth: %v", errRegister) + } + + // Initial model registration + reg.RegisterClient(authID, "codex", []*registry.ModelInfo{ + {ID: "gpt-6-astra", DisplayName: "GPT-6 Astra"}, + }) + + h := NewHandlerWithoutConfigFilePath(&config.Config{AuthDir: authDir}, manager) + // Wire postAuthPersistHook to synchronize registry models as runtimeAuthSyncHook does + h.SetPostAuthPersistHook(func(_ context.Context, a *coreauth.Auth) error { + if a == nil || a.ID == "" { + return nil + } + if a.Disabled { + reg.UnregisterClient(a.ID) + } else { + reg.RegisterClient(a.ID, a.Provider, []*registry.ModelInfo{ + {ID: "gpt-6-astra", DisplayName: "GPT-6 Astra"}, + }) + } + return nil + }) + + // Step 1: Confirm models exist before disabling + rec := httptest.NewRecorder() + ctx, _ := gin.CreateTestContext(rec) + ctx.Request = httptest.NewRequest(http.MethodGet, "/v0/management/auth-files/models?name="+fileName, nil) + h.GetAuthFileModels(ctx) + if rec.Code != http.StatusOK { + t.Fatalf("GetAuthFileModels initial status = %d, body = %s", rec.Code, rec.Body.String()) + } + var resp struct { + Models []map[string]any `json:"models"` + } + if err := json.Unmarshal(rec.Body.Bytes(), &resp); err != nil || len(resp.Models) != 1 { + t.Fatalf("expected 1 initial model, got: %s", rec.Body.String()) + } + + // Step 2: Disable the credential + rec = httptest.NewRecorder() + ctx, _ = gin.CreateTestContext(rec) + req := httptest.NewRequest(http.MethodPatch, "/v0/management/auth-files/status", strings.NewReader(`{"name":"codex-models.json","disabled":true}`)) + req.Header.Set("Content-Type", "application/json") + ctx.Request = req + h.PatchAuthFileStatus(ctx) + if rec.Code != http.StatusOK { + t.Fatalf("PatchAuthFileStatus disable failed: %s", rec.Body.String()) + } + + rec = httptest.NewRecorder() + ctx, _ = gin.CreateTestContext(rec) + ctx.Request = httptest.NewRequest(http.MethodGet, "/v0/management/auth-files/models?name="+fileName, nil) + h.GetAuthFileModels(ctx) + resp.Models = nil + if err := json.Unmarshal(rec.Body.Bytes(), &resp); err != nil || len(resp.Models) != 0 { + t.Fatalf("expected 0 models after disable, got: %s", rec.Body.String()) + } + + // Step 3: Re-enable the credential + rec = httptest.NewRecorder() + ctx, _ = gin.CreateTestContext(rec) + req = httptest.NewRequest(http.MethodPatch, "/v0/management/auth-files/status", strings.NewReader(`{"name":"codex-models.json","disabled":false}`)) + req.Header.Set("Content-Type", "application/json") + ctx.Request = req + h.PatchAuthFileStatus(ctx) + if rec.Code != http.StatusOK { + t.Fatalf("PatchAuthFileStatus re-enable failed: %s", rec.Body.String()) + } + + // Step 4: Confirm models are restored without restart + rec = httptest.NewRecorder() + ctx, _ = gin.CreateTestContext(rec) + ctx.Request = httptest.NewRequest(http.MethodGet, "/v0/management/auth-files/models?name="+fileName, nil) + h.GetAuthFileModels(ctx) + resp.Models = nil + if err := json.Unmarshal(rec.Body.Bytes(), &resp); err != nil || len(resp.Models) != 1 { + t.Fatalf("expected 1 model after re-enable, got: %s", rec.Body.String()) + } +} + +func TestPatchPluginVirtualSourceStatusInvokesPostAuthPersistHook(t *testing.T) { + t.Setenv("MANAGEMENT_PASSWORD", "") + + authDir := t.TempDir() + fileName := "source-sync.json" + filePath := filepath.Join(authDir, fileName) + if errWrite := os.WriteFile(filePath, []byte(`{"type":"gemini-cli","disabled":false}`), 0o600); errWrite != nil { + t.Fatalf("write source auth file: %v", errWrite) + } + + manager := coreauth.NewManager(nil, nil, nil) + for _, id := range []string{"source-sync.json", "virtual-project-a", "virtual-project-b"} { + auth := pluginVirtualAuthForTest(authDir, fileName, id) + if _, errRegister := manager.Register(context.Background(), auth); errRegister != nil { + t.Fatalf("register virtual auth %s: %v", id, errRegister) + } + } + + h := NewHandlerWithoutConfigFilePath(&config.Config{AuthDir: authDir}, manager) + + var hookCalls []*coreauth.Auth + h.SetPostAuthPersistHook(func(_ context.Context, updated *coreauth.Auth) error { + if updated != nil { + hookCalls = append(hookCalls, updated.Clone()) + } + return nil + }) + + // Disable the source file + rec := httptest.NewRecorder() + ctx, _ := gin.CreateTestContext(rec) + req := httptest.NewRequest(http.MethodPatch, "/v0/management/auth-files/status", strings.NewReader(`{"name":"source-sync.json","disabled":true}`)) + req.Header.Set("Content-Type", "application/json") + ctx.Request = req + + h.PatchAuthFileStatus(ctx) + if rec.Code != http.StatusOK { + t.Fatalf("status = %d, want %d body=%s", rec.Code, http.StatusOK, rec.Body.String()) + } + + if len(hookCalls) != 3 { + t.Fatalf("expected 3 hook calls for 3 virtual auths on disable, got %d", len(hookCalls)) + } + for _, call := range hookCalls { + if !call.Disabled { + t.Fatalf("expected virtual auth %s to be disabled in hook call", call.ID) + } + } + + // Re-enable the source file + hookCalls = nil + rec = httptest.NewRecorder() + ctx, _ = gin.CreateTestContext(rec) + req = httptest.NewRequest(http.MethodPatch, "/v0/management/auth-files/status", strings.NewReader(`{"name":"source-sync.json","disabled":false}`)) + req.Header.Set("Content-Type", "application/json") + ctx.Request = req + + h.PatchAuthFileStatus(ctx) + if rec.Code != http.StatusOK { + t.Fatalf("status = %d, want %d body=%s", rec.Code, http.StatusOK, rec.Body.String()) + } + + if len(hookCalls) != 3 { + t.Fatalf("expected 3 hook calls for 3 virtual auths on re-enable, got %d", len(hookCalls)) + } + for _, call := range hookCalls { + if call.Disabled { + t.Fatalf("expected virtual auth %s to be enabled in hook call", call.ID) + } + } +} + +func TestPatchAuthFileStatusHookErrorReturns500(t *testing.T) { + t.Setenv("MANAGEMENT_PASSWORD", "") + + authDir := t.TempDir() + fileName := "codex-hook-err.json" + filePath := filepath.Join(authDir, fileName) + if errWrite := os.WriteFile(filePath, []byte(`{"type":"codex","disabled":false}`), 0o600); errWrite != nil { + t.Fatalf("write auth file: %v", errWrite) + } + + manager := coreauth.NewManager(nil, nil, nil) + auth := &coreauth.Auth{ + ID: fileName, + FileName: fileName, + Provider: "codex", + Status: coreauth.StatusActive, + Attributes: map[string]string{ + "path": filePath, + }, + } + if _, errRegister := manager.Register(context.Background(), auth); errRegister != nil { + t.Fatalf("register auth: %v", errRegister) + } + + h := NewHandlerWithoutConfigFilePath(&config.Config{AuthDir: authDir}, manager) + h.SetPostAuthPersistHook(func(_ context.Context, _ *coreauth.Auth) error { + return errors.New("simulated sync hook failure") + }) + + rec := httptest.NewRecorder() + ctx, _ := gin.CreateTestContext(rec) + req := httptest.NewRequest(http.MethodPatch, "/v0/management/auth-files/status", strings.NewReader(`{"name":"codex-hook-err.json","disabled":true}`)) + req.Header.Set("Content-Type", "application/json") + ctx.Request = req + + h.PatchAuthFileStatus(ctx) + if rec.Code != http.StatusInternalServerError { + t.Fatalf("status = %d, want %d body=%s", rec.Code, http.StatusInternalServerError, rec.Body.String()) + } + if !strings.Contains(rec.Body.String(), "simulated sync hook failure") { + t.Fatalf("body = %s, want simulated sync hook failure error message", rec.Body.String()) + } +} + +func TestPatchPluginVirtualSourceStatusHookErrorReturnsError(t *testing.T) { + t.Setenv("MANAGEMENT_PASSWORD", "") + + authDir := t.TempDir() + fileName := "source-hook-err.json" + filePath := filepath.Join(authDir, fileName) + if errWrite := os.WriteFile(filePath, []byte(`{"type":"gemini-cli","disabled":false}`), 0o600); errWrite != nil { + t.Fatalf("write source auth file: %v", errWrite) + } + + manager := coreauth.NewManager(nil, nil, nil) + auth := pluginVirtualAuthForTest(authDir, fileName, "source-hook-err.json") + if _, errRegister := manager.Register(context.Background(), auth); errRegister != nil { + t.Fatalf("register virtual auth: %v", errRegister) + } + + h := NewHandlerWithoutConfigFilePath(&config.Config{AuthDir: authDir}, manager) + h.SetPostAuthPersistHook(func(_ context.Context, _ *coreauth.Auth) error { + return errors.New("plugin sync failed") + }) + + rec := httptest.NewRecorder() + ctx, _ := gin.CreateTestContext(rec) + req := httptest.NewRequest(http.MethodPatch, "/v0/management/auth-files/status", strings.NewReader(`{"name":"source-hook-err.json","disabled":true}`)) + req.Header.Set("Content-Type", "application/json") + ctx.Request = req + + h.PatchAuthFileStatus(ctx) + if rec.Code != http.StatusInternalServerError { + t.Fatalf("status = %d, want %d body=%s", rec.Code, http.StatusInternalServerError, rec.Body.String()) + } + if !strings.Contains(rec.Body.String(), "plugin sync failed") { + t.Fatalf("body = %s, want plugin sync failed error message", rec.Body.String()) + } +} diff --git a/internal/api/handlers/management/handler.go b/internal/api/handlers/management/handler.go index 0426bb667..883d532ee 100644 --- a/internal/api/handlers/management/handler.go +++ b/internal/api/handlers/management/handler.go @@ -41,6 +41,7 @@ type Handler struct { cfg *config.Config configFilePath string mu sync.Mutex + authStatusMu sync.Mutex reloadMu sync.Mutex reloadGeneration uint64 appliedReloadGeneration uint64 diff --git a/internal/api/server.go b/internal/api/server.go index f12ff4bf4..8108d5bf9 100644 --- a/internal/api/server.go +++ b/internal/api/server.go @@ -254,6 +254,14 @@ func NewServer(cfg *config.Config, authManager *auth.Manager, accessManager *sdk return s } +// Handler returns the HTTP handler used by the server. +func (s *Server) Handler() http.Handler { + if s == nil || s.server == nil { + return nil + } + return s.server.Handler +} + // Start begins listening for and serving HTTP or HTTPS requests. // It's a blocking call and will only return on an unrecoverable error. // diff --git a/internal/watcher/dispatcher.go b/internal/watcher/dispatcher.go index 2bf3f20a0..89b4aae5d 100644 --- a/internal/watcher/dispatcher.go +++ b/internal/watcher/dispatcher.go @@ -81,11 +81,16 @@ func (w *Watcher) dispatchRuntimeAuthUpdate(update AuthUpdate) bool { } func (w *Watcher) dispatchPersistedAuthUpdate(update AuthUpdate) bool { - if w == nil { - return false + ok, _ := w.dispatchPersistedAuthUpdateWithRevision(&update) + return ok +} + +func (w *Watcher) dispatchPersistedAuthUpdateWithRevision(update *AuthUpdate) (bool, uint64) { + if w == nil || update == nil { + return false, 0 } if update.Auth == nil || update.Auth.ID == "" { - return false + return false, 0 } path := "" if update.Auth.Attributes != nil { @@ -96,7 +101,7 @@ func (w *Watcher) dispatchPersistedAuthUpdate(update AuthUpdate) bool { } normalized := w.normalizeAuthPath(path) if normalized == "" { - return false + return false, 0 } clone := update.Auth.Clone() w.clientsMutex.Lock() @@ -116,15 +121,18 @@ func (w *Watcher) dispatchPersistedAuthUpdate(update AuthUpdate) bool { if update.ID == "" { update.ID = clone.ID } - update.Auth = clone.Clone() - updates := []AuthUpdate{update} + updateCopy := *update + updateCopy.Auth = clone.Clone() + updates := []AuthUpdate{updateCopy} w.stampAuthUpdatesLocked(updates) + rev := updates[0].revision + update.revision = rev w.clientsMutex.Unlock() if w.getAuthQueue() == nil { - return false + return false, rev } w.dispatchAuthUpdates(updates) - return true + return true, rev } func (w *Watcher) refreshAuthState(force bool) { diff --git a/internal/watcher/watcher.go b/internal/watcher/watcher.go index 7ce83f782..9dfb8617d 100644 --- a/internal/watcher/watcher.go +++ b/internal/watcher/watcher.go @@ -84,6 +84,18 @@ type AuthUpdate struct { revision uint64 // Watcher-local ordering, independent of runtime auth generations. } +// Revision returns the monotonic watcher revision assigned to this update. +func (u AuthUpdate) Revision() uint64 { + return u.revision +} + +// SetRevision updates the revision counter for this update. +func (u *AuthUpdate) SetRevision(rev uint64) { + if u != nil { + u.revision = rev + } +} + const ( // replaceCheckDelay is a short delay to allow atomic replace (rename) to settle // before deciding whether a Remove event indicates a real deletion. @@ -170,6 +182,12 @@ func (w *Watcher) DispatchPersistedAuthUpdate(update AuthUpdate) bool { return w.dispatchPersistedAuthUpdate(update) } +// DispatchPersistedAuthUpdateWithRevision pushes already-persisted file auth updates through the watcher queue +// and returns the stamped monotonic revision. +func (w *Watcher) DispatchPersistedAuthUpdateWithRevision(update *AuthUpdate) (bool, uint64) { + return w.dispatchPersistedAuthUpdateWithRevision(update) +} + // SnapshotCoreAuths converts current clients snapshot into core auth entries. func (w *Watcher) SnapshotCoreAuths() []*coreauth.Auth { w.clientsMutex.RLock() diff --git a/sdk/cliproxy/builder.go b/sdk/cliproxy/builder.go index f9b2084e3..cd4cc5b9f 100644 --- a/sdk/cliproxy/builder.go +++ b/sdk/cliproxy/builder.go @@ -320,10 +320,16 @@ func (s *Service) runtimeAuthSyncHook() coreauth.PostAuthHook { ID: auth.ID, Auth: auth, } - if s.watcher != nil && s.watcher.DispatchPersistedAuthUpdate(update) { - return nil + if s.watcher != nil { + _, rev := s.watcher.DispatchPersistedAuthUpdateWithRevision(&update) + if rev > 0 { + update.SetRevision(rev) + } } - s.handleAuthUpdate(coreauth.WithSkipPersist(ctx), update) + // Detach from request cancellation so runtime model registration always completes + // once the credential has been persisted to disk. + syncCtx := coreauth.WithSkipPersist(context.Background()) + s.handleAuthUpdate(syncCtx, update) return nil } } diff --git a/sdk/cliproxy/service.go b/sdk/cliproxy/service.go index bb1e146a6..839be4c6d 100644 --- a/sdk/cliproxy/service.go +++ b/sdk/cliproxy/service.go @@ -38,6 +38,8 @@ type Service struct { // configRuntimeMu orders side-effecting runtime application after config commits. configRuntimeMu sync.Mutex executorRegistrationMu sync.Mutex + authUpdateMu sync.Mutex + authRevisions map[string]uint64 configSequence uint64 appliedRoutingState *routingRuntimeState diff --git a/sdk/cliproxy/service_auth.go b/sdk/cliproxy/service_auth.go index 11b1e1d1b..20ca66fbc 100644 --- a/sdk/cliproxy/service_auth.go +++ b/sdk/cliproxy/service_auth.go @@ -94,7 +94,34 @@ func (s *Service) handleAuthUpdates(ctx context.Context, updates []watcher.AuthU if s == nil { return } - updates = coalesceAuthUpdates(updates) + s.authUpdateMu.Lock() + defer s.authUpdateMu.Unlock() + + if s.authRevisions == nil { + s.authRevisions = make(map[string]uint64) + } + + filtered := make([]watcher.AuthUpdate, 0, len(updates)) + for _, update := range updates { + id := authUpdateID(update) + if id == "" { + filtered = append(filtered, update) + continue + } + rev := update.Revision() + if rev > 0 { + if prevRev, exists := s.authRevisions[id]; exists && rev <= prevRev { + log.Debugf("skipping stale auth update for %s: rev %d <= processed %d", id, rev, prevRev) + continue + } + s.authRevisions[id] = rev + } + filtered = append(filtered, update) + } + if len(filtered) == 0 { + return + } + updates = coalesceAuthUpdates(filtered) s.cfgMu.RLock() cfg := s.cfg s.cfgMu.RUnlock() @@ -134,6 +161,12 @@ func (s *Service) handleAuthUpdates(ctx context.Context, updates []watcher.AuthU if id == "" { continue } + if existing, ok := s.coreManager.GetByID(id); ok && existing != nil && update.Auth != nil { + if isStaleCoreAuth(existing, update.Auth) { + log.Debugf("skipping stale auth delete for %s: incoming gen=%d, existing gen=%d", id, update.Auth.Generation, existing.Generation) + continue + } + } s.applyCoreAuthRemoval(registrationCtx, id) needsAliasRebuild = true default: @@ -282,6 +315,10 @@ func (s *Service) prepareCoreAuthForModelRegistration(ctx context.Context, auth op := "register" var err error if existing, ok := s.coreManager.GetByID(auth.ID); ok { + if isStaleCoreAuth(existing, auth) { + log.Debugf("skipping stale auth update for %s: incoming gen=%d, existing gen=%d", auth.ID, auth.Generation, existing.Generation) + return existing + } auth.CreatedAt = existing.CreatedAt if !existing.Disabled && existing.Status != coreauth.StatusDisabled && !auth.Disabled && auth.Status != coreauth.StatusDisabled { auth.LastRefreshedAt = existing.LastRefreshedAt @@ -307,6 +344,19 @@ func (s *Service) prepareCoreAuthForModelRegistration(ctx context.Context, auth return auth } +// isStaleCoreAuth reports whether an incoming auth update is older than the current +// state in coreManager, based on registration epoch. +func isStaleCoreAuth(existing, incoming *coreauth.Auth) bool { + if existing == nil || incoming == nil { + return false + } + // If incoming has an explicit registration epoch that is older than existing, it's stale. + if incoming.RegistrationEpoch > 0 && incoming.RegistrationEpoch < existing.RegistrationEpoch { + return true + } + return false +} + func (s *Service) completeModelRegistrationForAuth(ctx context.Context, auth *coreauth.Auth) { s.completeModelRegistrationForAuthWithCache(ctx, auth, nil) } diff --git a/sdk/cliproxy/service_auth_sync_test.go b/sdk/cliproxy/service_auth_sync_test.go new file mode 100644 index 000000000..e5d7f6ca4 --- /dev/null +++ b/sdk/cliproxy/service_auth_sync_test.go @@ -0,0 +1,710 @@ +package cliproxy + +import ( + "context" + "encoding/json" + "net/http" + "net/http/httptest" + "os" + "path/filepath" + "strings" + "testing" + "time" + + "golang.org/x/crypto/bcrypt" + + "github.com/router-for-me/CLIProxyAPI/v7/internal/api" + internalregistry "github.com/router-for-me/CLIProxyAPI/v7/internal/registry" + "github.com/router-for-me/CLIProxyAPI/v7/internal/watcher" + coreauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth" + cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor" + "github.com/router-for-me/CLIProxyAPI/v7/sdk/config" +) + +func TestRuntimeAuthSyncHook_SynchronousModelAndSchedulerRestoration(t *testing.T) { + reg := internalregistry.GetGlobalRegistry() + authID := "codex-sync-test-auth" + reg.UnregisterClient(authID) + t.Cleanup(func() { + reg.UnregisterClient(authID) + }) + + manager := coreauth.NewManager(nil, nil, nil) + manager.RegisterExecutor(&syncTestExecutor{}) + auth := &coreauth.Auth{ + ID: authID, + Provider: "codex", + Status: coreauth.StatusActive, + Attributes: map[string]string{ + "plan_type": "pro", + "path": "/path/to/codex-sync.json", + }, + } + if _, errRegister := manager.Register(context.Background(), auth); errRegister != nil { + t.Fatalf("register auth: %v", errRegister) + } + + // Set up a mock watcher wrapper where dispatchPersistedAuth records updates + // but DOES NOT consume or execute anything asynchronously, verifying that + // the synchronous restoration does not depend on the watcher background queue. + var persistedUpdates []watcher.AuthUpdate + watcherWrapper := &WatcherWrapper{ + dispatchPersistedAuth: func(update watcher.AuthUpdate) bool { + persistedUpdates = append(persistedUpdates, update) + return true + }, + } + + service := &Service{ + cfg: &config.Config{}, + coreManager: manager, + watcher: watcherWrapper, + } + + hook := service.runtimeAuthSyncHook() + if hook == nil { + t.Fatal("runtimeAuthSyncHook() returned nil") + } + + // 1. Initial sync (enabled -> registers models) + if err := hook(context.Background(), auth); err != nil { + t.Fatalf("hook failed for initial sync: %v", err) + } + manager.RegisterExecutor(&syncTestExecutor{}) + models := reg.GetModelsForClient(authID) + if len(models) == 0 { + t.Fatal("expected models registered for active auth, got none") + } + + // Verify scheduler selection works for enabled auth + resp, errExec := manager.Execute(context.Background(), []string{"codex"}, cliproxyexecutor.Request{Model: "gpt-6-astra"}, cliproxyexecutor.Options{}) + if errExec != nil || string(resp.Payload) != authID { + t.Fatalf("expected scheduler to select active auth, got resp=%q, err=%v", string(resp.Payload), errExec) + } + + // 2. Disable the credential + disabledAuth := auth.Clone() + disabledAuth.Disabled = true + disabledAuth.Status = coreauth.StatusDisabled + disabledAuth.StatusMessage = "disabled via management API" + disabledAuth.UpdatedAt = time.Now() + + if _, errUpdate := manager.Update(context.Background(), disabledAuth); errUpdate != nil { + t.Fatalf("manager update disable failed: %v", errUpdate) + } + if err := hook(context.Background(), disabledAuth); err != nil { + t.Fatalf("hook failed for disable: %v", err) + } + manager.RegisterExecutor(&syncTestExecutor{}) + + // Verify immediately: models must be unregistered from global registry + modelsAfterDisable := reg.GetModelsForClient(authID) + if len(modelsAfterDisable) != 0 { + t.Fatalf("expected 0 models for disabled client, got %d", len(modelsAfterDisable)) + } + + // Verify scheduler selection fails for disabled auth + if _, errExecDisabled := manager.Execute(context.Background(), []string{"codex"}, cliproxyexecutor.Request{Model: "gpt-6-astra"}, cliproxyexecutor.Options{}); errExecDisabled == nil { + t.Fatal("expected scheduler execution to fail for disabled auth, but got nil error") + } + + // 3. Re-enable the credential without quota + enabledAuth := disabledAuth.Clone() + enabledAuth.Disabled = false + enabledAuth.Status = coreauth.StatusActive + enabledAuth.StatusMessage = "" + enabledAuth.UpdatedAt = time.Now() + + if _, errUpdate := manager.Update(context.Background(), enabledAuth); errUpdate != nil { + t.Fatalf("manager update enable failed: %v", errUpdate) + } + + // 4. Re-enable hook invocation + if err := hook(context.Background(), enabledAuth); err != nil { + t.Fatalf("hook failed for re-enable: %v", err) + } + manager.RegisterExecutor(&syncTestExecutor{}) + + // Verify immediately: models must be restored in the global registry without restart + modelsAfterEnable := reg.GetModelsForClient(authID) + if len(modelsAfterEnable) == 0 { + t.Fatal("expected models to be restored synchronously upon re-enable, got none") + } + + hasAstra := false + for _, m := range modelsAfterEnable { + if m != nil && m.ID == "gpt-6-astra" { + hasAstra = true + break + } + } + if !hasAstra { + t.Fatalf("expected gpt-6-astra in restored models: %+v", modelsAfterEnable) + } + + // Verify scheduler selection is restored for the re-enabled auth + respAfterEnable, errExecEnable := manager.Execute(context.Background(), []string{"codex"}, cliproxyexecutor.Request{Model: "gpt-6-astra"}, cliproxyexecutor.Options{}) + if errExecEnable != nil || string(respAfterEnable.Payload) != authID { + t.Fatalf("expected scheduler to select restored auth, got resp=%q, err=%v", string(respAfterEnable.Payload), errExecEnable) + } + + // 5. Verify quota cooldown separation: set quota exceeded on the enabled auth + quotaAuth := enabledAuth.Clone() + quotaAuth.Quota = coreauth.QuotaState{ + Exceeded: true, + Reason: "credential_quota", + NextRecoverAt: time.Now().Add(10 * time.Minute), + } + if _, errUpdate := manager.Update(context.Background(), quotaAuth); errUpdate != nil { + t.Fatalf("manager update quota failed: %v", errUpdate) + } + if err := hook(context.Background(), quotaAuth); err != nil { + t.Fatalf("hook failed for quota update: %v", err) + } + manager.RegisterExecutor(&syncTestExecutor{}) + + // Verify quota separation: models remain registered in global registry + modelsWithQuota := reg.GetModelsForClient(authID) + if len(modelsWithQuota) == 0 { + t.Fatal("expected models to remain registered when auth is in quota cooldown") + } + // But scheduler blocks execution due to quota exhaustion + if _, errQuotaExec := manager.Execute(context.Background(), []string{"codex"}, cliproxyexecutor.Request{Model: "gpt-6-astra"}, cliproxyexecutor.Options{}); errQuotaExec == nil { + t.Fatal("expected scheduler to block execution when quota exceeded") + } + + // Verify watcher notification occurred + if len(persistedUpdates) == 0 { + t.Fatal("expected watcher to be notified of persisted auth updates") + } +} + +type syncTestExecutor struct{} + +func (e *syncTestExecutor) Identifier() string { return "codex" } + +func (e *syncTestExecutor) Execute(ctx context.Context, a *coreauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { + return cliproxyexecutor.Response{Payload: []byte(a.ID)}, nil +} + +func (e *syncTestExecutor) ExecuteStream(context.Context, *coreauth.Auth, cliproxyexecutor.Request, cliproxyexecutor.Options) (*cliproxyexecutor.StreamResult, error) { + return nil, &coreauth.Error{HTTPStatus: http.StatusNotImplemented, Message: "not implemented"} +} + +func (e *syncTestExecutor) Refresh(_ context.Context, auth *coreauth.Auth) (*coreauth.Auth, error) { + return auth, nil +} + +func (e *syncTestExecutor) CountTokens(context.Context, *coreauth.Auth, cliproxyexecutor.Request, cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { + return cliproxyexecutor.Response{}, &coreauth.Error{HTTPStatus: http.StatusNotImplemented, Message: "not implemented"} +} + +func (e *syncTestExecutor) HttpRequest(context.Context, *coreauth.Auth, *http.Request) (*http.Response, error) { + return nil, &coreauth.Error{HTTPStatus: http.StatusNotImplemented, Message: "not implemented"} +} + +func TestRuntimeAuthSyncHook_NilAndEmptyAuth(t *testing.T) { + service := &Service{} + hook := service.runtimeAuthSyncHook() + if hook == nil { + t.Fatal("runtimeAuthSyncHook() returned nil") + } + + if err := hook(context.Background(), nil); err != nil { + t.Fatalf("expected nil error for nil auth, got: %v", err) + } + + if err := hook(context.Background(), &coreauth.Auth{}); err != nil { + t.Fatalf("expected nil error for empty auth, got: %v", err) + } +} + +func TestEndToEndStatusPatch_RestoresModelsThroughServerPipeline(t *testing.T) { + authDir := t.TempDir() + fileName := "codex-e2e.json" + filePath := filepath.Join(authDir, fileName) + if errWrite := os.WriteFile(filePath, []byte(`{"type":"codex","disabled":false}`), 0o600); errWrite != nil { + t.Fatalf("write auth file: %v", errWrite) + } + + reg := internalregistry.GetGlobalRegistry() + authID := fileName + reg.UnregisterClient(authID) + t.Cleanup(func() { + reg.UnregisterClient(authID) + }) + + secretKey := "test-secret" + hashedSecret, errHash := bcrypt.GenerateFromPassword([]byte(secretKey), bcrypt.DefaultCost) + if errHash != nil { + t.Fatalf("bcrypt hash failed: %v", errHash) + } + + cfg := &config.Config{ + AuthDir: authDir, + RemoteManagement: config.RemoteManagement{ + SecretKey: string(hashedSecret), + AllowRemote: true, + }, + } + + service, errBuild := NewBuilder(). + WithConfig(cfg). + WithConfigPath(filepath.Join(authDir, "config.yaml")). + Build() + if errBuild != nil { + t.Fatalf("Build() failed: %v", errBuild) + } + + server := api.NewServer(service.cfg, service.coreManager, service.accessManager, service.configPath, service.serverOptions...) + if server == nil { + t.Fatal("NewServer() returned nil") + } + + auth := &coreauth.Auth{ + ID: authID, + FileName: fileName, + Provider: "codex", + Status: coreauth.StatusActive, + Attributes: map[string]string{ + "plan_type": "pro", + "path": filePath, + }, + } + if _, errRegister := service.coreManager.Register(context.Background(), auth); errRegister != nil { + t.Fatalf("register auth: %v", errRegister) + } + + // Initial model sync + syncHook := service.runtimeAuthSyncHook() + if err := syncHook(context.Background(), auth); err != nil { + t.Fatalf("initial sync failed: %v", err) + } + + handler := server.Handler() + + // 1. Confirm models exist before disabling + rec := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodGet, "/v0/management/auth-files/models?name="+fileName, nil) + req.Header.Set("Authorization", "Bearer test-secret") + handler.ServeHTTP(rec, req) + if rec.Code != http.StatusOK { + t.Fatalf("GetAuthFileModels initial status = %d, body = %s", rec.Code, rec.Body.String()) + } + var resp struct { + Models []map[string]any `json:"models"` + } + if err := json.Unmarshal(rec.Body.Bytes(), &resp); err != nil || len(resp.Models) == 0 { + t.Fatalf("expected initial models, got %s", rec.Body.String()) + } + + // 2. Disable credential via PATCH /v0/management/auth-files/status + rec = httptest.NewRecorder() + req = httptest.NewRequest(http.MethodPatch, "/v0/management/auth-files/status", strings.NewReader(`{"name":"`+fileName+`","disabled":true}`)) + req.Header.Set("Authorization", "Bearer test-secret") + req.Header.Set("Content-Type", "application/json") + handler.ServeHTTP(rec, req) + if rec.Code != http.StatusOK { + t.Fatalf("PATCH status disable failed: %s", rec.Body.String()) + } + + // 3. Confirm GET /v0/management/auth-files/models returns empty list immediately + rec = httptest.NewRecorder() + req = httptest.NewRequest(http.MethodGet, "/v0/management/auth-files/models?name="+fileName, nil) + req.Header.Set("Authorization", "Bearer test-secret") + handler.ServeHTTP(rec, req) + if rec.Code != http.StatusOK { + t.Fatalf("GetAuthFileModels after disable status = %d, body = %s", rec.Code, rec.Body.String()) + } + resp.Models = nil + if err := json.Unmarshal(rec.Body.Bytes(), &resp); err != nil || len(resp.Models) != 0 { + t.Fatalf("expected 0 models after disable, got: %s", rec.Body.String()) + } + + // 4. Re-enable credential via PATCH /v0/management/auth-files/status + rec = httptest.NewRecorder() + req = httptest.NewRequest(http.MethodPatch, "/v0/management/auth-files/status", strings.NewReader(`{"name":"`+fileName+`","disabled":false}`)) + req.Header.Set("Authorization", "Bearer test-secret") + req.Header.Set("Content-Type", "application/json") + handler.ServeHTTP(rec, req) + if rec.Code != http.StatusOK { + t.Fatalf("PATCH status re-enable failed: %s", rec.Body.String()) + } + + // 5. Confirm models are synchronously restored immediately without restart + rec = httptest.NewRecorder() + req = httptest.NewRequest(http.MethodGet, "/v0/management/auth-files/models?name="+fileName, nil) + req.Header.Set("Authorization", "Bearer test-secret") + handler.ServeHTTP(rec, req) + if rec.Code != http.StatusOK { + t.Fatalf("GetAuthFileModels after re-enable status = %d, body = %s", rec.Code, rec.Body.String()) + } + resp.Models = nil + if err := json.Unmarshal(rec.Body.Bytes(), &resp); err != nil || len(resp.Models) == 0 { + t.Fatalf("expected models to be restored after re-enable, got: %s", rec.Body.String()) + } +} + +func TestRuntimeAuthSync_StaleAsyncEventDoesNotOverwriteSynchronousEnable(t *testing.T) { + reg := internalregistry.GetGlobalRegistry() + authID := "codex-stale-test-auth" + reg.UnregisterClient(authID) + t.Cleanup(func() { + reg.UnregisterClient(authID) + }) + + manager := coreauth.NewManager(nil, nil, nil) + manager.RegisterExecutor(&syncTestExecutor{}) + + var currentRevision uint64 + watcherWrapper := &WatcherWrapper{ + dispatchPersistedAuthWithRev: func(update *watcher.AuthUpdate) (bool, uint64) { + currentRevision++ + return true, currentRevision + }, + } + + service := &Service{ + cfg: &config.Config{}, + coreManager: manager, + watcher: watcherWrapper, + } + + auth := &coreauth.Auth{ + ID: authID, + Provider: "codex", + Status: coreauth.StatusActive, + Attributes: map[string]string{ + "plan_type": "pro", + }, + } + registeredAuth, errRegister := manager.Register(context.Background(), auth) + if errRegister != nil { + t.Fatalf("register auth: %v", errRegister) + } + + hook := service.runtimeAuthSyncHook() + + // Initial sync + if err := hook(context.Background(), registeredAuth); err != nil { + t.Fatalf("initial hook sync failed: %v", err) + } + if len(reg.GetModelsForClient(authID)) == 0 { + t.Fatal("expected models after initial sync") + } + + // Step 1: Credential is disabled (Generation advances) + disabledAuth := registeredAuth.Clone() + disabledAuth.Disabled = true + disabledAuth.Status = coreauth.StatusDisabled + disabledAuth.UpdatedAt = time.Now() + + updatedDisabled, errDisable := manager.Update(context.Background(), disabledAuth) + if errDisable != nil { + t.Fatalf("update disable failed: %v", errDisable) + } + + // Capture the stale async event that would normally sit in an async queue + staleAsyncDisableUpdate := watcher.AuthUpdate{ + Action: watcher.AuthUpdateActionModify, + ID: authID, + Auth: updatedDisabled.Clone(), + } + staleAsyncDisableUpdate.SetRevision(2) + + // Step 2: Synchronous disable hook runs (assigns revision 2) + if err := hook(context.Background(), updatedDisabled); err != nil { + t.Fatalf("hook disable failed: %v", err) + } + if len(reg.GetModelsForClient(authID)) != 0 { + t.Fatal("expected 0 models after disable") + } + + // Step 3: Credential is immediately re-enabled (advances revision to 3) + enabledAuth := updatedDisabled.Clone() + enabledAuth.Disabled = false + enabledAuth.Status = coreauth.StatusActive + enabledAuth.UpdatedAt = time.Now() + + updatedEnabled, errEnable := manager.Update(context.Background(), enabledAuth) + if errEnable != nil { + t.Fatalf("update enable failed: %v", errEnable) + } + + // Step 4: Synchronous enable hook runs (assigns revision 3) + if err := hook(context.Background(), updatedEnabled); err != nil { + t.Fatalf("hook enable failed: %v", err) + } + manager.RegisterExecutor(&syncTestExecutor{}) + + // Models must be restored immediately + if len(reg.GetModelsForClient(authID)) == 0 { + t.Fatal("expected models restored immediately upon synchronous enable") + } + + // Step 5: The delayed stale disable update from Step 1 is delivered now via consumer + service.handleAuthUpdate(context.Background(), staleAsyncDisableUpdate) + manager.RegisterExecutor(&syncTestExecutor{}) + + // Step 6: Verify the stale async event was prevented from overwriting the newer enable state! + currentAuth, ok := manager.GetByID(authID) + if !ok || currentAuth == nil { + t.Fatal("expected auth to exist in manager") + } + if currentAuth.Disabled { + t.Fatal("stale async event overwrote the synchronous enable state in manager!") + } + modelsAfterStale := reg.GetModelsForClient(authID) + if len(modelsAfterStale) == 0 { + t.Fatal("stale async event unregistered models from global registry!") + } + + // Scheduler execution must still succeed and select the enabled auth + resp, errExec := manager.Execute(context.Background(), []string{"codex"}, cliproxyexecutor.Request{Model: "gpt-6-astra"}, cliproxyexecutor.Options{}) + if errExec != nil || string(resp.Payload) != authID { + t.Fatalf("expected scheduler to select active auth, got resp=%q err=%v", string(resp.Payload), errExec) + } +} + +func TestRuntimeAuthSync_StaleWatcherUnversionedEventDoesNotOverwriteSynchronousEnable(t *testing.T) { + reg := internalregistry.GetGlobalRegistry() + authID := "codex-unversioned-stale-test-auth" + reg.UnregisterClient(authID) + t.Cleanup(func() { + reg.UnregisterClient(authID) + }) + + manager := coreauth.NewManager(nil, nil, nil) + manager.RegisterExecutor(&syncTestExecutor{}) + + var currentRevision uint64 + watcherWrapper := &WatcherWrapper{ + dispatchPersistedAuthWithRev: func(update *watcher.AuthUpdate) (bool, uint64) { + currentRevision++ + return true, currentRevision + }, + } + + service := &Service{ + cfg: &config.Config{}, + coreManager: manager, + watcher: watcherWrapper, + } + + auth := &coreauth.Auth{ + ID: authID, + Provider: "codex", + Status: coreauth.StatusActive, + Attributes: map[string]string{ + "plan_type": "pro", + }, + } + registeredAuth, errRegister := manager.Register(context.Background(), auth) + if errRegister != nil { + t.Fatalf("register auth: %v", errRegister) + } + + hook := service.runtimeAuthSyncHook() + + // Initial sync + if err := hook(context.Background(), registeredAuth); err != nil { + t.Fatalf("initial hook sync failed: %v", err) + } + manager.RegisterExecutor(&syncTestExecutor{}) + if len(reg.GetModelsForClient(authID)) == 0 { + t.Fatal("expected models after initial sync") + } + + // Step 1: Simulate the watcher generating an unversioned Auth snapshot during a disable + // event on disk (as synthesizer.SynthesizeAuthFile produces: RegistrationEpoch=0, Generation=0). + // Timestamp is recorded at observation time T1. + unversionedDisableTime := time.Now().Add(-50 * time.Millisecond) + unversionedDisableAuth := &coreauth.Auth{ + ID: authID, + Provider: "codex", + Disabled: true, + Status: coreauth.StatusDisabled, + CreatedAt: unversionedDisableTime, + UpdatedAt: unversionedDisableTime, + RegistrationEpoch: 0, // Unversioned from file watcher synthesizer + Generation: 0, // Unversioned from file watcher synthesizer + Attributes: map[string]string{ + "plan_type": "pro", + }, + } + delayedWatcherUpdate := watcher.AuthUpdate{ + Action: watcher.AuthUpdateActionModify, + ID: authID, + Auth: unversionedDisableAuth, + } + delayedWatcherUpdate.SetRevision(2) + + // Step 2: In the meantime, the credential was synchronously re-enabled via management API + // at T2 > T1, updating coreManager with RegistrationEpoch >= 1, Generation >= 2, UpdatedAt = T2. + enabledAuth := registeredAuth.Clone() + enabledAuth.Disabled = false + enabledAuth.Status = coreauth.StatusActive + enabledAuth.UpdatedAt = time.Now() + + updatedEnabled, errEnable := manager.Update(context.Background(), enabledAuth) + if errEnable != nil { + t.Fatalf("update enable failed: %v", errEnable) + } + + // Synchronous enable hook runs and restores models in GlobalModelRegistry + if err := hook(context.Background(), updatedEnabled); err != nil { + t.Fatalf("hook enable failed: %v", err) + } + manager.RegisterExecutor(&syncTestExecutor{}) + + if len(reg.GetModelsForClient(authID)) == 0 { + t.Fatal("expected models restored immediately upon synchronous enable") + } + + // Step 3: Now the delayed unversioned watcher event from Step 1 arrives via the queue consumer + service.handleAuthUpdate(context.Background(), delayedWatcherUpdate) + manager.RegisterExecutor(&syncTestExecutor{}) + + // Scheduler execution must still succeed and select the enabled auth + resp, errExec := manager.Execute(context.Background(), []string{"codex"}, cliproxyexecutor.Request{Model: "gpt-6-astra"}, cliproxyexecutor.Options{}) + if errExec != nil || string(resp.Payload) != authID { + t.Fatalf("expected scheduler to select active auth, got resp=%q err=%v", string(resp.Payload), errExec) + } +} + +func TestRuntimeAuthSync_FileSnapshotEnqueuedThenRequestExecutionThenConsume(t *testing.T) { + reg := internalregistry.GetGlobalRegistry() + authID := "codex-interleaved-request-auth" + reg.UnregisterClient(authID) + t.Cleanup(func() { + reg.UnregisterClient(authID) + }) + + manager := coreauth.NewManager(nil, nil, nil) + manager.RegisterExecutor(&syncTestExecutor{}) + + service := &Service{ + cfg: &config.Config{}, + coreManager: manager, + } + + auth := &coreauth.Auth{ + ID: authID, + Provider: "codex", + Status: coreauth.StatusActive, + Attributes: map[string]string{ + "plan_type": "pro", + }, + } + registeredAuth, errRegister := manager.Register(context.Background(), auth) + if errRegister != nil { + t.Fatalf("register auth: %v", errRegister) + } + + hook := service.runtimeAuthSyncHook() + + // Initial sync (Revision 1) + if err := hook(context.Background(), registeredAuth); err != nil { + t.Fatalf("initial hook sync failed: %v", err) + } + manager.RegisterExecutor(&syncTestExecutor{}) + if len(reg.GetModelsForClient(authID)) == 0 { + t.Fatal("expected models after initial sync") + } + + // Step 1: A file change occurs on disk (e.g. user disables credential in auth file). + // The watcher generates an AuthUpdate with revision 2 and enqueues it. + fileSnapshotAuth := registeredAuth.Clone() + fileSnapshotAuth.Disabled = true + fileSnapshotAuth.Status = coreauth.StatusDisabled + fileSnapshotAuth.UpdatedAt = time.Now() + + fileUpdate := watcher.AuthUpdate{ + Action: watcher.AuthUpdateActionModify, + ID: authID, + Auth: fileSnapshotAuth, + } + fileUpdate.SetRevision(2) + + // Step 2: Before the queue consumer can process fileUpdate, an API request completes. + // MarkResult updates the auth in coreManager: auth.Generation++ and auth.UpdatedAt = now. + manager.MarkResult(context.Background(), coreauth.Result{ + AuthID: authID, + Provider: "codex", + Model: "gpt-6-astra", + Success: true, + }) + + // Step 3: Now the queue consumer processes the fileUpdate (revision 2). + // Because revision 2 > revision 1, the file update MUST NOT be discarded due to + // the intermediate request's UpdatedAt timestamp! + service.handleAuthUpdate(context.Background(), fileUpdate) + manager.RegisterExecutor(&syncTestExecutor{}) + + // Step 4: Verify the file update successfully took effect! + currentAuth, ok := manager.GetByID(authID) + if !ok || currentAuth == nil { + t.Fatal("expected auth to exist in manager") + } + if !currentAuth.Disabled { + t.Fatal("file update was incorrectly dropped after intermediate request execution!") + } + modelsAfterDisable := reg.GetModelsForClient(authID) + if len(modelsAfterDisable) != 0 { + t.Fatalf("expected 0 models after file update disable, got %d", len(modelsAfterDisable)) + } +} + +func TestRuntimeAuthSyncHook_ContextCancellationDoesNotAbortSync(t *testing.T) { + reg := internalregistry.GetGlobalRegistry() + authID := "codex-canceled-ctx-auth" + reg.UnregisterClient(authID) + t.Cleanup(func() { + reg.UnregisterClient(authID) + }) + + manager := coreauth.NewManager(nil, nil, nil) + manager.RegisterExecutor(&syncTestExecutor{}) + + service := &Service{ + cfg: &config.Config{}, + coreManager: manager, + } + + auth := &coreauth.Auth{ + ID: authID, + Provider: "codex", + Status: coreauth.StatusActive, + Attributes: map[string]string{ + "plan_type": "pro", + }, + } + registeredAuth, errRegister := manager.Register(context.Background(), auth) + if errRegister != nil { + t.Fatalf("register auth: %v", errRegister) + } + + hook := service.runtimeAuthSyncHook() + + // Create a context that is ALREADY canceled (e.g. client aborted request right after save) + canceledCtx, cancel := context.WithCancel(context.Background()) + cancel() + + // Calling hook with canceled context must NOT abort registration + if err := hook(canceledCtx, registeredAuth); err != nil { + t.Fatalf("hook failed with canceled context: %v", err) + } + manager.RegisterExecutor(&syncTestExecutor{}) + + // Verify models were registered successfully despite canceled incoming context + models := reg.GetModelsForClient(authID) + if len(models) == 0 { + t.Fatal("expected models to be registered even when incoming context is canceled") + } + + // Verify scheduler can execute + resp, errExec := manager.Execute(context.Background(), []string{"codex"}, cliproxyexecutor.Request{Model: "gpt-6-astra"}, cliproxyexecutor.Options{}) + if errExec != nil || string(resp.Payload) != authID { + t.Fatalf("expected scheduler execution to succeed, got resp=%q err=%v", string(resp.Payload), errExec) + } +} diff --git a/sdk/cliproxy/types.go b/sdk/cliproxy/types.go index dfde6d9c1..89a1a4845 100644 --- a/sdk/cliproxy/types.go +++ b/sdk/cliproxy/types.go @@ -101,13 +101,14 @@ type WatcherWrapper struct { start func(ctx context.Context) error stop func() error - setConfig func(cfg *config.Config) - snapshotAuths func() []*coreauth.Auth - setUpdateQueue func(queue chan<- watcher.AuthUpdate) - dispatchRuntimeUpdate func(update watcher.AuthUpdate) bool - dispatchPersistedAuth func(update watcher.AuthUpdate) bool - setPluginAuthParser func(parser PluginAuthParser) - reloadConfigIfChanged func() + setConfig func(cfg *config.Config) + snapshotAuths func() []*coreauth.Auth + setUpdateQueue func(queue chan<- watcher.AuthUpdate) + dispatchRuntimeUpdate func(update watcher.AuthUpdate) bool + dispatchPersistedAuth func(update watcher.AuthUpdate) bool + dispatchPersistedAuthWithRev func(update *watcher.AuthUpdate) (bool, uint64) + setPluginAuthParser func(parser PluginAuthParser) + reloadConfigIfChanged func() } // Start proxies to the underlying watcher Start implementation. @@ -163,10 +164,33 @@ func (w *WatcherWrapper) DispatchRuntimeAuthUpdate(update watcher.AuthUpdate) bo // DispatchPersistedAuthUpdate forwards already-persisted file auth updates. func (w *WatcherWrapper) DispatchPersistedAuthUpdate(update watcher.AuthUpdate) bool { - if w == nil || w.dispatchPersistedAuth == nil { + if w == nil { return false } - return w.dispatchPersistedAuth(update) + if w.dispatchPersistedAuthWithRev != nil { + ok, _ := w.dispatchPersistedAuthWithRev(&update) + return ok + } + if w.dispatchPersistedAuth != nil { + return w.dispatchPersistedAuth(update) + } + return false +} + +// DispatchPersistedAuthUpdateWithRevision forwards already-persisted file auth updates +// and returns whether it was enqueued along with its assigned watcher revision. +func (w *WatcherWrapper) DispatchPersistedAuthUpdateWithRevision(update *watcher.AuthUpdate) (bool, uint64) { + if w == nil { + return false, 0 + } + if w.dispatchPersistedAuthWithRev != nil { + return w.dispatchPersistedAuthWithRev(update) + } + if w.dispatchPersistedAuth != nil && update != nil { + ok := w.dispatchPersistedAuth(*update) + return ok, update.Revision() + } + return false, 0 } // SetClients updates the watcher file-backed clients registry. diff --git a/sdk/cliproxy/watcher.go b/sdk/cliproxy/watcher.go index 886b55646..b75d116aa 100644 --- a/sdk/cliproxy/watcher.go +++ b/sdk/cliproxy/watcher.go @@ -34,6 +34,9 @@ func defaultWatcherFactory(configPath, authDir string, reload func(*config.Confi dispatchPersistedAuth: func(update watcher.AuthUpdate) bool { return w.DispatchPersistedAuthUpdate(update) }, + dispatchPersistedAuthWithRev: func(update *watcher.AuthUpdate) (bool, uint64) { + return w.DispatchPersistedAuthUpdateWithRevision(update) + }, setPluginAuthParser: func(parser PluginAuthParser) { w.SetPluginAuthParser(parser) },