feat(api): integrate TokenStore for improved auth entry management

Replaced file-based auth entry counting with `TokenStore`-backed implementation, enhancing flexibility and context-aware token management. Updated related logic to reflect this change.
This commit is contained in:
Luis Pater
2026-01-03 04:53:47 +08:00
parent e02ceecd35
commit ebec293497
2 changed files with 25 additions and 33 deletions

View File

@@ -33,6 +33,7 @@ import (
"github.com/router-for-me/CLIProxyAPI/v6/sdk/api/handlers/claude" "github.com/router-for-me/CLIProxyAPI/v6/sdk/api/handlers/claude"
"github.com/router-for-me/CLIProxyAPI/v6/sdk/api/handlers/gemini" "github.com/router-for-me/CLIProxyAPI/v6/sdk/api/handlers/gemini"
"github.com/router-for-me/CLIProxyAPI/v6/sdk/api/handlers/openai" "github.com/router-for-me/CLIProxyAPI/v6/sdk/api/handlers/openai"
sdkAuth "github.com/router-for-me/CLIProxyAPI/v6/sdk/auth"
"github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/auth" "github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/auth"
log "github.com/sirupsen/logrus" log "github.com/sirupsen/logrus"
"gopkg.in/yaml.v3" "gopkg.in/yaml.v3"
@@ -966,8 +967,12 @@ func (s *Server) UpdateClients(cfg *config.Config) {
log.Warnf("amp module is nil, skipping config update") log.Warnf("amp module is nil, skipping config update")
} }
// Count client sources from configuration and auth directory // Count client sources from configuration and auth store.
authFiles := util.CountAuthFiles(cfg.AuthDir) tokenStore := sdkAuth.GetTokenStore()
if dirSetter, ok := tokenStore.(interface{ SetBaseDir(string) }); ok {
dirSetter.SetBaseDir(cfg.AuthDir)
}
authEntries := util.CountAuthFiles(context.Background(), tokenStore)
geminiAPIKeyCount := len(cfg.GeminiKey) geminiAPIKeyCount := len(cfg.GeminiKey)
claudeAPIKeyCount := len(cfg.ClaudeKey) claudeAPIKeyCount := len(cfg.ClaudeKey)
codexAPIKeyCount := len(cfg.CodexKey) codexAPIKeyCount := len(cfg.CodexKey)
@@ -978,10 +983,10 @@ func (s *Server) UpdateClients(cfg *config.Config) {
openAICompatCount += len(entry.APIKeyEntries) openAICompatCount += len(entry.APIKeyEntries)
} }
total := authFiles + geminiAPIKeyCount + claudeAPIKeyCount + codexAPIKeyCount + vertexAICompatCount + openAICompatCount total := authEntries + geminiAPIKeyCount + claudeAPIKeyCount + codexAPIKeyCount + vertexAICompatCount + openAICompatCount
fmt.Printf("server clients and configuration updated: %d clients (%d auth files + %d Gemini API keys + %d Claude API keys + %d Codex keys + %d Vertex-compat + %d OpenAI-compat)\n", fmt.Printf("server clients and configuration updated: %d clients (%d auth entries + %d Gemini API keys + %d Claude API keys + %d Codex keys + %d Vertex-compat + %d OpenAI-compat)\n",
total, total,
authFiles, authEntries,
geminiAPIKeyCount, geminiAPIKeyCount,
claudeAPIKeyCount, claudeAPIKeyCount,
codexAPIKeyCount, codexAPIKeyCount,

View File

@@ -4,8 +4,8 @@
package util package util
import ( import (
"context"
"fmt" "fmt"
"io/fs"
"os" "os"
"path/filepath" "path/filepath"
"regexp" "regexp"
@@ -93,36 +93,23 @@ func ResolveAuthDir(authDir string) (string, error) {
return filepath.Clean(authDir), nil return filepath.Clean(authDir), nil
} }
// CountAuthFiles returns the number of JSON auth files located under the provided directory. // CountAuthFiles returns the number of auth records available through the provided Store.
// The function resolves leading tildes to the user's home directory and performs a case-insensitive // For filesystem-backed stores, this reflects the number of JSON auth files under the configured directory.
// match on the ".json" suffix so that files saved with uppercase extensions are also counted. func CountAuthFiles[T any](ctx context.Context, store interface {
func CountAuthFiles(authDir string) int { List(context.Context) ([]T, error)
dir, err := ResolveAuthDir(authDir) }) int {
if err != nil { if store == nil {
log.Debugf("countAuthFiles: failed to resolve auth directory: %v", err)
return 0 return 0
} }
if dir == "" { if ctx == nil {
ctx = context.Background()
}
entries, err := store.List(ctx)
if err != nil {
log.Debugf("countAuthFiles: failed to list auth records: %v", err)
return 0 return 0
} }
count := 0 return len(entries)
walkErr := filepath.WalkDir(dir, func(path string, d fs.DirEntry, err error) error {
if err != nil {
log.Debugf("countAuthFiles: error accessing %s: %v", path, err)
return nil
}
if d.IsDir() {
return nil
}
if strings.HasSuffix(strings.ToLower(d.Name()), ".json") {
count++
}
return nil
})
if walkErr != nil {
log.Debugf("countAuthFiles: walk error: %v", walkErr)
}
return count
} }
// WritablePath returns the cleaned WRITABLE_PATH environment variable when it is set. // WritablePath returns the cleaned WRITABLE_PATH environment variable when it is set.