Files
workbuddy-gateway/main.go
T
CangShui 9fcd99f152 feat: auto-discover all workbuddy*.json in working directory
- When neither -auth nor -auth-dir is specified, scan the working
  directory for all workbuddy*.json files and build the account pool
  from them (multi-account without any flags)
- Example: /opt/workbuddy-gateway with workbuddy.json + workbuddy2.json
  is automatically loaded as a 2-account pool under systemd
- loadDisabledMarkers also auto-discovers *.json.disabled markers
- Fallback to default workbuddy.json when no credential files found
- Unit tests for auto-discovery and fallback behavior
2026-09-04 08:13:38 +08:00

1632 lines
51 KiB
Go
Raw Permalink Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
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.0"
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 <ip> 网关监听地址 (默认: 127.0.0.1)
-port <port> 网关监听端口 (默认: 8317)
-auth <path> 凭据文件路径;支持逗号分隔多个文件实现多账号
(默认: 自动发现当前目录下所有 workbuddy*.json)
-auth-dir <dir> 凭据目录:自动加载目录下所有 workbuddy*.json 作为账号池
-api-key <key> 设置后,调用网关必须携带 Bearer <key> 鉴权
-proxy <url> 设置上游转发代理 (例如 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
// -----------------------------------------------------------------------------
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.Printf("\n--- 账号 #%d ---\n", i+1)
fmt.Printf("凭据文件: %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
}
fmt.Printf("用户昵称: %s\n", ifEmpty(nickname, "(未知)"))
fmt.Printf("用户 UID: %s\n", ifEmpty(uid, "(未知)"))
fmt.Printf("账号状态: ❌ 授权失效(禁止调度)\n")
fmt.Printf("失效原因: %s\n", truncate(acc.DisabledReason, 200))
fmt.Printf("处理建议: ⚠️ 凭据文件已删除,请重新执行: workbuddy-gateway login -auth %s\n", acc.Path)
continue
}
if acc.Auth == nil {
continue
}
expTime := time.Unix(acc.Auth.Auth.ExpiresAt, 0)
remaining := time.Until(expTime)
statusStr := "有效"
if remaining <= 0 {
statusStr = "已过期"
}
fmt.Printf("用户昵称: %s\n", acc.Auth.Account.Nickname)
fmt.Printf("用户 UID: %s\n", acc.Auth.Account.UID)
fmt.Printf("企业 ID: %s\n", ifEmpty(acc.Auth.Account.EnterpriseID, "(个人账号)"))
fmt.Printf("认证域名: %s\n", ifEmpty(acc.Auth.Auth.Domain, "www.codebuddy.cn"))
if acc.CooldownUntil.After(now) {
fmt.Printf("冷却状态: 🔒 冷却中 (解封: %s, 剩余 %v)\n",
acc.CooldownUntil.Format("2006-01-02 15:04:05"),
time.Until(acc.CooldownUntil).Round(time.Minute))
fmt.Printf("冷却原因: %s\n", truncate(acc.CooldownMsg, 120))
} else {
fmt.Printf("冷却状态: ✅ 可用\n")
}
fmt.Printf("Token 状态: %s\n", statusStr)
fmt.Printf("过期时间: %s (剩余 %v)\n", expTime.Format("2006-01-02 15:04:05"), remaining.Round(time.Minute))
}
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()
accCount := len(accounts)
activeCount := 0
disabledCount := 0
for _, acc := range accounts {
if acc.Disabled || acc.Auth == nil {
disabledCount++
continue
}
activeCount++
_ = ensureValidTokenFor(acc)
}
accountMu.Unlock()
fmt.Printf("已就绪账号池: %d 个账号 (有效 %d, 失效 %d)\n", accCount, activeCount, disabledCount)
for i, acc := range accounts {
if acc.Disabled || acc.Auth == nil {
fmt.Printf(" #%d ❌ %s (授权失效,凭据文件已删除,请重新登录)\n", i+1, acc.Path)
continue
}
fmt.Printf(" #%d ✅ %s (UID: %s, 有效期至: %s)\n",
i+1, acc.Auth.Account.Nickname, acc.Auth.Account.UID,
time.Unix(acc.Auth.Auth.ExpiresAt, 0).Format("2006-01-02 15:04:05"))
}
}
// 启动后台自动刷新协程
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)
}
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] + "..."
}