From f17ec7ffd8c0561cd83b92016a78142c711f9dd5 Mon Sep 17 00:00:00 2001
From: hkfires <10558748+hkfires@users.noreply.github.com>
Date: Thu, 18 Sep 2025 22:43:33 +0800
Subject: [PATCH] fix(client): Prevent overwriting auth file on update

---
 internal/client/gemini-web_client.go | 18 ++++++++++++------
 internal/client/qwen_client.go       | 28 ++++++++++++++++++----------
 internal/interfaces/client.go        |  2 ++
 internal/watcher/watcher.go          |  2 +-
 4 files changed, 33 insertions(+), 17 deletions(-)

diff --git a/internal/client/gemini-web_client.go b/internal/client/gemini-web_client.go
index 0c764f9b..ea002380 100644
--- a/internal/client/gemini-web_client.go
+++ b/internal/client/gemini-web_client.go
@@ -61,25 +61,31 @@ type GeminiWebClient struct {
 	modelsRegistered bool
 }
 
-func (c *GeminiWebClient) UnregisterClient() { c.unregisterClient(false) }
+func (c *GeminiWebClient) UnregisterClient() { c.unregisterClient(interfaces.UnregisterReasonReload) }
 
 // UnregisterClientWithReason allows the watcher to avoid recreating deleted auth files.
 func (c *GeminiWebClient) UnregisterClientWithReason(reason interfaces.UnregisterReason) {
-	skipPersist := reason == interfaces.UnregisterReasonAuthFileRemoved
-	c.unregisterClient(skipPersist)
+	c.unregisterClient(reason)
 }
 
-func (c *GeminiWebClient) unregisterClient(skipPersist bool) {
+func (c *GeminiWebClient) unregisterClient(reason interfaces.UnregisterReason) {
 	if c.cookiePersistCancel != nil {
 		c.cookiePersistCancel()
 		c.cookiePersistCancel = nil
 	}
-	if skipPersist {
+	switch reason {
+	case interfaces.UnregisterReasonAuthFileRemoved:
 		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 {
+	case interfaces.UnregisterReasonAuthFileUpdated:
+		if c.snapshotManager != nil {
+			if err := c.snapshotManager.Persist(); err != nil {
+				log.Errorf("Failed to persist Gemini Web cookies for %s: %v", filepath.Base(c.tokenFilePath), err)
+			}
+		}
+	default:
 		// Flush cookie snapshot to main token file and remove snapshot
 		c.flushCookieSnapshotToMain()
 	}
diff --git a/internal/client/qwen_client.go b/internal/client/qwen_client.go
index 8e775753..a6c4e156 100644
--- a/internal/client/qwen_client.go
+++ b/internal/client/qwen_client.go
@@ -507,26 +507,34 @@ func (c *QwenClient) SetUnavailable() {
 }
 
 // UnregisterClient flushes cookie snapshot back into the main token file.
-func (c *QwenClient) UnregisterClient() { c.unregisterClient(false) }
+func (c *QwenClient) UnregisterClient() { c.unregisterClient(interfaces.UnregisterReasonReload) }
 
-// UnregisterClientWithReason allows the watcher to skip persistence when the auth file is removed.
+// UnregisterClientWithReason allows the watcher to adjust persistence behaviour.
 func (c *QwenClient) UnregisterClientWithReason(reason interfaces.UnregisterReason) {
-	skipPersist := reason == interfaces.UnregisterReasonAuthFileRemoved
-	c.unregisterClient(skipPersist)
+	c.unregisterClient(reason)
 }
 
-func (c *QwenClient) unregisterClient(skipPersist bool) {
+func (c *QwenClient) unregisterClient(reason interfaces.UnregisterReason) {
 	if c.snapshotManager == nil {
 		return
 	}
-	if skipPersist {
+	switch reason {
+	case interfaces.UnregisterReasonAuthFileRemoved:
 		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)
+	case interfaces.UnregisterReasonAuthFileUpdated:
+		if err := c.snapshotManager.Persist(); err != nil {
+			log.Errorf("Failed to persist Qwen cookies for %s: %v", filepath.Base(c.tokenFilePath), err)
+		}
+	case interfaces.UnregisterReasonShutdown, interfaces.UnregisterReasonReload:
+		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)
+		}
+	default:
+		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/interfaces/client.go b/internal/interfaces/client.go
index 1beed2fa..2600f6b1 100644
--- a/internal/interfaces/client.go
+++ b/internal/interfaces/client.go
@@ -72,4 +72,6 @@ const (
 	UnregisterReasonShutdown UnregisterReason = "shutdown"
 	// UnregisterReasonAuthFileRemoved indicates the underlying auth file was deleted.
 	UnregisterReasonAuthFileRemoved UnregisterReason = "auth-file-removed"
+	// UnregisterReasonAuthFileUpdated indicates the auth file content was modified.
+	UnregisterReasonAuthFileUpdated UnregisterReason = "auth-file-updated"
 )
diff --git a/internal/watcher/watcher.go b/internal/watcher/watcher.go
index 34373311..6cd5e54a 100644
--- a/internal/watcher/watcher.go
+++ b/internal/watcher/watcher.go
@@ -465,7 +465,7 @@ func (w *Watcher) addOrUpdateClient(path string) {
 		if _, canUnregister := any(oldClient).(interface{ UnregisterClient() }); canUnregister {
 			log.Debugf("unregistering old client for updated file: %s", filepath.Base(path))
 		}
-		unregisterClientWithReason(oldClient, interfaces.UnregisterReasonReload)
+		unregisterClientWithReason(oldClient, interfaces.UnregisterReasonAuthFileUpdated)
 	}
 
 	// Create new client (reads the file again internally; this is acceptable as the files are small and it keeps the change minimal)