Skip to content

Commit 55440f0

Browse files
luispaterluispater
committed
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: router-for-me#3690
1 parent 9c02454 commit 55440f0

7 files changed

Lines changed: 268 additions & 70 deletions

File tree

‎internal/api/handlers/management/auth_files.go‎

Lines changed: 7 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -770,7 +770,7 @@ func (h *Handler) DeleteAuthFile(c *gin.Context) {
770770
return
771771
}
772772
deleted++
773-
h.disableAuth(ctx, full)
773+
h.removeAuth(ctx, full)
774774
}
775775
}
776776
c.JSON(200, gin.H{"status": "ok", "deleted": deleted})
@@ -976,9 +976,9 @@ func (h *Handler) deleteAuthFileByName(ctx context.Context, name string) (string
976976
return filepath.Base(name), http.StatusInternalServerError, errDeleteRecord
977977
}
978978
if targetID != "" {
979-
h.disableAuth(ctx, targetID)
979+
h.removeAuth(ctx, targetID)
980980
} else {
981-
h.disableAuth(ctx, targetPath)
981+
h.removeAuth(ctx, targetPath)
982982
}
983983
return filepath.Base(name), http.StatusOK, nil
984984
}
@@ -1558,33 +1558,23 @@ func syncAuthFileDisabledState(auth *coreauth.Auth) {
15581558
auth.StatusMessage = ""
15591559
}
15601560

