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
3 changed files with 204 additions and 48 deletions
@@ -61,13 +61,19 @@ workbuddy-gateway serve -api-key sk-localsecret

 ### 配置多个账号

-网关支持同时挂载多个 CodeBuddy 账号，请求按**轮询（round-robin）**方式均匀分发，保持各账号额度消耗一致。两种配置方式：
+网关支持同时挂载多个 CodeBuddy 账号，请求按**轮询（round-robin）**方式均匀分发，保持各账号额度消耗一致。三种配置方式：

 ```bash
-# 方式一：-auth 逗号分隔多个凭据文件
+# 方式一（推荐）：自动发现 — 把多个凭据文件放进工作目录即可，无需任何参数
+# 网关启动时自动加载当前目录下所有 workbuddy*.json（如 workbuddy.json、workbuddy2.json…）
+cd /opt/workbuddy-gateway
+workbuddy-gateway serve -addr 0.0.0.0 -port 8317
+# 启动日志会显示: 已就绪账号池: 2 个账号 (有效 2, 失效 0)
+
+# 方式二：-auth 逗号分隔多个凭据文件
 workbuddy-gateway serve -auth workbuddy.json,workbuddy-2.json,workbuddy-3.json

-# 方式二：-auth-dir 指定凭据目录（自动加载目录下所有 workbuddy*.json）
+# 方式三：-auth-dir 指定凭据目录（自动加载目录下所有 workbuddy*.json）
 mkdir -p auths
 workbuddy-gateway login -auth auths/workbuddy-1.json   # 依次为每个账号扫码登录
 workbuddy-gateway login -auth auths/workbuddy-2.json
@@ -126,7 +132,7 @@ workbuddy-gateway status   # 该账号恢复为 ✅ 可用

 ### 与单账号模式的兼容性

- 不传 `-auth` 时默认使用 `./workbuddy.json`，行为与旧版完全一致；
+- 不传 `-auth` / `-auth-dir` 时，自动发现当前工作目录下所有 `workbuddy*.json`；目录中只有一个凭据文件时行为与旧版完全一致；
 - 只有一个账号时，请求始终使用该账号，429 冷却逻辑同样生效（冷却期间请求将返回 429 提示）；
 - 账号之间使用独立的串行锁：同一账号请求严格排队，不同账号可并行，兼顾风控与吞吐。

@@ -230,8 +236,18 @@ sudo mkdir -p /opt/workbuddy-gateway
 sudo cp workbuddy-gateway /opt/workbuddy-gateway/      # 对应架构的二进制
 # 首次登录（会生成 workbuddy.json）
 sudo /opt/workbuddy-gateway/workbuddy-gateway login
+# 登录第二个账号（可选）：生成 workbuddy2.json
+sudo /opt/workbuddy-gateway/workbuddy-gateway login -auth /opt/workbuddy-gateway/workbuddy2.json
 ```

