From 55440f0a3907f9085edbe179de71877c9fde9369 Mon Sep 17 00:00:00 2001 From: Luis Pater Date: Wed, 3 Jun 2026 11:52:27 +0800 Subject: [PATCH] feat(auth): add runtime auth removal and unscheduling logic - Introduced `Manager.Remove` to delete runtime auth and unschedule associated tasks. - Updated handler logic to directly remove auth instead of marking as disabled. - Added tests to validate removal, unscheduling, and runtime state handling. - Added a test to validate `skipPersist` behavior during registration. - Enhanced `Remove` test to verify auto-refresh loop state before and after removal. Closes: #3690 --- .../api/handlers/management/auth_files.go | 24 ++-- .../management/auth_files_delete_test.go | 46 ++++++++ sdk/cliproxy/auth/conductor.go | 99 ++++++++++++++-- sdk/cliproxy/auth/conductor_remove_test.go | 111 ++++++++++++++++++ sdk/cliproxy/auth/persist_policy_test.go | 7 ++ sdk/cliproxy/service.go | 18 ++- sdk/cliproxy/service_stale_state_test.go | 33 +----- 7 files changed, 268 insertions(+), 70 deletions(-) create mode 100644 sdk/cliproxy/auth/conductor_remove_test.go diff --git a/internal/api/handlers/management/auth_files.go b/internal/api/handlers/management/auth_files.go index c32f41a71..b26bea753 100644 --- a/internal/api/handlers/management/auth_files.go +++ b/internal/api/handlers/management/auth_files.go @@ -770,7 +770,7 @@ func (h *Handler) DeleteAuthFile(c *gin.Context) { return } deleted++ - h.disableAuth(ctx, full) + h.removeAuth(ctx, full) } } c.JSON(200, gin.H{"status": "ok", "deleted": deleted}) @@ -976,9 +976,9 @@ func (h *Handler) deleteAuthFileByName(ctx context.Context, name string) (string return filepath.Base(name), http.StatusInternalServerError, errDeleteRecord } if targetID != "" { - h.disableAuth(ctx, targetID) + h.removeAuth(ctx, targetID) } else { - h.disableAuth(ctx, targetPath) + h.removeAuth(ctx, targetPath) } return filepath.Base(name), http.StatusOK, nil } @@ -1558,7 +1558,7 @@ func syncAuthFileDisabledState(auth *coreauth.Auth) { auth.StatusMessage = "" } -func (h *Handler) disableAuth(ctx context.Context, id string) { +func (h *Handler) removeAuth(ctx context.Context, id string) { if h == nil || h.authManager == nil { return } @@ -1566,25 +1566,15 @@ func (h *Handler) disableAuth(ctx context.Context, id string) { if id == "" { return } - if auth, ok := h.authManager.GetByID(id); ok { - auth.Disabled = true - auth.Status = coreauth.StatusDisabled - auth.StatusMessage = "removed via management API" - auth.UpdatedAt = time.Now() - _, _ = h.authManager.Update(ctx, auth) + if _, ok := h.authManager.GetByID(id); ok { + h.authManager.Remove(ctx, id) return } authID := h.authIDForPath(id) if authID == "" { return } - if auth, ok := h.authManager.GetByID(authID); ok { - auth.Disabled = true - auth.Status = coreauth.StatusDisabled - auth.StatusMessage = "removed via management API" - auth.UpdatedAt = time.Now() - _, _ = h.authManager.Update(ctx, auth) - } + h.authManager.Remove(ctx, authID) } func (h *Handler) deleteTokenRecord(ctx context.Context, path string) error { diff --git a/internal/api/handlers/management/auth_files_delete_test.go b/internal/api/handlers/management/auth_files_delete_test.go index a57c9993a..b67f1f66c 100644 --- a/internal/api/handlers/management/auth_files_delete_test.go +++ b/internal/api/handlers/management/auth_files_delete_test.go @@ -127,3 +127,49 @@ func TestDeleteAuthFile_FallbackToAuthDirPath(t *testing.T) { t.Fatalf("expected auth file to be removed from auth dir, stat err: %v", errStat) } } + +func TestDeleteAuthFile_RemovesRuntimeAuth(t *testing.T) { + t.Setenv("MANAGEMENT_PASSWORD", "") + gin.SetMode(gin.TestMode) + + authDir := t.TempDir() + fileName := "runtime-remove-user.json" + filePath := filepath.Join(authDir, fileName) + if errWrite := os.WriteFile(filePath, []byte(`{"type":"codex","email":"runtime@example.com"}`), 0o600); errWrite != nil { + t.Fatalf("failed to write auth file: %v", errWrite) + } + + manager := coreauth.NewManager(nil, nil, nil) + record := &coreauth.Auth{ + ID: "runtime-remove-auth", + FileName: fileName, + Provider: "codex", + Status: coreauth.StatusActive, + Attributes: map[string]string{ + "path": filePath, + }, + Metadata: map[string]any{ + "type": "codex", + "email": "runtime@example.com", + }, + } + if _, errRegister := manager.Register(context.Background(), record); errRegister != nil { + t.Fatalf("failed to register auth record: %v", errRegister) + } + + h := NewHandlerWithoutConfigFilePath(&config.Config{AuthDir: authDir}, manager) + h.tokenStore = &memoryAuthStore{} + + deleteRec := httptest.NewRecorder() + deleteCtx, _ := gin.CreateTestContext(deleteRec) + deleteReq := httptest.NewRequest(http.MethodDelete, "/v0/management/auth-files?name="+url.QueryEscape(fileName), nil) + deleteCtx.Request = deleteReq + h.DeleteAuthFile(deleteCtx) + + if deleteRec.Code != http.StatusOK { + t.Fatalf("expected delete status %d, got %d with body %s", http.StatusOK, deleteRec.Code, deleteRec.Body.String()) + } + if _, ok := manager.GetByID(record.ID); ok { + t.Fatalf("expected runtime auth %q to be removed", record.ID) + } +} diff --git a/sdk/cliproxy/auth/conductor.go b/sdk/cliproxy/auth/conductor.go index 2d355d48a..8c8effcdd 100644 --- a/sdk/cliproxy/auth/conductor.go +++ b/sdk/cliproxy/auth/conductor.go @@ -1164,18 +1164,21 @@ func (m *Manager) Update(ctx context.Context, auth *Auth) (*Auth, error) { return nil, nil } m.mu.Lock() - if existing, ok := m.auths[auth.ID]; ok && existing != nil { - if !auth.indexAssigned && auth.Index == "" { - auth.Index = existing.Index - auth.indexAssigned = existing.indexAssigned - } - auth.Success = existing.Success - auth.Failed = existing.Failed - auth.recentRequests = existing.recentRequests - if !existing.Disabled && existing.Status != StatusDisabled && !auth.Disabled && auth.Status != StatusDisabled { - if len(auth.ModelStates) == 0 && len(existing.ModelStates) > 0 { - auth.ModelStates = existing.ModelStates - } + existing, ok := m.auths[auth.ID] + if !ok || existing == nil { + m.mu.Unlock() + return nil, nil + } + if !auth.indexAssigned && auth.Index == "" { + auth.Index = existing.Index + auth.indexAssigned = existing.indexAssigned + } + auth.Success = existing.Success + auth.Failed = existing.Failed + auth.recentRequests = existing.recentRequests + if !existing.Disabled && existing.Status != StatusDisabled && !auth.Disabled && auth.Status != StatusDisabled { + if len(auth.ModelStates) == 0 && len(existing.ModelStates) > 0 { + auth.ModelStates = existing.ModelStates } } auth.EnsureIndex() @@ -1192,6 +1195,65 @@ func (m *Manager) Update(ctx context.Context, auth *Auth) (*Auth, error) { return auth.Clone(), nil } +// Remove deletes an auth from runtime state without persisting. +// Disk and token-store deletion must be handled by the caller. +func (m *Manager) Remove(ctx context.Context, id string) { + if m == nil { + return + } + id = strings.TrimSpace(id) + if id == "" { + return + } + _ = ctx + + m.mu.Lock() + existing := m.auths[id] + if existing == nil { + m.mu.Unlock() + return + } + provider := strings.TrimSpace(existing.Provider) + delete(m.auths, id) + if m.modelPoolOffsets != nil { + delete(m.modelPoolOffsets, id) + } + for sessionID, sessionAuths := range m.homeRuntimeAuths { + if sessionAuths == nil { + continue + } + delete(sessionAuths, id) + if len(sessionAuths) == 0 { + delete(m.homeRuntimeAuths, sessionID) + } + } + m.mu.Unlock() + + m.rebuildAPIKeyModelAliasFromRuntimeConfig() + if m.scheduler != nil { + m.scheduler.removeAuth(id) + } + m.queueRefreshUnschedule(id) + m.invalidateSessionAffinity(id) + + if provider != "" { + if exec, ok := m.Executor(provider); ok && exec != nil { + if closer, okCloser := exec.(ExecutionSessionCloser); okCloser { + closer.CloseExecutionSession(CloseAllExecutionSessionsID) + } + } + } +} + +func (m *Manager) invalidateSessionAffinity(authID string) { + if m == nil || authID == "" { + return + } + if invalidator, ok := m.selector.(interface{ InvalidateAuth(string) }); ok && invalidator != nil { + invalidator.InvalidateAuth(authID) + } +} + // Load resets manager state from the backing store. func (m *Manager) Load(ctx context.Context) error { m.mu.Lock() @@ -4041,6 +4103,19 @@ func (m *Manager) queueRefreshReschedule(authID string) { loop.queueReschedule(authID) } +func (m *Manager) queueRefreshUnschedule(authID string) { + if m == nil || authID == "" { + return + } + m.mu.RLock() + loop := m.refreshLoop + m.mu.RUnlock() + if loop == nil { + return + } + loop.remove(authID) +} + func (m *Manager) shouldRefresh(a *Auth, now time.Time) bool { if a == nil { return false diff --git a/sdk/cliproxy/auth/conductor_remove_test.go b/sdk/cliproxy/auth/conductor_remove_test.go new file mode 100644 index 000000000..1ada1d74f --- /dev/null +++ b/sdk/cliproxy/auth/conductor_remove_test.go @@ -0,0 +1,111 @@ +package auth + +import ( + "context" + "testing" + "time" +) + +func TestManager_Remove_DeletesRuntimeAuth(t *testing.T) { + manager := NewManager(nil, nil, nil) + ctx := context.Background() + + auth := &Auth{ + ID: "remove-runtime-auth", + Provider: "claude", + Status: StatusActive, + Metadata: map[string]any{"email": "x@example.com"}, + } + if _, errRegister := manager.Register(ctx, auth); errRegister != nil { + t.Fatalf("register auth: %v", errRegister) + } + + manager.Remove(ctx, auth.ID) + + if _, ok := manager.GetByID(auth.ID); ok { + t.Fatalf("expected auth %q to be removed", auth.ID) + } +} + +func TestManager_Update_MissingAuthIsNoOp(t *testing.T) { + manager := NewManager(nil, nil, nil) + ctx := context.Background() + + auth := &Auth{ + ID: "missing-update-auth", + Provider: "claude", + Status: StatusActive, + } + if _, errRegister := manager.Register(ctx, auth); errRegister != nil { + t.Fatalf("register auth: %v", errRegister) + } + manager.Remove(ctx, auth.ID) + + updated, errUpdate := manager.Update(ctx, &Auth{ + ID: auth.ID, + Provider: "claude", + Status: StatusDisabled, + Disabled: true, + }) + if errUpdate != nil { + t.Fatalf("update removed auth: %v", errUpdate) + } + if updated != nil { + t.Fatalf("expected update on removed auth to be no-op, got %#v", updated) + } + if _, ok := manager.GetByID(auth.ID); ok { + t.Fatalf("expected removed auth to stay absent after late update") + } +} + +func TestManager_Remove_UnschedulesAutoRefresh(t *testing.T) { + ctx := context.Background() + + manager := NewManager(nil, nil, nil) + loop := newAuthAutoRefreshLoop(manager, time.Second, 1) + manager.mu.Lock() + manager.refreshLoop = loop + manager.mu.Unlock() + + lead := 10 * time.Minute + setRefreshLeadFactory(t, "provider-lead-expiry", func() *time.Duration { + d := lead + return &d + }) + + auth := &Auth{ + ID: "remove-refresh-auth", + Provider: "provider-lead-expiry", + Metadata: map[string]any{ + "email": "x@example.com", + "expires_at": time.Now().Add(time.Hour).Format(time.RFC3339), + }, + } + if _, errRegister := manager.Register(ctx, auth); errRegister != nil { + t.Fatalf("register auth: %v", errRegister) + } + + now := time.Now() + if _, ok := nextRefreshCheckAt(now, auth, time.Second); !ok { + t.Fatalf("expected auth to be scheduled before removal") + } + loop.applyDirty(now) + loop.mu.Lock() + if _, ok := loop.index[auth.ID]; !ok { + loop.mu.Unlock() + t.Fatalf("expected auth %q to be present in auto-refresh index before removal", auth.ID) + } + loop.mu.Unlock() + + manager.Remove(ctx, auth.ID) + + if _, ok := manager.GetByID(auth.ID); ok { + t.Fatalf("expected auth to be removed") + } + loop.mu.Lock() + if _, ok := loop.index[auth.ID]; ok { + loop.mu.Unlock() + t.Fatalf("expected auth %q to be removed from auto-refresh index", auth.ID) + } + loop.mu.Unlock() +} diff --git a/sdk/cliproxy/auth/persist_policy_test.go b/sdk/cliproxy/auth/persist_policy_test.go index f408c872d..6ec4aaf2f 100644 --- a/sdk/cliproxy/auth/persist_policy_test.go +++ b/sdk/cliproxy/auth/persist_policy_test.go @@ -28,6 +28,13 @@ func TestWithSkipPersist_DisablesUpdatePersistence(t *testing.T) { Metadata: map[string]any{"type": "antigravity"}, } + if _, err := mgr.Register(WithSkipPersist(context.Background()), auth); err != nil { + t.Fatalf("Register(skipPersist) returned error: %v", err) + } + if got := store.saveCount.Load(); got != 0 { + t.Fatalf("expected 0 Save calls, got %d", got) + } + if _, err := mgr.Update(context.Background(), auth); err != nil { t.Fatalf("Update returned error: %v", err) } diff --git a/sdk/cliproxy/service.go b/sdk/cliproxy/service.go index 10c3d0dd9..ff30ad372 100644 --- a/sdk/cliproxy/service.go +++ b/sdk/cliproxy/service.go @@ -339,17 +339,15 @@ func (s *Service) applyCoreAuthRemoval(ctx context.Context, id string) { if s.coreManager == nil { return } - GlobalModelRegistry().UnregisterClient(id) + id = strings.TrimSpace(id) + var provider string if existing, ok := s.coreManager.GetByID(id); ok && existing != nil { - existing.Disabled = true - existing.Status = coreauth.StatusDisabled - if _, err := s.coreManager.Update(ctx, existing); err != nil { - log.Errorf("failed to disable auth %s: %v", id, err) - } - if strings.EqualFold(strings.TrimSpace(existing.Provider), "codex") { - executor.CloseCodexWebsocketSessionsForAuthID(existing.ID, "auth_removed") - s.ensureExecutorsForAuth(existing) - } + provider = strings.TrimSpace(existing.Provider) + } + GlobalModelRegistry().UnregisterClient(id) + s.coreManager.Remove(ctx, id) + if strings.EqualFold(provider, "codex") { + executor.CloseCodexWebsocketSessionsForAuthID(id, "auth_removed") } } diff --git a/sdk/cliproxy/service_stale_state_test.go b/sdk/cliproxy/service_stale_state_test.go index 53849eb34..f5f72e7ec 100644 --- a/sdk/cliproxy/service_stale_state_test.go +++ b/sdk/cliproxy/service_stale_state_test.go @@ -40,37 +40,8 @@ func TestServiceApplyCoreAuthAddOrUpdate_DeleteReAddDoesNotInheritStaleRuntimeSt service.applyCoreAuthRemoval(context.Background(), authID) - disabled, ok := service.coreManager.GetByID(authID) - if !ok || disabled == nil { - t.Fatalf("expected disabled auth after removal") - } - if !disabled.Disabled || disabled.Status != coreauth.StatusDisabled { - t.Fatalf("expected disabled auth after removal, got disabled=%v status=%v", disabled.Disabled, disabled.Status) - } - if disabled.LastRefreshedAt.IsZero() { - t.Fatalf("expected disabled auth to still carry prior LastRefreshedAt for regression setup") - } - if disabled.NextRefreshAfter.IsZero() { - t.Fatalf("expected disabled auth to still carry prior NextRefreshAfter for regression setup") - } - - // Reconcile prunes unsupported model state during registration, so seed the - // disabled snapshot explicitly before exercising delete -> re-add behavior. - disabled.ModelStates = map[string]*coreauth.ModelState{ - modelID: { - Quota: coreauth.QuotaState{BackoffLevel: 7}, - }, - } - if _, err := service.coreManager.Update(context.Background(), disabled); err != nil { - t.Fatalf("seed disabled auth stale ModelStates: %v", err) - } - - disabled, ok = service.coreManager.GetByID(authID) - if !ok || disabled == nil { - t.Fatalf("expected disabled auth after stale state seeding") - } - if len(disabled.ModelStates) == 0 { - t.Fatalf("expected disabled auth to carry seeded ModelStates for regression setup") + if _, ok := service.coreManager.GetByID(authID); ok { + t.Fatalf("expected auth %q to be removed from runtime state", authID) } service.applyCoreAuthAddOrUpdate(context.Background(), &coreauth.Auth{