clear-session-bindings.ts 30 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911912913914915916917918919920921922923924925926927928929930931932933934935936937938939940941942943944945946947948949950951952953954955956957958959960961962963964965966967968969970971972973974975976977978979980981982983984985986987988989990991992993994995996997998999100010011002100310041005
  1. #!/usr/bin/env bun
  2. /**
  3. * 清除指定供应商的会话绑定
  4. *
  5. * 用法:
  6. * # 交互式模式(推荐)
  7. * bun run scripts/clear-session-bindings.ts
  8. *
  9. * # 按优先级筛选
  10. * bun run scripts/clear-session-bindings.ts --priority <number>
  11. *
  12. * # 指定供应商 ID
  13. * bun run scripts/clear-session-bindings.ts --id 1,2,3
  14. *
  15. * # 指定供应商名称(模糊匹配)
  16. * bun run scripts/clear-session-bindings.ts --name "cubence"
  17. *
  18. * 选项:
  19. * --priority, -p <number> 优先级阈值(清除 priority < 该值的供应商绑定)
  20. * --id <ids> 指定供应商 ID(逗号分隔)
  21. * --name <pattern> 按名称模糊匹配供应商
  22. * --type <type> 供应商类型筛选(claude, claude-auth, codex, 默认全部)
  23. * --yes, -y 跳过确认提示
  24. * --dry-run 仅显示将要清理的内容,不实际执行
  25. * --help, -h 显示帮助信息
  26. */
  27. import { createInterface, type Interface as ReadlineInterface } from "node:readline";
  28. import { and, asc, eq, ilike, inArray, isNull, lt, or } from "drizzle-orm";
  29. import { drizzle, type PostgresJsDatabase } from "drizzle-orm/postgres-js";
  30. import Redis, { type RedisOptions } from "ioredis";
  31. import postgres from "postgres";
  32. import * as schema from "@/drizzle/schema";
  33. import { getGlobalActiveSessionsKey, getKeyActiveSessionsKey } from "@/lib/redis/active-session-keys";
  34. // ============================================================================
  35. // 常量配置
  36. // ============================================================================
  37. const SCAN_BATCH_SIZE = 500;
  38. const PIPELINE_BATCH_SIZE = 200;
  39. // Session 相关的 Redis Key 后缀
  40. const SESSION_KEY_SUFFIXES = [
  41. "provider",
  42. "info",
  43. "usage",
  44. "key",
  45. "last_seen",
  46. "messages",
  47. "response",
  48. "concurrent_count",
  49. ];
  50. // ============================================================================
  51. // 类型定义
  52. // ============================================================================
  53. type Database = PostgresJsDatabase<typeof schema>;
  54. type PostgresClient = ReturnType<typeof postgres>;
  55. type ProviderType =
  56. | "claude"
  57. | "claude-auth"
  58. | "codex"
  59. | "gemini-cli"
  60. | "gemini"
  61. | "openai-compatible";
  62. const VALID_PROVIDER_TYPES: ReadonlyArray<ProviderType | "all"> = [
  63. "claude",
  64. "claude-auth",
  65. "codex",
  66. "gemini-cli",
  67. "gemini",
  68. "openai-compatible",
  69. "all",
  70. ];
  71. function isValidProviderType(value: string): value is ProviderType | "all" {
  72. return VALID_PROVIDER_TYPES.includes(value as ProviderType | "all");
  73. }
  74. interface ProviderRecord {
  75. id: number;
  76. name: string;
  77. priority: number;
  78. providerType: string;
  79. isEnabled: boolean;
  80. }
  81. interface SessionBinding {
  82. providerId: number;
  83. keyId?: number | null;
  84. }
  85. interface CliOptions {
  86. mode: "interactive" | "priority" | "id" | "name";
  87. priorityThreshold?: number;
  88. providerIds?: number[];
  89. namePattern?: string;
  90. providerType?: ProviderType | "all";
  91. assumeYes: boolean;
  92. dryRun: boolean;
  93. }
  94. interface CleanupResult {
  95. sessionCount: number;
  96. deletedKeys: number;
  97. zsetRemovals: number;
  98. missingKeyRefs: number;
  99. }
  100. // ============================================================================
  101. // CLI 参数解析
  102. // ============================================================================
  103. function printUsage(): void {
  104. console.log(`
  105. 用法: bun run scripts/clear-session-bindings.ts [选项]
  106. 清除指定供应商的会话绑定。支持交互式选择或通过参数指定。
  107. 模式选择(互斥,不指定则进入交互模式):
  108. --priority, -p <number> 按优先级筛选(清除 priority < 该值的供应商)
  109. --id <ids> 指定供应商 ID(逗号分隔,如 1,2,3)
  110. --name <pattern> 按名称模糊匹配供应商
  111. 筛选选项:
  112. --type <type> 供应商类型筛选(claude, claude-auth, codex, all)
  113. 默认: all
  114. 执行选项:
  115. --yes, -y 跳过确认提示,直接执行
  116. --dry-run 仅显示将要清理的内容,不实际执行删除操作
  117. --help, -h 显示此帮助信息
  118. 示例:
  119. # 交互式模式(推荐新手使用)
  120. bun run scripts/clear-session-bindings.ts
  121. # 清除优先级小于 10 的 Claude 供应商
  122. bun run scripts/clear-session-bindings.ts --priority 10 --type claude
  123. # 清除指定 ID 的供应商绑定
  124. bun run scripts/clear-session-bindings.ts --id 1,2,3
  125. # 按名称模糊匹配
  126. bun run scripts/clear-session-bindings.ts --name "cubence"
  127. # 预览模式
  128. bun run scripts/clear-session-bindings.ts --priority 5 --dry-run
  129. `);
  130. }
  131. function parseCliArgs(args: string[]): CliOptions {
  132. let priorityValue: number | null = null;
  133. let providerIds: number[] | null = null;
  134. let namePattern: string | null = null;
  135. let providerType: ProviderType | "all" | undefined;
  136. let assumeYes = false;
  137. let dryRun = false;
  138. for (let i = 0; i < args.length; i++) {
  139. const arg = args[i];
  140. if (arg === "--priority" || arg === "-p") {
  141. const nextValue = args[++i];
  142. if (!nextValue) throw new Error("--priority 需要一个数值参数");
  143. priorityValue = Number.parseInt(nextValue, 10);
  144. if (Number.isNaN(priorityValue)) throw new Error("--priority 必须是整数");
  145. } else if (arg.startsWith("--priority=")) {
  146. priorityValue = Number.parseInt(arg.split("=")[1], 10);
  147. if (Number.isNaN(priorityValue)) throw new Error("--priority 必须是整数");
  148. } else if (arg === "--id") {
  149. const nextValue = args[++i];
  150. if (!nextValue) throw new Error("--id 需要供应商 ID 列表");
  151. providerIds = nextValue.split(",").map((s) => {
  152. const id = Number.parseInt(s.trim(), 10);
  153. if (Number.isNaN(id)) throw new Error(`无效的供应商 ID: ${s}`);
  154. return id;
  155. });
  156. } else if (arg.startsWith("--id=")) {
  157. providerIds = arg
  158. .split("=")[1]
  159. .split(",")
  160. .map((s) => {
  161. const id = Number.parseInt(s.trim(), 10);
  162. if (Number.isNaN(id)) throw new Error(`无效的供应商 ID: ${s}`);
  163. return id;
  164. });
  165. } else if (arg === "--name") {
  166. namePattern = args[++i];
  167. if (!namePattern) throw new Error("--name 需要一个匹配模式");
  168. } else if (arg.startsWith("--name=")) {
  169. namePattern = arg.split("=")[1];
  170. } else if (arg === "--type") {
  171. const typeValue = args[++i];
  172. if (!typeValue || !isValidProviderType(typeValue)) {
  173. throw new Error(
  174. `无效的供应商类型: ${typeValue}。有效值: ${VALID_PROVIDER_TYPES.join(", ")}`
  175. );
  176. }
  177. providerType = typeValue;
  178. } else if (arg.startsWith("--type=")) {
  179. const typeValue = arg.split("=")[1];
  180. if (!isValidProviderType(typeValue)) {
  181. throw new Error(
  182. `无效的供应商类型: ${typeValue}。有效值: ${VALID_PROVIDER_TYPES.join(", ")}`
  183. );
  184. }
  185. providerType = typeValue;
  186. } else if (arg === "--yes" || arg === "-y") {
  187. assumeYes = true;
  188. } else if (arg === "--dry-run") {
  189. dryRun = true;
  190. } else if (arg === "--help" || arg === "-h") {
  191. printUsage();
  192. process.exit(0);
  193. } else {
  194. throw new Error(`未知参数: ${arg}`);
  195. }
  196. }
  197. // 确定模式
  198. const modeCount = [priorityValue !== null, providerIds !== null, namePattern !== null].filter(
  199. Boolean
  200. ).length;
  201. if (modeCount > 1) {
  202. throw new Error("--priority, --id, --name 不能同时使用,请选择一种筛选方式");
  203. }
  204. let mode: CliOptions["mode"] = "interactive";
  205. if (priorityValue !== null) mode = "priority";
  206. else if (providerIds !== null) mode = "id";
  207. else if (namePattern !== null) mode = "name";
  208. return {
  209. mode,
  210. priorityThreshold: priorityValue ?? undefined,
  211. providerIds: providerIds ?? undefined,
  212. namePattern: namePattern ?? undefined,
  213. providerType,
  214. assumeYes,
  215. dryRun,
  216. };
  217. }
  218. // ============================================================================
  219. // 数据库和 Redis 连接
  220. // ============================================================================
  221. function createDatabaseConnection(connectionString: string): {
  222. client: PostgresClient;
  223. db: Database;
  224. } {
  225. const client = postgres(connectionString);
  226. const db = drizzle(client, { schema });
  227. return { client, db };
  228. }
  229. async function createRedisClient(redisUrl: string): Promise<Redis> {
  230. const options: RedisOptions = {
  231. enableOfflineQueue: false,
  232. maxRetriesPerRequest: 3,
  233. retryStrategy(times) {
  234. if (times > 5) {
  235. console.error("[Redis] 重试次数已达上限,放弃连接");
  236. return null;
  237. }
  238. return Math.min(times * 200, 2000);
  239. },
  240. lazyConnect: true,
  241. };
  242. if (redisUrl.startsWith("rediss://")) {
  243. const rejectUnauthorized = process.env.REDIS_TLS_REJECT_UNAUTHORIZED !== "false";
  244. try {
  245. const url = new URL(redisUrl);
  246. options.tls = {
  247. host: url.hostname,
  248. servername: url.hostname, // SNI support for cloud Redis providers
  249. rejectUnauthorized,
  250. };
  251. } catch {
  252. options.tls = { rejectUnauthorized };
  253. }
  254. }
  255. const redis = new Redis(redisUrl, options);
  256. // 等待连接就绪
  257. await redis.connect();
  258. return redis;
  259. }
  260. async function safeCloseRedis(redis: Redis | null): Promise<void> {
  261. if (!redis) return;
  262. try {
  263. await redis.quit();
  264. } catch (error) {
  265. console.warn("关闭 Redis 连接时发生错误:", error);
  266. }
  267. }
  268. async function safeClosePostgres(client: PostgresClient | null): Promise<void> {
  269. if (!client) return;
  270. try {
  271. await client.end({ timeout: 5 });
  272. } catch (error) {
  273. console.warn("关闭数据库连接时发生错误:", error);
  274. }
  275. }
  276. // ============================================================================
  277. // 数据库查询
  278. // ============================================================================
  279. async function fetchAllProviders(db: Database): Promise<ProviderRecord[]> {
  280. const rows = await db
  281. .select({
  282. id: schema.providers.id,
  283. name: schema.providers.name,
  284. priority: schema.providers.priority,
  285. providerType: schema.providers.providerType,
  286. isEnabled: schema.providers.isEnabled,
  287. })
  288. .from(schema.providers)
  289. .where(isNull(schema.providers.deletedAt))
  290. .orderBy(asc(schema.providers.priority), asc(schema.providers.id));
  291. return rows.map((row) => ({
  292. id: row.id,
  293. name: row.name,
  294. priority: row.priority ?? 0,
  295. providerType: row.providerType,
  296. isEnabled: row.isEnabled,
  297. }));
  298. }
  299. async function fetchProvidersByPriority(
  300. db: Database,
  301. threshold: number,
  302. providerType?: ProviderType | "all"
  303. ): Promise<ProviderRecord[]> {
  304. const conditions = [isNull(schema.providers.deletedAt), lt(schema.providers.priority, threshold)];
  305. if (providerType && providerType !== "all") {
  306. if (providerType === "claude") {
  307. conditions.push(
  308. or(
  309. eq(schema.providers.providerType, "claude"),
  310. eq(schema.providers.providerType, "claude-auth")
  311. )!
  312. );
  313. } else {
  314. conditions.push(eq(schema.providers.providerType, providerType));
  315. }
  316. }
  317. const rows = await db
  318. .select({
  319. id: schema.providers.id,
  320. name: schema.providers.name,
  321. priority: schema.providers.priority,
  322. providerType: schema.providers.providerType,
  323. isEnabled: schema.providers.isEnabled,
  324. })
  325. .from(schema.providers)
  326. .where(and(...conditions))
  327. .orderBy(asc(schema.providers.priority), asc(schema.providers.id));
  328. return rows.map((row) => ({
  329. id: row.id,
  330. name: row.name,
  331. priority: row.priority ?? 0,
  332. providerType: row.providerType,
  333. isEnabled: row.isEnabled,
  334. }));
  335. }
  336. async function fetchProvidersByIds(db: Database, ids: number[]): Promise<ProviderRecord[]> {
  337. const rows = await db
  338. .select({
  339. id: schema.providers.id,
  340. name: schema.providers.name,
  341. priority: schema.providers.priority,
  342. providerType: schema.providers.providerType,
  343. isEnabled: schema.providers.isEnabled,
  344. })
  345. .from(schema.providers)
  346. .where(and(isNull(schema.providers.deletedAt), inArray(schema.providers.id, ids)))
  347. .orderBy(asc(schema.providers.priority), asc(schema.providers.id));
  348. return rows.map((row) => ({
  349. id: row.id,
  350. name: row.name,
  351. priority: row.priority ?? 0,
  352. providerType: row.providerType,
  353. isEnabled: row.isEnabled,
  354. }));
  355. }
  356. async function fetchProvidersByName(db: Database, pattern: string): Promise<ProviderRecord[]> {
  357. const rows = await db
  358. .select({
  359. id: schema.providers.id,
  360. name: schema.providers.name,
  361. priority: schema.providers.priority,
  362. providerType: schema.providers.providerType,
  363. isEnabled: schema.providers.isEnabled,
  364. })
  365. .from(schema.providers)
  366. .where(and(isNull(schema.providers.deletedAt), ilike(schema.providers.name, `%${pattern}%`)))
  367. .orderBy(asc(schema.providers.priority), asc(schema.providers.id));
  368. return rows.map((row) => ({
  369. id: row.id,
  370. name: row.name,
  371. priority: row.priority ?? 0,
  372. providerType: row.providerType,
  373. isEnabled: row.isEnabled,
  374. }));
  375. }
  376. // ============================================================================
  377. // Redis 操作
  378. // ============================================================================
  379. function extractSessionIdFromKey(redisKey: string): string | null {
  380. const match = /^session:(.+):provider$/.exec(redisKey);
  381. return match ? match[1] : null;
  382. }
  383. async function findSessionsBoundToProviders(
  384. redis: Redis,
  385. providerIds: Set<number>
  386. ): Promise<Map<string, SessionBinding>> {
  387. const sessionMap = new Map<string, SessionBinding>();
  388. if (providerIds.size === 0) return sessionMap;
  389. let cursor = "0";
  390. let scannedKeys = 0;
  391. do {
  392. const [nextCursor, keys] = await redis.scan(
  393. cursor,
  394. "MATCH",
  395. "session:*:provider",
  396. "COUNT",
  397. SCAN_BATCH_SIZE
  398. );
  399. cursor = nextCursor;
  400. if (!keys || keys.length === 0) continue;
  401. scannedKeys += keys.length;
  402. const pipeline = redis.pipeline();
  403. for (const key of keys) {
  404. pipeline.get(key);
  405. }
  406. const results = await pipeline.exec();
  407. keys.forEach((key, index) => {
  408. const result = results?.[index];
  409. if (!result) return;
  410. const [error, value] = result;
  411. if (error || typeof value !== "string") return;
  412. const providerId = Number.parseInt(value, 10);
  413. if (!Number.isFinite(providerId) || !providerIds.has(providerId)) return;
  414. const sessionId = extractSessionIdFromKey(key);
  415. if (!sessionId) return;
  416. sessionMap.set(sessionId, { providerId });
  417. });
  418. } while (cursor !== "0");
  419. console.log(` 扫描了 ${scannedKeys} 个 session:*:provider key`);
  420. return sessionMap;
  421. }
  422. async function populateSessionKeyBindings(
  423. redis: Redis,
  424. sessionMap: Map<string, SessionBinding>
  425. ): Promise<void> {
  426. if (sessionMap.size === 0) return;
  427. const sessionIds = Array.from(sessionMap.keys());
  428. // 第一轮:尝试从 session:${sessionId}:key 获取
  429. for (let i = 0; i < sessionIds.length; i += PIPELINE_BATCH_SIZE) {
  430. const chunk = sessionIds.slice(i, i + PIPELINE_BATCH_SIZE);
  431. const pipeline = redis.pipeline();
  432. for (const sessionId of chunk) {
  433. pipeline.get(`session:${sessionId}:key`);
  434. }
  435. const results = await pipeline.exec();
  436. chunk.forEach((sessionId, index) => {
  437. const result = results?.[index];
  438. if (!result) return;
  439. const [error, value] = result;
  440. if (error || typeof value !== "string") return;
  441. const keyId = Number.parseInt(value, 10);
  442. if (Number.isFinite(keyId)) {
  443. const binding = sessionMap.get(sessionId);
  444. if (binding) binding.keyId = keyId;
  445. }
  446. });
  447. }
  448. // 第二轮:从 info hash 补充缺失的 keyId
  449. const missingKeyIds = sessionIds.filter((id) => sessionMap.get(id)?.keyId == null);
  450. if (missingKeyIds.length === 0) return;
  451. console.log(` ${missingKeyIds.length} 个 session 缺少 key 绑定,尝试从 info hash 获取...`);
  452. for (let i = 0; i < missingKeyIds.length; i += PIPELINE_BATCH_SIZE) {
  453. const chunk = missingKeyIds.slice(i, i + PIPELINE_BATCH_SIZE);
  454. const pipeline = redis.pipeline();
  455. for (const sessionId of chunk) {
  456. pipeline.hget(`session:${sessionId}:info`, "keyId");
  457. }
  458. const results = await pipeline.exec();
  459. chunk.forEach((sessionId, index) => {
  460. const result = results?.[index];
  461. if (!result) return;
  462. const [error, value] = result;
  463. if (error || typeof value !== "string") return;
  464. const keyId = Number.parseInt(value, 10);
  465. if (Number.isFinite(keyId)) {
  466. const binding = sessionMap.get(sessionId);
  467. if (binding) binding.keyId = keyId;
  468. }
  469. });
  470. }
  471. }
  472. async function clearSessionBindings(
  473. redis: Redis,
  474. sessionMap: Map<string, SessionBinding>,
  475. dryRun: boolean
  476. ): Promise<CleanupResult> {
  477. const entries = Array.from(sessionMap.entries());
  478. let deletedKeys = 0;
  479. let zsetRemovals = 0;
  480. let missingKeyRefs = 0;
  481. if (dryRun) {
  482. for (const [, binding] of entries) {
  483. deletedKeys += SESSION_KEY_SUFFIXES.length;
  484. zsetRemovals += 2;
  485. if (binding.keyId != null) zsetRemovals += 1;
  486. else missingKeyRefs += 1;
  487. }
  488. return { sessionCount: entries.length, deletedKeys, zsetRemovals, missingKeyRefs };
  489. }
  490. for (let i = 0; i < entries.length; i += PIPELINE_BATCH_SIZE) {
  491. const chunk = entries.slice(i, i + PIPELINE_BATCH_SIZE);
  492. const pipeline = redis.pipeline();
  493. const commandTypes: Array<"del" | "zrem"> = [];
  494. for (const [sessionId, binding] of chunk) {
  495. const keysToDelete = SESSION_KEY_SUFFIXES.map((suffix) => `session:${sessionId}:${suffix}`);
  496. pipeline.del(...keysToDelete);
  497. commandTypes.push("del");
  498. pipeline.zrem(getGlobalActiveSessionsKey(), sessionId);
  499. commandTypes.push("zrem");
  500. pipeline.zrem(`provider:${binding.providerId}:active_sessions`, sessionId);
  501. commandTypes.push("zrem");
  502. if (binding.keyId != null) {
  503. pipeline.zrem(getKeyActiveSessionsKey(binding.keyId), sessionId);
  504. commandTypes.push("zrem");
  505. } else {
  506. missingKeyRefs += 1;
  507. }
  508. }
  509. const results = await pipeline.exec();
  510. results?.forEach(([error, value], index) => {
  511. if (error) return;
  512. const type = commandTypes[index];
  513. if (type === "del" && typeof value === "number") deletedKeys += value;
  514. else if (type === "zrem" && typeof value === "number") zsetRemovals += value;
  515. });
  516. const processed = Math.min(i + PIPELINE_BATCH_SIZE, entries.length);
  517. process.stdout.write(`\r 清理进度: ${processed}/${entries.length}`);
  518. }
  519. console.log();
  520. return { sessionCount: entries.length, deletedKeys, zsetRemovals, missingKeyRefs };
  521. }
  522. // ============================================================================
  523. // 交互式界面
  524. // ============================================================================
  525. class InteractiveMenu {
  526. private rl: ReadlineInterface;
  527. constructor() {
  528. this.rl = createInterface({
  529. input: process.stdin,
  530. output: process.stdout,
  531. });
  532. }
  533. async question(prompt: string): Promise<string> {
  534. return new Promise((resolve) => {
  535. this.rl.question(prompt, (answer) => resolve(answer.trim()));
  536. });
  537. }
  538. close(): void {
  539. this.rl.close();
  540. }
  541. displayProviderList(providers: ProviderRecord[]): void {
  542. console.log("\n可用供应商列表:\n");
  543. console.log(" 序号 ID 名称 优先级 类型 状态");
  544. console.log(` ${"-".repeat(75)}`);
  545. providers.forEach((p, index) => {
  546. const status = p.isEnabled ? "启用" : "禁用";
  547. const name = p.name.length > 28 ? `${p.name.substring(0, 25)}...` : p.name.padEnd(28);
  548. console.log(
  549. ` ${String(index + 1).padStart(4)} ${String(p.id).padStart(4)} ${name} ${String(p.priority).padStart(6)} ${p.providerType.padEnd(13)} ${status}`
  550. );
  551. });
  552. console.log();
  553. }
  554. async selectProviders(providers: ProviderRecord[]): Promise<ProviderRecord[]> {
  555. this.displayProviderList(providers);
  556. console.log("选择方式:");
  557. console.log(" - 输入序号(逗号分隔): 1,2,3");
  558. console.log(" - 输入范围: 1-5");
  559. console.log(" - 输入 'all' 选择全部");
  560. console.log(" - 输入 'q' 退出\n");
  561. const input = await this.question("请选择要清理的供应商: ");
  562. if (input.toLowerCase() === "q") {
  563. return [];
  564. }
  565. if (input.toLowerCase() === "all") {
  566. return providers;
  567. }
  568. const selectedIndices = new Set<number>();
  569. // 解析输入
  570. const parts = input.split(",").map((s) => s.trim());
  571. for (const part of parts) {
  572. if (part.includes("-")) {
  573. const [start, end] = part.split("-").map((s) => Number.parseInt(s.trim(), 10));
  574. if (!Number.isNaN(start) && !Number.isNaN(end)) {
  575. for (let i = start; i <= end; i++) {
  576. if (i >= 1 && i <= providers.length) {
  577. selectedIndices.add(i - 1);
  578. }
  579. }
  580. }
  581. } else {
  582. const index = Number.parseInt(part, 10);
  583. if (!Number.isNaN(index) && index >= 1 && index <= providers.length) {
  584. selectedIndices.add(index - 1);
  585. }
  586. }
  587. }
  588. return Array.from(selectedIndices)
  589. .sort((a, b) => a - b)
  590. .map((i) => providers[i]);
  591. }
  592. async confirm(message: string): Promise<boolean> {
  593. const answer = await this.question(`${message} [y/N]: `);
  594. return answer.toLowerCase() === "y" || answer.toLowerCase() === "yes";
  595. }
  596. async selectMainAction(): Promise<"priority" | "select" | "type" | "name" | "quit"> {
  597. console.log("\n请选择操作:");
  598. console.log(" 1. 按优先级筛选");
  599. console.log(" 2. 手动选择供应商");
  600. console.log(" 3. 按类型筛选");
  601. console.log(" 4. 按名称搜索");
  602. console.log(" q. 退出\n");
  603. const input = await this.question("请输入选项: ");
  604. switch (input.toLowerCase()) {
  605. case "1":
  606. return "priority";
  607. case "2":
  608. return "select";
  609. case "3":
  610. return "type";
  611. case "4":
  612. return "name";
  613. case "q":
  614. case "quit":
  615. case "exit":
  616. return "quit";
  617. default:
  618. console.log("无效选项,请重新输入");
  619. return this.selectMainAction();
  620. }
  621. }
  622. async selectProviderType(providers: ProviderRecord[]): Promise<string | null> {
  623. // 统计每种类型的数量
  624. const typeCounts = new Map<string, number>();
  625. for (const p of providers) {
  626. typeCounts.set(p.providerType, (typeCounts.get(p.providerType) || 0) + 1);
  627. }
  628. const types = Array.from(typeCounts.keys()).sort();
  629. console.log("\n可用的供应商类型:\n");
  630. types.forEach((type, index) => {
  631. console.log(` ${index + 1}. ${type} (${typeCounts.get(type)} 个)`);
  632. });
  633. console.log(` a. 全部类型`);
  634. console.log(` q. 返回上级菜单\n`);
  635. const input = await this.question("请选择类型: ");
  636. if (input.toLowerCase() === "q") {
  637. return null;
  638. }
  639. if (input.toLowerCase() === "a" || input.toLowerCase() === "all") {
  640. return "all";
  641. }
  642. const index = Number.parseInt(input, 10);
  643. if (!Number.isNaN(index) && index >= 1 && index <= types.length) {
  644. return types[index - 1];
  645. }
  646. // 也支持直接输入类型名称
  647. if (types.includes(input)) {
  648. return input;
  649. }
  650. console.log("无效选项,请重新输入");
  651. return this.selectProviderType(providers);
  652. }
  653. async inputPriority(): Promise<number | null> {
  654. const input = await this.question("请输入优先级阈值(清除 priority < 该值的供应商): ");
  655. const value = Number.parseInt(input, 10);
  656. if (Number.isNaN(value)) {
  657. console.log("无效的数字,请重新输入");
  658. return this.inputPriority();
  659. }
  660. return value;
  661. }
  662. async inputNamePattern(): Promise<string | null> {
  663. const input = await this.question("请输入供应商名称(模糊匹配): ");
  664. return input || null;
  665. }
  666. }
  667. // ============================================================================
  668. // 主流程
  669. // ============================================================================
  670. function displaySelectedProviders(providers: ProviderRecord[]): void {
  671. console.log(`\n已选择 ${providers.length} 个供应商:\n`);
  672. console.table(
  673. providers.map((p) => ({
  674. ID: p.id,
  675. 名称: p.name,
  676. 优先级: p.priority,
  677. 类型: p.providerType,
  678. 状态: p.isEnabled ? "启用" : "禁用",
  679. }))
  680. );
  681. }
  682. function displayResult(result: CleanupResult, dryRun: boolean): void {
  683. console.log(`\n${"=".repeat(60)}`);
  684. console.log(dryRun ? " Dry-Run 结果摘要" : " 清理完成");
  685. console.log("=".repeat(60));
  686. console.log(` Session 数量: ${result.sessionCount}`);
  687. console.log(` 删除的 Key 数量: ${result.deletedKeys}`);
  688. console.log(` ZSET 移除数量: ${result.zsetRemovals}`);
  689. if (result.missingKeyRefs > 0) {
  690. console.log(` 缺少 key 绑定: ${result.missingKeyRefs}`);
  691. }
  692. if (dryRun) {
  693. console.log("\n[Dry-Run] 以上为预计操作,实际未执行任何删除。");
  694. }
  695. }
  696. async function runCleanup(
  697. redis: Redis,
  698. providers: ProviderRecord[],
  699. dryRun: boolean
  700. ): Promise<CleanupResult | null> {
  701. console.log("\n正在扫描 session 绑定...");
  702. const providerIds = new Set(providers.map((p) => p.id));
  703. const sessions = await findSessionsBoundToProviders(redis, providerIds);
  704. if (sessions.size === 0) {
  705. console.log("\n没有找到需要清理的 session 绑定。");
  706. return null;
  707. }
  708. console.log(` 匹配到 ${sessions.size} 个 session`);
  709. console.log("正在获取 session 的 key 绑定信息...");
  710. await populateSessionKeyBindings(redis, sessions);
  711. console.log(dryRun ? "\n[Dry-Run] 预计清理:" : "\n正在清理...");
  712. return clearSessionBindings(redis, sessions, dryRun);
  713. }
  714. async function main(): Promise<void> {
  715. console.log("=".repeat(60));
  716. console.log(" Claude Code Hub - Session 绑定清理工具");
  717. console.log("=".repeat(60));
  718. const options = parseCliArgs(process.argv.slice(2));
  719. if (options.dryRun) {
  720. console.log("\n[Dry-Run 模式] 仅显示将要执行的操作,不实际删除数据\n");
  721. }
  722. const dsn = process.env.DSN;
  723. if (!dsn) throw new Error("DSN 环境变量未设置");
  724. const redisUrl = process.env.REDIS_URL;
  725. if (!redisUrl) throw new Error("REDIS_URL 环境变量未设置");
  726. console.log("正在连接数据库...");
  727. const { client, db } = createDatabaseConnection(dsn);
  728. let redis: Redis | null = null;
  729. let menu: InteractiveMenu | null = null;
  730. try {
  731. let targetProviders: ProviderRecord[] = [];
  732. if (options.mode === "interactive") {
  733. // 交互式模式
  734. menu = new InteractiveMenu();
  735. const allProviders = await fetchAllProviders(db);
  736. if (allProviders.length === 0) {
  737. console.log("\n数据库中没有供应商。");
  738. return;
  739. }
  740. const action = await menu.selectMainAction();
  741. if (action === "quit") {
  742. console.log("\n已退出。");
  743. return;
  744. }
  745. if (action === "priority") {
  746. const threshold = await menu.inputPriority();
  747. if (threshold === null) return;
  748. targetProviders = allProviders.filter((p) => p.priority < threshold);
  749. } else if (action === "type") {
  750. const selectedType = await menu.selectProviderType(allProviders);
  751. if (selectedType === null) {
  752. // 返回主菜单
  753. const retryAction = await menu.selectMainAction();
  754. if (retryAction === "quit") {
  755. console.log("\n已退出。");
  756. return;
  757. }
  758. // 简单处理:重新开始
  759. console.log("\n请重新运行脚本。");
  760. return;
  761. }
  762. const filteredProviders =
  763. selectedType === "all"
  764. ? allProviders
  765. : allProviders.filter((p) => p.providerType === selectedType);
  766. if (filteredProviders.length === 0) {
  767. console.log(`\n没有找到类型为 "${selectedType}" 的供应商。`);
  768. return;
  769. }
  770. // 从筛选后的列表中选择
  771. targetProviders = await menu.selectProviders(filteredProviders);
  772. } else if (action === "name") {
  773. const pattern = await menu.inputNamePattern();
  774. if (!pattern) return;
  775. targetProviders = allProviders.filter((p) =>
  776. p.name.toLowerCase().includes(pattern.toLowerCase())
  777. );
  778. } else {
  779. targetProviders = await menu.selectProviders(allProviders);
  780. }
  781. if (targetProviders.length === 0) {
  782. console.log("\n未选择任何供应商。");
  783. return;
  784. }
  785. displaySelectedProviders(targetProviders);
  786. if (!options.dryRun) {
  787. const confirmed = await menu.confirm("\n确认清理这些供应商的 session 绑定?");
  788. if (!confirmed) {
  789. console.log("\n操作已取消。");
  790. return;
  791. }
  792. }
  793. } else {
  794. // 命令行模式
  795. if (options.mode === "priority") {
  796. targetProviders = await fetchProvidersByPriority(
  797. db,
  798. options.priorityThreshold!,
  799. options.providerType
  800. );
  801. } else if (options.mode === "id") {
  802. targetProviders = await fetchProvidersByIds(db, options.providerIds!);
  803. } else if (options.mode === "name") {
  804. targetProviders = await fetchProvidersByName(db, options.namePattern!);
  805. }
  806. if (targetProviders.length === 0) {
  807. console.log("\n未找到符合条件的供应商。");
  808. return;
  809. }
  810. displaySelectedProviders(targetProviders);
  811. if (!options.assumeYes && !options.dryRun) {
  812. menu = new InteractiveMenu();
  813. const confirmed = await menu.confirm("\n确认清理这些供应商的 session 绑定?");
  814. if (!confirmed) {
  815. console.log("\n操作已取消。");
  816. return;
  817. }
  818. }
  819. }
  820. // 连接 Redis 并执行清理
  821. console.log("\n正在连接 Redis...");
  822. redis = await createRedisClient(redisUrl);
  823. console.log("Redis 连接成功");
  824. const result = await runCleanup(redis, targetProviders, options.dryRun);
  825. if (result) {
  826. displayResult(result, options.dryRun);
  827. }
  828. console.log();
  829. } finally {
  830. menu?.close();
  831. await safeCloseRedis(redis);
  832. await safeClosePostgres(client);
  833. }
  834. }
  835. main()
  836. .catch((error) => {
  837. console.error("\n发生错误:", error);
  838. process.exitCode = 1;
  839. })
  840. .finally(() => {
  841. process.exit(process.exitCode ?? 0);
  842. });