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] 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 +}