Spaces:

bg5
/

zai2api2

Sleeping

部署 zai2api 到 HF Space

9aab84d verified 5 months ago

12.6 kB

package internal

import (
	"bufio"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"os"

	fhttp "github.com/bogdanfinn/fhttp"
	"path/filepath"
	"strings"
	"sync"
	"sync/atomic"
	"time"

	"github.com/fsnotify/fsnotify"
)

// TokenInfo 存储单个 token 的信息
type TokenInfo struct {
	Token       string    `json:"token"`
	Email       string    `json:"email"`
	UserID      string    `json:"user_id"`
	Valid       bool      `json:"valid"`
	LastChecked time.Time `json:"last_checked"`
	UseCount    int64     `json:"use_count"`
}

// TokenManager 管理所有用户 token
type TokenManager struct {
	mu              sync.RWMutex
	tokens          map[string]*TokenInfo // token -> TokenInfo
	validTokens     []string              // 有效 token 列表
	currentIndex    int                   // 轮询索引
	dataDir         string
	watcher         *fsnotify.Watcher
	checkInterval   time.Duration
	stopChan        chan struct{}
	multimodalCount int64 // 多模态请求计数
	totalCalls      int64 // 累计调用次数
	successCalls    int64 // 成功调用次数
}

var (
	tokenManager *TokenManager
	tokenOnce    sync.Once
)

// GetTokenManager 获取单例 TokenManager
func GetTokenManager() *TokenManager {
	tokenOnce.Do(func() {
		tokenManager = &TokenManager{
			tokens:        make(map[string]*TokenInfo),
			validTokens:   make([]string, 0),
			dataDir:       "data",
			checkInterval: 5 * time.Minute, // 每5分钟检查一次
			stopChan:      make(chan struct{}),
		}
	})
	return tokenManager
}

// Start 启动 token 管理器
func (tm *TokenManager) Start() error {
	// 确保 data 目录存在
	if err := os.MkdirAll(tm.dataDir, 0755); err != nil {
		return fmt.Errorf("创建 data 目录失败: %v", err)
	}

	// 初始加载 token
	if err := tm.loadTokens(); err != nil {
		LogWarn("初始加载 token 失败: %v", err)
	}

	// 启动文件监听
	if err := tm.startWatcher(); err != nil {
		LogWarn("启动文件监听失败: %v", err)
	}

	// 启动定期验证
	go tm.startValidator()

	LogInfo("TokenManager 已启动，当前有效 token 数: %d", len(tm.validTokens))
	return nil
}

// Stop 停止 token 管理器
func (tm *TokenManager) Stop() {
	close(tm.stopChan)
	if tm.watcher != nil {
		tm.watcher.Close()
	}
}

// loadTokens 从 data 目录加载所有 token
func (tm *TokenManager) loadTokens() error {
	tokenFile := filepath.Join(tm.dataDir, "tokens.txt")

	file, err := os.Open(tokenFile)
	if err != nil {
		if os.IsNotExist(err) {
			// 创建示例文件
			tm.createExampleTokenFile(tokenFile)
			return nil
		}
		return err
	}
	defer file.Close()

	tm.mu.Lock()
	defer tm.mu.Unlock()

	// 保留旧的统计数据
	oldTokens := tm.tokens
	tm.tokens = make(map[string]*TokenInfo)
	tm.validTokens = make([]string, 0)

	scanner := bufio.NewScanner(file)
	for scanner.Scan() {
		line := strings.TrimSpace(scanner.Text())
		if line == "" || strings.HasPrefix(line, "#") {
			continue
		}

		token := line
		// 支持 token=xxx 格式
		if strings.HasPrefix(line, "token=") {
			token = strings.TrimPrefix(line, "token=")
		}

		if token == "" {
			continue
		}

		// 复用旧的 TokenInfo 如果存在
		if oldInfo, exists := oldTokens[token]; exists {
			tm.tokens[token] = oldInfo
			if oldInfo.Valid {
				tm.validTokens = append(tm.validTokens, token)
			}
		} else {
			// 新 token，解析并标记为待验证
			info := &TokenInfo{
				Token: token,
				Valid: true, // 初始假设有效，验证时会更新
			}
			// 尝试解析 JWT 获取信息
			if payload, err := DecodeJWTPayload(token); err == nil && payload != nil {
				info.Email = payload.Email
				info.UserID = payload.ID
			}
			tm.tokens[token] = info
			tm.validTokens = append(tm.validTokens, token)
		}
	}

	validN := len(tm.validTokens)
	LogInfo("已加载 %d 个 token", validN)
	scanErr := scanner.Err()
	if validN > 0 {
		invalidateAnonymousPoolSlots()
	}
	return scanErr
}