1561-
func (h *Handler) disableAuth(ctx context.Context, id string) {
1561+
func (h *Handler) removeAuth(ctx context.Context, id string) {
15621562
if h == nil || h.authManager == nil {
15631563
return
15641564
}
15651565
id = strings.TrimSpace(id)
15661566
if id == "" {
15671567
return
15681568
}
1569-
if auth, ok := h.authManager.GetByID(id); ok {
1570-
auth.Disabled = true
1571-
auth.Status = coreauth.StatusDisabled
1572-
auth.StatusMessage = "removed via management API"
1573-
auth.UpdatedAt = time.Now()
1574-
_, _ = h.authManager.Update(ctx, auth)
1569+
if _, ok := h.authManager.GetByID(id); ok {
1570+
h.authManager.Remove(ctx, id)
15751571
return
15761572
}
15771573
authID := h.authIDForPath(id)
15781574
if authID == "" {
15791575
return
15801576
}
1581-
if auth, ok := h.authManager.GetByID(authID); ok {
1582-
auth.Disabled = true
1583-
auth.Status = coreauth.StatusDisabled
1584-
auth.StatusMessage = "removed via management API"
1585-
auth.UpdatedAt = time.Now()
1586-
_, _ = h.authManager.Update(ctx, auth)
1587-
}
1577+
h.authManager.Remove(ctx, authID)
15881578
}
15891579

15901580
func (h *Handler) deleteTokenRecord(ctx context.Context, path string) error {

‎internal/api/handlers/management/auth_files_delete_test.go‎

Lines changed: 46 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -127,3 +127,49 @@ func TestDeleteAuthFile_FallbackToAuthDirPath(t *testing.T) {
127127
t.Fatalf("expected auth file to be removed from auth dir, stat err: %v", errStat)
128128
}
129129
}
130+
131+
func TestDeleteAuthFile_RemovesRuntimeAuth(t *testing.T) {
132+
t.Setenv("MANAGEMENT_PASSWORD", "")
133+
gin.SetMode(gin.TestMode)
134+
135+
authDir := t.TempDir()
136+
fileName := "runtime-remove-user.json"
137+
filePath := filepath.Join(authDir, fileName)
138+
if errWrite := os.WriteFile(filePath, []byte(`{"type":"codex","email":"runtime@example.com"}`), 0o600); errWrite != nil {
139+
t.Fatalf("failed to write auth file: %v", errWrite)
140+
}
141+
142+
manager := coreauth.NewManager(nil, nil, nil)
143+
record := &coreauth.Auth{
144+
ID: "runtime-remove-auth",
145+
FileName: fileName,
146+
Provider: "codex",
147+
Status: coreauth.StatusActive,
148+
Attributes: map[string]string{
149+
"path": filePath,
150+
},
151+
Metadata: map[string]any{
152+
"type": "codex",
153+
"email": "runtime@example.com",
154+
},
155+
}
156+
if _, errRegister := manager.Register(context.Background(), record); errRegister != nil {
157+
t.Fatalf("failed to register auth record: %v", errRegister)
158+
}
159+
160+
h := NewHandlerWithoutConfigFilePath(&config.Config{AuthDir: authDir}, manager)
161+
h.tokenStore = &memoryAuthStore{}
162+
163+
deleteRec := httptest.NewRecorder()
164+
deleteCtx, _ := gin.CreateTestContext(deleteRec)
165+
deleteReq := httptest.NewRequest(http.MethodDelete, "/v0/management/auth-files?name="+url.QueryEscape(fileName), nil)
166+
deleteCtx.Request = deleteReq
167+
h.DeleteAuthFile(deleteCtx)
168+
169+
if deleteRec.Code != http.StatusOK {
170+
t.Fatalf("expected delete status %d, got %d with body %s", http.StatusOK, deleteRec.Code, deleteRec.Body.String())
171+
}
172+
if _, ok := manager.GetByID(record.ID); ok {
173+
t.Fatalf("expected runtime auth %q to be removed", record.ID)
174+
}
175+
}

‎sdk/cliproxy/auth/conductor.go‎

Lines changed: 87 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -1164,18 +1164,21 @@ func (m *Manager) Update(ctx context.Context, auth *Auth) (*Auth, error) {
11641164
return nil, nil
11651165
}
11661166
m.mu.Lock()
1167-
if existing, ok := m.auths[auth.ID]; ok && existing != nil {
1168-
if !auth.indexAssigned && auth.Index == "" {
1169-
auth.Index = existing.Index
1170-
auth.indexAssigned = existing.indexAssigned
1171-
}
1172-
auth.Success = existing.Success
1173-
auth.Failed = existing.Failed
1174-
auth.recentRequests = existing.recentRequests
1175-
if !existing.Disabled && existing.Status != StatusDisabled && !auth.Disabled && auth.Status != StatusDisabled {
1176-
if len(auth.ModelStates) == 0 && len(existing.ModelStates) > 0 {
1177-
auth.ModelStates = existing.ModelStates
1178-
}
1167+
existing, ok := m.auths[auth.ID]
1168+
if !ok || existing == nil {
1169+
m.mu.Unlock()
1170+
return nil, nil
1171+
}
1172+
if !auth.indexAssigned && auth.Index == "" {
1173+
auth.Index = existing.Index
1174+
auth.indexAssigned = existing.indexAssigned
1175+
}
1176+
auth.Success = existing.Success
1177+
auth.Failed = existing.Failed
1178+
auth.recentRequests = existing.recentRequests
1179+
if !existing.Disabled && existing.Status != StatusDisabled && !auth.Disabled && auth.Status != StatusDisabled {
1180+
if len(auth.ModelStates) == 0 && len(existing.ModelStates) > 0 {
1181+
auth.ModelStates = existing.ModelStates
11791182
}
11801183
}
11811184
auth.EnsureIndex()
@@ -1192,6 +1195,65 @@ func (m *Manager) Update(ctx context.Context, auth *Auth) (*Auth, error) {
11921195
return auth.Clone(), nil
11931196
}
11941197

1198+
// Remove deletes an auth from runtime state without persisting.
1199+
// Disk and token-store deletion must be handled by the caller.
1200+
func (m *Manager) Remove(ctx context.Context, id string) {
1201+
if m == nil {
1202+
return
1203+
}
1204+
id = strings.TrimSpace(id)
1205+
if id == "" {
1206+
return
1207+
}
1208+
_ = ctx
1209+
1210+
m.mu.Lock()
1211+
existing := m.auths[id]
1212+
if existing == nil {
1213+
m.mu.Unlock()
1214+
return
1215+
}
1216+
provider := strings.TrimSpace(existing.Provider)
1217+
delete(m.auths, id)
1218+
if m.modelPoolOffsets != nil {
1219+
delete(m.modelPoolOffsets, id)
1220+
}
1221+
for sessionID, sessionAuths := range m.homeRuntimeAuths {
1222+
if sessionAuths == nil {
1223+
continue
1224+
}
1225+
delete(sessionAuths, id)
1226+
if len(sessionAuths) == 0 {
1227+
delete(m.homeRuntimeAuths, sessionID)
1228+
}
1229+
}
1230+
m.mu.Unlock()
1231+
1232+
m.rebuildAPIKeyModelAliasFromRuntimeConfig()
1233+
if m.scheduler != nil {
1234+
m.scheduler.removeAuth(id)
1235+
}
1236+
m.queueRefreshUnschedule(id)
1237+
m.invalidateSessionAffinity(id)
1238+
1239+
if provider != "" {
1240+
if exec, ok := m.Executor(provider); ok && exec != nil {
1241+
if closer, okCloser := exec.(ExecutionSessionCloser); okCloser {
1242+
closer.CloseExecutionSession(CloseAllExecutionSessionsID)
1243+
}
1244+
}
1245+
}
1246+
}
1247+
1248+
func (m *Manager) invalidateSessionAffinity(authID string) {
1249+
if m == nil || authID == "" {
1250+
return
1251+
}
1252+
if invalidator, ok := m.selector.(interface{ InvalidateAuth(string) }); ok && invalidator != nil {
1253+
invalidator.InvalidateAuth(authID)
1254+
}
1255+
}
1256+
11951257
// Load resets manager state from the backing store.
11961258
func (m *Manager) Load(ctx context.Context) error {
11971259
m.mu.Lock()
@@ -4041,6 +4103,19 @@ func (m *Manager) queueRefreshReschedule(authID string) {
40414103
loop.queueReschedule(authID)
40424104
}
40434105

4106+
func (m *Manager) queueRefreshUnschedule(authID string) {
4107+
if m == nil || authID == "" {
4108+
return
4109+
}
4110+
m.mu.RLock()
4111+
loop := m.refreshLoop
4112+
m.mu.RUnlock()
4113+
if loop == nil {
4114+
return
4115+
}
4116+
loop.remove(authID)
4117+
}
4118+
40444119
func (m *Manager) shouldRefresh(a *Auth, now time.Time) bool {
40454120
if a == nil {
40464121
return false

‎sdk/cliproxy/auth/conductor_remove_test.go‎

Lines changed: 111 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,111 @@
1+
package auth
2+
3+
import (
4+
"context"
5+
"testing"
6+
"time"
7+
)
8+
9+
func TestManager_Remove_DeletesRuntimeAuth(t *testing.T) {
10+
manager := NewManager(nil, nil, nil)
11+
ctx := context.Background()
12+
13+
auth := &Auth{
14+
ID: "remove-runtime-auth",
15+
Provider: "claude",
16+
Status: StatusActive,
17+
Metadata: map[string]any{"email": "x@example.com"},
18+
}
19+
if _, errRegister := manager.Register(ctx, auth); errRegister != nil {
20+
t.Fatalf("register auth: %v", errRegister)
21+
}
22+
23+
manager.Remove(ctx, auth.ID)
24+
25+
if _, ok := manager.GetByID(auth.ID); ok {
26+
t.Fatalf("expected auth %q to be removed", auth.ID)
27+
}
28+
}
29+
30+
func TestManager_Update_MissingAuthIsNoOp(t *testing.T) {
31+
manager := NewManager(nil, nil, nil)
32+
ctx := context.Background()
33+
34+
auth := &Auth{
35+
ID: "missing-update-auth",
36+
Provider: "claude",
37+
Status: StatusActive,
38+
}
39+
if _, errRegister := manager.Register(ctx, auth); errRegister != nil {
40+
t.Fatalf("register auth: %v", errRegister)
41+
}
42+
manager.Remove(ctx, auth.ID)
43+
44+
updated, errUpdate := manager.Update(ctx, &Auth{
45+
ID: auth.ID,
46+
Provider: "claude",
47+
Status: StatusDisabled,
48+
Disabled: true,
49+
})
50+
if errUpdate != nil {
51+
t.Fatalf("update removed auth: %v", errUpdate)
52+
}
53+
if updated != nil {
54+
t.Fatalf("expected update on removed auth to be no-op, got %#v", updated)
55+
}
56+
if _, ok := manager.GetByID(auth.ID); ok {
57+
t.Fatalf("expected removed auth to stay absent after late update")
58+
}
59+
}
60+
61+
func TestManager_Remove_UnschedulesAutoRefresh(t *testing.T) {
62+
ctx := context.Background()
63+
64+
manager := NewManager(nil, nil, nil)
65+
loop := newAuthAutoRefreshLoop(manager, time.Second, 1)
66+
manager.mu.Lock()
67+
manager.refreshLoop = loop
68+
manager.mu.Unlock()
69+
70+
lead := 10 * time.Minute
71+
setRefreshLeadFactory(t, "provider-lead-expiry", func() *time.Duration {
72+
d := lead
73+
return &d
74+
})
75+
76+
auth := &Auth{
77+
ID: "remove-refresh-auth",
78+
Provider: "provider-lead-expiry",
79+
Metadata: map[string]any{
80+
"email": "x@example.com",
81+
"expires_at": time.Now().Add(time.Hour).Format(time.RFC3339),
82+
},
83+
}
84+
if _, errRegister := manager.Register(ctx, auth); errRegister != nil {
85+
t.Fatalf("register auth: %v", errRegister)
86+
}
87+
88+
now := time.Now()
89+
if _, ok := nextRefreshCheckAt(now, auth, time.Second); !ok {
90+
t.Fatalf("expected auth to be scheduled before removal")
91+
}
92+
loop.applyDirty(now)
93+
loop.mu.Lock()
94+
if _, ok := loop.index[auth.ID]; !ok {
95+
loop.mu.Unlock()
96+
t.Fatalf("expected auth %q to be present in auto-refresh index before removal", auth.ID)
97+
}
98+
loop.mu.Unlock()
99+
100+
manager.Remove(ctx, auth.ID)
101+
102+
if _, ok := manager.GetByID(auth.ID); ok {
103+
t.Fatalf("expected auth to be removed")
104+
}
105+
loop.mu.Lock()
106+
if _, ok := loop.index[auth.ID]; ok {
107+
loop.mu.Unlock()
108+
t.Fatalf("expected auth %q to be removed from auto-refresh index", auth.ID)
109+
}
110+
loop.mu.Unlock()
111+
}

‎sdk/cliproxy/auth/persist_policy_test.go‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,13 @@ func TestWithSkipPersist_DisablesUpdatePersistence(t *testing.T) {
2828
Metadata: map[string]any{"type": "antigravity"},
2929
}
3030

31+
if _, err := mgr.Register(WithSkipPersist(context.Background()), auth); err != nil {
32+
t.Fatalf("Register(skipPersist) returned error: %v", err)
33+
}
34+
if got := store.saveCount.Load(); got != 0 {
35+
t.Fatalf("expected 0 Save calls, got %d", got)
36+
}
37+
3138
if _, err := mgr.Update(context.Background(), auth); err != nil {
3239
t.Fatalf("Update returned error: %v", err)
3340
}

‎sdk/cliproxy/service.go‎

Lines changed: 8 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -339,17 +339,15 @@ func (s *Service) applyCoreAuthRemoval(ctx context.Context, id string) {
339339
if s.coreManager == nil {
340340
return
341341
}
342-
GlobalModelRegistry().UnregisterClient(id)
342+
id = strings.TrimSpace(id)
343+
var provider string
343344
if existing, ok := s.coreManager.GetByID(id); ok && existing != nil {
344-
existing.Disabled = true
345-
existing.Status = coreauth.StatusDisabled
346-
if _, err := s.coreManager.Update(ctx, existing); err != nil {
347-
log.Errorf("failed to disable auth %s: %v", id, err)
348-
}
349-
if strings.EqualFold(strings.TrimSpace(existing.Provider), "codex") {
350-
executor.CloseCodexWebsocketSessionsForAuthID(existing.ID, "auth_removed")
351-
s.ensureExecutorsForAuth(existing)
352-
}
345+
provider = strings.TrimSpace(existing.Provider)
346+
}
347+
GlobalModelRegistry().UnregisterClient(id)
348+
s.coreManager.Remove(ctx, id)
349+
if strings.EqualFold(provider, "codex") {
350+
executor.CloseCodexWebsocketSessionsForAuthID(id, "auth_removed")
353351
}
354352
}
355353

0 commit comments

Comments
 (0)