package main import ( "bufio" "bytes" "context" "crypto/tls" "encoding/json" "flag" "fmt" "io" "log" "net/http" "net/http/cookiejar" "net/url" "os" "os/signal" "path/filepath" "regexp" "sort" "strings" "sync" "sync/atomic" "syscall" "time" "github.com/google/uuid" "github.com/skip2/go-qrcode" ) const ( version = "1.4.1" upstreamBase = "https://copilot.tencent.com" clientUA = "CLI/2.143.1 CodeBuddy/2.143.1" originReferer = "https://www.codebuddy.cn" endpointAuthState = upstreamBase + "/v2/plugin/auth/state?platform=VSCode" endpointLoginAcct = upstreamBase + "/v2/plugin/login/account?state=" endpointAuthToken = upstreamBase + "/v2/plugin/auth/token?state=" endpointTokenRefresh = upstreamBase + "/v2/plugin/auth/token/refresh" endpointChat = upstreamBase + "/v2/chat/completions" loginTTL = 5 * time.Minute ) // ----------------------------------------------------------------------------- // 数据结构定义 // ----------------------------------------------------------------------------- type StoredAuth struct { Auth StoredTokens `json:"auth"` Account StoredAccount `json:"account"` } type StoredTokens struct { AccessToken string `json:"accessToken"` RefreshToken string `json:"refreshToken"` ExpiresAt int64 `json:"expiresAt"` Domain string `json:"domain"` } type StoredAccount struct { UID string `json:"uid"` EnterpriseID string `json:"enterpriseId"` Nickname string `json:"nickname"` } type apiEnvelope struct { Code int `json:"code"` Msg string `json:"msg"` Data json.RawMessage `json:"data"` } type tokenData struct { AccessToken string `json:"accessToken"` RefreshToken string `json:"refreshToken"` ExpiresIn int64 `json:"expiresIn"` RefreshExpiresIn int64 `json:"refreshExpiresIn"` Domain string `json:"domain"` } type accountData struct { UID string `json:"uid"` EnterpriseID string `json:"enterpriseId"` Nickname string `json:"nickname"` } type authStateData struct { State string `json:"state"` AuthURL string `json:"authUrl"` } // ----------------------------------------------------------------------------- // 全局状态与配置 // ----------------------------------------------------------------------------- type Config struct { Addr string Port int AuthFile string AuthDir string AuthExplicit bool // 用户是否显式指定了 -auth(未指定时自动扫描目录下所有 workbuddy*.json) APIKey string ProxyURL string Verbose bool HttpClient *http.Client } // Account 表示一个 CodeBuddy 账号凭据及其运行时状态。 type Account struct { Path string // 凭据文件路径 Auth *StoredAuth // 凭据内容(失效标记恢复时可能为 nil) CooldownUntil time.Time // 冷却截止时间(429 频率限制后自动屏蔽到该时间) CooldownMsg string // 触发冷却的原因/上游提示 Disabled bool // 授权失效(token 过期且无法刷新/已撤销):禁止调度 DisabledReason string // 失效原因,用于控制台展示 Nickname string // 失效标记持久化字段:昵称(Auth 为 nil 时使用) UID string // 失效标记持久化字段:UID lock sync.Mutex // 单账号串行锁(防止同账号并发触发 11128) } // disabledMarker 是授权失效账号的持久化标记(凭据文件删除后用于控制台提示重新登录)。 type disabledMarker struct { Path string `json:"path"` Reason string `json:"reason"` DisabledAt int64 `json:"disabledAt"` Nickname string `json:"nickname,omitempty"` UID string `json:"uid,omitempty"` } // markerPath 返回与凭据文件同目录的失效标记文件路径。 func markerPath(authPath string) string { return authPath + ".disabled" } var ( cfg Config authLock sync.RWMutex currAuth *StoredAuth reqCounter uint64 accountMu sync.Mutex // 账号池保护锁 accounts []*Account // 多账号池(单账号时长度为 1,行为与旧版完全一致) rrIndex int // 轮询游标 ) // ----------------------------------------------------------------------------- // 主入口与命令行控制 // ----------------------------------------------------------------------------- func main() { if len(os.Args) > 1 { first := os.Args[1] if first == "help" || first == "-h" || first == "--help" { printHelp() return } if first == "version" || first == "-v" || first == "--version" { fmt.Printf("WorkBuddy Local Gateway v%s\n", version) return } } command := "serve" args := os.Args[1:] if len(os.Args) > 1 && !strings.HasPrefix(os.Args[1], "-") { command = os.Args[1] args = os.Args[2:] } fs := flag.NewFlagSet(command, flag.ExitOnError) fs.StringVar(&cfg.Addr, "addr", "127.0.0.1", "网关监听地址") fs.IntVar(&cfg.Port, "port", 8317, "网关监听端口") fs.StringVar(&cfg.AuthFile, "auth", "workbuddy.json", "凭据存储文件路径(支持逗号分隔多个文件实现多账号)") fs.StringVar(&cfg.AuthDir, "auth-dir", "", "凭据目录:自动加载目录下所有 workbuddy*.json 作为多账号池") fs.StringVar(&cfg.APIKey, "api-key", "", "可选:访问网关所需的 API Key (客户端 Bearer 校验)") fs.StringVar(&cfg.ProxyURL, "proxy", "", "可选:上游请求代理 (如 http://127.0.0.1:7890)") fs.BoolVar(&cfg.Verbose, "verbose", false, "输出详细调试日志") _ = fs.Parse(args) // 检测 -auth 是否被显式指定: // 若未指定 -auth 且未指定 -auth-dir,则自动扫描当前目录下所有 workbuddy*.json 组成账号池, // 这样把多个凭据文件放进工作目录即可自动多账号,无需手写参数。 fs.Visit(func(f *flag.Flag) { if f.Name == "auth" { cfg.AuthExplicit = true } }) // 初始化 HTTP 客户端 initHTTPClient() switch command { case "serve", "run", "start": runServe() case "login": runLogin() case "status": runStatus() case "refresh": runRefresh() case "version", "-v", "--version": fmt.Printf("WorkBuddy Local Gateway v%s\n", version) default: fmt.Printf("未知命令: %s\n\n", command) printHelp() os.Exit(1) } } func printHelp() { fmt.Println(`WorkBuddy Local Gateway - 轻量化跨平台本地 AI 网关 将腾讯 CodeBuddy / 混元反代为标准 OpenAI 协议,支持直接透传任意模型到上游。 支持多账号池:轮询使用多个账号均衡额度;账号触发 429 频率限制后自动冷却屏蔽, 由其余账号代偿,冷却到期自动恢复。 用法: workbuddy-gateway [command] [options] 命令: serve 启动本地网关 (默认操作) login 微信/企业微信扫码登录,获取/更新凭据 status 查看账号池状态(含冷却状态与过期时间) refresh 手动立即刷新所有账号访问令牌 (Access Token) version 查看版本信息 help 查看帮助说明 参数选项: -addr 网关监听地址 (默认: 127.0.0.1) -port 网关监听端口 (默认: 8317) -auth 凭据文件路径;支持逗号分隔多个文件实现多账号 (默认: 自动发现当前目录下所有 workbuddy*.json) -auth-dir 凭据目录:自动加载目录下所有 workbuddy*.json 作为账号池 -api-key 设置后,调用网关必须携带 Bearer 鉴权 -proxy 设置上游转发代理 (例如 http://127.0.0.1:7890 或 socks5://...) -verbose 输出详细调试日志 (请求/响应体) 多账号说明: # 登录第二个账号(保存到不同文件) workbuddy-gateway login -auth workbuddy2.json # 自动发现:把多个凭据文件放进工作目录即可自动多账号(无需任何参数) workbuddy-gateway serve # 自动加载 ./workbuddy*.json # 启动时指定多个凭据文件(轮询 + 429 自动冷却代偿) workbuddy-gateway serve -auth workbuddy.json,workbuddy2.json # 或使用目录模式:目录内所有 workbuddy*.json 自动组成账号池 workbuddy-gateway serve -auth-dir ./auths 示例: # 首次使用扫码登录 workbuddy-gateway login # 启动本地网关 (监听 127.0.0.1:8317) workbuddy-gateway serve # 启动网关并指定端口和代理 workbuddy-gateway serve -port 9000 -proxy http://127.0.0.1:7890`) } func initHTTPClient() { jar, _ := cookiejar.New(nil) transport := &http.Transport{ MaxIdleConns: 50, IdleConnTimeout: 90 * time.Second, MaxIdleConnsPerHost: 10, TLSClientConfig: &tls.Config{InsecureSkipVerify: false}, } if cfg.ProxyURL != "" { pURL, err := url.Parse(cfg.ProxyURL) if err != nil { log.Fatalf("错误: 无效的代理地址 %s: %v", cfg.ProxyURL, err) } transport.Proxy = http.ProxyURL(pURL) } cfg.HttpClient = &http.Client{ Timeout: 180 * time.Second, Transport: transport, Jar: jar, } } // ----------------------------------------------------------------------------- // 凭据加载、保存与自动续期 // ----------------------------------------------------------------------------- func loadAuth() (*StoredAuth, error) { authLock.RLock() if currAuth != nil { defer authLock.RUnlock() return currAuth, nil } authLock.RUnlock() data, err := os.ReadFile(cfg.AuthFile) if err != nil { return nil, fmt.Errorf("读取凭据文件失败 (%s): %w", cfg.AuthFile, err) } var sa StoredAuth if err := json.Unmarshal(data, &sa); err != nil { return nil, fmt.Errorf("解析凭据文件失败: %w", err) } if sa.Auth.AccessToken == "" { return nil, fmt.Errorf("凭据文件中缺少 AccessToken") } authLock.Lock() currAuth = &sa authLock.Unlock() return &sa, nil } func saveAuth(sa *StoredAuth) error { authLock.Lock() currAuth = sa authLock.Unlock() dir := filepath.Dir(cfg.AuthFile) if dir != "" && dir != "." { _ = os.MkdirAll(dir, 0755) } data, err := json.MarshalIndent(sa, "", " ") if err != nil { return err } if err := os.WriteFile(cfg.AuthFile, data, 0600); err != nil { return err } clearDisabledMarker(cfg.AuthFile) return nil } // saveAuthTo 将凭据写入指定路径(多账号模式使用)。 // 写入成功后清除该路径的失效标记(表示账号已重新登录)。 func saveAuthTo(path string, sa *StoredAuth) error { dir := filepath.Dir(path) if dir != "" && dir != "." { _ = os.MkdirAll(dir, 0755) } data, err := json.MarshalIndent(sa, "", " ") if err != nil { return err } if err := os.WriteFile(path, data, 0600); err != nil { return err } clearDisabledMarker(path) return nil } // loadAccountFile 从指定路径读取并解析凭据(不修改全局缓存)。 func loadAccountFile(path string) (*StoredAuth, error) { data, err := os.ReadFile(path) if err != nil { return nil, fmt.Errorf("读取凭据文件失败 (%s): %w", path, err) } var sa StoredAuth if err := json.Unmarshal(data, &sa); err != nil { return nil, fmt.Errorf("解析凭据文件失败 (%s): %w", path, err) } if sa.Auth.AccessToken == "" { return nil, fmt.Errorf("凭据文件缺少 AccessToken (%s)", path) } return &sa, nil } // collectConfiguredAuthPaths 解析应加载的凭据路径列表: // 1) -auth-dir 指定目录 → 目录下所有 workbuddy*.json // 2) 显式 -auth → 逗号分隔的凭据文件列表(保持用户顺序) // 3) 均未指定(自动发现模式)→ 扫描当前目录下所有 workbuddy*.json, // 这样把多个凭据文件放进工作目录即可自动组成多账号池;若一个都没有则回退默认 cfg.AuthFile func collectConfiguredAuthPaths() []string { var paths []string if cfg.AuthDir != "" { entries, err := os.ReadDir(cfg.AuthDir) if err == nil { for _, e := range entries { if e.IsDir() { continue } name := e.Name() if strings.HasPrefix(name, "workbuddy") && strings.HasSuffix(name, ".json") { paths = append(paths, filepath.Join(cfg.AuthDir, name)) } } sort.Strings(paths) } return paths } if cfg.AuthExplicit { for _, p := range strings.Split(cfg.AuthFile, ",") { p = strings.TrimSpace(p) if p != "" { paths = append(paths, p) } } return paths } // 自动发现模式:扫描当前工作目录(systemd WorkingDirectory 即凭据目录) entries, err := os.ReadDir(".") if err == nil { for _, e := range entries { if e.IsDir() { continue } name := e.Name() if strings.HasPrefix(name, "workbuddy") && strings.HasSuffix(name, ".json") { paths = append(paths, name) } } sort.Strings(paths) } if len(paths) == 0 { // 无任何凭据文件时回退默认路径,便于给出"请先 login"的友好提示 paths = append(paths, cfg.AuthFile) } return paths } // loadAccounts 构建账号池。 // -auth 支持逗号分隔多个凭据文件;-auth-dir 自动加载目录下所有 workbuddy*.json; // 两者都未指定时自动发现当前目录下所有 workbuddy*.json(多账号免参数)。 func loadAccounts() error { accountMu.Lock() defer accountMu.Unlock() accounts = nil rrIndex = 0 for _, p := range collectConfiguredAuthPaths() { sa, err := loadAccountFile(p) if err != nil { log.Printf("[Auth] 跳过无效凭据文件 %s: %v", p, err) continue } accounts = append(accounts, &Account{Path: p, Auth: sa}) } // 恢复持久化的失效账号标记(凭据文件已删除,仅供控制台提示重新登录) loadDisabledMarkers() if len(accounts) == 0 { return fmt.Errorf("未找到有效凭据,请先执行 login 命令扫码登录") } return nil } // nextAccount 轮询选择下一个未处于冷却状态且未被禁用的账号。 // 若全部账号不可用,返回错误并附带说明。 func nextAccount() (*Account, error) { accountMu.Lock() defer accountMu.Unlock() if len(accounts) == 0 { return nil, fmt.Errorf("账号池为空") } now := time.Now() disabledCount := 0 cooldownCount := 0 for i := 0; i < len(accounts); i++ { idx := (rrIndex + i) % len(accounts) acc := accounts[idx] if acc.Disabled { disabledCount++ continue } if acc.CooldownUntil.After(now) { cooldownCount++ continue } rrIndex = (idx + 1) % len(accounts) return acc, nil } // 全部不可用:优先报告失效账号 for _, a := range accounts { if a.Disabled { return nil, fmt.Errorf("账号 %s 已失效(授权过期或已撤销,凭据文件已删除),请重新执行 login 登录", a.Path) } } // 全部冷却:返回最早解封时间 earliest := accounts[0].CooldownUntil for _, a := range accounts { if a.CooldownUntil.Before(earliest) { earliest = a.CooldownUntil } } return nil, fmt.Errorf("所有账号均处于冷却状态,最早解封时间: %s", earliest.Format("2006-01-02 15:04:05")) } // markCooldown 将账号屏蔽至指定时间。 func markCooldown(acc *Account, until time.Time, msg string) { accountMu.Lock() acc.CooldownUntil = until acc.CooldownMsg = msg accountMu.Unlock() log.Printf("[Cooldown] 账号 %s 触发频率限制,自动屏蔽至 %s (提示: %s)", acc.Path, until.Format("2006-01-02 15:04:05"), msg) } // disableAccount 将账号标记为失效(授权过期/撤销),禁止调度并删除凭据文件, // 同时写入持久化失效标记,便于控制台提示用户重新登录。 func disableAccount(acc *Account, reason string) { accountMu.Lock() acc.Disabled = true acc.DisabledReason = reason acc.CooldownUntil = time.Time{} path := acc.Path nickname := "" uid := "" if acc.Auth != nil { nickname = acc.Auth.Account.Nickname uid = acc.Auth.Account.UID acc.Nickname = nickname acc.UID = uid } accountMu.Unlock() // 写持久化失效标记(凭据删除后仍能在 status 中提示重新登录) marker := disabledMarker{ Path: path, Reason: reason, DisabledAt: time.Now().Unix(), Nickname: nickname, UID: uid, } if data, err := json.MarshalIndent(marker, "", " "); err == nil { if werr := os.WriteFile(markerPath(path), data, 0600); werr != nil { log.Printf("[Auth] 账号 %s 失效标记写入失败: %v", path, werr) } } // 删除失效的凭据文件,方便用户下次重新登录 if err := os.Remove(path); err != nil && !os.IsNotExist(err) { log.Printf("[Auth] 账号 %s 凭据文件删除失败: %v", path, err) } log.Printf("[Auth] ⚠️ 账号 %s 授权失效,已禁止调度并删除凭据文件: %s", path, reason) } // loadDisabledMarkers 扫描失效标记文件,将其恢复为账号池中的失效账号(Auth 为 nil)。 func loadDisabledMarkers() { // 收集标记文件路径: // - AuthDir 模式:目录下 *.json.disabled 即为标记文件 // - 显式 -auth 模式:<凭据路径>.disabled 为标记文件 // - 自动发现模式:扫描当前目录下所有 *.json.disabled 标记文件 var markerPaths []string if cfg.AuthDir != "" { entries, err := os.ReadDir(cfg.AuthDir) if err != nil { return } for _, e := range entries { if e.IsDir() { continue } name := e.Name() if strings.HasSuffix(name, ".json.disabled") { markerPaths = append(markerPaths, filepath.Join(cfg.AuthDir, name)) } } } else if cfg.AuthExplicit { for _, p := range strings.Split(cfg.AuthFile, ",") { p = strings.TrimSpace(p) if p != "" { markerPaths = append(markerPaths, markerPath(p)) } } } else { // 自动发现模式:扫描当前工作目录 entries, err := os.ReadDir(".") if err == nil { for _, e := range entries { if e.IsDir() { continue } name := e.Name() if strings.HasPrefix(name, "workbuddy") && strings.HasSuffix(name, ".json.disabled") { markerPaths = append(markerPaths, name) } } sort.Strings(markerPaths) } } for _, mp := range markerPaths { data, err := os.ReadFile(mp) if err != nil { continue } var m disabledMarker if err := json.Unmarshal(data, &m); err != nil { continue } // 避免与已加载的有效账号重复 dup := false for _, a := range accounts { if a.Path == m.Path { dup = true break } } if dup { continue } accounts = append(accounts, &Account{ Path: m.Path, Disabled: true, DisabledReason: m.Reason, Nickname: m.Nickname, UID: m.UID, }) } } // clearDisabledMarker 在重新登录成功后清除失效标记。 func clearDisabledMarker(path string) { mp := markerPath(path) if err := os.Remove(mp); err != nil && !os.IsNotExist(err) { log.Printf("[Auth] 失效标记清除失败 (%s): %v", mp, err) } } // isAuthFailure 判断上游响应是否为授权失效(401/403 / invalid token / 登录过期等)。 func isAuthFailure(statusCode int, body string) bool { if statusCode == http.StatusUnauthorized || statusCode == http.StatusForbidden { return true } low := strings.ToLower(body) if strings.Contains(low, "invalid token") || strings.Contains(low, "unauthorized") || strings.Contains(low, "登录已过期") || strings.Contains(low, "登录失效") || strings.Contains(low, "token 已失效") || strings.Contains(low, "authentication required") { return true } return false } var resetTimeRe = regexp.MustCompile(`将在\s*(\d{4}-\d{2}-\d{2})\s+(\d{2}:\d{2}:\d{2})\s*(UTC[+-]\d+(?::\d{2})?)?`) // parseResetTime 从上游 429 错误消息中解析频率限制重置时间。 // 示例: "您的使用量已超出频率限制,将在 2026-09-04 07:48:15 UTC+8 重置" func parseResetTime(s string) (time.Time, bool) { m := resetTimeRe.FindStringSubmatch(s) if m == nil { return time.Time{}, false } offset := 8 * 3600 // 默认按 UTC+8 解析 if m[3] != "" { zone := m[3] zone = strings.TrimPrefix(zone, "UTC") zone = strings.TrimPrefix(zone, "utc") var h, mi int if _, err := fmt.Sscanf(zone, "%d:%d", &h, &mi); err == nil { offset = h*3600 + mi*60 } else if _, err := fmt.Sscanf(zone, "%d", &h); err == nil { offset = h * 3600 } } loc := time.FixedZone("UTC", offset) t, err := time.ParseInLocation("2006-01-02 15:04:05", m[1]+" "+m[2], loc) if err != nil { return time.Time{}, false } return t, true } // isRateLimited 判断上游响应是否属于频率限制(429 或 code 6004 / 频率限制提示)。 func isRateLimited(statusCode int, body string) bool { if statusCode == http.StatusTooManyRequests { return true } if strings.Contains(body, `"code":6004`) || strings.Contains(body, "频率限制") || strings.Contains(body, "frequency limit") { return true } return false } // ensureValidTokenFor 对指定账号检查并刷新令牌(多账号版)。 // 失效账号(Disabled 或 Auth 为 nil)直接跳过,不参与调度。 func ensureValidTokenFor(acc *Account) error { if acc == nil || acc.Disabled || acc.Auth == nil { return nil } if time.Now().Unix() > acc.Auth.Auth.ExpiresAt-900 { log.Printf("[Auth] 账号 %s 访问令牌即将/已经过期,正在自动刷新...", acc.Path) return doRefreshTokenFor(acc) } return nil } func ensureValidToken() error { sa, err := loadAuth() if err != nil { return err } // 如果距过期不足 15 分钟,则自动刷新 if time.Now().Unix() > sa.Auth.ExpiresAt-900 { log.Printf("[Auth] 访问令牌即将或已经过期 (ExpiresAt=%s),正在自动刷新...", time.Unix(sa.Auth.ExpiresAt, 0).Format("2006-01-02 15:04:05")) return doRefreshToken(sa) } return nil } func doRefreshToken(sa *StoredAuth) error { if sa == nil || sa.Auth.RefreshToken == "" { return fmt.Errorf("无法刷新:缺少 RefreshToken") } if _, err := refreshTokenPayload(sa); err != nil { return err } if err := saveAuth(sa); err != nil { return fmt.Errorf("写回凭据失败: %w", err) } log.Printf("[Auth] Token 刷新成功!新过期时间: %s", time.Unix(sa.Auth.ExpiresAt, 0).Format("2006-01-02 15:04:05")) return nil } // doRefreshTokenFor 刷新指定账号的令牌并保存回其凭据文件(多账号版)。 // 若刷新因授权失效失败(401/403/refresh token 无效),自动禁用该账号并删除凭据文件。 func doRefreshTokenFor(acc *Account) error { if acc == nil || acc.Disabled || acc.Auth == nil || acc.Auth.Auth.RefreshToken == "" { return fmt.Errorf("无法刷新:账号缺少 RefreshToken 或已失效") } status, err := refreshTokenPayload(acc.Auth) if err != nil { if isAuthFailure(status, err.Error()) { disableAccount(acc, fmt.Sprintf("令牌刷新失败 (HTTP %d): %v", status, err)) } return err } if err := saveAuthTo(acc.Path, acc.Auth); err != nil { return fmt.Errorf("写回凭据失败: %w", err) } log.Printf("[Auth] 账号 %s Token 刷新成功!新过期时间: %s", acc.Path, time.Unix(acc.Auth.Auth.ExpiresAt, 0).Format("2006-01-02 15:04:05")) return nil } // refreshTokenPayload 调用上游刷新接口并更新内存中的令牌字段(不落盘)。 // 返回上游 HTTP 状态码(成功或失败时均为实际状态;网络错误为 0)。 func refreshTokenPayload(sa *StoredAuth) (int, error) { headers := func(r *http.Request) { commonHeaders(r) r.Header.Set("X-Refresh-Token", sa.Auth.RefreshToken) if sa.Account.EnterpriseID != "" { r.Header.Set("X-Enterprise-Id", sa.Account.EnterpriseID) } r.Header.Set("X-Auth-Refresh-Source", "workbuddy") } data, status, err := doJSON(cfg.HttpClient, http.MethodPost, endpointTokenRefresh, headers, nil) if err != nil { return status, fmt.Errorf("上游刷新拒绝 (HTTP %d): %w", status, err) } var tok tokenData if err := json.Unmarshal(data, &tok); err != nil || tok.AccessToken == "" { return status, fmt.Errorf("解析新 Token 失败: %w", err) } sa.Auth.AccessToken = tok.AccessToken if tok.RefreshToken != "" { sa.Auth.RefreshToken = tok.RefreshToken } if tok.Domain != "" { sa.Auth.Domain = tok.Domain } if tok.ExpiresIn > 0 { sa.Auth.ExpiresAt = time.Now().Add(time.Duration(tok.ExpiresIn) * time.Second).Unix() } return status, nil } // ----------------------------------------------------------------------------- // CLI 子命令实现: login, status, refresh // ----------------------------------------------------------------------------- // formatAccountStatus 生成单个账号的完整状态文本(status 命令与 serve 启动横幅共用)。 // idx 从 1 开始的账号序号;now 为当前时间(用于冷却/过期判定)。 func formatAccountStatus(acc *Account, idx int, now time.Time) string { var sb strings.Builder sb.WriteString(fmt.Sprintf("\n--- 账号 #%d ---\n", idx)) sb.WriteString(fmt.Sprintf("凭据文件: %s\n", acc.Path)) if acc.Disabled { // 失效账号(Auth 可能为 nil:凭据文件已删除,信息来自失效标记) nickname := acc.Nickname uid := acc.UID if acc.Auth != nil { nickname = acc.Auth.Account.Nickname uid = acc.Auth.Account.UID } sb.WriteString(fmt.Sprintf("用户昵称: %s\n", ifEmpty(nickname, "(未知)"))) sb.WriteString(fmt.Sprintf("用户 UID: %s\n", ifEmpty(uid, "(未知)"))) sb.WriteString("账号状态: ❌ 授权失效(禁止调度)\n") sb.WriteString(fmt.Sprintf("失效原因: %s\n", truncate(acc.DisabledReason, 200))) sb.WriteString(fmt.Sprintf("处理建议: ⚠️ 凭据文件已删除,请重新执行: workbuddy-gateway login -auth %s\n", acc.Path)) return sb.String() } if acc.Auth == nil { return sb.String() } expTime := time.Unix(acc.Auth.Auth.ExpiresAt, 0) remaining := time.Until(expTime) statusStr := "有效" if remaining <= 0 { statusStr = "已过期" } sb.WriteString(fmt.Sprintf("用户昵称: %s\n", acc.Auth.Account.Nickname)) sb.WriteString(fmt.Sprintf("用户 UID: %s\n", acc.Auth.Account.UID)) sb.WriteString(fmt.Sprintf("企业 ID: %s\n", ifEmpty(acc.Auth.Account.EnterpriseID, "(个人账号)"))) sb.WriteString(fmt.Sprintf("认证域名: %s\n", ifEmpty(acc.Auth.Auth.Domain, "www.codebuddy.cn"))) if acc.CooldownUntil.After(now) { sb.WriteString(fmt.Sprintf("冷却状态: 🔒 冷却中 (解封: %s, 剩余 %v)\n", acc.CooldownUntil.Format("2006-01-02 15:04:05"), time.Until(acc.CooldownUntil).Round(time.Minute))) sb.WriteString(fmt.Sprintf("冷却原因: %s\n", truncate(acc.CooldownMsg, 120))) } else { sb.WriteString("冷却状态: ✅ 可用\n") } sb.WriteString(fmt.Sprintf("Token 状态: %s\n", statusStr)) sb.WriteString(fmt.Sprintf("过期时间: %s (剩余 %v)\n", expTime.Format("2006-01-02 15:04:05"), remaining.Round(time.Minute))) return sb.String() } func runStatus() { if err := loadAccounts(); err != nil { fmt.Printf("未找到有效凭据: %v\n请先执行: workbuddy-gateway login 扫码登录。\n", err) return } accountMu.Lock() defer accountMu.Unlock() fmt.Println("================== WorkBuddy 账号池状态 ==================") fmt.Printf("账号总数: %d\n", len(accounts)) now := time.Now() for i, acc := range accounts { fmt.Print(formatAccountStatus(acc, i+1, now)) } fmt.Println("\n=======================================================") } func runRefresh() { if err := loadAccounts(); err != nil { fmt.Printf("读取凭据失败: %v\n", err) return } accountMu.Lock() defer accountMu.Unlock() if len(accounts) == 0 { fmt.Println("账号池为空") return } ok := 0 fail := 0 skipped := 0 for i, acc := range accounts { fmt.Printf("正在刷新账号 #%d (%s)... ", i+1, acc.Path) if acc.Disabled || acc.Auth == nil { fmt.Printf("⏭️ 跳过(授权失效,请重新登录)\n") skipped++ continue } if err := doRefreshTokenFor(acc); err != nil { fmt.Printf("❌ %v\n", err) fail++ } else { fmt.Println("✅") ok++ } } fmt.Printf("\n刷新完成: 成功 %d 个,失败 %d 个,跳过 %d 个(授权失效)\n", ok, fail, skipped) } func runLogin() { fmt.Println("================ WorkBuddy 扫码登录 ================") fmt.Println("正在生成登录凭据与二维码...") loginClient := &http.Client{ Timeout: 30 * time.Second, Jar: cfg.HttpClient.Jar, } data, _, err := doJSON(loginClient, http.MethodPost, endpointAuthState, nil, bytes.NewReader([]byte("{}"))) if err != nil { fmt.Printf("获取登录状态失败: %v\n", err) return } var st authStateData _ = json.Unmarshal(data, &st) if st.State == "" || st.AuthURL == "" { fmt.Println("上游返回的登录状态信息异常,请重试。") return } // 终端字符二维码 qr, err := qrcode.New(st.AuthURL, qrcode.Medium) if err == nil { fmt.Println("\n请使用 微信 或 企业微信 扫描下方二维码登录:") fmt.Println(qr.ToSmallString(false)) } fmt.Println("如无法扫码,也可在浏览器中直接打开以下链接:") fmt.Printf("👉 %s\n\n", st.AuthURL) fmt.Println("等待扫码授权中 (按 Ctrl+C 可取消)...") deadline := time.Now().Add(loginTTL) ticker := time.NewTicker(2 * time.Second) defer ticker.Stop() for { select { case <-ticker.C: if time.Now().After(deadline) { fmt.Println("\n登录已超时,请重新执行 login 命令。") return } tokRaw, _, errTok := doJSON(loginClient, http.MethodGet, endpointAuthToken+st.State, nil, nil) if errTok != nil { continue } var tok tokenData if err := json.Unmarshal(tokRaw, &tok); err != nil || tok.AccessToken == "" { continue } // 登录成功,拉取账号信息 var acct accountData acctHeaders := func(r *http.Request) { commonHeaders(r) r.Header.Set("Authorization", "Bearer "+tok.AccessToken) } if acctRaw, _, errAcct := doJSON(loginClient, http.MethodGet, endpointLoginAcct+st.State, acctHeaders, nil); errAcct == nil { _ = json.Unmarshal(acctRaw, &acct) } sa := &StoredAuth{ Auth: StoredTokens{ AccessToken: tok.AccessToken, RefreshToken: tok.RefreshToken, ExpiresAt: time.Now().Add(time.Duration(tok.ExpiresIn) * time.Second).Unix(), Domain: tok.Domain, }, Account: StoredAccount{ UID: acct.UID, EnterpriseID: acct.EnterpriseID, Nickname: acct.Nickname, }, } if err := saveAuth(sa); err != nil { fmt.Printf("\n登录成功但保存凭据失败: %v\n", err) return } fmt.Println("\n🎉 登录成功!") fmt.Printf("欢迎,%s (UID: %s)\n", sa.Account.Nickname, sa.Account.UID) fmt.Printf("凭据已成功保存至: %s\n", cfg.AuthFile) fmt.Printf("令牌有效期至: %s\n", time.Unix(sa.Auth.ExpiresAt, 0).Format("2006-01-02 15:04:05")) fmt.Println("\n现在您可以运行以下命令启动网关服务:") fmt.Println(" workbuddy-gateway serve") return } } } // ----------------------------------------------------------------------------- // 本地 HTTP API 网关服务 (OpenAI 协议兼容) // ----------------------------------------------------------------------------- func runServe() { if err := loadAccounts(); err != nil { fmt.Printf("警告: 未检测到有效凭据 (%v)。\n请先执行: workbuddy-gateway login 扫码登录,或确保凭据文件存在。\n\n", err) } else { accountMu.Lock() for _, acc := range accounts { if acc.Disabled || acc.Auth == nil { continue } _ = ensureValidTokenFor(acc) } accountMu.Unlock() } // 启动后台自动刷新协程 go backgroundTokenRefresher() mux := http.NewServeMux() mux.HandleFunc("/v1/chat/completions", handleChatCompletions) mux.HandleFunc("/chat/completions", handleChatCompletions) mux.HandleFunc("/v1/models", handleModels) mux.HandleFunc("/models", handleModels) mux.HandleFunc("/health", handleHealth) mux.HandleFunc("/ping", handleHealth) mux.HandleFunc("/", handleIndex) listenAddr := fmt.Sprintf("%s:%d", cfg.Addr, cfg.Port) server := &http.Server{ Addr: listenAddr, Handler: corsMiddleware(authMiddleware(mux)), ReadTimeout: 120 * time.Second, WriteTimeout: 300 * time.Second, } fmt.Println("================================================================") fmt.Printf("🚀 WorkBuddy 本地网关已启动!\n") fmt.Printf(" 服务监听地址: http://%s\n", listenAddr) fmt.Printf(" Chat 接口地址: http://%s/v1/chat/completions\n", listenAddr) fmt.Printf(" Models 接口: http://%s/v1/models\n", listenAddr) fmt.Printf(" 模型转发策略: 【完全透传】客户端请求的任意 model 原样中继至上游\n") if cfg.APIKey != "" { fmt.Printf(" API 鉴权: 已启用 (Bearer %s)\n", cfg.APIKey) } else { fmt.Printf(" API 鉴权: 未启用 (任何客户端均可直连)\n") } if cfg.ProxyURL != "" { fmt.Printf(" 上游出口代理: %s\n", cfg.ProxyURL) } // 启动时展示所有账号状态(与 status 命令一致) accountMu.Lock() accCount := len(accounts) activeCount := 0 disabledCount := 0 now := time.Now() for _, acc := range accounts { if acc.Disabled || acc.Auth == nil { disabledCount++ } else { activeCount++ } } fmt.Printf(" 账号池: %d 个账号 (有效 %d, 失效 %d)\n", accCount, activeCount, disabledCount) if accCount > 0 { fmt.Println(" ----------------------------------------------------------") for i, acc := range accounts { fmt.Print(formatAccountStatus(acc, i+1, now)) } fmt.Println(" ----------------------------------------------------------") } accountMu.Unlock() fmt.Println("================================================================") fmt.Println("等待客户端请求中 (按 Ctrl+C 安全停止)...") stopChan := make(chan os.Signal, 1) signal.Notify(stopChan, os.Interrupt, syscall.SIGTERM) go func() { if err := server.ListenAndServe(); err != nil && err != http.ErrServerClosed { log.Fatalf("HTTP 服务异常退出: %v", err) } }() <-stopChan fmt.Println("\n正在关闭网关服务...") ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) defer cancel() _ = server.Shutdown(ctx) fmt.Println("网关已安全停止。") } func backgroundTokenRefresher() { ticker := time.NewTicker(5 * time.Minute) defer ticker.Stop() for range ticker.C { accountMu.Lock() accs := make([]*Account, len(accounts)) copy(accs, accounts) accountMu.Unlock() for _, acc := range accs { if err := ensureValidTokenFor(acc); err != nil { log.Printf("[BackgroundAuth] 账号 %s 自动检查/续期令牌异常: %v", acc.Path, err) } } } } // ----------------------------------------------------------------------------- // 中间件: CORS 与可选 APIKey 校验 // ----------------------------------------------------------------------------- func corsMiddleware(next http.Handler) http.Handler { return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { w.Header().Set("Access-Control-Allow-Origin", "*") w.Header().Set("Access-Control-Allow-Methods", "GET, POST, PUT, DELETE, OPTIONS") w.Header().Set("Access-Control-Allow-Headers", "*") if r.Method == http.MethodOptions { w.WriteHeader(http.StatusOK) return } next.ServeHTTP(w, r) }) } func authMiddleware(next http.Handler) http.Handler { return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { if cfg.APIKey != "" && r.URL.Path != "/health" && r.URL.Path != "/ping" && r.URL.Path != "/" { authHeader := r.Header.Get("Authorization") token := strings.TrimPrefix(authHeader, "Bearer ") if token != cfg.APIKey { writeOpenAIError(w, http.StatusUnauthorized, "invalid_api_key", "未提供有效 API 密钥") return } } next.ServeHTTP(w, r) }) } // ----------------------------------------------------------------------------- // 路由处理: /v1/chat/completions (支持任意 model 透传) // ----------------------------------------------------------------------------- func handleChatCompletions(w http.ResponseWriter, r *http.Request) { if r.Method != http.MethodPost { writeOpenAIError(w, http.StatusMethodNotAllowed, "method_not_allowed", "仅支持 POST 请求") return } reqID := atomic.AddUint64(&reqCounter, 1) startTime := time.Now() bodyBytes, err := io.ReadAll(r.Body) if err != nil { writeOpenAIError(w, http.StatusBadRequest, "read_error", "读取请求体失败") return } defer r.Body.Close() var reqObj map[string]any if err := json.Unmarshal(bodyBytes, &reqObj); err != nil { writeOpenAIError(w, http.StatusBadRequest, "invalid_json", "无效的 JSON 请求体") return } // 核心特性:完全透传 model 字段 // 客户端传什么 model,我们就透传什么 model 给上游,不做任何硬编码限制! modelName, _ := reqObj["model"].(string) if modelName == "" { modelName = "hy4-preview" // 默认保底 reqObj["model"] = modelName } isStream, _ := reqObj["stream"].(bool) // 腾讯上游强制要求 stream 必须为 true,非流式会被拦截 (code 11101) reqObj["stream"] = true // 深度思考 (Thinking) 自动适配:混元系列如果未关闭思考,自动赋予 high 档位保证深度思考输出 applyThinkingRules(reqObj, modelName) // 模板净化:改写 Claude Code 等框架被腾讯官方逐字拉黑的固定 prompt 语句 sanitizeMessages(reqObj) upstreamBytes, err := json.Marshal(reqObj) if err != nil { writeOpenAIError(w, http.StatusInternalServerError, "encode_error", "序列化请求失败") return } if cfg.Verbose { log.Printf("[#%d][Req] Model: %s | Stream: %v | BodyLen: %d", reqID, modelName, isStream, len(upstreamBytes)) } else { log.Printf("[#%d] POST /v1/chat/completions -> Upstream [Model: %s, Stream: %v]", reqID, modelName, isStream) } // 多账号轮询 + 429 自动冷却代偿: // 按轮询顺序尝试账号;命中 429 频率限制时立即将当前账号屏蔽到重置时间, // 并自动改用下一个可用账号重试(代偿),直至成功或所有账号均不可用。 accountMu.Lock() poolSize := len(accounts) accountMu.Unlock() if poolSize == 0 { writeOpenAIError(w, http.StatusUnauthorized, "no_auth", "未找到有效登录凭据,请先执行 login 命令扫码登录") return } var lastRateErr string var lastAuthErr string for attempt := 0; attempt < poolSize; attempt++ { acc, err := nextAccount() if err != nil { // 所有账号均不可用(冷却或失效) msg := fmt.Sprintf("无可用账号: %v", err) if lastAuthErr != "" { msg += " | 最近一次授权失效: " + truncate(lastAuthErr, 200) } if lastRateErr != "" { msg += " | 最近一次频率限制: " + truncate(lastRateErr, 200) } log.Printf("[#%d] %s", reqID, msg) writeOpenAIError(w, http.StatusServiceUnavailable, "no_available_account", msg) return } _ = ensureValidTokenFor(acc) // 令牌刷新时若发现授权失效会禁用账号;若被禁用则跳过换下一个 accountMu.Lock() disabled := acc.Disabled accountMu.Unlock() if disabled { lastAuthErr = acc.DisabledReason continue } upstreamReq, err := http.NewRequestWithContext(r.Context(), http.MethodPost, endpointChat, bytes.NewReader(upstreamBytes)) if err != nil { writeOpenAIError(w, http.StatusInternalServerError, "req_create_error", err.Error()) return } // 注入 CodeBuddy 凭据与指纹 Header backendHeaders(upstreamReq, acc.Auth) // 限制同一账号向腾讯上游的请求严格单并发串行排队,防止 DSH 同时发起标题生成+主对话触发腾讯 11128 风控 acc.lock.Lock() resp, err := cfg.HttpClient.Do(upstreamReq) acc.lock.Unlock() if err != nil { log.Printf("[#%d] 账号 %s 上游请求失败: %v", reqID, acc.Path, err) writeOpenAIError(w, http.StatusBadGateway, "upstream_network_error", fmt.Sprintf("网络转发失败: %v", err)) return } // 上游非 200 响应处理 if resp.StatusCode >= 400 { errBody, _ := io.ReadAll(resp.Body) resp.Body.Close() errStr := string(errBody) log.Printf("[#%d] 账号 %s 上游返回 HTTP %d: %s (耗时 %v)", reqID, acc.Path, resp.StatusCode, errStr, time.Since(startTime)) if isRateLimited(resp.StatusCode, errStr) { // 429 频率限制:解析重置时间并屏蔽该账号,交由其他账号代偿 until, ok := parseResetTime(errStr) if !ok { until = time.Now().Add(60 * time.Second) // 无法解析时默认冷却 60 秒 } markCooldown(acc, until, errStr) lastRateErr = errStr continue // 尝试下一个账号 } if isAuthFailure(resp.StatusCode, errStr) { // 授权失效(401/403 / token 无效 / 登录过期):禁用该账号并删除凭据文件, // 自动改用下一个可用账号,控制台提示用户重新登录 disableAccount(acc, fmt.Sprintf("上游鉴权失败 (HTTP %d): %s", resp.StatusCode, truncate(errStr, 200))) lastAuthErr = errStr continue // 尝试下一个账号 } writeOpenAIError(w, resp.StatusCode, "upstream_error", fmt.Sprintf("upstream %d: %s", resp.StatusCode, errStr)) return } if isStream { // 客户端需要流式响应 w.Header().Set("Content-Type", "text/event-stream") w.Header().Set("Cache-Control", "no-cache") w.Header().Set("Connection", "keep-alive") w.Header().Set("X-Accel-Buffering", "no") flusher, ok := w.(http.Flusher) if !ok { resp.Body.Close() writeOpenAIError(w, http.StatusInternalServerError, "streaming_unsupported", "服务器不支持流式响应 Flush") return } scanner := bufio.NewScanner(resp.Body) scanner.Buffer(make([]byte, 64*1024), 4*1024*1024) for scanner.Scan() { line := scanner.Text() cleanData := stripDataPrefix(line) if cleanData == "" { continue } if cleanData == "[DONE]" { _, _ = fmt.Fprintf(w, "data: [DONE]\n\n") flusher.Flush() break } cleanedChunk := cleanChunkJSON(cleanData) if cleanedChunk != "" { _, _ = fmt.Fprintf(w, "data: %s\n\n", cleanedChunk) flusher.Flush() } } resp.Body.Close() log.Printf("[#%d] 流式输出完成 (账号 %s, 耗时 %v)", reqID, acc.Path, time.Since(startTime)) } else { // 客户端需要完整 JSON 响应,聚合 SSE 数据流 completionJSON, err := aggregateCompletion(resp.Body, modelName) resp.Body.Close() if err != nil { log.Printf("[#%d] 聚合响应失败: %v", reqID, err) writeOpenAIError(w, http.StatusInternalServerError, "aggregate_error", "聚合上游流式响应失败: "+err.Error()) return } w.Header().Set("Content-Type", "application/json") w.WriteHeader(http.StatusOK) _, _ = w.Write(completionJSON) log.Printf("[#%d] 非流式响应完成 (账号 %s, 耗时 %v)", reqID, acc.Path, time.Since(startTime)) } return } // 理论上不可达(poolSize 次尝试后未成功即已在循环内返回) writeOpenAIError(w, http.StatusTooManyRequests, "all_accounts_cooldown", "所有账号均处于冷却状态") } // ----------------------------------------------------------------------------- // 路由处理: /v1/models (列出常见模型供客户端自动补全) // ----------------------------------------------------------------------------- func handleModels(w http.ResponseWriter, r *http.Request) { // 动态内置常见模型,用户直接传入任何未列出的 model 也会直接透传到上游 modelsList := []map[string]any{ {"id": "hy4-preview", "object": "model", "owned_by": "workbuddy", "permission": []any{}}, {"id": "hy3-preview-agent", "object": "model", "owned_by": "workbuddy", "permission": []any{}}, {"id": "hy3-preview", "object": "model", "owned_by": "workbuddy", "permission": []any{}}, {"id": "hy3", "object": "model", "owned_by": "workbuddy", "permission": []any{}}, {"id": "glm-5.2", "object": "model", "owned_by": "workbuddy", "permission": []any{}}, {"id": "glm-5.1", "object": "model", "owned_by": "workbuddy", "permission": []any{}}, {"id": "kimi-k2.7", "object": "model", "owned_by": "workbuddy", "permission": []any{}}, {"id": "deepseek-v4-pro", "object": "model", "owned_by": "workbuddy", "permission": []any{}}, {"id": "deepseek-v4-flash", "object": "model", "owned_by": "workbuddy", "permission": []any{}}, {"id": "minimax-m3-pay", "object": "model", "owned_by": "workbuddy", "permission": []any{}}, } resp := map[string]any{ "object": "list", "data": modelsList, } w.Header().Set("Content-Type", "application/json") _ = json.NewEncoder(w).Encode(resp) } func handleHealth(w http.ResponseWriter, r *http.Request) { w.Header().Set("Content-Type", "application/json") _ = json.NewEncoder(w).Encode(map[string]any{ "status": "healthy", "timestamp": time.Now().Unix(), "version": version, }) } func handleIndex(w http.ResponseWriter, r *http.Request) { if r.URL.Path != "/" { http.NotFound(w, r) return } w.Header().Set("Content-Type", "text/plain; charset=utf-8") _, _ = fmt.Fprintf(w, "WorkBuddy Local Gateway v%s is running.\n\nEndpoints:\n- POST /v1/chat/completions\n- GET /v1/models\n- GET /health\n", version) } // ----------------------------------------------------------------------------- // 请求转译与净化工具函数 // ----------------------------------------------------------------------------- func applyThinkingRules(obj map[string]any, modelName string) { // 遵循 CodeBuddy 规范:仅当客户端显式设置了 reasoning_effort 时才传递与规范化 // 绝不可强行对普通请求注入 reasoning_effort,否则极易触发腾讯内容与安全策略拦截 (code 11128) currEff, exists := obj["reasoning_effort"].(string) if !exists || currEff == "" || currEff == "off" || currEff == "none" { delete(obj, "reasoning_effort") delete(obj, "reasoning_summary") return } // 客户端显式请求思考时,设置 auto obj["reasoning_summary"] = "auto" } func sanitizeMessages(obj map[string]any) { messages, ok := obj["messages"].([]any) if !ok { return } for _, m := range messages { msg, ok := m.(map[string]any) if !ok { continue } switch c := msg["content"].(type) { case string: msg["content"] = sanitizeBlockedTemplates(c) case []any: for _, partAny := range c { part, ok := partAny.(map[string]any) if !ok { continue } if t, ok := part["text"].(string); ok { part["text"] = sanitizeBlockedTemplates(t) } } } } } func sanitizeBlockedTemplates(s string) string { s = strings.ReplaceAll(s, "You are Claude Code, Anthropic's official CLI for Claude.", "You are Claude Code, Anthropic's official CLI tool for Claude.") s = strings.ReplaceAll(s, "Main branch (you will usually use this for PRs)", "Default branch (you will usually use this for PRs)") return s } // backendHeaders 设置 CodeBuddy 上游专用指纹与鉴权 Header func backendHeaders(req *http.Request, sa *StoredAuth) { commonHeaders(req) reqID := uuid.New().String() req.Header.Set("X-Request-ID", reqID) req.Header.Set("X-Trace-ID", reqID) req.Header.Set("X-Client-ID", "codebuddy-cli") req.Header.Set("X-Client-Version", "2.143.1") if sa != nil && sa.Auth.AccessToken != "" { req.Header.Set("Authorization", "Bearer "+sa.Auth.AccessToken) } else { req.Header.Set("X-No-Authorization", "1") } if sa != nil && sa.Account.UID != "" { req.Header.Set("X-User-Id", sa.Account.UID) } else { req.Header.Set("X-No-User-Id", "1") } if sa != nil && sa.Account.EnterpriseID != "" { req.Header.Set("X-Enterprise-Id", sa.Account.EnterpriseID) } else { req.Header.Set("X-No-Enterprise-Id", "1") } if sa != nil && sa.Auth.RefreshToken != "" { req.Header.Set("X-Refresh-Token", sa.Auth.RefreshToken) } if sa != nil && sa.Auth.Domain != "" { req.Header.Set("X-Domain", sa.Auth.Domain) } else { req.Header.Set("X-No-Department-Info", "1") } req.Header.Set("X-Product", "SaaS") } func commonHeaders(req *http.Request) { req.Header.Set("Content-Type", "application/json") req.Header.Set("Accept", "application/json, text/plain, */*") req.Header.Set("X-Requested-With", "XMLHttpRequest") req.Header.Set("Origin", originReferer) req.Header.Set("Referer", originReferer+"/") req.Header.Set("User-Agent", clientUA) } func doJSON(client *http.Client, method, fullURL string, headers func(*http.Request), body io.Reader) (json.RawMessage, int, error) { req, err := http.NewRequest(method, fullURL, body) if err != nil { return nil, 0, err } if headers != nil { headers(req) } else { commonHeaders(req) } resp, err := client.Do(req) if err != nil { return nil, 0, err } defer resp.Body.Close() raw, _ := io.ReadAll(resp.Body) if resp.StatusCode >= 400 { return nil, resp.StatusCode, fmt.Errorf("http_error: upstream %d: %s", resp.StatusCode, string(raw)) } var env apiEnvelope if err := json.Unmarshal(raw, &env); err != nil { return nil, resp.StatusCode, fmt.Errorf("解析上游 JSON 失败: %w", err) } if env.Code != 0 { return nil, resp.StatusCode, fmt.Errorf("code=%d msg=%s", env.Code, env.Msg) } return env.Data, resp.StatusCode, nil } // ----------------------------------------------------------------------------- // 流式聚合与格式清理 // ----------------------------------------------------------------------------- func aggregateCompletion(r io.Reader, model string) ([]byte, error) { var content, reasoning, role, respModel, respID, finish string var created int64 var usage map[string]any var toolCalls []map[string]any scanner := bufio.NewScanner(r) scanner.Buffer(make([]byte, 64*1024), 4*1024*1024) for scanner.Scan() { data := stripDataPrefix(scanner.Text()) if data == "" || data == "[DONE]" { continue } var chunk map[string]any if json.Unmarshal([]byte(data), &chunk) != nil { continue } if v, ok := chunk["id"].(string); ok && v != "" { respID = v } if v, ok := chunk["model"].(string); ok && v != "" { respModel = v } if v, ok := chunk["created"].(float64); ok { created = int64(v) } if v, ok := chunk["usage"].(map[string]any); ok { usage = v } choices, _ := chunk["choices"].([]any) for _, c := range choices { choice, _ := c.(map[string]any) if delta, ok := choice["delta"].(map[string]any); ok { if v, ok := delta["role"].(string); ok && v != "" { role = v } if v, ok := delta["content"].(string); ok { content += v } if v, ok := delta["reasoning_content"].(string); ok { reasoning += v } if tcs, ok := delta["tool_calls"].([]any); ok { for _, tc := range tcs { if call, ok := tc.(map[string]any); ok { toolCalls = append(toolCalls, call) } } } } if v, ok := choice["finish_reason"].(string); ok && v != "" { finish = v } } } message := map[string]any{"role": ifEmpty(role, "assistant"), "content": content} if reasoning != "" { message["reasoning_content"] = reasoning } if len(toolCalls) > 0 { message["tool_calls"] = toolCalls } if created == 0 { created = time.Now().Unix() } result := map[string]any{ "id": ifEmpty(respID, fmt.Sprintf("chatcmpl-%d", time.Now().UnixNano())), "object": "chat.completion", "created": created, "model": ifEmpty(respModel, model), "choices": []map[string]any{{ "index": 0, "message": message, "finish_reason": ifEmpty(finish, "stop"), }}, } if usage != nil { result["usage"] = usage } return json.Marshal(result) } func cleanChunkJSON(s string) string { var obj map[string]any if json.Unmarshal([]byte(s), &obj) != nil { return s } if choices, ok := obj["choices"].([]any); ok { for _, c := range choices { choice, ok := c.(map[string]any) if !ok { continue } if delta, ok := choice["delta"].(map[string]any); ok { for k, v := range delta { if isEmptyValue(v) { delete(delta, k) } } } } } out, err := json.Marshal(obj) if err != nil { return s } return string(out) } func isEmptyValue(v any) bool { switch x := v.(type) { case nil: return true case string: return x == "" case []any: return len(x) == 0 case map[string]any: return len(x) == 0 } return false } func stripDataPrefix(s string) string { s = strings.TrimSpace(s) for strings.HasPrefix(s, "data:") { s = strings.TrimSpace(strings.TrimPrefix(s, "data:")) } return s } func writeOpenAIError(w http.ResponseWriter, statusCode int, errType, message string) { w.Header().Set("Content-Type", "application/json") w.WriteHeader(statusCode) _ = json.NewEncoder(w).Encode(map[string]any{ "error": map[string]any{ "message": message, "type": errType, "code": statusCode, }, }) } func ifEmpty(val, fallback string) string { if strings.TrimSpace(val) == "" { return fallback } return val } func truncate(s string, n int) string { if len(s) <= n { return s } return s[:n] + "..." }