+> 💡 **多账号自动发现**：服务通过 `WorkingDirectory` 固定在 `/opt/workbuddy-gateway`，
+> 无需修改 ExecStart——把多个凭据文件（`workbuddy.json`、`workbuddy2.json`…）放进该目录，
+> 重启服务即自动加载全部账号组成轮询池。确认方式：
+> ```bash
+> sudo journalctl -u workbuddy-gateway | grep 账号池
+> # 输出示例: 已就绪账号池: 2 个账号 (有效 2, 失效 0)
+> ```
+
 启用并启动服务：

 ```bash
@@ -29,7 +29,7 @@ import (
 )

 const (
-	version       = "1.3.0"
+	version       = "1.4.0"
 	upstreamBase  = "https://copilot.tencent.com"
 	clientUA      = "CLI/2.143.1 CodeBuddy/2.143.1"
 	originReferer = "https://www.codebuddy.cn"
@@ -95,14 +95,15 @@ type authStateData struct {
 // -----------------------------------------------------------------------------

 type Config struct {
-	Addr       string
-	Port       int
-	AuthFile   string
-	AuthDir    string
-	APIKey     string
-	ProxyURL   string
-	Verbose    bool
-	HttpClient *http.Client
+	Addr         string
+	Port         int
+	AuthFile     string
+	AuthDir      string
+	AuthExplicit bool // 用户是否显式指定了 -auth（未指定时自动扫描目录下所有 workbuddy*.json）
+	APIKey       string
+	ProxyURL     string
+	Verbose      bool
+	HttpClient   *http.Client
 }

 // Account 表示一个 CodeBuddy 账号凭据及其运行时状态。
@@ -177,6 +178,15 @@ func main() {
 	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()

@@ -219,7 +229,7 @@ func printHelp() {
  -addr <ip>        网关监听地址 (默认: 127.0.0.1)
  -port <port>      网关监听端口 (默认: 8317)
  -auth <path>      凭据文件路径；支持逗号分隔多个文件实现多账号
-                    (默认: ./workbuddy.json)
+                    (默认: 自动发现当前目录下所有 workbuddy*.json)
  -auth-dir <dir>   凭据目录：自动加载目录下所有 workbuddy*.json 作为账号池
  -api-key <key>    设置后，调用网关必须携带 Bearer <key> 鉴权
  -proxy <url>      设置上游转发代理 (例如 http://127.0.0.1:7890 或 socks5://...)
@@ -227,10 +237,13 @@ func printHelp() {

 多账号说明:
  # 登录第二个账号（保存到不同文件）
-  workbuddy-gateway login -auth workbuddy-2.json
+  workbuddy-gateway login -auth workbuddy2.json
+
+  # 自动发现：把多个凭据文件放进工作目录即可自动多账号（无需任何参数）
+  workbuddy-gateway serve        # 自动加载 ./workbuddy*.json

  # 启动时指定多个凭据文件（轮询 + 429 自动冷却代偿）
-  workbuddy-gateway serve -auth workbuddy.json,workbuddy-2.json
+  workbuddy-gateway serve -auth workbuddy.json,workbuddy2.json

  # 或使用目录模式：目录内所有 workbuddy*.json 自动组成账号池
  workbuddy-gateway serve -auth-dir ./auths
@@ -352,9 +365,65 @@ func loadAccountFile(path string) (*StoredAuth, error) {
 	return &sa, nil
 }

-// loadAccounts 根据 -auth / -auth-dir 配置构建账号池。
-// -auth 支持逗号分隔多个凭据文件；-auth-dir 自动加载目录下所有 workbuddy*.json。
-// 单账号时池长度为 1，行为与旧版完全一致。
+// 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()
@@ -362,32 +431,7 @@ func loadAccounts() error {
 	accounts = nil
 	rrIndex = 0

-	var paths []string
-	if cfg.AuthDir != "" {
-		entries, err := os.ReadDir(cfg.AuthDir)
-		if err != nil {
-			return fmt.Errorf("读取凭据目录失败 (%s): %w", cfg.AuthDir, err)
-		}
-		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)
-	} else {
-		for _, p := range strings.Split(cfg.AuthFile, ",") {
-			p = strings.TrimSpace(p)
-			if p != "" {
-				paths = append(paths, p)
-			}
-		}
-	}
-
-	for _, p := range paths {
+	for _, p := range collectConfiguredAuthPaths() {
 		sa, err := loadAccountFile(p)
 		if err != nil {
 			log.Printf("[Auth] 跳过无效凭据文件 %s: %v", p, err)
@@ -500,7 +544,8 @@ func disableAccount(acc *Account, reason string) {
 func loadDisabledMarkers() {
 	// 收集标记文件路径：
 	// - AuthDir 模式：目录下 *.json.disabled 即为标记文件
-	// - 单文件模式：<凭据路径>.disabled 为标记文件
+	// - 显式 -auth 模式：<凭据路径>.disabled 为标记文件
+	// - 自动发现模式：扫描当前目录下所有 *.json.disabled 标记文件
 	var markerPaths []string
 	if cfg.AuthDir != "" {
 		entries, err := os.ReadDir(cfg.AuthDir)
@@ -516,13 +561,28 @@ func loadDisabledMarkers() {
 				markerPaths = append(markerPaths, filepath.Join(cfg.AuthDir, name))
 			}
 		}
-	} else {
+	} 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 {
@@ -232,15 +232,18 @@ func TestLoadDisabledMarkers(t *testing.T) {
 		t.Fatal(err)
 	}

-	// 模拟 -auth 指定该路径
+	// 模拟 -auth 显式指定该路径
 	oldAuthFile := cfg.AuthFile
 	oldAuthDir := cfg.AuthDir
+	oldAuthExplicit := cfg.AuthExplicit
 	defer func() {
 		cfg.AuthFile = oldAuthFile
 		cfg.AuthDir = oldAuthDir
+		cfg.AuthExplicit = oldAuthExplicit
 	}()
 	cfg.AuthFile = authPath
 	cfg.AuthDir = ""
+	cfg.AuthExplicit = true

 	accounts = nil
 	loadDisabledMarkers()
@@ -274,3 +277,80 @@ func TestClearDisabledMarker(t *testing.T) {
 	}
 	t.Log("clearDisabledMarker works")
 }
+
+// 验证自动发现模式：未指定 -auth/-auth-dir 时，扫描当前目录下所有 workbuddy*.json
+func TestCollectConfiguredAuthPathsAutoDiscover(t *testing.T) {
+	oldDir, err := os.Getwd()
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer os.Chdir(oldDir)
+
+	dir := t.TempDir()
+	if err := os.Chdir(dir); err != nil {
+		t.Fatal(err)
+	}
+	// 模拟目录内多个凭据文件 + 非凭据文件 + 失效标记
+	for _, name := range []string{"workbuddy.json", "workbuddy2.json", "workbuddy-3.json", "other.json", "workbuddy.json.disabled"} {
+		if err := os.WriteFile(name, []byte(`{"auth":{"accessToken":"x"}}`), 0600); err != nil {
+			t.Fatal(err)
+		}
+	}
+
+	oldAuthFile := cfg.AuthFile
+	oldAuthDir := cfg.AuthDir
+	oldAuthExplicit := cfg.AuthExplicit
+	defer func() {
+		cfg.AuthFile = oldAuthFile
+		cfg.AuthDir = oldAuthDir
+		cfg.AuthExplicit = oldAuthExplicit
+	}()
+	cfg.AuthFile = "workbuddy.json"
+	cfg.AuthDir = ""
+	cfg.AuthExplicit = false
+
+	paths := collectConfiguredAuthPaths()
+	if len(paths) != 3 {
+		t.Fatalf("expected 3 workbuddy json files, got %d: %v", len(paths), paths)
+	}
+	want := []string{"workbuddy-3.json", "workbuddy.json", "workbuddy2.json"} // sort.Strings 排序结果
+	for i, p := range want {
+		if paths[i] != p {
+			t.Fatalf("paths[%d] = %s, want %s (all: %v)", i, paths[i], p, paths)
+		}
+	}
+	t.Logf("auto-discover paths: %v", paths)
+}
+
+// 验证自动发现模式：目录内没有任何凭据文件时回退到默认路径
+func TestCollectConfiguredAuthPathsAutoDiscoverFallback(t *testing.T) {
+	oldDir, err := os.Getwd()
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer os.Chdir(oldDir)
+
+	dir := t.TempDir()
+	if err := os.Chdir(dir); err != nil {
+		t.Fatal(err)
+	}
+	// 目录为空
+
+	oldAuthFile := cfg.AuthFile
+	oldAuthDir := cfg.AuthDir
+	oldAuthExplicit := cfg.AuthExplicit
+	defer func() {
+		cfg.AuthFile = oldAuthFile
+		cfg.AuthDir = oldAuthDir
+		cfg.AuthExplicit = oldAuthExplicit
+	}()
+	cfg.AuthFile = "workbuddy.json"
+	cfg.AuthDir = ""
+	cfg.AuthExplicit = false
+
+	paths := collectConfiguredAuthPaths()
+	if len(paths) != 1 || paths[0] != "workbuddy.json" {
+		t.Fatalf("expected fallback to default workbuddy.json, got %v", paths)
+	}
+	t.Logf("fallback path: %v", paths)
+}