// createExampleTokenFile 创建示例 token 文件
func (tm *TokenManager) createExampleTokenFile(path string) {
	content := `# 用户 Token 文件
# 每行一个 token，支持以下格式：
# 1. 直接写 token
# 2. token=xxx 格式
# 以 # 开头的行为注释

# 示例:
# eyJhbGciOiJFUzI1NiIsInR5cCI6IkpXVCJ9.xxxxx
`
	os.WriteFile(path, []byte(content), 0644)
	LogInfo("已创建示例 token 文件: %s", path)
}

// startWatcher 启动文件变化监听
func (tm *TokenManager) startWatcher() error {
	watcher, err := fsnotify.NewWatcher()
	if err != nil {
		return err
	}
	tm.watcher = watcher

	go func() {
		for {
			select {
			case event, ok := <-watcher.Events:
				if !ok {
					return
				}
				if event.Op&(fsnotify.Write|fsnotify.Create) != 0 {
					if strings.HasSuffix(event.Name, "tokens.txt") {
						LogInfo("检测到 token 文件变化，重新加载...")
						time.Sleep(100 * time.Millisecond) // 等待文件写入完成
						tm.loadTokens()
					}
				}
			case err, ok := <-watcher.Errors:
				if !ok {
					return
				}
				LogError("文件监听错误: %v", err)
			case <-tm.stopChan:
				return
			}
		}
	}()

	return watcher.Add(tm.dataDir)
}

// startValidator 启动定期验证
func (tm *TokenManager) startValidator() {
	// 首次延迟验证
	time.Sleep(10 * time.Second)
	tm.validateAllTokens()

	ticker := time.NewTicker(tm.checkInterval)
	defer ticker.Stop()

	for {
		select {
		case <-ticker.C:
			tm.validateAllTokens()
		case <-tm.stopChan:
			return
		}
	}
}

// validateAllTokens 验证所有 token
func (tm *TokenManager) validateAllTokens() {
	tm.mu.RLock()
	tokens := make([]string, 0, len(tm.tokens))
	for token := range tm.tokens {
		tokens = append(tokens, token)
	}
	tm.mu.RUnlock()

	LogInfo("开始验证 %d 个 token...", len(tokens))
	invalidCount := 0

	for _, token := range tokens {
		valid := tm.validateToken(token)
		tm.mu.Lock()
		if info, exists := tm.tokens[token]; exists {
			info.Valid = valid
			info.LastChecked = time.Now()
			if !valid {
				invalidCount++
			}
		}
		tm.mu.Unlock()
		time.Sleep(500 * time.Millisecond) // 避免请求过快
	}

	// 更新有效 token 列表
	tm.rebuildValidTokens()
	LogInfo("Token 验证完成，失效 %d 个，剩余有效 %d 个", invalidCount, len(tm.validTokens))

	// 自动删除失效 token
	if invalidCount > 0 {
		tm.removeInvalidTokens()
	}
}

