From 906aceff6e2824e7d74fc916335e1742f8ca835e Mon Sep 17 00:00:00 2001
From: kyx236 <kyxjames23@gmail.com>
Date: Tue, 20 Jan 2026 21:15:50 +0800
Subject: [PATCH] update

---
 team_service.py | 102 ++++++++++++++++++++++++------------------------
 1 file changed, 50 insertions(+), 52 deletions(-)

diff --git a/team_service.py b/team_service.py
index 78b9a2a..200a8bf 100644
--- a/team_service.py
+++ b/team_service.py
@@ -104,8 +104,17 @@ def fetch_account_id(team: dict, silent: bool = False) -> str:
     return ""
 
 
+def _is_shutdown_requested() -> bool:
+    """检查是否收到停止请求"""
+    try:
+        import run
+        return run._shutdown_requested
+    except Exception:
+        return False
+
+
 def preload_all_account_ids() -> tuple[int, int]:
-    """预加载所有 Team 的 account_id
+    """预加载所有 Team 的 account_id (并行处理)
 
     在程序启动时调用，避免后续重复获取
     只处理有 token 的 Team，没有 token 的跳过
@@ -114,6 +123,7 @@ def preload_all_account_ids() -> tuple[int, int]:
         tuple: (success_count, fail_count)
     """
     import sys
+    from concurrent.futures import ThreadPoolExecutor, as_completed
 
     success_count = 0
     fail_count = 0
@@ -122,83 +132,71 @@ def preload_all_account_ids() -> tuple[int, int]:
     teams_with_token = [t for t in TEAMS if t.get("auth_token")]
     teams_need_fetch = [t for t in teams_with_token if not t.get("account_id")]
 
+    # 已缓存的数量
+    cached_count = len(teams_with_token) - len(teams_need_fetch)
+
     if not teams_need_fetch:
         if teams_with_token:
             log.success(f"所有 Team account_id 已缓存 ({len(teams_with_token)} 个)")
         return len(teams_with_token), 0
 
-    total = len(teams_with_token)
-    log.info(f"预加载 {len(teams_need_fetch)} 个 Team 的 account_id...", icon="sync")
+    total = len(teams_need_fetch)
+    log.info(f"并行预加载 {total} 个 Team 的 account_id...", icon="sync")
 
     need_save = False
     failed_teams = []
+    stopped = False
 
-    # 检测是否为 TTY 环境
-    is_tty = sys.stdout.isatty()
+    def fetch_one(team):
+        """获取单个 Team 的 account_id"""
+        if _is_shutdown_requested():
+            return team, None, "stopped"
+        account_id = fetch_account_id(team, silent=True)
+        return team, account_id, "ok" if account_id else "failed"
 
-    if is_tty:
-        # TTY 环境: 使用 rich 进度条
-        from rich.progress import Progress, SpinnerColumn, TextColumn, BarColumn, TaskProgressColumn
+    # 使用线程池并行获取 (最多 10 个并发)
+    max_workers = min(10, total)
+    completed = 0
 
-        with Progress(
-            SpinnerColumn(),
-            TextColumn("[progress.description]{task.description}"),
-            BarColumn(),
-            TaskProgressColumn(),
-            TextColumn("{task.fields[status]}"),
-        ) as progress:
-            task = progress.add_task("加载中", total=total, status="")
+    with ThreadPoolExecutor(max_workers=max_workers) as executor:
+        futures = {executor.submit(fetch_one, team): team for team in teams_need_fetch}
 
-            for team in teams_with_token:
-                progress.update(task, description=f"[cyan]{team['name']}", status="")
+        for future in as_completed(futures):
+            # 检查停止标志
+            if _is_shutdown_requested():
+                stopped = True
+                # 取消剩余任务
+                for f in futures:
+                    f.cancel()
+                break
 
-                if team.get("account_id"):
-                    success_count += 1
-                    progress.update(task, advance=1, status="[green]✓ 已缓存")
-                    continue
+            team, account_id, status = future.result()
+            completed += 1
 
-                account_id = fetch_account_id(team, silent=True)
-                if account_id:
-                    success_count += 1
-                    progress.update(task, advance=1, status="[green]✓")
-                    if team.get("format") == "new":
-                        need_save = True
-                else:
-                    fail_count += 1
-                    failed_teams.append(team['name'])
-                    progress.update(task, advance=1, status="[red]✗")
-    else:
-        # 非 TTY 环境 (systemd/journalctl): 使用普通日志输出
-        for idx, team in enumerate(teams_with_token, 1):
-            team_name = team['name']
-
-            if team.get("account_id"):
+            if status == "stopped":
+                stopped = True
+                break
+            elif status == "ok":
                 success_count += 1
-                log.info(f"预加载 [{idx}/{total}] {team_name}: ✓ 已缓存")
-                continue
-
-            account_id = fetch_account_id(team, silent=True)
-            if account_id:
-                success_count += 1
-                log.info(f"预加载 [{idx}/{total}] {team_name}: ✓ {account_id}")
                 if team.get("format") == "new":
                     need_save = True
+                log.info(f"预加载 [{completed}/{total}] {team['name']}: ✓ {account_id}")
             else:
                 fail_count += 1
-                failed_teams.append(team_name)
-                log.warning(f"预加载 [{idx}/{total}] {team_name}: ✗ 失败")
+                failed_teams.append(team['name'])
+                log.warning(f"预加载 [{completed}/{total}] {team['name']}: ✗ 失败")
 
-    # 输出失败的 team (仅 TTY 环境，非 TTY 已在循环中输出)
-    if is_tty:
-        for name in failed_teams:
-            log.warning(f"Team {name}: 获取 account_id 失败")
+    # 加上已缓存的数量
+    success_count += cached_count
 
     # 持久化到 team.json
     if need_save:
         if save_team_json():
             log.success(f"account_id 已保存到 team.json")
 
-    if fail_count == 0 and success_count > 0:
+    if stopped:
+        log.warning(f"预加载已停止 (已完成 {completed} 个)")
+    elif fail_count == 0 and success_count > 0:
         log.success(f"所有 Team account_id 加载完成 ({success_count} 个)")
     elif fail_count > 0:
         log.warning(f"account_id 加载: 成功 {success_count}, 失败 {fail_count}")