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

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

internal/client/gemini-web_client.go

@@ -40,10 +40,11 @@ const (
 
 type GeminiWebClient struct {
 	ClientBase
-	gwc           *geminiWeb.GeminiClient
+	gwc             *geminiWeb.GeminiClient
-	tokenFilePath string
+	tokenFilePath   string
-	convStore     map[string][]string
+	snapshotManager *util.Manager[gemini.GeminiWebTokenStorage]
-	convMutex     sync.RWMutex
+	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)
+	if tokenFilePath != "" {
-
+		client.snapshotManager = util.NewManager[gemini.GeminiWebTokenStorage](
-	// Prefer cookie snapshot at startup if present
+			tokenFilePath,
-	if ok, err := geminiWeb.ApplyCookieSnapshotToTokenStorage(tokenFilePath, ts); err == nil && ok {
+			ts,
-		log.Debugf("Loaded Gemini Web cookie snapshot: %s", filepath.Base(util.CookieSnapshotPath(tokenFilePath)))
+			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)
+	ts := c.tokenStorage.(*gemini.GeminiWebTokenStorage)
-	if err := geminiWeb.FlushCookieSnapshotToMain(c.tokenFilePath, c.gwc.Cookies, base); err != nil {
+	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)
 	}
 }

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
+	qwenAuth        *qwen.QwenAuth
-	tokenFilePath string
+	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.
+	if client.tokenFilePath != "" {
-	_ = client.applyCookieSnapshot()
+		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 != "" {
+	if c.snapshotManager != nil {
-		return c.saveCookieSnapshot(ts)
+		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.snapshotManager.Flush(); err != nil {
-	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
-}

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