fix(client): Add reason to unregistration to skip persistence

This commit is contained in:
hkfires
2025-09-18 20:58:43 +08:00
parent 8f0a345e2a
commit d9f8129a32
5 changed files with 70 additions and 15 deletions

View File

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

View File

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

View File

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

View File

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

View File

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