From 76322049668aeabea7a817d860e33c897ee6e2cc Mon Sep 17 00:00:00 2001 From: hkfires <10558748+hkfires@users.noreply.github.com> Date: Thu, 18 Sep 2025 19:04:41 +0800 Subject: [PATCH 1/9] refactor(cookie): Extract cookie snapshot logic to util package The logic for managing cookie persistence files was previously implemented directly within the `gemini-web` client's persistence layer. This approach was not reusable and led to duplicated helper functions. This commit refactors the cookie persistence mechanism by: - Renaming the concept from "sidecar" to "snapshot" for clarity. - Extracting file I/O and path manipulation logic into a new, generic `internal/util/cookie_snapshot.go` file. - Creating reusable utility functions: `WriteCookieSnapshot`, `TryReadCookieSnapshotInto`, and `RemoveCookieSnapshot`. - Updating the `gemini-web` persistence code to use these new centralized utility functions. This change improves code organization, reduces duplication, and makes the cookie snapshot functionality easier to maintain and potentially reuse across other clients. --- internal/client/gemini-web/persistence.go | 101 +++++++--------------- internal/client/gemini-web_client.go | 32 +++---- internal/client/qwen_client.go | 96 +++++++++++++++++++- internal/util/cookie_snapshot.go | 89 +++++++++++++++++++ 4 files changed, 225 insertions(+), 93 deletions(-) create mode 100644 internal/util/cookie_snapshot.go diff --git a/internal/client/gemini-web/persistence.go b/internal/client/gemini-web/persistence.go index e9631da7..114d89ad 100644 --- a/internal/client/gemini-web/persistence.go +++ b/internal/client/gemini-web/persistence.go @@ -11,6 +11,7 @@ import ( "time" "github.com/luispater/CLIProxyAPI/v5/internal/auth/gemini" + "github.com/luispater/CLIProxyAPI/v5/internal/util" ) // StoredMessage represents a single message in a conversation record. @@ -268,42 +269,17 @@ func FindReusableSessionIn(items map[string]ConversationRecord, index map[string return nil, nil } -// CookiesSidecarPath derives the sidecar cookie file path from the main token JSON path. -func CookiesSidecarPath(mainPath string) string { - if strings.HasSuffix(mainPath, ".json") { - return strings.TrimSuffix(mainPath, ".json") + ".cookies" - } - return mainPath + ".cookies" -} - -// FileExists reports whether the given path exists and is a regular file. -func FileExists(path string) bool { - if path == "" { - return false - } - if st, err := os.Stat(path); err == nil && !st.IsDir() { - return true - } - return false -} - -// ApplyCookiesSidecarToTokenStorage loads cookies from sidecar into the provided token storage. -// Returns true when a sidecar was found and applied. -func ApplyCookiesSidecarToTokenStorage(tokenFilePath string, ts *gemini.GeminiWebTokenStorage) (bool, error) { +// ApplyCookieSnapshotToTokenStorage loads cookies from cookie snapshot into the provided token storage. +// Returns true when a snapshot was found and applied. +func ApplyCookieSnapshotToTokenStorage(tokenFilePath string, ts *gemini.GeminiWebTokenStorage) (bool, error) { if ts == nil { return false, nil } - side := CookiesSidecarPath(tokenFilePath) - if !FileExists(side) { - return false, nil - } - data, err := os.ReadFile(side) - if err != nil || len(data) == 0 { - return false, err - } var latest gemini.GeminiWebTokenStorage - if err := json.Unmarshal(data, &latest); err != nil { + if ok, err := util.TryReadCookieSnapshotInto(tokenFilePath, &latest); err != nil { return false, err + } else if !ok { + return false, nil } if latest.Secure1PSID != "" { ts.Secure1PSID = latest.Secure1PSID @@ -314,10 +290,9 @@ func ApplyCookiesSidecarToTokenStorage(tokenFilePath string, ts *gemini.GeminiWe return true, nil } -// SaveCookiesSidecar writes the current cookies into a sidecar file next to the token file. +// SaveCookieSnapshot writes the current cookies into a snapshot file next to the token file. // This keeps the main token JSON stable until an orderly flush. -func SaveCookiesSidecar(tokenFilePath string, cookies map[string]string) error { - side := CookiesSidecarPath(tokenFilePath) +func SaveCookieSnapshot(tokenFilePath string, cookies map[string]string) error { ts := &gemini.GeminiWebTokenStorage{Type: "gemini-web"} if v := cookies["__Secure-1PSID"]; v != "" { ts.Secure1PSID = v @@ -325,51 +300,35 @@ func SaveCookiesSidecar(tokenFilePath string, cookies map[string]string) error { if v := cookies["__Secure-1PSIDTS"]; v != "" { ts.Secure1PSIDTS = v } - if err := os.MkdirAll(filepath.Dir(side), 0o700); err != nil { - return err - } - return ts.SaveTokenToFile(side) + return util.WriteCookieSnapshot(tokenFilePath, ts) } -// FlushCookiesSidecarToMain merges the sidecar cookies into the main token file and removes the sidecar. -// If sidecar is missing, it will combine the provided base token storage with the latest cookies. -func FlushCookiesSidecarToMain(tokenFilePath string, cookies map[string]string, base *gemini.GeminiWebTokenStorage) error { +// FlushCookieSnapshotToMain merges the cookie snapshot into the main token file and removes the snapshot. +// If snapshot is missing, it will combine the provided base token storage with the latest cookies. +func FlushCookieSnapshotToMain(tokenFilePath string, cookies map[string]string, base *gemini.GeminiWebTokenStorage) error { if tokenFilePath == "" { return nil } - side := CookiesSidecarPath(tokenFilePath) - var merged gemini.GeminiWebTokenStorage - var fromSidecar bool - if FileExists(side) { - if data, err := os.ReadFile(side); err == nil && len(data) > 0 { - if err2 := json.Unmarshal(data, &merged); err2 == nil { - fromSidecar = true - } - } - } - if !fromSidecar { - if base != nil { - merged = *base - } - if v := cookies["__Secure-1PSID"]; v != "" { - merged.Secure1PSID = v - } - if v := cookies["__Secure-1PSIDTS"]; v != "" { - merged.Secure1PSIDTS = v - } - } + var merged gemini.GeminiWebTokenStorage + var fromSnapshot bool + if ok, _ := util.TryReadCookieSnapshotInto(tokenFilePath, &merged); ok { + fromSnapshot = true + } + if !fromSnapshot { + if base != nil { + merged = *base + } + if v := cookies["__Secure-1PSID"]; v != "" { + merged.Secure1PSID = v + } + if v := cookies["__Secure-1PSIDTS"]; v != "" { + merged.Secure1PSIDTS = v + } + } merged.Type = "gemini-web" - if err := os.MkdirAll(filepath.Dir(tokenFilePath), 0o700); err != nil { - return err - } if err := merged.SaveTokenToFile(tokenFilePath); err != nil { return err } - if FileExists(side) { - _ = os.Remove(side) - } + util.RemoveCookieSnapshots(tokenFilePath) return nil } - -// IsSelfPersistedToken compares provided token storage with current cookies. -// Removed: IsSelfPersistedToken (client-side no longer needs self-originated write detection) diff --git a/internal/client/gemini-web_client.go b/internal/client/gemini-web_client.go index 308473a9..3f7c9960 100644 --- a/internal/client/gemini-web_client.go +++ b/internal/client/gemini-web_client.go @@ -65,8 +65,8 @@ func (c *GeminiWebClient) UnregisterClient() { c.cookiePersistCancel() c.cookiePersistCancel = nil } - // Flush sidecar cookies to main token file and remove sidecar - c.flushCookiesSidecarToMain() + // Flush cookie snapshot to main token file and remove snapshot + c.flushCookieSnapshotToMain() if c.gwc != nil { c.gwc.Close(0) c.gwc = nil @@ -115,9 +115,9 @@ func NewGeminiWebClient(cfg *config.Config, ts *gemini.GeminiWebTokenStorage, to client.InitializeModelRegistry(clientID) - // Prefer sidecar cookies at startup if present - if ok, err := geminiWeb.ApplyCookiesSidecarToTokenStorage(tokenFilePath, ts); err == nil && ok { - log.Debugf("Loaded Gemini Web cookies from sidecar: %s", filepath.Base(geminiWeb.CookiesSidecarPath(tokenFilePath))) + // Prefer cookie snapshot at startup if present + if ok, err := geminiWeb.ApplyCookieSnapshotToTokenStorage(tokenFilePath, ts); err == nil && ok { + log.Debugf("Loaded Gemini Web cookie snapshot: %s", filepath.Base(util.CookieSnapshotPath(tokenFilePath))) } client.gwc = geminiWeb.NewGeminiClient(ts.Secure1PSID, ts.Secure1PSIDTS, cfg.ProxyURL, geminiWeb.WithAccountLabel(strings.TrimSuffix(filepath.Base(tokenFilePath), ".json"))) @@ -783,7 +783,7 @@ func (c *GeminiWebClient) SendRawTokenCount(ctx context.Context, modelName strin return []byte(fmt.Sprintf(`{"totalTokens":%d}`, est)), nil } -// SaveTokenToFile persists current cookies to a sidecar file via gemini-web helpers. +// SaveTokenToFile persists current cookies to a cookie snapshot via gemini-web helpers. func (c *GeminiWebClient) SaveTokenToFile() error { ts := c.tokenStorage.(*gemini.GeminiWebTokenStorage) if c.gwc != nil && c.gwc.Cookies != nil { @@ -794,11 +794,11 @@ func (c *GeminiWebClient) SaveTokenToFile() error { ts.Secure1PSIDTS = v } } - log.Debugf("Saving Gemini Web cookies sidecar to %s", filepath.Base(geminiWeb.CookiesSidecarPath(c.tokenFilePath))) - return geminiWeb.SaveCookiesSidecar(c.tokenFilePath, c.gwc.Cookies) + log.Debugf("Saving Gemini Web cookie snapshot to %s", filepath.Base(util.CookieSnapshotPath(c.tokenFilePath))) + return geminiWeb.SaveCookieSnapshot(c.tokenFilePath, c.gwc.Cookies) } -// startCookiePersist periodically writes refreshed cookies into the sidecar file. +// startCookiePersist periodically writes refreshed cookies into the cookie snapshot file. func (c *GeminiWebClient) startCookiePersist() { if c.gwc == nil { return @@ -827,7 +827,7 @@ func (c *GeminiWebClient) startCookiePersist() { case <-ticker.C: if c.gwc != nil && c.gwc.Cookies != nil { if err := c.SaveTokenToFile(); err != nil { - log.Errorf("Failed to persist cookies sidecar for %s: %v", c.GetEmail(), err) + log.Errorf("Failed to persist cookie snapshot for %s: %v", c.GetEmail(), err) } } } @@ -1020,22 +1020,18 @@ func (c *GeminiWebClient) backgroundInitRetry() { } } -// IsSelfPersistedToken compares provided token storage with currently active cookies. -// Removed: IsSelfPersistedToken (no longer needed with sidecar-only periodic persistence) - -// flushCookiesSidecarToMain merges sidecar cookies into the main token file. -func (c *GeminiWebClient) flushCookiesSidecarToMain() { +// flushCookieSnapshotToMain merges snapshot cookies into the main token file. +func (c *GeminiWebClient) flushCookieSnapshotToMain() { if c.tokenFilePath == "" { return } base := c.tokenStorage.(*gemini.GeminiWebTokenStorage) - if err := geminiWeb.FlushCookiesSidecarToMain(c.tokenFilePath, c.gwc.Cookies, base); err != nil { - log.Errorf("Failed to flush cookies sidecar to main for %s: %v", filepath.Base(c.tokenFilePath), err) + if err := geminiWeb.FlushCookieSnapshotToMain(c.tokenFilePath, c.gwc.Cookies, base); err != nil { + log.Errorf("Failed to flush cookie snapshot to main for %s: %v", filepath.Base(c.tokenFilePath), err) } } // findReusableSession and storeConversationJSON live here as client bridges; hashing/records in gemini-web - func (c *GeminiWebClient) getConfiguredGem() *geminiWeb.Gem { if c.cfg.GeminiWeb.CodeMode { return &geminiWeb.Gem{ID: "coding-partner", Name: "Coding partner", Predefined: true} diff --git a/internal/client/qwen_client.go b/internal/client/qwen_client.go index 59b2670b..d09da8af 100644 --- a/internal/client/qwen_client.go +++ b/internal/client/qwen_client.go @@ -37,7 +37,8 @@ const ( // QwenClient implements the Client interface for OpenAI API type QwenClient struct { ClientBase - qwenAuth *qwen.QwenAuth + qwenAuth *qwen.QwenAuth + tokenFilePath string } // NewQwenClient creates a new OpenAI client instance @@ -48,7 +49,7 @@ type QwenClient struct { // // Returns: // - *QwenClient: A new Qwen client instance. -func NewQwenClient(cfg *config.Config, ts *qwen.QwenTokenStorage) *QwenClient { +func NewQwenClient(cfg *config.Config, ts *qwen.QwenTokenStorage, tokenFilePath ...string) *QwenClient { httpClient := util.SetProxy(cfg, &http.Client{}) // Generate unique client ID @@ -66,6 +67,19 @@ func NewQwenClient(cfg *config.Config, ts *qwen.QwenTokenStorage) *QwenClient { qwenAuth: qwen.NewQwenAuth(cfg), } + // If created with a known token file path, record it. + if len(tokenFilePath) > 0 && tokenFilePath[0] != "" { + client.tokenFilePath = tokenFilePath[0] + } + + // If no explicit path provided but email exists, derive the canonical path. + if client.tokenFilePath == "" && ts != nil && ts.Email != "" { + client.tokenFilePath = filepath.Join(cfg.AuthDir, fmt.Sprintf("qwen-%s.json", ts.Email)) + } + + // Prefer cookie snapshot at startup if present. + _ = client.applyCookieSnapshot() + // Initialize model registry and register Qwen models client.InitializeModelRegistry(clientID) client.RegisterModels("qwen", registry.GetQwenModels()) @@ -275,7 +289,13 @@ func (c *QwenClient) SendRawTokenCount(_ context.Context, _ string, _ []byte, _ // Returns: // - error: An error if the save operation fails, nil otherwise. func (c *QwenClient) SaveTokenToFile() error { - fileName := filepath.Join(c.cfg.AuthDir, fmt.Sprintf("qwen-%s.json", c.tokenStorage.(*qwen.QwenTokenStorage).Email)) + ts := c.tokenStorage.(*qwen.QwenTokenStorage) + // When the client was created from an auth file, persist via cookie snapshot + if c.tokenFilePath != "" { + return c.saveCookieSnapshot(ts) + } + // Initial bootstrap (e.g., during OAuth flow) writes the main token file + fileName := filepath.Join(c.cfg.AuthDir, fmt.Sprintf("qwen-%s.json", ts.Email)) return c.tokenStorage.SaveTokenToFile(fileName) } @@ -347,7 +367,7 @@ func (c *QwenClient) APIRequest(ctx context.Context, modelName, endpoint string, } var url string - if c.tokenStorage.(*qwen.QwenTokenStorage).ResourceURL == "" { + if c.tokenStorage.(*qwen.QwenTokenStorage).ResourceURL != "" { url = fmt.Sprintf("https://%s/v1%s", c.tokenStorage.(*qwen.QwenTokenStorage).ResourceURL, endpoint) } else { url = fmt.Sprintf("%s%s", qwenEndpoint, endpoint) @@ -458,3 +478,71 @@ func (c *QwenClient) IsAvailable() bool { func (c *QwenClient) SetUnavailable() { c.isAvailable = false } + +// UnregisterClient flushes cookie snapshot back into the main token file. +func (c *QwenClient) UnregisterClient() { + if c.tokenFilePath == "" { + return + } + base := c.tokenStorage.(*qwen.QwenTokenStorage) + if err := c.flushCookieSnapshotToMain(base); err != nil { + log.Errorf("Failed to flush Qwen cookie snapshot to main for %s: %v", filepath.Base(c.tokenFilePath), err) + } +} + +// applyCookieSnapshot loads latest tokens from cookie snapshot if present. +func (c *QwenClient) applyCookieSnapshot() error { + if c.tokenFilePath == "" { + return nil + } + ts := c.tokenStorage.(*qwen.QwenTokenStorage) + var latest qwen.QwenTokenStorage + if ok, err := util.TryReadCookieSnapshotInto(c.tokenFilePath, &latest); err != nil { + return err + } else if !ok { + return nil + } + if latest.AccessToken != "" { + ts.AccessToken = latest.AccessToken + } + if latest.RefreshToken != "" { + ts.RefreshToken = latest.RefreshToken + } + if latest.ResourceURL != "" { + ts.ResourceURL = latest.ResourceURL + } + if latest.Expire != "" { + ts.Expire = latest.Expire + } + return nil +} + +// saveCookieSnapshot writes the token storage into the snapshot file next to the token file. +func (c *QwenClient) saveCookieSnapshot(ts *qwen.QwenTokenStorage) error { + if c.tokenFilePath == "" || ts == nil { + return nil + } + ts.Type = "qwen" + return util.WriteCookieSnapshot(c.tokenFilePath, ts) +} + +// flushCookieSnapshotToMain merges snapshot tokens into the main token file and removes the snapshot. +func (c *QwenClient) flushCookieSnapshotToMain(base *qwen.QwenTokenStorage) error { + if c.tokenFilePath == "" { + return nil + } + var merged qwen.QwenTokenStorage + var fromSnapshot bool + if ok, _ := util.TryReadCookieSnapshotInto(c.tokenFilePath, &merged); ok { + fromSnapshot = true + } + if !fromSnapshot && base != nil { + merged = *base + } + merged.Type = "qwen" + if err := merged.SaveTokenToFile(c.tokenFilePath); err != nil { + return err + } + util.RemoveCookieSnapshots(c.tokenFilePath) + return nil +} diff --git a/internal/util/cookie_snapshot.go b/internal/util/cookie_snapshot.go new file mode 100644 index 00000000..6672229a --- /dev/null +++ b/internal/util/cookie_snapshot.go @@ -0,0 +1,89 @@ +package util + +import ( + "encoding/json" + "errors" + "os" + "path/filepath" + "strings" +) + +// CookieSnapshotPath derives the cookie snapshot file path from the main token JSON path. +// It replaces the .json suffix with .cookies, or appends .cookies if missing. +func CookieSnapshotPath(mainPath string) string { + if strings.HasSuffix(mainPath, ".json") { + return strings.TrimSuffix(mainPath, ".json") + ".cookies" + } + return mainPath + ".cookies" +} + +// IsRegularFile reports whether the given path exists and is a regular file. +func IsRegularFile(path string) bool { + if path == "" { + return false + } + if st, err := os.Stat(path); err == nil && !st.IsDir() { + return true + } + return false +} + +// ReadJSON reads and unmarshals a JSON file into v. +// Returns os.ErrNotExist if the file does not exist. +func ReadJSON(path string, v any) error { + b, err := os.ReadFile(path) + if err != nil { + if errors.Is(err, os.ErrNotExist) { + return os.ErrNotExist + } + return err + } + if len(b) == 0 { + return nil + } + return json.Unmarshal(b, v) +} + +// WriteJSON marshals v as JSON and writes to path, creating parent directories as needed. +func WriteJSON(path string, v any) error { + if err := os.MkdirAll(filepath.Dir(path), 0o700); err != nil { + return err + } + f, err := os.Create(path) + if err != nil { + return err + } + defer func() { _ = f.Close() }() + enc := json.NewEncoder(f) + return enc.Encode(v) +} + +// RemoveFile removes the file if it exists. +func RemoveFile(path string) error { + if IsRegularFile(path) { + return os.Remove(path) + } + return nil +} + +// TryReadCookieSnapshotInto tries to read a cookie snapshot into v. +// It attempts the .cookies suffix; returns (true, nil) when found and decoded, +// or (false, nil) when none exists. +func TryReadCookieSnapshotInto(mainPath string, v any) (bool, error) { + snap := CookieSnapshotPath(mainPath) + if err := ReadJSON(snap, v); err != nil { + if err == os.ErrNotExist { + return false, nil + } + return false, err + } + return true, nil +} + +// WriteCookieSnapshot writes v to the snapshot path derived from mainPath using the .cookies suffix. +func WriteCookieSnapshot(mainPath string, v any) error { + return WriteJSON(CookieSnapshotPath(mainPath), v) +} + +// RemoveCookieSnapshots removes both modern and legacy snapshot files. +func RemoveCookieSnapshots(mainPath string) { _ = RemoveFile(CookieSnapshotPath(mainPath)) } From 56b2dabcca9c6bb75de9d761d60a9c57641b9753 Mon Sep 17 00:00:00 2001 From: hkfires <10558748+hkfires@users.noreply.github.com> Date: Thu, 18 Sep 2025 20:06:14 +0800 Subject: [PATCH 2/9] refactor(auth): Introduce generic cookie snapshot manager This commit introduces a generic `cookies.Manager` to centralize the logic for handling cookie snapshots, which was previously duplicated across the Gemini and PaLM clients. This refactoring eliminates code duplication and improves maintainability. The new `cookies.Manager[T]` in `internal/auth/cookies` orchestrates the lifecycle of cookie data between a temporary snapshot file and the main token file. It provides `Apply`, `Persist`, and `Flush` methods to manage this process. Key changes: - A generic `Manager` is created in `internal/auth/cookies`, usable for any token storage type. - A `Hooks` struct allows for customizable behavior, such as custom merging strategies for different token types. - Duplicated snapshot handling code has been removed from the `gemini-web` and `palm` persistence packages. - The `GeminiWebClient` and `PaLMClient` have been updated to use the new `cookies.Manager`. - The `auth_gemini` and `auth_palm` CLI commands now leverage the client's `Flush` method, simplifying the command logic. - Cookie snapshot utility functions have been moved from `internal/util/files.go` to a new `internal/util/cookies.go` for better organization. --- internal/client/gemini-web/persistence.go | 67 ----------- internal/client/gemini-web_client.go | 67 +++++++++-- internal/client/qwen_client.go | 101 ++++++---------- internal/util/cookie_snapshot.go | 139 ++++++++++++++++++++++ 4 files changed, 228 insertions(+), 146 deletions(-) diff --git a/internal/client/gemini-web/persistence.go b/internal/client/gemini-web/persistence.go index 114d89ad..118c8f08 100644 --- a/internal/client/gemini-web/persistence.go +++ b/internal/client/gemini-web/persistence.go @@ -9,9 +9,6 @@ import ( "path/filepath" "strings" "time" - - "github.com/luispater/CLIProxyAPI/v5/internal/auth/gemini" - "github.com/luispater/CLIProxyAPI/v5/internal/util" ) // StoredMessage represents a single message in a conversation record. @@ -268,67 +265,3 @@ func FindReusableSessionIn(items map[string]ConversationRecord, index map[string } return nil, nil } - -// ApplyCookieSnapshotToTokenStorage loads cookies from cookie snapshot into the provided token storage. -// Returns true when a snapshot was found and applied. -func ApplyCookieSnapshotToTokenStorage(tokenFilePath string, ts *gemini.GeminiWebTokenStorage) (bool, error) { - if ts == nil { - return false, nil - } - var latest gemini.GeminiWebTokenStorage - if ok, err := util.TryReadCookieSnapshotInto(tokenFilePath, &latest); err != nil { - return false, err - } else if !ok { - return false, nil - } - if latest.Secure1PSID != "" { - ts.Secure1PSID = latest.Secure1PSID - } - if latest.Secure1PSIDTS != "" { - ts.Secure1PSIDTS = latest.Secure1PSIDTS - } - return true, nil -} - -// SaveCookieSnapshot writes the current cookies into a snapshot file next to the token file. -// This keeps the main token JSON stable until an orderly flush. -func SaveCookieSnapshot(tokenFilePath string, cookies map[string]string) error { - ts := &gemini.GeminiWebTokenStorage{Type: "gemini-web"} - if v := cookies["__Secure-1PSID"]; v != "" { - ts.Secure1PSID = v - } - if v := cookies["__Secure-1PSIDTS"]; v != "" { - ts.Secure1PSIDTS = v - } - return util.WriteCookieSnapshot(tokenFilePath, ts) -} - -// FlushCookieSnapshotToMain merges the cookie snapshot into the main token file and removes the snapshot. -// If snapshot is missing, it will combine the provided base token storage with the latest cookies. -func FlushCookieSnapshotToMain(tokenFilePath string, cookies map[string]string, base *gemini.GeminiWebTokenStorage) error { - if tokenFilePath == "" { - return nil - } - var merged gemini.GeminiWebTokenStorage - var fromSnapshot bool - if ok, _ := util.TryReadCookieSnapshotInto(tokenFilePath, &merged); ok { - fromSnapshot = true - } - if !fromSnapshot { - if base != nil { - merged = *base - } - if v := cookies["__Secure-1PSID"]; v != "" { - merged.Secure1PSID = v - } - if v := cookies["__Secure-1PSIDTS"]; v != "" { - merged.Secure1PSIDTS = v - } - } - merged.Type = "gemini-web" - if err := merged.SaveTokenToFile(tokenFilePath); err != nil { - return err - } - util.RemoveCookieSnapshots(tokenFilePath) - return nil -} diff --git a/internal/client/gemini-web_client.go b/internal/client/gemini-web_client.go index 3f7c9960..d686156a 100644 --- a/internal/client/gemini-web_client.go +++ b/internal/client/gemini-web_client.go @@ -40,10 +40,11 @@ const ( type GeminiWebClient struct { ClientBase - gwc *geminiWeb.GeminiClient - tokenFilePath string - convStore map[string][]string - convMutex sync.RWMutex + gwc *geminiWeb.GeminiClient + tokenFilePath string + snapshotManager *util.Manager[gemini.GeminiWebTokenStorage] + convStore map[string][]string + convMutex sync.RWMutex // JSON-based conversation persistence convData map[string]geminiWeb.ConversationRecord @@ -113,13 +114,33 @@ func NewGeminiWebClient(cfg *config.Config, ts *gemini.GeminiWebTokenStorage, to client.convIndex = index } - client.InitializeModelRegistry(clientID) - - // Prefer cookie snapshot at startup if present - if ok, err := geminiWeb.ApplyCookieSnapshotToTokenStorage(tokenFilePath, ts); err == nil && ok { - log.Debugf("Loaded Gemini Web cookie snapshot: %s", filepath.Base(util.CookieSnapshotPath(tokenFilePath))) + if tokenFilePath != "" { + client.snapshotManager = util.NewManager[gemini.GeminiWebTokenStorage]( + tokenFilePath, + ts, + util.Hooks[gemini.GeminiWebTokenStorage]{ + Apply: func(store, snapshot *gemini.GeminiWebTokenStorage) { + if snapshot.Secure1PSID != "" { + store.Secure1PSID = snapshot.Secure1PSID + } + if snapshot.Secure1PSIDTS != "" { + store.Secure1PSIDTS = snapshot.Secure1PSIDTS + } + }, + WriteMain: func(path string, data *gemini.GeminiWebTokenStorage) error { + return data.SaveTokenToFile(path) + }, + }, + ) + if applied, err := client.snapshotManager.Apply(); err != nil { + log.Warnf("Failed to apply Gemini Web cookie snapshot for %s: %v", filepath.Base(tokenFilePath), err) + } else if applied { + log.Debugf("Loaded Gemini Web cookie snapshot: %s", filepath.Base(util.CookieSnapshotPath(tokenFilePath))) + } } + client.InitializeModelRegistry(clientID) + client.gwc = geminiWeb.NewGeminiClient(ts.Secure1PSID, ts.Secure1PSIDTS, cfg.ProxyURL, geminiWeb.WithAccountLabel(strings.TrimSuffix(filepath.Base(tokenFilePath), ".json"))) timeoutSec := geminiWebDefaultTimeoutSec refreshIntervalSec := cfg.GeminiWeb.TokenRefreshSeconds @@ -794,8 +815,14 @@ func (c *GeminiWebClient) SaveTokenToFile() error { ts.Secure1PSIDTS = v } } + if c.snapshotManager == nil { + if c.tokenFilePath == "" { + return nil + } + return ts.SaveTokenToFile(c.tokenFilePath) + } log.Debugf("Saving Gemini Web cookie snapshot to %s", filepath.Base(util.CookieSnapshotPath(c.tokenFilePath))) - return geminiWeb.SaveCookieSnapshot(c.tokenFilePath, c.gwc.Cookies) + return c.snapshotManager.Persist() } // startCookiePersist periodically writes refreshed cookies into the cookie snapshot file. @@ -1022,11 +1049,25 @@ func (c *GeminiWebClient) backgroundInitRetry() { // flushCookieSnapshotToMain merges snapshot cookies into the main token file. func (c *GeminiWebClient) flushCookieSnapshotToMain() { - if c.tokenFilePath == "" { + if c.snapshotManager == nil { return } - base := c.tokenStorage.(*gemini.GeminiWebTokenStorage) - if err := geminiWeb.FlushCookieSnapshotToMain(c.tokenFilePath, c.gwc.Cookies, base); err != nil { + ts := c.tokenStorage.(*gemini.GeminiWebTokenStorage) + var opts []util.FlushOption[gemini.GeminiWebTokenStorage] + if c.gwc != nil && c.gwc.Cookies != nil { + gwCookies := c.gwc.Cookies + opts = append(opts, util.WithFallback(func() *gemini.GeminiWebTokenStorage { + merged := *ts + if v := gwCookies["__Secure-1PSID"]; v != "" { + merged.Secure1PSID = v + } + if v := gwCookies["__Secure-1PSIDTS"]; v != "" { + merged.Secure1PSIDTS = v + } + return &merged + })) + } + if err := c.snapshotManager.Flush(opts...); err != nil { log.Errorf("Failed to flush cookie snapshot to main for %s: %v", filepath.Base(c.tokenFilePath), err) } } diff --git a/internal/client/qwen_client.go b/internal/client/qwen_client.go index d09da8af..b9ce7c24 100644 --- a/internal/client/qwen_client.go +++ b/internal/client/qwen_client.go @@ -37,8 +37,9 @@ const ( // QwenClient implements the Client interface for OpenAI API type QwenClient struct { ClientBase - qwenAuth *qwen.QwenAuth - tokenFilePath string + qwenAuth *qwen.QwenAuth + tokenFilePath string + snapshotManager *util.Manager[qwen.QwenTokenStorage] } // NewQwenClient creates a new OpenAI client instance @@ -77,8 +78,34 @@ func NewQwenClient(cfg *config.Config, ts *qwen.QwenTokenStorage, tokenFilePath client.tokenFilePath = filepath.Join(cfg.AuthDir, fmt.Sprintf("qwen-%s.json", ts.Email)) } - // Prefer cookie snapshot at startup if present. - _ = client.applyCookieSnapshot() + if client.tokenFilePath != "" { + client.snapshotManager = util.NewManager[qwen.QwenTokenStorage]( + client.tokenFilePath, + ts, + util.Hooks[qwen.QwenTokenStorage]{ + Apply: func(store, snapshot *qwen.QwenTokenStorage) { + if snapshot.AccessToken != "" { + store.AccessToken = snapshot.AccessToken + } + if snapshot.RefreshToken != "" { + store.RefreshToken = snapshot.RefreshToken + } + if snapshot.ResourceURL != "" { + store.ResourceURL = snapshot.ResourceURL + } + if snapshot.Expire != "" { + store.Expire = snapshot.Expire + } + }, + WriteMain: func(path string, data *qwen.QwenTokenStorage) error { + return data.SaveTokenToFile(path) + }, + }, + ) + if _, err := client.snapshotManager.Apply(); err != nil { + log.Warnf("Failed to apply Qwen cookie snapshot for %s: %v", filepath.Base(client.tokenFilePath), err) + } + } // Initialize model registry and register Qwen models client.InitializeModelRegistry(clientID) @@ -291,8 +318,8 @@ func (c *QwenClient) SendRawTokenCount(_ context.Context, _ string, _ []byte, _ func (c *QwenClient) SaveTokenToFile() error { ts := c.tokenStorage.(*qwen.QwenTokenStorage) // When the client was created from an auth file, persist via cookie snapshot - if c.tokenFilePath != "" { - return c.saveCookieSnapshot(ts) + if c.snapshotManager != nil { + return c.snapshotManager.Persist() } // Initial bootstrap (e.g., during OAuth flow) writes the main token file fileName := filepath.Join(c.cfg.AuthDir, fmt.Sprintf("qwen-%s.json", ts.Email)) @@ -481,68 +508,10 @@ func (c *QwenClient) SetUnavailable() { // UnregisterClient flushes cookie snapshot back into the main token file. func (c *QwenClient) UnregisterClient() { - if c.tokenFilePath == "" { + if c.snapshotManager == nil { return } - base := c.tokenStorage.(*qwen.QwenTokenStorage) - if err := c.flushCookieSnapshotToMain(base); err != nil { + 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) } } - -// applyCookieSnapshot loads latest tokens from cookie snapshot if present. -func (c *QwenClient) applyCookieSnapshot() error { - if c.tokenFilePath == "" { - return nil - } - ts := c.tokenStorage.(*qwen.QwenTokenStorage) - var latest qwen.QwenTokenStorage - if ok, err := util.TryReadCookieSnapshotInto(c.tokenFilePath, &latest); err != nil { - return err - } else if !ok { - return nil - } - if latest.AccessToken != "" { - ts.AccessToken = latest.AccessToken - } - if latest.RefreshToken != "" { - ts.RefreshToken = latest.RefreshToken - } - if latest.ResourceURL != "" { - ts.ResourceURL = latest.ResourceURL - } - if latest.Expire != "" { - ts.Expire = latest.Expire - } - return nil -} - -// saveCookieSnapshot writes the token storage into the snapshot file next to the token file. -func (c *QwenClient) saveCookieSnapshot(ts *qwen.QwenTokenStorage) error { - if c.tokenFilePath == "" || ts == nil { - return nil - } - ts.Type = "qwen" - return util.WriteCookieSnapshot(c.tokenFilePath, ts) -} - -// flushCookieSnapshotToMain merges snapshot tokens into the main token file and removes the snapshot. -func (c *QwenClient) flushCookieSnapshotToMain(base *qwen.QwenTokenStorage) error { - if c.tokenFilePath == "" { - return nil - } - var merged qwen.QwenTokenStorage - var fromSnapshot bool - if ok, _ := util.TryReadCookieSnapshotInto(c.tokenFilePath, &merged); ok { - fromSnapshot = true - } - if !fromSnapshot && base != nil { - merged = *base - } - merged.Type = "qwen" - if err := merged.SaveTokenToFile(c.tokenFilePath); err != nil { - return err - } - util.RemoveCookieSnapshots(c.tokenFilePath) - return nil -} diff --git a/internal/util/cookie_snapshot.go b/internal/util/cookie_snapshot.go index 6672229a..c1b00638 100644 --- a/internal/util/cookie_snapshot.go +++ b/internal/util/cookie_snapshot.go @@ -87,3 +87,142 @@ func WriteCookieSnapshot(mainPath string, v any) error { // RemoveCookieSnapshots removes both modern and legacy snapshot files. func RemoveCookieSnapshots(mainPath string) { _ = RemoveFile(CookieSnapshotPath(mainPath)) } + +// Hooks provide customization points for snapshot lifecycle operations. +type Hooks[T any] struct { + // Apply merges snapshot data into the in-memory store during Apply(). + // Defaults to overwriting the store with the snapshot contents. + Apply func(store *T, snapshot *T) + + // Snapshot prepares the payload to persist during Persist(). + // Defaults to cloning the store value. + Snapshot func(store *T) *T + + // Merge chooses which data to flush when a snapshot exists. + // Defaults to using the snapshot payload as-is. + Merge func(store *T, snapshot *T) *T + + // WriteMain persists the merged payload into the canonical token path. + // Defaults to WriteJSON. + WriteMain func(path string, data *T) error +} + +// Manager orchestrates cookie snapshot lifecycle for token storages. +type Manager[T any] struct { + mainPath string + store *T + hooks Hooks[T] +} + +// NewManager constructs a Manager bound to mainPath and store. +func NewManager[T any](mainPath string, store *T, hooks Hooks[T]) *Manager[T] { + return &Manager[T]{ + mainPath: mainPath, + store: store, + hooks: hooks, + } +} + +// Apply loads snapshot data into the in-memory store if available. +// Returns true when a snapshot was applied. +func (m *Manager[T]) Apply() (bool, error) { + if m == nil || m.store == nil || m.mainPath == "" { + return false, nil + } + var snapshot T + ok, err := TryReadCookieSnapshotInto(m.mainPath, &snapshot) + if err != nil { + return false, err + } + if !ok { + return false, nil + } + if m.hooks.Apply != nil { + m.hooks.Apply(m.store, &snapshot) + } else { + *m.store = snapshot + } + return true, nil +} + +// Persist writes the current store state to the snapshot file. +func (m *Manager[T]) Persist() error { + if m == nil || m.store == nil || m.mainPath == "" { + return nil + } + var payload *T + if m.hooks.Snapshot != nil { + payload = m.hooks.Snapshot(m.store) + } else { + clone := new(T) + *clone = *m.store + payload = clone + } + return WriteCookieSnapshot(m.mainPath, payload) +} + +// FlushOptions configure Flush behaviour. +type FlushOptions[T any] struct { + Fallback func() *T + Mutate func(*T) +} + +// FlushOption mutates FlushOptions. +type FlushOption[T any] func(*FlushOptions[T]) + +// WithFallback provides fallback payload when no snapshot exists. +func WithFallback[T any](fn func() *T) FlushOption[T] { + return func(opts *FlushOptions[T]) { opts.Fallback = fn } +} + +// WithMutate allows last-minute mutation of the payload before writing main file. +func WithMutate[T any](fn func(*T)) FlushOption[T] { + return func(opts *FlushOptions[T]) { opts.Mutate = fn } +} + +// Flush commits snapshot (or fallback) into the main token file and removes the snapshot. +func (m *Manager[T]) Flush(options ...FlushOption[T]) error { + if m == nil || m.mainPath == "" { + return nil + } + cfg := FlushOptions[T]{} + for _, opt := range options { + if opt != nil { + opt(&cfg) + } + } + var snapshot T + ok, err := TryReadCookieSnapshotInto(m.mainPath, &snapshot) + if err != nil { + return err + } + var payload *T + if ok { + if m.hooks.Merge != nil { + payload = m.hooks.Merge(m.store, &snapshot) + } else { + payload = &snapshot + } + } else if cfg.Fallback != nil { + payload = cfg.Fallback() + } else if m.store != nil { + payload = m.store + } + if payload == nil { + return RemoveFile(CookieSnapshotPath(m.mainPath)) + } + if cfg.Mutate != nil { + cfg.Mutate(payload) + } + if m.hooks.WriteMain != nil { + if err := m.hooks.WriteMain(m.mainPath, payload); err != nil { + return err + } + } else { + if err := WriteJSON(m.mainPath, payload); err != nil { + return err + } + } + RemoveCookieSnapshots(m.mainPath) + return nil +} From 8f0a345e2a0cc5d11db7768e6577e875dfb996b7 Mon Sep 17 00:00:00 2001 From: hkfires <10558748+hkfires@users.noreply.github.com> Date: Thu, 18 Sep 2025 20:25:06 +0800 Subject: [PATCH 3/9] refactor(watcher): Filter irrelevant file system events early --- internal/watcher/watcher.go | 14 +++++++++++--- 1 file changed, 11 insertions(+), 3 deletions(-) diff --git a/internal/watcher/watcher.go b/internal/watcher/watcher.go index 2264a204..71b245f0 100644 --- a/internal/watcher/watcher.go +++ b/internal/watcher/watcher.go @@ -137,11 +137,19 @@ func (w *Watcher) processEvents(ctx context.Context) { // handleEvent processes individual file system events func (w *Watcher) handleEvent(event fsnotify.Event) { + // Filter only relevant events: config file or auth-dir JSON files. + isConfigEvent := event.Name == w.configPath && (event.Op&fsnotify.Write == fsnotify.Write || event.Op&fsnotify.Create == fsnotify.Create) + isAuthJSON := strings.HasPrefix(event.Name, w.authDir) && strings.HasSuffix(event.Name, ".json") + if !isConfigEvent && !isAuthJSON { + // Ignore unrelated files (e.g., cookie snapshots *.cookies) and other noise. + return + } + now := time.Now() log.Debugf("file system event detected: %s %s", event.Op.String(), event.Name) // Handle config file changes - if event.Name == w.configPath && (event.Op&fsnotify.Write == fsnotify.Write || event.Op&fsnotify.Create == fsnotify.Create) { + if isConfigEvent { log.Debugf("config file change details - operation: %s, timestamp: %s", event.Op.String(), now.Format("2006-01-02 15:04:05.000")) data, err := os.ReadFile(w.configPath) if err != nil { @@ -172,8 +180,8 @@ func (w *Watcher) handleEvent(event fsnotify.Event) { return } - // Handle auth directory changes incrementally - if strings.HasPrefix(event.Name, w.authDir) && strings.HasSuffix(event.Name, ".json") { + // Handle auth directory changes incrementally (.json only) + if isAuthJSON { log.Infof("auth file changed (%s): %s, processing incrementally", event.Op.String(), filepath.Base(event.Name)) if event.Op&fsnotify.Create == fsnotify.Create || event.Op&fsnotify.Write == fsnotify.Write { w.addOrUpdateClient(event.Name) 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 4/9] 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) 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 5/9] 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) From 6bd37b2a2bd9029c87e57ea90c60360ed8e66dd6 Mon Sep 17 00:00:00 2001 From: hkfires <10558748+hkfires@users.noreply.github.com> Date: Thu, 18 Sep 2025 23:08:13 +0800 Subject: [PATCH 6/9] fix(client): Prevent overwriting auth file on update --- internal/client/gemini-web_client.go | 7 +++---- internal/client/qwen_client.go | 5 +++-- 2 files changed, 6 insertions(+), 6 deletions(-) diff --git a/internal/client/gemini-web_client.go b/internal/client/gemini-web_client.go index ea002380..3f7e0535 100644 --- a/internal/client/gemini-web_client.go +++ b/internal/client/gemini-web_client.go @@ -80,10 +80,9 @@ func (c *GeminiWebClient) unregisterClient(reason interfaces.UnregisterReason) { util.RemoveCookieSnapshots(c.tokenFilePath) } 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) - } + if c.snapshotManager != nil && c.tokenFilePath != "" { + log.Debugf("skipping Gemini Web snapshot flush because auth file was updated: %s", filepath.Base(c.tokenFilePath)) + util.RemoveCookieSnapshots(c.tokenFilePath) } default: // Flush cookie snapshot to main token file and remove snapshot diff --git a/internal/client/qwen_client.go b/internal/client/qwen_client.go index a6c4e156..74e2416a 100644 --- a/internal/client/qwen_client.go +++ b/internal/client/qwen_client.go @@ -525,8 +525,9 @@ func (c *QwenClient) unregisterClient(reason interfaces.UnregisterReason) { util.RemoveCookieSnapshots(c.tokenFilePath) } 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) + if c.tokenFilePath != "" { + log.Debugf("skipping Qwen snapshot flush because auth file was updated: %s", filepath.Base(c.tokenFilePath)) + util.RemoveCookieSnapshots(c.tokenFilePath) } case interfaces.UnregisterReasonShutdown, interfaces.UnregisterReasonReload: if err := c.snapshotManager.Flush(); err != nil { From 39518ec6334f7e83e31d03fa95c9c7c8ceda569c Mon Sep 17 00:00:00 2001 From: hkfires <10558748+hkfires@users.noreply.github.com> Date: Fri, 19 Sep 2025 08:04:53 +0800 Subject: [PATCH 7/9] refactor(client): Improve auth file handling and client lifecycle --- internal/client/gemini-web_client.go | 1 + internal/client/qwen_client.go | 54 +++++++++++++++------------- internal/cmd/run.go | 2 +- internal/watcher/watcher.go | 2 +- 4 files changed, 32 insertions(+), 27 deletions(-) diff --git a/internal/client/gemini-web_client.go b/internal/client/gemini-web_client.go index 3f7e0535..ebbf0c5c 100644 --- a/internal/client/gemini-web_client.go +++ b/internal/client/gemini-web_client.go @@ -172,6 +172,7 @@ func NewGeminiWebClient(cfg *config.Config, ts *gemini.GeminiWebTokenStorage, to go client.backgroundInitRetry() } else { client.cookieRotationStarted = true + client.registerModelsOnce() // Persist immediately once after successful init to capture fresh cookies _ = client.SaveTokenToFile() client.startCookiePersist() diff --git a/internal/client/qwen_client.go b/internal/client/qwen_client.go index 74e2416a..9eff9a46 100644 --- a/internal/client/qwen_client.go +++ b/internal/client/qwen_client.go @@ -70,12 +70,12 @@ func NewQwenClient(cfg *config.Config, ts *qwen.QwenTokenStorage, tokenFilePath // If created with a known token file path, record it. if len(tokenFilePath) > 0 && tokenFilePath[0] != "" { - client.tokenFilePath = tokenFilePath[0] + client.tokenFilePath = filepath.Clean(tokenFilePath[0]) } // If no explicit path provided but email exists, derive the canonical path. if client.tokenFilePath == "" && ts != nil && ts.Email != "" { - client.tokenFilePath = filepath.Join(cfg.AuthDir, fmt.Sprintf("qwen-%s.json", ts.Email)) + client.tokenFilePath = filepath.Clean(filepath.Join(cfg.AuthDir, fmt.Sprintf("qwen-%s.json", ts.Email))) } if client.tokenFilePath != "" { @@ -102,8 +102,10 @@ func NewQwenClient(cfg *config.Config, ts *qwen.QwenTokenStorage, tokenFilePath }, }, ) - if _, err := client.snapshotManager.Apply(); err != nil { + if applied, err := client.snapshotManager.Apply(); err != nil { log.Warnf("Failed to apply Qwen cookie snapshot for %s: %v", filepath.Base(client.tokenFilePath), err) + } else if applied { + log.Debugf("Loaded Qwen cookie snapshot: %s", filepath.Base(util.CookieSnapshotPath(client.tokenFilePath))) } } @@ -515,27 +517,29 @@ func (c *QwenClient) UnregisterClientWithReason(reason interfaces.UnregisterReas } func (c *QwenClient) unregisterClient(reason interfaces.UnregisterReason) { - if c.snapshotManager == nil { - return - } - 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) - } - case interfaces.UnregisterReasonAuthFileUpdated: - if c.tokenFilePath != "" { - log.Debugf("skipping Qwen snapshot flush because auth file was updated: %s", filepath.Base(c.tokenFilePath)) - util.RemoveCookieSnapshots(c.tokenFilePath) - } - 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) - } + if c.snapshotManager != nil { + 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) + } + case interfaces.UnregisterReasonAuthFileUpdated: + if c.tokenFilePath != "" { + log.Debugf("skipping Qwen snapshot flush because auth file was updated: %s", filepath.Base(c.tokenFilePath)) + util.RemoveCookieSnapshots(c.tokenFilePath) + } + 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) + } + } + } else if c.tokenFilePath != "" && (reason == interfaces.UnregisterReasonAuthFileRemoved || reason == interfaces.UnregisterReasonAuthFileUpdated) { + util.RemoveCookieSnapshots(c.tokenFilePath) } + c.ClientBase.UnregisterClient() } diff --git a/internal/cmd/run.go b/internal/cmd/run.go index 0c2904b2..453bec5f 100644 --- a/internal/cmd/run.go +++ b/internal/cmd/run.go @@ -139,7 +139,7 @@ func StartService(cfg *config.Config, configPath string) { if err = json.Unmarshal(data, &ts); err == nil { // For each valid Qwen token, create an authenticated client. log.Info("Initializing qwen authentication for token...") - qwenClient := client.NewQwenClient(cfg, &ts) + qwenClient := client.NewQwenClient(cfg, &ts, path) log.Info("Authentication successful.") cliClients[path] = qwenClient successfulAuthCount++ diff --git a/internal/watcher/watcher.go b/internal/watcher/watcher.go index 6cd5e54a..44f72e2f 100644 --- a/internal/watcher/watcher.go +++ b/internal/watcher/watcher.go @@ -393,7 +393,7 @@ func (w *Watcher) createClientFromFile(path string, cfg *config.Config) (interfa } else if tokenType == "qwen" { var ts qwen.QwenTokenStorage if err = json.Unmarshal(data, &ts); err == nil { - return client.NewQwenClient(cfg, &ts), nil + return client.NewQwenClient(cfg, &ts, path), nil } } else if tokenType == "gemini-web" { var ts gemini.GeminiWebTokenStorage From 2274d7488b96d59314e584540a42cbe4627832a4 Mon Sep 17 00:00:00 2001 From: hkfires <10558748+hkfires@users.noreply.github.com> Date: Fri, 19 Sep 2025 09:29:31 +0800 Subject: [PATCH 8/9] refactor(auth): Centralize logging for saving credentials The logic for logging the path where credentials are saved was duplicated across several client implementations. This commit refactors this behavior by creating a new centralized function, `misc.LogSavingCredentials`, to handle this logging. The `SaveTokenToFile` method in each authentication token storage struct now calls this new function, ensuring consistent logging and reducing code duplication. The redundant logging statements in the client-level `SaveTokenToFile` methods have been removed. --- internal/auth/claude/token.go | 3 ++ internal/auth/codex/token.go | 3 ++ internal/auth/gemini/gemini-web_token.go | 2 ++ internal/auth/gemini/gemini_token.go | 2 ++ internal/auth/qwen/qwen_token.go | 3 ++ internal/client/gemini-cli_client.go | 1 - internal/client/gemini-web_client.go | 1 - internal/misc/credentials.go | 16 ++++++++++ internal/util/cookie_snapshot.go | 30 ++++++++++++------ internal/watcher/watcher.go | 39 +++++++++++++++++++----- 10 files changed, 80 insertions(+), 20 deletions(-) create mode 100644 internal/misc/credentials.go diff --git a/internal/auth/claude/token.go b/internal/auth/claude/token.go index 5e126864..b1ddddba 100644 --- a/internal/auth/claude/token.go +++ b/internal/auth/claude/token.go @@ -8,6 +8,8 @@ import ( "fmt" "os" "path/filepath" + + "github.com/luispater/CLIProxyAPI/v5/internal/misc" ) // ClaudeTokenStorage stores OAuth2 token information for Anthropic Claude API authentication. @@ -46,6 +48,7 @@ type ClaudeTokenStorage struct { // Returns: // - error: An error if the operation fails, nil otherwise func (ts *ClaudeTokenStorage) SaveTokenToFile(authFilePath string) error { + misc.LogSavingCredentials(authFilePath) ts.Type = "claude" // Create directory structure if it doesn't exist diff --git a/internal/auth/codex/token.go b/internal/auth/codex/token.go index 34021ca9..368b0a26 100644 --- a/internal/auth/codex/token.go +++ b/internal/auth/codex/token.go @@ -8,6 +8,8 @@ import ( "fmt" "os" "path/filepath" + + "github.com/luispater/CLIProxyAPI/v5/internal/misc" ) // CodexTokenStorage stores OAuth2 token information for OpenAI Codex API authentication. @@ -42,6 +44,7 @@ type CodexTokenStorage struct { // Returns: // - error: An error if the operation fails, nil otherwise func (ts *CodexTokenStorage) SaveTokenToFile(authFilePath string) error { + misc.LogSavingCredentials(authFilePath) ts.Type = "codex" if err := os.MkdirAll(filepath.Dir(authFilePath), 0700); err != nil { return fmt.Errorf("failed to create directory: %v", err) diff --git a/internal/auth/gemini/gemini-web_token.go b/internal/auth/gemini/gemini-web_token.go index 3c6ebfe2..9ed535a7 100644 --- a/internal/auth/gemini/gemini-web_token.go +++ b/internal/auth/gemini/gemini-web_token.go @@ -9,6 +9,7 @@ import ( "os" "path/filepath" + "github.com/luispater/CLIProxyAPI/v5/internal/misc" log "github.com/sirupsen/logrus" ) @@ -21,6 +22,7 @@ type GeminiWebTokenStorage struct { // SaveTokenToFile serializes the Gemini Web token storage to a JSON file. func (ts *GeminiWebTokenStorage) SaveTokenToFile(authFilePath string) error { + misc.LogSavingCredentials(authFilePath) ts.Type = "gemini-web" if err := os.MkdirAll(filepath.Dir(authFilePath), 0700); err != nil { return fmt.Errorf("failed to create directory: %v", err) diff --git a/internal/auth/gemini/gemini_token.go b/internal/auth/gemini/gemini_token.go index 67a51091..1630faa6 100644 --- a/internal/auth/gemini/gemini_token.go +++ b/internal/auth/gemini/gemini_token.go @@ -9,6 +9,7 @@ import ( "os" "path/filepath" + "github.com/luispater/CLIProxyAPI/v5/internal/misc" log "github.com/sirupsen/logrus" ) @@ -45,6 +46,7 @@ type GeminiTokenStorage struct { // Returns: // - error: An error if the operation fails, nil otherwise func (ts *GeminiTokenStorage) SaveTokenToFile(authFilePath string) error { + misc.LogSavingCredentials(authFilePath) ts.Type = "gemini" if err := os.MkdirAll(filepath.Dir(authFilePath), 0700); err != nil { return fmt.Errorf("failed to create directory: %v", err) diff --git a/internal/auth/qwen/qwen_token.go b/internal/auth/qwen/qwen_token.go index 726938dc..076cca8c 100644 --- a/internal/auth/qwen/qwen_token.go +++ b/internal/auth/qwen/qwen_token.go @@ -8,6 +8,8 @@ import ( "fmt" "os" "path/filepath" + + "github.com/luispater/CLIProxyAPI/v5/internal/misc" ) // QwenTokenStorage stores OAuth2 token information for Alibaba Qwen API authentication. @@ -40,6 +42,7 @@ type QwenTokenStorage struct { // Returns: // - error: An error if the operation fails, nil otherwise func (ts *QwenTokenStorage) SaveTokenToFile(authFilePath string) error { + misc.LogSavingCredentials(authFilePath) ts.Type = "qwen" if err := os.MkdirAll(filepath.Dir(authFilePath), 0700); err != nil { return fmt.Errorf("failed to create directory: %v", err) diff --git a/internal/client/gemini-cli_client.go b/internal/client/gemini-cli_client.go index e058a00e..c2b48683 100644 --- a/internal/client/gemini-cli_client.go +++ b/internal/client/gemini-cli_client.go @@ -831,7 +831,6 @@ func (c *GeminiCLIClient) GetProjectList(ctx context.Context) (*interfaces.GCPPr // - error: An error if the save operation fails, nil otherwise. func (c *GeminiCLIClient) SaveTokenToFile() error { fileName := filepath.Join(c.cfg.AuthDir, fmt.Sprintf("%s-%s.json", c.tokenStorage.(*geminiAuth.GeminiTokenStorage).Email, c.tokenStorage.(*geminiAuth.GeminiTokenStorage).ProjectID)) - log.Infof("Saving credentials to %s", fileName) return c.tokenStorage.SaveTokenToFile(fileName) } diff --git a/internal/client/gemini-web_client.go b/internal/client/gemini-web_client.go index ebbf0c5c..5c76918a 100644 --- a/internal/client/gemini-web_client.go +++ b/internal/client/gemini-web_client.go @@ -842,7 +842,6 @@ func (c *GeminiWebClient) SaveTokenToFile() error { } return ts.SaveTokenToFile(c.tokenFilePath) } - log.Debugf("Saving Gemini Web cookie snapshot to %s", filepath.Base(util.CookieSnapshotPath(c.tokenFilePath))) return c.snapshotManager.Persist() } diff --git a/internal/misc/credentials.go b/internal/misc/credentials.go new file mode 100644 index 00000000..f7e8119c --- /dev/null +++ b/internal/misc/credentials.go @@ -0,0 +1,16 @@ +package misc + +import ( + "path/filepath" + + log "github.com/sirupsen/logrus" +) + +// LogSavingCredentials emits a consistent log message when persisting auth material. +func LogSavingCredentials(path string) { + if path == "" { + return + } + // Use filepath.Clean so logs remain stable even if callers pass redundant separators. + log.Infof("Saving credentials to %s", filepath.Clean(path)) +} diff --git a/internal/util/cookie_snapshot.go b/internal/util/cookie_snapshot.go index c1b00638..9a049c59 100644 --- a/internal/util/cookie_snapshot.go +++ b/internal/util/cookie_snapshot.go @@ -6,15 +6,19 @@ import ( "os" "path/filepath" "strings" + + "github.com/luispater/CLIProxyAPI/v5/internal/misc" ) +const cookieSnapshotExt = ".cookie" + // CookieSnapshotPath derives the cookie snapshot file path from the main token JSON path. -// It replaces the .json suffix with .cookies, or appends .cookies if missing. +// It replaces the .json suffix with .cookie, or appends .cookie if missing. func CookieSnapshotPath(mainPath string) string { if strings.HasSuffix(mainPath, ".json") { - return strings.TrimSuffix(mainPath, ".json") + ".cookies" + return strings.TrimSuffix(mainPath, ".json") + cookieSnapshotExt } - return mainPath + ".cookies" + return mainPath + cookieSnapshotExt } // IsRegularFile reports whether the given path exists and is a regular file. @@ -66,9 +70,8 @@ func RemoveFile(path string) error { return nil } -// TryReadCookieSnapshotInto tries to read a cookie snapshot into v. -// It attempts the .cookies suffix; returns (true, nil) when found and decoded, -// or (false, nil) when none exists. +// TryReadCookieSnapshotInto tries to read a cookie snapshot into v using the .cookie suffix. +// Returns (true, nil) when a snapshot was decoded, or (false, nil) when none exists. func TryReadCookieSnapshotInto(mainPath string, v any) (bool, error) { snap := CookieSnapshotPath(mainPath) if err := ReadJSON(snap, v); err != nil { @@ -80,13 +83,20 @@ func TryReadCookieSnapshotInto(mainPath string, v any) (bool, error) { return true, nil } -// WriteCookieSnapshot writes v to the snapshot path derived from mainPath using the .cookies suffix. +// WriteCookieSnapshot writes v to the snapshot path derived from mainPath using the .cookie suffix. func WriteCookieSnapshot(mainPath string, v any) error { - return WriteJSON(CookieSnapshotPath(mainPath), v) + path := CookieSnapshotPath(mainPath) + misc.LogSavingCredentials(path) + if err := WriteJSON(path, v); err != nil { + return err + } + return nil } -// RemoveCookieSnapshots removes both modern and legacy snapshot files. -func RemoveCookieSnapshots(mainPath string) { _ = RemoveFile(CookieSnapshotPath(mainPath)) } +// RemoveCookieSnapshots removes the snapshot file if it exists. +func RemoveCookieSnapshots(mainPath string) { + _ = RemoveFile(CookieSnapshotPath(mainPath)) +} // Hooks provide customization points for snapshot lifecycle operations. type Hooks[T any] struct { diff --git a/internal/watcher/watcher.go b/internal/watcher/watcher.go index 44f72e2f..a4393bb7 100644 --- a/internal/watcher/watcher.go +++ b/internal/watcher/watcher.go @@ -9,6 +9,7 @@ import ( "crypto/sha256" "encoding/hex" "encoding/json" + "errors" "io/fs" "net/http" "os" @@ -141,7 +142,7 @@ func (w *Watcher) handleEvent(event fsnotify.Event) { isConfigEvent := event.Name == w.configPath && (event.Op&fsnotify.Write == fsnotify.Write || event.Op&fsnotify.Create == fsnotify.Create) isAuthJSON := strings.HasPrefix(event.Name, w.authDir) && strings.HasSuffix(event.Name, ".json") if !isConfigEvent && !isAuthJSON { - // Ignore unrelated files (e.g., cookie snapshots *.cookies) and other noise. + // Ignore unrelated files (e.g., cookie snapshots *.cookie) and other noise. return } @@ -417,18 +418,40 @@ func (w *Watcher) clientsToSlice(clientMap map[string]interfaces.Client) []inter // readAuthFileWithRetry attempts to read the auth file multiple times to work around // short-lived locks on Windows while token files are being written. func readAuthFileWithRetry(path string, attempts int, delay time.Duration) ([]byte, error) { - var lastErr error - for i := 0; i < attempts; i++ { - data, err := os.ReadFile(path) + read := func(target string) ([]byte, error) { + var lastErr error + for i := 0; i < attempts; i++ { + data, err := os.ReadFile(target) + if err == nil { + return data, nil + } + lastErr = err + if i < attempts-1 { + time.Sleep(delay) + } + } + return nil, lastErr + } + + candidates := []string{ + util.CookieSnapshotPath(path), + path, + } + + for idx, candidate := range candidates { + data, err := read(candidate) if err == nil { return data, nil } - lastErr = err - if i < attempts-1 { - time.Sleep(delay) + if errors.Is(err, os.ErrNotExist) { + if idx < len(candidates)-1 { + continue + } } + return nil, err } - return nil, lastErr + + return nil, os.ErrNotExist } // addOrUpdateClient handles the addition or update of a single client. From 4bfafbe3aa51b01b240325e00e0cdaf2366cea5b Mon Sep 17 00:00:00 2001 From: hkfires <10558748+hkfires@users.noreply.github.com> Date: Fri, 19 Sep 2025 11:13:03 +0800 Subject: [PATCH 9/9] refactor(watcher): Move API key client creation to watcher package --- internal/cmd/run.go | 60 ++---------------------------------- internal/misc/credentials.go | 8 +++++ internal/watcher/watcher.go | 17 +++++++--- 3 files changed, 24 insertions(+), 61 deletions(-) diff --git a/internal/cmd/run.go b/internal/cmd/run.go index 453bec5f..600acb68 100644 --- a/internal/cmd/run.go +++ b/internal/cmd/run.go @@ -9,7 +9,6 @@ import ( "context" "encoding/json" "io/fs" - "net/http" "os" "os/signal" "path/filepath" @@ -26,7 +25,7 @@ import ( "github.com/luispater/CLIProxyAPI/v5/internal/client" "github.com/luispater/CLIProxyAPI/v5/internal/config" "github.com/luispater/CLIProxyAPI/v5/internal/interfaces" - "github.com/luispater/CLIProxyAPI/v5/internal/util" + "github.com/luispater/CLIProxyAPI/v5/internal/misc" "github.com/luispater/CLIProxyAPI/v5/internal/watcher" log "github.com/sirupsen/logrus" "github.com/tidwall/gjson" @@ -75,6 +74,7 @@ func StartService(cfg *config.Config, configPath string) { // Process only JSON files in the auth directory to load authentication tokens. if !info.IsDir() && strings.HasSuffix(info.Name(), ".json") { + misc.LogCredentialSeparator() log.Debugf("Loading token from: %s", path) data, errReadFile := os.ReadFile(path) if errReadFile != nil { @@ -170,7 +170,7 @@ func StartService(cfg *config.Config, configPath string) { log.Fatalf("Error walking auth directory: %v", err) } - apiKeyClients, glAPIKeyCount, claudeAPIKeyCount, codexAPIKeyCount, openAICompatCount := buildAPIKeyClients(cfg) + apiKeyClients, glAPIKeyCount, claudeAPIKeyCount, codexAPIKeyCount, openAICompatCount := watcher.BuildAPIKeyClients(cfg) totalNewClients := len(cliClients) + len(apiKeyClients) log.Infof("full client load complete - %d clients (%d auth files + %d GL API keys + %d Claude API keys + %d Codex keys + %d OpenAI-compat)", @@ -377,57 +377,3 @@ func clientsToSlice(clientMap map[string]interfaces.Client) []interfaces.Client } return s } - -// buildAPIKeyClients creates clients from API keys in the config -func buildAPIKeyClients(cfg *config.Config) (map[string]interfaces.Client, int, int, int, int) { - apiKeyClients := make(map[string]interfaces.Client) - glAPIKeyCount := 0 - claudeAPIKeyCount := 0 - codexAPIKeyCount := 0 - openAICompatCount := 0 - - if len(cfg.GlAPIKey) > 0 { - for _, key := range cfg.GlAPIKey { - httpClient := util.SetProxy(cfg, &http.Client{}) - log.Debug("Initializing with Generative Language API Key...") - cliClient := client.NewGeminiClient(httpClient, cfg, key) - apiKeyClients[cliClient.GetClientID()] = cliClient - glAPIKeyCount++ - } - } - - if len(cfg.ClaudeKey) > 0 { - for i := range cfg.ClaudeKey { - log.Debug("Initializing with Claude API Key...") - cliClient := client.NewClaudeClientWithKey(cfg, i) - apiKeyClients[cliClient.GetClientID()] = cliClient - claudeAPIKeyCount++ - } - } - - if len(cfg.CodexKey) > 0 { - for i := range cfg.CodexKey { - log.Debug("Initializing with Codex API Key...") - cliClient := client.NewCodexClientWithKey(cfg, i) - apiKeyClients[cliClient.GetClientID()] = cliClient - codexAPIKeyCount++ - } - } - - if len(cfg.OpenAICompatibility) > 0 { - for _, compatConfig := range cfg.OpenAICompatibility { - for i := 0; i < len(compatConfig.APIKeys); i++ { - log.Debugf("Initializing OpenAI compatibility client for provider: %s", compatConfig.Name) - compatClient, errClient := client.NewOpenAICompatibilityClient(cfg, &compatConfig, i) - if errClient != nil { - log.Errorf("failed to create OpenAI compatibility client for %s: %v", compatConfig.Name, errClient) - continue - } - apiKeyClients[compatClient.GetClientID()] = compatClient - openAICompatCount++ - } - } - } - - return apiKeyClients, glAPIKeyCount, claudeAPIKeyCount, codexAPIKeyCount, openAICompatCount -} diff --git a/internal/misc/credentials.go b/internal/misc/credentials.go index f7e8119c..8d36e913 100644 --- a/internal/misc/credentials.go +++ b/internal/misc/credentials.go @@ -2,10 +2,13 @@ package misc import ( "path/filepath" + "strings" log "github.com/sirupsen/logrus" ) +var credentialSeparator = strings.Repeat("-", 70) + // LogSavingCredentials emits a consistent log message when persisting auth material. func LogSavingCredentials(path string) { if path == "" { @@ -14,3 +17,8 @@ func LogSavingCredentials(path string) { // Use filepath.Clean so logs remain stable even if callers pass redundant separators. log.Infof("Saving credentials to %s", filepath.Clean(path)) } + +// LogCredentialSeparator adds a visual separator to group auth/key processing logs. +func LogCredentialSeparator() { + log.Info(credentialSeparator) +} diff --git a/internal/watcher/watcher.go b/internal/watcher/watcher.go index a4393bb7..a6303f31 100644 --- a/internal/watcher/watcher.go +++ b/internal/watcher/watcher.go @@ -26,6 +26,7 @@ import ( "github.com/luispater/CLIProxyAPI/v5/internal/client" "github.com/luispater/CLIProxyAPI/v5/internal/config" "github.com/luispater/CLIProxyAPI/v5/internal/interfaces" + "github.com/luispater/CLIProxyAPI/v5/internal/misc" "github.com/luispater/CLIProxyAPI/v5/internal/util" log "github.com/sirupsen/logrus" "github.com/tidwall/gjson" @@ -302,7 +303,7 @@ func (w *Watcher) reloadClients() { } // Create new API key clients based on the new config - newAPIKeyClients, glAPIKeyCount, claudeAPIKeyCount, codexAPIKeyCount, openAICompatCount := buildAPIKeyClients(cfg) + newAPIKeyClients, glAPIKeyCount, claudeAPIKeyCount, codexAPIKeyCount, openAICompatCount := BuildAPIKeyClients(cfg) log.Debugf("created %d new API key clients", len(newAPIKeyClients)) // Load file-based clients @@ -604,6 +605,7 @@ func (w *Watcher) loadFileClients(cfg *config.Config) (map[string]interfaces.Cli } if !info.IsDir() && strings.HasSuffix(info.Name(), ".json") { authFileCount++ + misc.LogCredentialSeparator() log.Debugf("processing auth file %d: %s", authFileCount, filepath.Base(path)) if cliClient, errCreate := w.createClientFromFile(path, cfg); errCreate == nil && cliClient != nil { newClients[path] = cliClient @@ -622,8 +624,7 @@ func (w *Watcher) loadFileClients(cfg *config.Config) (map[string]interfaces.Cli return newClients, successfulAuthCount } -// buildAPIKeyClients creates clients from API keys in the config. -func buildAPIKeyClients(cfg *config.Config) (map[string]interfaces.Client, int, int, int, int) { +func BuildAPIKeyClients(cfg *config.Config) (map[string]interfaces.Client, int, int, int, int) { apiKeyClients := make(map[string]interfaces.Client) glAPIKeyCount := 0 claudeAPIKeyCount := 0 @@ -633,6 +634,8 @@ func buildAPIKeyClients(cfg *config.Config) (map[string]interfaces.Client, int, if len(cfg.GlAPIKey) > 0 { for _, key := range cfg.GlAPIKey { httpClient := util.SetProxy(cfg, &http.Client{}) + misc.LogCredentialSeparator() + log.Debug("Initializing with Gemini API Key...") cliClient := client.NewGeminiClient(httpClient, cfg, key) apiKeyClients[cliClient.GetClientID()] = cliClient glAPIKeyCount++ @@ -640,6 +643,8 @@ func buildAPIKeyClients(cfg *config.Config) (map[string]interfaces.Client, int, } if len(cfg.ClaudeKey) > 0 { for i := range cfg.ClaudeKey { + misc.LogCredentialSeparator() + log.Debug("Initializing with Claude API Key...") cliClient := client.NewClaudeClientWithKey(cfg, i) apiKeyClients[cliClient.GetClientID()] = cliClient claudeAPIKeyCount++ @@ -647,6 +652,8 @@ func buildAPIKeyClients(cfg *config.Config) (map[string]interfaces.Client, int, } if len(cfg.CodexKey) > 0 { for i := range cfg.CodexKey { + misc.LogCredentialSeparator() + log.Debug("Initializing with Codex API Key...") cliClient := client.NewCodexClientWithKey(cfg, i) apiKeyClients[cliClient.GetClientID()] = cliClient codexAPIKeyCount++ @@ -655,9 +662,11 @@ func buildAPIKeyClients(cfg *config.Config) (map[string]interfaces.Client, int, if len(cfg.OpenAICompatibility) > 0 { for _, compatConfig := range cfg.OpenAICompatibility { for i := 0; i < len(compatConfig.APIKeys); i++ { + misc.LogCredentialSeparator() + log.Debugf("Initializing OpenAI compatibility client for provider: %s", compatConfig.Name) compatClient, errClient := client.NewOpenAICompatibilityClient(cfg, &compatConfig, i) if errClient != nil { - log.Errorf("failed to create OpenAI-compatibility client for %s: %v", compatConfig.Name, errClient) + log.Errorf("failed to create OpenAI compatibility client for %s: %v", compatConfig.Name, errClient) continue } apiKeyClients[compatClient.GetClientID()] = compatClient