custom-provider-pi 0.1.0 → 0.1.2

package/README.md

@@ -81,6 +81,8 @@ Tab 键可自动补全子命令与 provider 名（大小写不敏感）。
 | `--ua <预设\|原始UA>` | 预设 User-Agent（见下），或直接给自定义字符串 |
 | `--profile <模板键>` | 应用完整请求头模板（`claude-code` / `codex` / `opencode` / `browser` 等，含 UA + 客户端典型头集合） |
 | `--proxy <URL>` | HTTP/SOCKS 代理地址（支持 `$ENV`） |
+| `--lb-keys "$K1,$K2"` | 多 Key 负载均衡（轮询 + 429 自动冷却） |
+| `--lb-cooldown 60` | 冷却秒数（默认 60） |
 
 ### JSON 参数
 
@@ -180,6 +182,35 @@ Tab 键可自动补全子命令与 provider 名（大小写不敏感）。
 ```
 
 > ⚠️ Node.js 的 `fetch` 需要 `NODE_USE_ENV_PROXY=1` 才会读取代理环境变量（Vite/Tauri 等桌面应用的内置 fetch 也类似）。传统 `http.request` / `https.request` 不读取这些变量；如需在 Node 底层走代理，请用 `global-agent` 等库或升级到 Node ≥ 22 并设置该环境变量。
+
+### 负载均衡（多 Key 轮询 + 自动冷却）
+
+应对中转服务的 RPM/TPM/RTM 限制：一个渠道配多个 API Key，轮询使用；某个 Key 触发限流（HTTP 429）后自动冷却，恢复后重新参与轮询。连续 429 会指数退避（×2，上限 10 分钟）；成功请求重置退避计数。
+
+**交互向导：** add 时选「需要多 Key 负载均衡？」→ 输入逗号分隔的 Keys → 设定冷却时间（默认 60s）
+
+**flags：**
+
+```
+/custom-provider add relay --base-url https://api.gw.example.com/v1 \
+    --lb-keys "$KEY_A,$KEY_B,sk-plain" --lb-cooldown 60 --models gpt-4o
+```
+
+**JSON 配置：**
+
+```json
+{
+  "name": "relay",
+  "baseUrl": "https://api.gw.example.com/v1",
+  "lbKeys": ["$KEY_A", "$KEY_B", "sk-plain-c"],
+  "lbCooldown": 30,
+  "models": ["gpt-4o"]
+}
+```
+
+- Key 值支持 `$ENV` / `!cmd` 引用，注册时解析一次存入内存池
+- `/list` 显示活跃 Key 数（如 `3 Key（2 活跃 / 60s 冷却）`）
+- 单 key 模式不受影响（`apiKey` 字符串仍向后兼容）
 ## prune：模型修剪
 
 添加时自动拉取后已提供关键字过滤；对已有 provider 可用 `prune` 事后修剪：

package/custom-provider.ts

@@ -319,6 +319,10 @@ interface IProvider {
    *  生效方式：注册时写入 HTTPS_PROXY/HTTP_PROXY/ALL_PROXY 环境变量，
    *  需请求库读取这些变量（Node 21+ 的 fetch 需启用 NODE_USE_ENV_PROXY=1） */
   proxy?: string;
+  /** 多 Key 负载均衡：逗号分隔的 API Key 列表（支持 $ENV / !cmd 引用） */
+  lbKeys?: string[];
+  /** 负载均衡默认冷却时间（秒），不填默认 60 */
+  lbCooldown?: number;
   /** false 表示已禁用（不注册、不出现在 /model）；缺失视为启用 */
   enabled?: boolean;
   models: (string | IModel)[];
@@ -505,16 +509,28 @@ function buildProviderConfig(provider: IProvider): ProviderConfig {
   const providerConfig: ProviderConfig = {
     name: provider.name,
     baseUrl: provider.baseUrl,
-    apiKey: provider.apiKey ?? "local",
+    // LB 模式：apiKey 用占位符（SDK 生成 Bearer <placeholder>，before_provider_headers 替换）
+    // LB 模式：apiKey 用第一个 key 的原始值作占位（保证 pi 鉴权检查通过，
+    // before_provider_headers 会在发送前替换为轮询到的真实 key）
+    apiKey: provider.lbKeys && provider.lbKeys.length > 0
+      ? resolveValue(provider.lbKeys[0]) || "lb-pool"
+      : (provider.apiKey ?? "local"),
     models: provider.models.map((m) => prepareModel(m, provider)),
   };
 
+  // LB 模式：注入 X-LB-POOL 标记头（before_provider_headers 检测此标记识别 LB 请求）
+  if (provider.lbKeys && provider.lbKeys.length > 0) {
+    providerHeaders["X-LB-POOL"] = provider.name;
+  }
+
   if (Object.keys(providerHeaders).length > 0) {
     providerConfig.headers = providerHeaders;
   }
 
-  // authHeader：非标准 API 需要 Authorization: Bearer 时由用户显式开启（pi 负责生成）
-  if (provider.authHeader) {
+  // authHeader：LB 模式强制 true（保证 Authorization 头存在供 hook 覆盖）
+  if (provider.lbKeys && provider.lbKeys.length > 0) {
+    providerConfig.authHeader = true;
+  } else if (provider.authHeader) {
     providerConfig.authHeader = true;
   }
 
@@ -767,6 +783,7 @@ export default function customProviderExtension(pi: ExtensionAPI) {
 
     try {
       saveConfig(config);
+      loadLBPools(); // 同步刷新 LB key 池
     } catch (error) {
       ctx.ui.notify(`保存配置失败: ${error instanceof Error ? error.message : String(error)}`, "error");
       return false;
@@ -845,6 +862,13 @@ export default function customProviderExtension(pi: ExtensionAPI) {
     if (typeof data.authHeader === "boolean") provider.authHeader = data.authHeader;
     if (typeof data.proxy === "string" && data.proxy.trim()) provider.proxy = data.proxy;
     if (typeof data.enabled === "boolean") provider.enabled = data.enabled;
+    // 多 Key 负载均衡（JSON 路径）：lbKeys: ["$KEY_A","sk-plain"], lbCooldown: 30
+    if (Array.isArray(data.lbKeys) && (data.lbKeys as string[]).length > 0) {
+      provider.lbKeys = (data.lbKeys as string[]).map(String).filter(Boolean);
+      if (typeof data.lbCooldown === "number" && data.lbCooldown > 0) {
+        provider.lbCooldown = data.lbCooldown;
+      }
+    }
     if (data.headers && typeof data.headers === "object") {
       // 与 flags/向导同一套校验：拒绝非法名称/值（含 CR/LF 注入）
       const bad = Object.entries(data.headers).find(([k, v]) => {
@@ -885,6 +909,127 @@ export default function customProviderExtension(pi: ExtensionAPI) {
   // 监听 session_start 以支持热重载
   pi.on("session_start", () => {
     registerProviders();
+    loadLBPools();
+  });
+
+  // ================= 负载均衡（LBKeyPool）=================
+  // 多 Key 轮询 + 429 自动冷却（默认 60s）
+  // 通过 before_provider_headers 替换 Authorization / after_provider_response 冷却
+  // 原理：LB 模式的 provider 在 headers 里注入 X-LB-POOL: <name> 标记头，
+  // before_provider_headers 检测此标记 → 识别为 LB 请求 → 替换 Authorization
+
+  class LBKeyPool {
+    keys: string[];           // 解析后的 key 值
+    cooldownMs: number;       // 默认冷却毫秒
+    cursor: number = 0;       // 轮询游标
+    cooldownEnd: number[];    // 每个 key 的冷却结束时间戳
+    consecutive429s: number[];// 指数退避计数
+    perKeyCooldowns: (number | null | undefined)[]; // 每 key 的冷却覆盖
+
+    constructor(keys: string[], defaultCooldownSec: number, perKey?: (number | null | undefined)[]) {
+      this.keys = keys;
+      this.cooldownMs = defaultCooldownSec * 1000;
+      this.cooldownEnd = new Array(keys.length).fill(0);
+      this.consecutive429s = new Array(keys.length).fill(0);
+      this.perKeyCooldowns = perKey ?? [];
+    }
+
+    pick(): string | null {
+      const n = this.keys.length;
+      const now = Date.now();
+      // 最多遍历一轮，找到未在冷却中的 key
+      for (let i = 0; i < n; i++) {
+        const idx = (this.cursor + i) % n;
+        if (this.cooldownEnd[idx] <= now) {
+          this.cursor = (idx + 1) % n; // 下次从下一个开始
+          return this.keys[idx];
+        }
+      }
+      // 全部冷却中，返回冷却结束最早的（允许微小过期）
+      let bestIdx = 0;
+      for (let i = 1; i < n; i++) {
+        if (this.cooldownEnd[i] < this.cooldownEnd[bestIdx]) bestIdx = i;
+      }
+      this.cursor = (bestIdx + 1) % n;
+      return this.keys[bestIdx];
+    }
+
+    on429(idx: number): void {
+      this.consecutive429s[idx]++;
+      const base = this.perKeyCooldowns[idx] ?? this.cooldownMs;
+      // 指数退避：连续429时 × 2^n，上限 10 分钟
+      const multiplier = Math.min(this.consecutive429s[idx], 10);
+      const cooldown = Math.min(base * Math.pow(2, multiplier - 1), 600000);
+      this.cooldownEnd[idx] = Date.now() + cooldown;
+    }
+
+    onSuccess(idx: number): void {
+      this.consecutive429s[idx] = 0;
+    }
+
+    activeCount(): number {
+      const now = Date.now();
+      return this.keys.filter((_, i) => this.cooldownEnd[i] <= now).length;
+    }
+  }
+
+  // 运行时 LB 状态：provider name → key pool
+  let lbPools: Map<string, LBKeyPool> = new Map();
+  let lastUsedPool: string | null = null; // 上一次请求的 provider name
+  let lastUsedKeyIdx: number = -1;
+
+  function loadLBPools(): void {
+    lbPools.clear();
+    const config = loadConfig();
+    for (const provider of config.providers) {
+      if (provider.enabled === false) continue;
+      // 优先从 IProvider.lbKeys 读取；兼容旧格式 compat.lb.keys
+      const keys: string[] = provider.lbKeys ?? (provider.compat?.lb?.keys as string[] | undefined) ?? [];
+      const defaultCooldown = provider.lbCooldown ?? (provider.compat?.lb?.cooldown as number | undefined) ?? 60;
+      if (keys.length > 0) {
+        const resolved = keys.map(resolveValue).filter(Boolean);
+        if (resolved.length === 0) continue;
+        lbPools.set(provider.name, new LBKeyPool(resolved, defaultCooldown));
+      }
+    }
+  }
+
+  // 加载 LB 池 + 注册事件（启动时执行一次）
+  loadLBPools();
+
+  pi.on("before_provider_headers", (event) => {
+    // 检测 X-LB-POOL 标记头 → 识别为 LB provider 的请求
+    const poolName = event.headers["X-LB-POOL"] as string | undefined;
+    if (!poolName) return;
+    const pool = lbPools.get(poolName);
+    if (!pool) {
+      delete event.headers["X-LB-POOL"];
+      return;
+    }
+    const key = pool.pick();
+    if (!key) {
+      delete event.headers["X-LB-POOL"];
+      return;
+    }
+    // 记录本次使用的 key 索引（用于 after_provider_response 冷却）
+    const keyIdx = pool.keys.indexOf(key);
+    lastUsedPool = poolName;
+    lastUsedKeyIdx = keyIdx;
+    // 替换 Authorization 为真实 key（SDK 已发送 Bearer $LB）
+    event.headers["Authorization"] = `Bearer ${key}`;
+    // 移除标记头（不转发给上游）
+    delete event.headers["X-LB-POOL"];
+  });
+
+  pi.on("after_provider_response", (event) => {
+    if (!lastUsedPool || lastUsedKeyIdx < 0) return;
+    const pool = lbPools.get(lastUsedPool);
+    if (!pool) return;
+    if (event.status === 429) {
+      pool.on429(lastUsedKeyIdx);
+    } else if (event.status >= 200 && event.status < 300) {
+      pool.onSuccess(lastUsedKeyIdx);
+    }
   });
 
 // ================= 子命令实现 =================
@@ -1031,6 +1176,32 @@ export default function customProviderExtension(pi: ExtensionAPI) {
         }
       }
 
+      // ---- 4.6 多 Key 负载均衡（应对 RPM/RTM 限制；选 Yes 输入多个 API Key）----
+      const needLB = await ctx.ui.confirm(
+        "需要多 Key 负载均衡？",
+        "应对 RPM/RTM 等限流：多个 Key 轮询，受限后自动冷却 60s 再恢复"
+      );
+      let lbKeys: string[] | undefined;
+      let lbCooldown: number | undefined;
+      if (needLB) {
+        const lbInput = await ctx.ui.input(
+          "API Keys（逗号分隔，支持 $ENV / !cmd 引用）",
+          "$KEY_A,$KEY_B,sk-plain-c"
+        );
+        if (lbInput && lbInput.trim()) {
+          lbKeys = lbInput.split(",").map((s: string) => s.trim()).filter(Boolean);
+          const cdInput = await ctx.ui.input(
+            "默认冷却时间（秒），受限后自动暂停该 Key，恢复后重新参与轮询",
+            "60"
+          );
+          const cd = Number(cdInput?.trim() || "60");
+          if (cd > 0) lbCooldown = cd;
+        }
+        if (lbKeys && lbKeys.length > 0) {
+          ctx.ui.notify(`已启用 ${lbKeys.length} Key 负载均衡，冷却 ${lbCooldown ?? 60}s`, "info");
+        }
+      }
+
       // ---- 5. 是否自动拉取模型 ----
       const autoFetch = await ctx.ui.confirm(
         "自动从 /v1/models 拉取模型列表？",
@@ -1270,6 +1441,10 @@ export default function customProviderExtension(pi: ExtensionAPI) {
         models,
       };
       if (proxyUrl) newProvider.proxy = proxyUrl;
+      if (lbKeys && lbKeys.length > 0) {
+        newProvider.lbKeys = lbKeys;
+        if (lbCooldown) newProvider.lbCooldown = lbCooldown;
+      }
       if (apiType !== "自动推断" || inferredApi !== "openai-completions") {
         // 显式选择的协议，或自动推断出的非默认协议，需存盘保证幂等
         newProvider.api = inferredApi;
@@ -1403,6 +1578,15 @@ export default function customProviderExtension(pi: ExtensionAPI) {
     if (apiRaw && apiRaw !== "auto") provider.api = apiRaw;
     else if (api !== "openai-completions") provider.api = api;
     if (getFlag(flags, "auth-header")) provider.authHeader = true;
+    const lbKeysFlag = getFlag(flags, "lb-keys");
+    if (lbKeysFlag) {
+      provider.lbKeys = lbKeysFlag.split(",").map((s) => s.trim()).filter(Boolean);
+      const lbCooldownFlag = getFlag(flags, "lb-cooldown");
+      if (lbCooldownFlag) {
+        const n = Number(lbCooldownFlag);
+        if (n > 0) provider.lbCooldown = n;
+      }
+    }
 
     const compatJson = getFlag(flags, "compat");
     if (compatJson) {
@@ -1576,6 +1760,7 @@ export default function customProviderExtension(pi: ExtensionAPI) {
     try {
       saveConfig(config);
       pi.unregisterProvider(name);
+      lbPools.delete(name); // 同步移除 LB 池
       ctx.ui.notify(`Provider "${name}" 已删除并注销`, "info");
     } catch (error) {
       ctx.ui.notify(`删除失败: ${error instanceof Error ? error.message : String(error)}`, "error");
@@ -1660,7 +1845,9 @@ export default function customProviderExtension(pi: ExtensionAPI) {
       const preview =
         ids.length <= 4 ? ids.join(", ") : `${ids.slice(0, 4).join(", ")}, …（共 ${ids.length} 个）`;
       const proxyLine = p.proxy ? `\n  代理: ${p.proxy}` : "";
-      return `• ${p.name}  [${state}] [${api}]\n  端点: ${p.baseUrl}${proxyLine}\n  模型: ${preview}`;
+      const lbPool = lbPools.get(p.name);
+      const lbLine = lbPool ? `\n  负载均衡: ${lbPool.keys.length} Key（${lbPool.activeCount()} 活跃 / ${lbPool.cooldownMs / 1000}s 冷却）` : "";
+      return `• ${p.name}  [${state}] [${api}]\n  端点: ${p.baseUrl}${proxyLine}${lbLine}\n  模型: ${preview}`;
     });
     ctx.ui.notify(`已配置 ${config.providers.length} 个 provider:\n\n${lines.join("\n\n")}\n\n启用/禁用: /custom-provider enable|disable <名称>` , "info");
   };
@@ -1854,7 +2041,9 @@ export default function customProviderExtension(pi: ExtensionAPI) {
         const api = p.api ?? inferApi(p.baseUrl);
         const state = p.enabled === false ? "✗ 禁用" : "✓ 启用";
         const proxyLine = p.proxy ? `\n  代理: ${p.proxy}` : "";
-        return `• ${p.name}  [${state}] [${api}]  ${p.models.length} 个模型\n  端点: ${p.baseUrl}${proxyLine}`;
+        const lbPool = lbPools.get(p.name);
+        const lbLine = lbPool ? `\n  负载均衡: ${lbPool.keys.length} Key（${lbPool.activeCount()} 活跃 / ${lbPool.cooldownMs / 1000}s 冷却）` : "";
+        return `• ${p.name}  [${state}] [${api}]  ${p.models.length} 个模型\n  端点: ${p.baseUrl}${proxyLine}${lbLine}`;
       });
       ctx.ui.notify(
         `配置文件: ${CONFIG_PATH}\n已配置 ${config.providers.length} 个 provider:\n\n${lines.join("\n\n")}\n\nconfig <name> 查看详情 · config edit 编辑 · config path 路径`,
@@ -1955,13 +2144,21 @@ export default function customProviderExtension(pi: ExtensionAPI) {
     "  --name · --base-url/--url · --api-key/--key · --api TYPE",
     "  --models \"m1,m2\" · --model m（可多次）",
     "  --header \"K: V\"（可多次）· --headers '{\"k\":\"v\"}'",
-    "  --auth-header · --compat '{...}' · --overrides '{\"modelId\":{...}}'",
+    "  --profile <模板键>（完整请求头模板：claude-code / codex / browser 等）",
+    "  --ua <预设键|原始UA>（如 claude-code / codex / browser）",
+    "  --auth-header · --proxy <URL>（HTTP/SOCKS 代理，$ENV 可用）",
+    "  --lb-keys \"$K1,$K2\"（多 Key 负载均衡）· --lb-cooldown 60（冷却秒数）",
+    "  --model-api 'id:协议'（可多次，如 claude-x:anthropic-messages）",
+    "  --model-base-url 'id:url'（和 --model-api 搭配混用双协议）",
+    "  --compat '{...}' · --overrides '{\"modelId\":{...}}'",
     "  --force（覆盖已存在）· --json '{...}'（完整配置）",
-        "  --ua <预设|原始UA>（如 claude-code / codex / opencode / cursor，或直接给自定义串）",
-"  --model-api 'id:协议'（可多次，如 claude-x:anthropic-messages）",
-    "  --model-base-url 'id:url'（可多次，和 --model-api 搭配混用双协议）",
-    "",
-    "示例:",
+    "  ",
+    "负载均衡示例（同渠道多 Key 轮询，429 后自动冷却 60s）:",
+    "  /custom-provider add relay --base-url https://api.gw.com/v1 \\",
+    "      --lb-keys \"$KEY_A,$KEY_B,sk-plain\" --lb-cooldown 60 --force",
+    "  JSON: --json '{\"name\":\"relay\",\"baseUrl\":\"...\",\"lbKeys\":[\"k1\",\"k2\"],\"lbCooldown\":30}'",
+    "  ",
+"示例:",
     "  /custom-provider add deepseek --base-url https://api.deepseek.com/v1 \\",
     "      --api-key $DEEPSEEK_API_KEY --models deepseek-chat,deepseek-reasoner",
     "  /custom-provider add --json '{\"name\":\"x\",\"baseUrl\":\"...\",\"models\":[\"a\"]}'",

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "custom-provider-pi",
-  "version": "0.1.0",
+  "version": "0.1.2",
   "description": "pi 扩展：统一管理第三方模型 Provider。子命令体系 /custom-provider add|remove|refresh|list|test|config|enable|disable|prune，支持交互向导、flags/JSON 非交互添加、双协议混用（OpenAI + Anthropic）、模型关键字过滤与修剪",
   "keywords": [
     "pi-package",