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