From d9f8129a327cacbdc881e5c7b764740fe01f3627 Mon Sep 17 00:00:00 2001 From: hkfires <10558748+hkfires@users.noreply.github.com> Date: Thu, 18 Sep 2025 20:58:43 +0800 Subject: [PATCH] fix(client): Add reason to unregistration to skip persistence --- internal/client/gemini-web_client.go | 21 ++++++++++++++++++--- internal/client/qwen_client.go | 17 ++++++++++++++++- internal/cmd/run.go | 7 ++++++- internal/interfaces/client.go | 12 ++++++++++++ internal/watcher/watcher.go | 28 ++++++++++++++++++---------- 5 files changed, 70 insertions(+), 15 deletions(-) diff --git a/internal/client/gemini-web_client.go b/internal/client/gemini-web_client.go index d686156a..0c764f9b 100644 --- a/internal/client/gemini-web_client.go +++ b/internal/client/gemini-web_client.go @@ -61,13 +61,28 @@ type GeminiWebClient struct { modelsRegistered bool } -func (c *GeminiWebClient) UnregisterClient() { +func (c *GeminiWebClient) UnregisterClient() { c.unregisterClient(false) } + +// UnregisterClientWithReason allows the watcher to avoid recreating deleted auth files. +func (c *GeminiWebClient) UnregisterClientWithReason(reason interfaces.UnregisterReason) { + skipPersist := reason == interfaces.UnregisterReasonAuthFileRemoved + c.unregisterClient(skipPersist) +} + +func (c *GeminiWebClient) unregisterClient(skipPersist bool) { if c.cookiePersistCancel != nil { c.cookiePersistCancel() c.cookiePersistCancel = nil } - // Flush cookie snapshot to main token file and remove snapshot - c.flushCookieSnapshotToMain() + if skipPersist { + if c.snapshotManager != nil && c.tokenFilePath != "" { + log.Debugf("skipping Gemini Web snapshot flush because auth file is missing: %s", filepath.Base(c.tokenFilePath)) + util.RemoveCookieSnapshots(c.tokenFilePath) + } + } else { + // Flush cookie snapshot to main token file and remove snapshot + c.flushCookieSnapshotToMain() + } if c.gwc != nil { c.gwc.Close(0) c.gwc = nil diff --git a/internal/client/qwen_client.go b/internal/client/qwen_client.go index b9ce7c24..8e775753 100644 --- a/internal/client/qwen_client.go +++ b/internal/client/qwen_client.go @@ -507,10 +507,25 @@ func (c *QwenClient) SetUnavailable() { } // UnregisterClient flushes cookie snapshot back into the main token file. -func (c *QwenClient) UnregisterClient() { +func (c *QwenClient) UnregisterClient() { c.unregisterClient(false) } + +// UnregisterClientWithReason allows the watcher to skip persistence when the auth file is removed. +func (c *QwenClient) UnregisterClientWithReason(reason interfaces.UnregisterReason) { + skipPersist := reason == interfaces.UnregisterReasonAuthFileRemoved + c.unregisterClient(skipPersist) +} + +func (c *QwenClient) unregisterClient(skipPersist bool) { if c.snapshotManager == nil { return } + if skipPersist { + if c.tokenFilePath != "" { + log.Debugf("skipping Qwen snapshot flush because auth file is missing: %s", filepath.Base(c.tokenFilePath)) + util.RemoveCookieSnapshots(c.tokenFilePath) + } + return + } if err := c.snapshotManager.Flush(); err != nil { log.Errorf("Failed to flush Qwen cookie snapshot to main for %s: %v", filepath.Base(c.tokenFilePath), err) } diff --git a/internal/cmd/run.go b/internal/cmd/run.go index 57743e7e..0c2904b2 100644 --- a/internal/cmd/run.go +++ b/internal/cmd/run.go @@ -346,7 +346,12 @@ func StartService(cfg *config.Config, configPath string) { for _, c := range snapshot { // Persist tokens/cookies then unregister/cleanup per client. _ = c.SaveTokenToFile() - if u, ok := any(c).(interface{ UnregisterClient() }); ok { + switch u := any(c).(type) { + case interface { + UnregisterClientWithReason(interfaces.UnregisterReason) + }: + u.UnregisterClientWithReason(interfaces.UnregisterReasonShutdown) + case interface{ UnregisterClient() }: u.UnregisterClient() } } diff --git a/internal/interfaces/client.go b/internal/interfaces/client.go index ae7a641f..1beed2fa 100644 --- a/internal/interfaces/client.go +++ b/internal/interfaces/client.go @@ -61,3 +61,15 @@ type Client interface { // SetUnavailable sets the client to unavailable. SetUnavailable() } + +// UnregisterReason describes the context for unregistering a client instance. +type UnregisterReason string + +const ( + // UnregisterReasonReload indicates a full reload is replacing the client. + UnregisterReasonReload UnregisterReason = "reload" + // UnregisterReasonShutdown indicates the service is shutting down. + UnregisterReasonShutdown UnregisterReason = "shutdown" + // UnregisterReasonAuthFileRemoved indicates the underlying auth file was deleted. + UnregisterReasonAuthFileRemoved UnregisterReason = "auth-file-removed" +) diff --git a/internal/watcher/watcher.go b/internal/watcher/watcher.go index 71b245f0..34373311 100644 --- a/internal/watcher/watcher.go +++ b/internal/watcher/watcher.go @@ -297,9 +297,7 @@ func (w *Watcher) reloadClients() { // Unregister all old API key clients before creating new ones log.Debugf("unregistering %d old API key clients", oldAPIKeyClientCount) for _, oldClient := range w.apiKeyClients { - if u, ok := oldClient.(interface{ UnregisterClient() }); ok { - u.UnregisterClient() - } + unregisterClientWithReason(oldClient, interfaces.UnregisterReasonReload) } // Create new API key clients based on the new config @@ -313,9 +311,7 @@ func (w *Watcher) reloadClients() { // Unregister all old file-based clients log.Debugf("unregistering %d old file-based clients", oldFileClientCount) for _, oldClient := range w.clients { - if u, ok := any(oldClient).(interface{ UnregisterClient() }); ok { - u.UnregisterClient() - } + unregisterClientWithReason(oldClient, interfaces.UnregisterReasonReload) } // Update client maps @@ -466,10 +462,10 @@ func (w *Watcher) addOrUpdateClient(path string) { // If an old client exists, unregister it first if oldClient, ok := w.clients[path]; ok { - if u, canUnregister := any(oldClient).(interface{ UnregisterClient() }); canUnregister { + if _, canUnregister := any(oldClient).(interface{ UnregisterClient() }); canUnregister { log.Debugf("unregistering old client for updated file: %s", filepath.Base(path)) - u.UnregisterClient() } + unregisterClientWithReason(oldClient, interfaces.UnregisterReasonReload) } // Create new client (reads the file again internally; this is acceptable as the files are small and it keeps the change minimal) @@ -511,10 +507,10 @@ func (w *Watcher) removeClient(path string) { // Unregister client if it exists if oldClient, ok := w.clients[path]; ok { - if u, canUnregister := any(oldClient).(interface{ UnregisterClient() }); canUnregister { + if _, canUnregister := any(oldClient).(interface{ UnregisterClient() }); canUnregister { log.Debugf("unregistering client for removed file: %s", filepath.Base(path)) - u.UnregisterClient() } + unregisterClientWithReason(oldClient, interfaces.UnregisterReasonAuthFileRemoved) delete(w.clients, path) delete(w.lastAuthHashes, path) log.Debugf("removed client for %s", filepath.Base(path)) @@ -550,6 +546,18 @@ func (w *Watcher) buildCombinedClientMap() map[string]interfaces.Client { return combined } +// unregisterClientWithReason attempts to call client-specific unregister hooks with context. +func unregisterClientWithReason(c interfaces.Client, reason interfaces.UnregisterReason) { + switch u := any(c).(type) { + case interface { + UnregisterClientWithReason(interfaces.UnregisterReason) + }: + u.UnregisterClientWithReason(reason) + case interface{ UnregisterClient() }: + u.UnregisterClient() + } +} + // loadFileClients scans the auth directory and creates clients from .json files. func (w *Watcher) loadFileClients(cfg *config.Config) (map[string]interfaces.Client, int) { newClients := make(map[string]interfaces.Client)