// validateToken 验证单个 token
func (tm *TokenManager) validateToken(token string) bool {
	req, err := fhttp.NewRequest("GET", "https://chat.z.ai/api/v1/auths/", nil)
	if err != nil {
		return false
	}

	ApplyBrowserFingerprintHeaders(req.Header)
	req.Header.Set("Accept", "application/json")
	req.Header.Set("Accept-Language", "zh-CN")
	req.Header.Set("Cache-Control", "no-cache")
	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("DNT", "1")
	req.Header.Set("Pragma", "no-cache")
	req.Header.Set("Referer", "https://chat.z.ai/")
	req.Header.Set("Sec-Fetch-Dest", "empty")
	req.Header.Set("Sec-Fetch-Mode", "cors")
	req.Header.Set("Sec-Fetch-Site", "same-origin")
	req.Header.Set("sec-gpc", "1")
	req.AddCookie(&fhttp.Cookie{Name: "token", Value: token})

	client, err := TLSHTTPClient(10 * time.Second)
	if err != nil {
		LogDebug("Token 验证 tls client: %v", err)
		return false
	}
	resp, err := client.Do(req)
	if err != nil {
		LogDebug("Token 验证请求失败: %v", err)
		return false
	}
	defer resp.Body.Close()

	// 读取响应
	body, _ := io.ReadAll(resp.Body)

	if resp.StatusCode != http.StatusOK {
		LogDebug("Token 验证失败，状态码: %d", resp.StatusCode)
		return false
	}

	// 尝试解析响应获取新 token
	var authResp struct {
		Token string `json:"token"`
		Email string `json:"email"`
		ID    string `json:"id"`
	}
	if err := json.Unmarshal(body, &authResp); err == nil && authResp.Token != "" {
		// 更新 token 信息
		tm.mu.Lock()
		if info, exists := tm.tokens[token]; exists {
			if authResp.Email != "" {
				info.Email = authResp.Email
			}
			if authResp.ID != "" {
				info.UserID = authResp.ID
			}
		}
		tm.mu.Unlock()
	}

	return true
}

// rebuildValidTokens 重建有效 token 列表
func (tm *TokenManager) rebuildValidTokens() {
	tm.mu.Lock()
	defer tm.mu.Unlock()

	tm.validTokens = make([]string, 0)
	for token, info := range tm.tokens {
		if info.Valid {
			tm.validTokens = append(tm.validTokens, token)
		}
	}
}

// removeInvalidTokens 从文件中移除失效 token
func (tm *TokenManager) removeInvalidTokens() {
	tm.mu.Lock()
	defer tm.mu.Unlock()

	tokenFile := filepath.Join(tm.dataDir, "tokens.txt")
	invalidFile := filepath.Join(tm.dataDir, "tokens_invalid.txt")

	// 收集失效 token
	var invalidTokens []string
	for token, info := range tm.tokens {
		if !info.Valid {
			invalidTokens = append(invalidTokens, token)
			delete(tm.tokens, token)
		}
	}

	if len(invalidTokens) == 0 {
		return
	}

	// 追加到失效文件
	f, err := os.OpenFile(invalidFile, os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0644)
	if err == nil {
		defer f.Close()
		timestamp := time.Now().Format("2006-01-02 15:04:05")
		for _, token := range invalidTokens {
			f.WriteString(fmt.Sprintf("# 失效于 %s\n%s\n", timestamp, token))
		}
	}

	// 重写有效 token 文件
	var validTokenLines []string
	for _, token := range tm.validTokens {
		validTokenLines = append(validTokenLines, token)
	}

	content := "# 用户 Token 文件（自动更新）\n"
	content += fmt.Sprintf("# 更新时间: %s\n", time.Now().Format("2006-01-02 15:04:05"))
	content += "# 失效 token 已移至 tokens_invalid.txt\n\n"
	content += strings.Join(validTokenLines, "\n")
	if len(validTokenLines) > 0 {
		content += "\n"
	}

	os.WriteFile(tokenFile, []byte(content), 0644)
	LogInfo("已移除 %d 个失效 token 到 %s", len(invalidTokens), invalidFile)
}

// HasValidUpstreamTokens 是否存在可用的 z.ai 上游 token（TokenManager 轮询用）
func (tm *TokenManager) HasValidUpstreamTokens() bool {
	tm.mu.RLock()
	defer tm.mu.RUnlock()
	return len(tm.validTokens) > 0
}

// GetToken 获取一个有效 token（轮询）
func (tm *TokenManager) GetToken() string {
	tm.mu.Lock()
	defer tm.mu.Unlock()

	if len(tm.validTokens) == 0 {
		return ""
	}

	token := tm.validTokens[tm.currentIndex%len(tm.validTokens)]
	tm.currentIndex++

	// 增加使用计数
	if info, exists := tm.tokens[token]; exists {
		info.UseCount++
	}

	return token
}

// RecordCall 记录调用
func (tm *TokenManager) RecordCall(success bool, isMultimodal bool) {
	atomic.AddInt64(&tm.totalCalls, 1)
	if success {
		atomic.AddInt64(&tm.successCalls, 1)
	}
	if isMultimodal {
		atomic.AddInt64(&tm.multimodalCount, 1)
	}
}

