- serve now periodically writes account-pool state (active/cooldown/ disabled, cooldown until, token expiry) to workbuddy-status.json, refreshed every 3s and immediately on cooldown/disable/refresh events - new 'monitor' command runs in foreground: refreshes a live panel of all accounts (usable/cooldown/expired + reasons) every N seconds, optionally showing recent logs via journalctl (-journal) or a log file (-logfile) - statusSnapshotFile excluded from git via .gitignore - unit tests for tailLines and status snapshot states
1941 lines
61 KiB
Go
1941 lines
61 KiB
Go
package main
|
||
|
||
import (
|
||
"bufio"
|
||
"bytes"
|
||
"context"
|
||
"crypto/tls"
|
||
"encoding/json"
|
||
"flag"
|
||
"fmt"
|
||
"io"
|
||
"log"
|
||
"net/http"
|
||
"net/http/cookiejar"
|
||
"net/url"
|
||
"os"
|
||
"os/exec"
|
||
"os/signal"
|
||
"path/filepath"
|
||
"regexp"
|
||
"sort"
|
||
"strconv"
|
||
"strings"
|
||
"sync"
|
||
"sync/atomic"
|
||
"syscall"
|
||
"time"
|
||
|
||
"github.com/google/uuid"
|
||
"github.com/skip2/go-qrcode"
|
||
)
|
||
|
||
const (
|
||
version = "1.5.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
|
||
|
||
// 状态快照文件名:serve 后台周期写入,monitor 前台命令实时读取展示
|
||
statusSnapshotFile = "workbuddy-status.json"
|
||
)
|
||
|
||
// -----------------------------------------------------------------------------
|
||
// 数据结构定义
|
||
// -----------------------------------------------------------------------------
|
||
|
||
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
|
||
MonitorInterval int // monitor 状态刷新间隔(秒)
|
||
LogFile string // monitor 附加展示的日志文件路径
|
||
JournalService string // monitor 附加展示的 systemd 服务名(journalctl -u)
|
||
LogLines int // monitor 展示的最近日志行数
|
||
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.IntVar(&cfg.MonitorInterval, "interval", 3, "monitor 状态刷新间隔(秒)")
|
||
fs.StringVar(&cfg.LogFile, "logfile", "", "monitor 附加跟随的日志文件路径(如 -logfile /var/log/workbuddy-gateway.log)")
|
||
fs.StringVar(&cfg.JournalService, "journal", "", "monitor 附加跟随的 systemd 服务名(Linux 下用 journalctl -u <服务> -f 跟随)")
|
||
fs.IntVar(&cfg.LogLines, "lines", 5, "monitor 每次刷新展示的最近日志行数")
|
||
_ = 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 "monitor":
|
||
runMonitor()
|
||
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)
|
||
monitor 前台实时监控:周期刷新展示账号状态 + 最近日志 (Ctrl+C 退出)
|
||
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 输出详细调试日志 (请求/响应体)
|
||
|
||
monitor 选项:
|
||
-interval <sec> 状态刷新间隔秒数 (默认: 3)
|
||
-journal <svc> 同时展示 systemd 服务最近日志 (Linux, 如 -journal workbuddy-gateway)
|
||
-logfile <path> 同时展示指定日志文件最近内容 (如 -logfile /var/log/wb.log)
|
||
-lines <n> 每次刷新展示的最近日志行数 (默认: 5)
|
||
|
||
多账号说明:
|
||
# 登录第二个账号(保存到不同文件)
|
||
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)
|
||
writeStatusSnapshot()
|
||
}
|
||
|
||
// 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)
|
||
writeStatusSnapshot()
|
||
}
|
||
|
||
// 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"))
|
||
writeStatusSnapshot()
|
||
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)
|
||
}
|
||
|
||
// -----------------------------------------------------------------------------
|
||
// 状态快照与前台实时监控 (monitor)
|
||
// -----------------------------------------------------------------------------
|
||
|
||
// accountSnapshot 是写入状态快照文件的单个账号状态。
|
||
type accountSnapshot struct {
|
||
Path string `json:"path"`
|
||
Nickname string `json:"nickname"`
|
||
UID string `json:"uid"`
|
||
State string `json:"state"` // active | cooldown | disabled
|
||
CooldownUntil int64 `json:"cooldownUntil,omitempty"`
|
||
CooldownMsg string `json:"cooldownMsg,omitempty"`
|
||
DisabledReason string `json:"disabledReason,omitempty"`
|
||
TokenExpiresAt int64 `json:"tokenExpiresAt,omitempty"`
|
||
}
|
||
|
||
// statusSnapshot 是写入 workbuddy-status.json 的完整快照。
|
||
type statusSnapshot struct {
|
||
UpdatedAt int64 `json:"updatedAt"`
|
||
Accounts []accountSnapshot `json:"accounts"`
|
||
}
|
||
|
||
// writeStatusSnapshot 将账号池实时状态(含冷却/失效)原子写入状态快照文件。
|
||
// serve 后台周期调用;monitor 命令前台读取展示。
|
||
func writeStatusSnapshot() {
|
||
accountMu.Lock()
|
||
snap := statusSnapshot{UpdatedAt: time.Now().Unix()}
|
||
now := time.Now()
|
||
for _, acc := range accounts {
|
||
as := accountSnapshot{Path: acc.Path}
|
||
switch {
|
||
case acc.Disabled:
|
||
as.State = "disabled"
|
||
as.DisabledReason = acc.DisabledReason
|
||
if acc.Auth != nil {
|
||
as.Nickname = acc.Auth.Account.Nickname
|
||
as.UID = acc.Auth.Account.UID
|
||
} else {
|
||
as.Nickname = acc.Nickname
|
||
as.UID = acc.UID
|
||
}
|
||
case acc.CooldownUntil.After(now):
|
||
as.State = "cooldown"
|
||
as.CooldownUntil = acc.CooldownUntil.Unix()
|
||
as.CooldownMsg = acc.CooldownMsg
|
||
if acc.Auth != nil {
|
||
as.Nickname = acc.Auth.Account.Nickname
|
||
as.UID = acc.Auth.Account.UID
|
||
as.TokenExpiresAt = acc.Auth.Auth.ExpiresAt
|
||
}
|
||
default:
|
||
as.State = "active"
|
||
if acc.Auth != nil {
|
||
as.Nickname = acc.Auth.Account.Nickname
|
||
as.UID = acc.Auth.Account.UID
|
||
as.TokenExpiresAt = acc.Auth.Auth.ExpiresAt
|
||
}
|
||
}
|
||
snap.Accounts = append(snap.Accounts, as)
|
||
}
|
||
accountMu.Unlock()
|
||
|
||
data, err := json.MarshalIndent(snap, "", " ")
|
||
if err != nil {
|
||
return
|
||
}
|
||
// 原子写:先写临时文件再改名,避免 monitor 读到半截内容
|
||
tmp := statusSnapshotFile + ".tmp"
|
||
if err := os.WriteFile(tmp, data, 0600); err != nil {
|
||
return
|
||
}
|
||
_ = os.Rename(tmp, statusSnapshotFile)
|
||
}
|
||
|
||
// statusSnapshotLoop serve 后台每 3 秒刷新一次状态快照。
|
||
func statusSnapshotLoop() {
|
||
writeStatusSnapshot()
|
||
ticker := time.NewTicker(3 * time.Second)
|
||
defer ticker.Stop()
|
||
for range ticker.C {
|
||
writeStatusSnapshot()
|
||
}
|
||
}
|
||
|
||
// tailLines 读取文件末尾 n 行(用于 monitor 展示最近日志)。
|
||
func tailLines(path string, n int) ([]string, error) {
|
||
f, err := os.Open(path)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
defer f.Close()
|
||
// 先定位到文件末尾,从后向前扫描 n 个换行符
|
||
const chunk = 4096
|
||
fi, err := f.Stat()
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
size := fi.Size()
|
||
var lines []string
|
||
buf := make([]byte, chunk)
|
||
pos := size
|
||
lineBuf := make([]byte, 0, chunk)
|
||
newlines := 0
|
||
for pos > 0 && newlines <= n {
|
||
read := int64(chunk)
|
||
if pos < chunk {
|
||
read = pos
|
||
}
|
||
pos -= read
|
||
if _, err := f.Seek(pos, io.SeekStart); err != nil {
|
||
break
|
||
}
|
||
rn, err := f.Read(buf[:read])
|
||
if err != nil && rn == 0 {
|
||
break
|
||
}
|
||
for i := rn - 1; i >= 0; i-- {
|
||
if buf[i] == '\n' {
|
||
if len(lineBuf) > 0 {
|
||
lines = append([]string{string(lineBuf)}, lines...)
|
||
lineBuf = lineBuf[:0]
|
||
newlines++
|
||
if newlines > n {
|
||
break
|
||
}
|
||
}
|
||
} else {
|
||
lineBuf = append([]byte{buf[i]}, lineBuf...)
|
||
}
|
||
}
|
||
if newlines > n {
|
||
break
|
||
}
|
||
}
|
||
if len(lineBuf) > 0 && newlines <= n {
|
||
lines = append([]string{string(lineBuf)}, lines...)
|
||
}
|
||
if len(lines) > n {
|
||
lines = lines[len(lines)-n:]
|
||
}
|
||
return lines, nil
|
||
}
|
||
|
||
// journalLines 获取 systemd 服务最近 n 行日志(Linux journalctl)。
|
||
func journalLines(service string, n int) ([]string, error) {
|
||
cmd := exec.Command("journalctl", "-u", service, "-n", strconv.Itoa(n), "--no-pager", "-o", "short")
|
||
out, err := cmd.Output()
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
s := strings.TrimSpace(string(out))
|
||
if s == "" {
|
||
return nil, nil
|
||
}
|
||
lines := strings.Split(s, "\n")
|
||
if len(lines) > n {
|
||
lines = lines[len(lines)-n:]
|
||
}
|
||
return lines, nil
|
||
}
|
||
|
||
// runMonitor 前台实时监控:周期刷新展示账号池状态 + 最近日志(Ctrl+C 退出)。
|
||
func runMonitor() {
|
||
interval := time.Duration(cfg.MonitorInterval) * time.Second
|
||
if interval <= 0 {
|
||
interval = 3 * time.Second
|
||
}
|
||
logN := cfg.LogLines
|
||
if logN <= 0 {
|
||
logN = 5
|
||
}
|
||
|
||
// 判断 stdout 是否为终端(是则用 ANSI 清屏重绘,否则滚动输出)
|
||
isTTY := false
|
||
if fi, err := os.Stdout.Stat(); err == nil {
|
||
isTTY = fi.Mode()&os.ModeCharDevice != 0
|
||
}
|
||
|
||
fmt.Println("================ WorkBuddy 实时监控 ================")
|
||
fmt.Println("按 Ctrl+C 退出 | 状态文件: " + statusSnapshotFile)
|
||
fmt.Println("---------------------------------------------------------------")
|
||
|
||
for {
|
||
if isTTY {
|
||
fmt.Print("\033[H\033[2J") // 清屏
|
||
} else {
|
||
fmt.Println()
|
||
}
|
||
fmt.Printf("🕐 更新时间: %s\n", time.Now().Format("2006-01-02 15:04:05"))
|
||
|
||
data, err := os.ReadFile(statusSnapshotFile)
|
||
if err != nil {
|
||
fmt.Printf("⚠️ 未找到状态文件 %s(服务是否在运行?请确认已在服务工作目录执行 monitor)\n", statusSnapshotFile)
|
||
} else {
|
||
var snap statusSnapshot
|
||
if err := json.Unmarshal(data, &snap); err == nil {
|
||
active, cooldown, disabled := 0, 0, 0
|
||
for _, a := range snap.Accounts {
|
||
switch a.State {
|
||
case "active":
|
||
active++
|
||
case "cooldown":
|
||
cooldown++
|
||
case "disabled":
|
||
disabled++
|
||
}
|
||
}
|
||
fmt.Printf("📊 账号池: 共 %d 个 | ✅ 可用 %d | 🔒 冷却 %d | ❌ 失效 %d\n",
|
||
len(snap.Accounts), active, cooldown, disabled)
|
||
for i, a := range snap.Accounts {
|
||
fmt.Printf(" #%d %s (%s)\n", i+1, a.Path, ifEmpty(a.Nickname, "?"))
|
||
switch a.State {
|
||
case "active":
|
||
fmt.Printf(" ✅ 可用 | Token 有效期至: %s\n",
|
||
time.Unix(a.TokenExpiresAt, 0).Format("01-02 15:04"))
|
||
case "cooldown":
|
||
fmt.Printf(" 🔒 冷却中 至 %s (剩余 %v)\n",
|
||
time.Unix(a.CooldownUntil, 0).Format("01-02 15:04:05"),
|
||
time.Until(time.Unix(a.CooldownUntil, 0)).Round(time.Minute))
|
||
if a.CooldownMsg != "" {
|
||
fmt.Printf(" 原因: %s\n", truncate(a.CooldownMsg, 100))
|
||
}
|
||
case "disabled":
|
||
fmt.Printf(" ❌ 授权失效: %s\n", truncate(a.DisabledReason, 100))
|
||
fmt.Printf(" ⚠️ 请重新执行: workbuddy-gateway login -auth %s\n", a.Path)
|
||
}
|
||
}
|
||
} else {
|
||
fmt.Println("⚠️ 状态文件解析失败")
|
||
}
|
||
}
|
||
|
||
// 最近日志
|
||
if cfg.JournalService != "" {
|
||
if lines, err := journalLines(cfg.JournalService, logN); err == nil && len(lines) > 0 {
|
||
fmt.Printf("\n📜 最近日志 (journalctl -u %s):\n", cfg.JournalService)
|
||
for _, l := range lines {
|
||
fmt.Println(" " + l)
|
||
}
|
||
}
|
||
} else if cfg.LogFile != "" {
|
||
if lines, err := tailLines(cfg.LogFile, logN); err == nil && len(lines) > 0 {
|
||
fmt.Printf("\n📜 最近日志 (%s):\n", cfg.LogFile)
|
||
for _, l := range lines {
|
||
fmt.Println(" " + l)
|
||
}
|
||
}
|
||
}
|
||
|
||
fmt.Println("---------------------------------------------------------------")
|
||
time.Sleep(interval)
|
||
}
|
||
}
|
||
|
||
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()
|
||
|
||
// 启动状态快照协程(monitor 命令实时读取展示)
|
||
go statusSnapshotLoop()
|
||
|
||
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(" ----------------------------------------------------------")
|
||
}
|
||
// 末尾汇总:加载到的凭据文件清单(即使上方日志被截断也能确认)
|
||
fileList := make([]string, 0, len(accounts))
|
||
for _, acc := range accounts {
|
||
fileList = append(fileList, acc.Path)
|
||
}
|
||
if len(fileList) > 0 {
|
||
fmt.Printf(" ✅ 已加载凭据文件 (%d): %s\n", len(fileList), strings.Join(fileList, ", "))
|
||
} else {
|
||
fmt.Println(" ⚠️ 未加载到任何凭据文件,请先执行 login 命令扫码登录")
|
||
}
|
||
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] + "..."
|
||
}
|