mirror of
https://github.com/router-for-me/CLIProxyAPI.git
synced 2026-02-18 20:30:51 +08:00
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.
This commit is contained in:
@@ -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
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user