// GetStats 获取统计数据
func (tm *TokenManager) GetStats() TokenManagerStats {
	tm.mu.RLock()
	defer tm.mu.RUnlock()

	total := atomic.LoadInt64(&tm.totalCalls)
	success := atomic.LoadInt64(&tm.successCalls)
	multimodal := atomic.LoadInt64(&tm.multimodalCount)

	var successRate float64
	if total > 0 {
		successRate = float64(success) / float64(total) * 100
	}

	return TokenManagerStats{
		ValidTokenCount: len(tm.validTokens),
		TotalTokenCount: len(tm.tokens),
		MultimodalCount: multimodal,
		TotalCalls:      total,
		SuccessCalls:    success,
		SuccessRate:     successRate,
	}
}

// TokenManagerStats token 管理器统计数据
type TokenManagerStats struct {
	ValidTokenCount int     `json:"valid_token_count"`
	TotalTokenCount int     `json:"total_token_count"`
	MultimodalCount int64   `json:"multimodal_count"`
	TotalCalls      int64   `json:"total_calls"`
	SuccessCalls    int64   `json:"success_calls"`
	SuccessRate     float64 `json:"success_rate"`
}

// GetClientIP 从请求中获取客户端 IP
func GetClientIP(r *http.Request) string {
	// 优先从 X-Forwarded-For 获取
	xff := r.Header.Get("X-Forwarded-For")
	if xff != "" {
		// X-Forwarded-For 可能包含多个 IP，取第一个
		ips := strings.Split(xff, ",")
		if len(ips) > 0 {
			ip := strings.TrimSpace(ips[0])
			if ip != "" {
				return ip
			}
		}
	}

	// 尝试 X-Real-IP
	xri := r.Header.Get("X-Real-IP")
	if xri != "" {
		return xri
	}

	// 最后使用 RemoteAddr
	ip := r.RemoteAddr
	// 去除端口
	if idx := strings.LastIndex(ip, ":"); idx != -1 {
		ip = ip[:idx]
	}
	return ip
}

	package internal

	import (
	"bufio"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"os"

	fhttp "github.com/bogdanfinn/fhttp"
	"path/filepath"
	"strings"
	"sync"
	"sync/atomic"
	"time"

	"github.com/fsnotify/fsnotify"
	)

	// TokenInfo 存储单个 token 的信息
	type TokenInfo struct {
	Token string `json:"token"`
	Email string `json:"email"`
	UserID string `json:"user_id"`
	Valid bool `json:"valid"`
	LastChecked time.Time `json:"last_checked"`
	UseCount int64 `json:"use_count"`
	}

	// TokenManager 管理所有用户 token
	type TokenManager struct {
	mu sync.RWMutex
	tokens map[string]*TokenInfo // token -> TokenInfo
	validTokens []string // 有效 token 列表
	currentIndex int // 轮询索引
	dataDir string
	watcher *fsnotify.Watcher
	checkInterval time.Duration
	stopChan chan struct{}
	multimodalCount int64 // 多模态请求计数
	totalCalls int64 // 累计调用次数
	successCalls int64 // 成功调用次数
	}

	var (
	tokenManager *TokenManager
	tokenOnce sync.Once
	)

	// GetTokenManager 获取单例 TokenManager
	func GetTokenManager() *TokenManager {
	tokenOnce.Do(func() {
	tokenManager = &TokenManager{
	tokens: make(map[string]*TokenInfo),
	validTokens: make([]string, 0),
	dataDir: "data",
	checkInterval: 5 * time.Minute, // 每5分钟检查一次
	stopChan: make(chan struct{}),
	}
	})
	return tokenManager
	}

	// Start 启动 token 管理器
	func (tm *TokenManager) Start() error {
	// 确保 data 目录存在
	if err := os.MkdirAll(tm.dataDir, 0755); err != nil {
	return fmt.Errorf("创建 data 目录失败: %v", err)
	}

	// 初始加载 token
	if err := tm.loadTokens(); err != nil {
	LogWarn("初始加载 token 失败: %v", err)
	}

	// 启动文件监听
	if err := tm.startWatcher(); err != nil {
	LogWarn("启动文件监听失败: %v", err)
	}

	// 启动定期验证
	go tm.startValidator()

	LogInfo("TokenManager 已启动，当前有效 token 数: %d", len(tm.validTokens))
	return nil
	}

	// Stop 停止 token 管理器
	func (tm *TokenManager) Stop() {
	close(tm.stopChan)
	if tm.watcher != nil {
	tm.watcher.Close()
	}
	}

	// loadTokens 从 data 目录加载所有 token
	func (tm *TokenManager) loadTokens() error {
	tokenFile := filepath.Join(tm.dataDir, "tokens.txt")

	file, err := os.Open(tokenFile)
	if err != nil {
	if os.IsNotExist(err) {
	// 创建示例文件
	tm.createExampleTokenFile(tokenFile)
	return nil
	}
	return err
	}
	defer file.Close()

	tm.mu.Lock()
	defer tm.mu.Unlock()

	// 保留旧的统计数据
	oldTokens := tm.tokens
	tm.tokens = make(map[string]*TokenInfo)
	tm.validTokens = make([]string, 0)

	scanner := bufio.NewScanner(file)
	for scanner.Scan() {
	line := strings.TrimSpace(scanner.Text())
	if line == "" \|\| strings.HasPrefix(line, "#") {
	continue
	}

	token := line
	// 支持 token=xxx 格式
	if strings.HasPrefix(line, "token=") {
	token = strings.TrimPrefix(line, "token=")
	}

	if token == "" {
	continue
	}

	// 复用旧的 TokenInfo 如果存在
	if oldInfo, exists := oldTokens[token]; exists {
	tm.tokens[token] = oldInfo
	if oldInfo.Valid {
	tm.validTokens = append(tm.validTokens, token)
	}
	} else {
	// 新 token，解析并标记为待验证
	info := &TokenInfo{
	Token: token,
	Valid: true, // 初始假设有效，验证时会更新
	}
	// 尝试解析 JWT 获取信息
	if payload, err := DecodeJWTPayload(token); err == nil && payload != nil {
	info.Email = payload.Email
	info.UserID = payload.ID
	}
	tm.tokens[token] = info
	tm.validTokens = append(tm.validTokens, token)
	}
	}

	validN := len(tm.validTokens)
	LogInfo("已加载 %d 个 token", validN)
	scanErr := scanner.Err()
	if validN > 0 {
	invalidateAnonymousPoolSlots()
	}
	return scanErr
	}

	// createExampleTokenFile 创建示例 token 文件
	func (tm *TokenManager) createExampleTokenFile(path string) {
	content := `# 用户 Token 文件
	# 每行一个 token，支持以下格式：
	# 1. 直接写 token
	# 2. token=xxx 格式
	# 以 # 开头的行为注释

	# 示例:
	# eyJhbGciOiJFUzI1NiIsInR5cCI6IkpXVCJ9.xxxxx
	`
	os.WriteFile(path, []byte(content), 0644)
	LogInfo("已创建示例 token 文件: %s", path)
	}

	// startWatcher 启动文件变化监听
	func (tm *TokenManager) startWatcher() error {
	watcher, err := fsnotify.NewWatcher()
	if err != nil {
	return err
	}
	tm.watcher = watcher

	go func() {
	for {
	select {
	case event, ok := <-watcher.Events:
	if !ok {
	return
	}
	if event.Op&(fsnotify.Write\|fsnotify.Create) != 0 {
	if strings.HasSuffix(event.Name, "tokens.txt") {
	LogInfo("检测到 token 文件变化，重新加载...")
	time.Sleep(100 * time.Millisecond) // 等待文件写入完成
	tm.loadTokens()
	}
	}
	case err, ok := <-watcher.Errors:
	if !ok {
	return
	}
	LogError("文件监听错误: %v", err)
	case <-tm.stopChan:
	return
	}
	}
	}()

	return watcher.Add(tm.dataDir)
	}

	// startValidator 启动定期验证
	func (tm *TokenManager) startValidator() {
	// 首次延迟验证
	time.Sleep(10 * time.Second)
	tm.validateAllTokens()

	ticker := time.NewTicker(tm.checkInterval)
	defer ticker.Stop()

	for {
	select {
	case <-ticker.C:
	tm.validateAllTokens()
	case <-tm.stopChan:
	return
	}
	}
	}

	// validateAllTokens 验证所有 token
	func (tm *TokenManager) validateAllTokens() {
	tm.mu.RLock()
	tokens := make([]string, 0, len(tm.tokens))
	for token := range tm.tokens {
	tokens = append(tokens, token)
	}
	tm.mu.RUnlock()

	LogInfo("开始验证 %d 个 token...", len(tokens))
	invalidCount := 0

	for _, token := range tokens {
	valid := tm.validateToken(token)
	tm.mu.Lock()
	if info, exists := tm.tokens[token]; exists {
	info.Valid = valid
	info.LastChecked = time.Now()
	if !valid {
	invalidCount++
	}
	}
	tm.mu.Unlock()
	time.Sleep(500 * time.Millisecond) // 避免请求过快
	}

	// 更新有效 token 列表
	tm.rebuildValidTokens()
	LogInfo("Token 验证完成，失效 %d 个，剩余有效 %d 个", invalidCount, len(tm.validTokens))

	// 自动删除失效 token
	if invalidCount > 0 {
	tm.removeInvalidTokens()
	}
	}

	// validateToken 验证单个 token
	func (tm *TokenManager) validateToken(token string) bool {
	req, err := fhttp.NewRequest("GET", "https://chat.z.ai/api/v1/auths/", nil)
	if err != nil {
	return false
	}

	ApplyBrowserFingerprintHeaders(req.Header)
	req.Header.Set("Accept", "application/json")
	req.Header.Set("Accept-Language", "zh-CN")
	req.Header.Set("Cache-Control", "no-cache")
	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("DNT", "1")
	req.Header.Set("Pragma", "no-cache")
	req.Header.Set("Referer", "https://chat.z.ai/")
	req.Header.Set("Sec-Fetch-Dest", "empty")
	req.Header.Set("Sec-Fetch-Mode", "cors")
	req.Header.Set("Sec-Fetch-Site", "same-origin")
	req.Header.Set("sec-gpc", "1")
	req.AddCookie(&fhttp.Cookie{Name: "token", Value: token})

	client, err := TLSHTTPClient(10 * time.Second)
	if err != nil {
	LogDebug("Token 验证 tls client: %v", err)
	return false
	}
	resp, err := client.Do(req)
	if err != nil {
	LogDebug("Token 验证请求失败: %v", err)
	return false
	}
	defer resp.Body.Close()

	// 读取响应
	body, _ := io.ReadAll(resp.Body)

	if resp.StatusCode != http.StatusOK {
	LogDebug("Token 验证失败，状态码: %d", resp.StatusCode)
	return false
	}

	// 尝试解析响应获取新 token
	var authResp struct {
	Token string `json:"token"`
	Email string `json:"email"`
	ID string `json:"id"`
	}
	if err := json.Unmarshal(body, &authResp); err == nil && authResp.Token != "" {
	// 更新 token 信息
	tm.mu.Lock()
	if info, exists := tm.tokens[token]; exists {
	if authResp.Email != "" {
	info.Email = authResp.Email
	}
	if authResp.ID != "" {
	info.UserID = authResp.ID
	}
	}
	tm.mu.Unlock()
	}

	return true
	}

	// rebuildValidTokens 重建有效 token 列表
	func (tm *TokenManager) rebuildValidTokens() {
	tm.mu.Lock()
	defer tm.mu.Unlock()

	tm.validTokens = make([]string, 0)
	for token, info := range tm.tokens {
	if info.Valid {
	tm.validTokens = append(tm.validTokens, token)
	}
	}
	}

	// removeInvalidTokens 从文件中移除失效 token
	func (tm *TokenManager) removeInvalidTokens() {
	tm.mu.Lock()
	defer tm.mu.Unlock()

	tokenFile := filepath.Join(tm.dataDir, "tokens.txt")
	invalidFile := filepath.Join(tm.dataDir, "tokens_invalid.txt")

	// 收集失效 token
	var invalidTokens []string
	for token, info := range tm.tokens {
	if !info.Valid {
	invalidTokens = append(invalidTokens, token)
	delete(tm.tokens, token)
	}
	}

	if len(invalidTokens) == 0 {
	return
	}

	// 追加到失效文件
	f, err := os.OpenFile(invalidFile, os.O_APPEND\|os.O_CREATE\|os.O_WRONLY, 0644)
	if err == nil {
	defer f.Close()
	timestamp := time.Now().Format("2006-01-02 15:04:05")
	for _, token := range invalidTokens {
	f.WriteString(fmt.Sprintf("# 失效于 %s\n%s\n", timestamp, token))
	}
	}

	// 重写有效 token 文件
	var validTokenLines []string
	for _, token := range tm.validTokens {
	validTokenLines = append(validTokenLines, token)
	}

	content := "# 用户 Token 文件（自动更新）\n"
	content += fmt.Sprintf("# 更新时间: %s\n", time.Now().Format("2006-01-02 15:04:05"))
	content += "# 失效 token 已移至 tokens_invalid.txt\n\n"
	content += strings.Join(validTokenLines, "\n")
	if len(validTokenLines) > 0 {
	content += "\n"
	}

	os.WriteFile(tokenFile, []byte(content), 0644)
	LogInfo("已移除 %d 个失效 token 到 %s", len(invalidTokens), invalidFile)
	}

	// HasValidUpstreamTokens 是否存在可用的 z.ai 上游 token（TokenManager 轮询用）
	func (tm *TokenManager) HasValidUpstreamTokens() bool {
	tm.mu.RLock()
	defer tm.mu.RUnlock()
	return len(tm.validTokens) > 0
	}

	// GetToken 获取一个有效 token（轮询）
	func (tm *TokenManager) GetToken() string {
	tm.mu.Lock()
	defer tm.mu.Unlock()

	if len(tm.validTokens) == 0 {
	return ""
	}

	token := tm.validTokens[tm.currentIndex%len(tm.validTokens)]
	tm.currentIndex++

	// 增加使用计数
	if info, exists := tm.tokens[token]; exists {
	info.UseCount++
	}

	return token
	}

	// RecordCall 记录调用
	func (tm *TokenManager) RecordCall(success bool, isMultimodal bool) {
	atomic.AddInt64(&tm.totalCalls, 1)
	if success {
	atomic.AddInt64(&tm.successCalls, 1)
	}
	if isMultimodal {
	atomic.AddInt64(&tm.multimodalCount, 1)
	}
	}

	// GetStats 获取统计数据
	func (tm *TokenManager) GetStats() TokenManagerStats {
	tm.mu.RLock()
	defer tm.mu.RUnlock()

	total := atomic.LoadInt64(&tm.totalCalls)
	success := atomic.LoadInt64(&tm.successCalls)
	multimodal := atomic.LoadInt64(&tm.multimodalCount)

	var successRate float64
	if total > 0 {
	successRate = float64(success) / float64(total) * 100
	}

	return TokenManagerStats{
	ValidTokenCount: len(tm.validTokens),
	TotalTokenCount: len(tm.tokens),
	MultimodalCount: multimodal,
	TotalCalls: total,
	SuccessCalls: success,
	SuccessRate: successRate,
	}
	}

	// TokenManagerStats token 管理器统计数据
	type TokenManagerStats struct {
	ValidTokenCount int `json:"valid_token_count"`
	TotalTokenCount int `json:"total_token_count"`
	MultimodalCount int64 `json:"multimodal_count"`
	TotalCalls int64 `json:"total_calls"`
	SuccessCalls int64 `json:"success_calls"`
	SuccessRate float64 `json:"success_rate"`
	}

	// GetClientIP 从请求中获取客户端 IP
	func GetClientIP(r *http.Request) string {
	// 优先从 X-Forwarded-For 获取
	xff := r.Header.Get("X-Forwarded-For")
	if xff != "" {
	// X-Forwarded-For 可能包含多个 IP，取第一个
	ips := strings.Split(xff, ",")
	if len(ips) > 0 {
	ip := strings.TrimSpace(ips[0])
	if ip != "" {
	return ip
	}
	}
	}

	// 尝试 X-Real-IP
	xri := r.Header.Get("X-Real-IP")
	if xri != "" {
	return xri
	}

	// 最后使用 RemoteAddr
	ip := r.RemoteAddr
	// 去除端口
	if idx := strings.LastIndex(ip, ":"); idx != -1 {
	ip = ip[:idx]
	}
	return ip
	}