Auto-sync from GitHub Actions

core/engine.py

@@ -18,6 +18,7 @@ import os
 import re
 import time
 from collections import OrderedDict
+from contextlib import asynccontextmanager
 from pathlib import Path
 from datetime import datetime
 from typing import Any, Optional
@@ -203,6 +204,9 @@ class DanbooruTagger:
     # 进程级 CPU 并发闸门：串行化所有 CPU 密集型操作（search / get_related /
     # get_group_candidates），避免并发抢占 CPU 而拖垮 asyncio 事件循环。
     _cpu_sem: Optional[asyncio.Semaphore] = None
+    _cpu_sem_limit: Optional[int] = None
+    _cpu_active_tasks: int = 0
+    _cpu_waiting_tasks: int = 0
 
     @classmethod
     def is_ready(cls) -> bool:
@@ -636,9 +640,40 @@ class DanbooruTagger:
     @classmethod
     def _get_cpu_sem(cls) -> asyncio.Semaphore:
         if cls._cpu_sem is None:
-            cls._cpu_sem = asyncio.Semaphore(_resolve_cpu_sem_limit())
+            cls._cpu_sem_limit = _resolve_cpu_sem_limit()
+            cls._cpu_sem = asyncio.Semaphore(cls._cpu_sem_limit)
         return cls._cpu_sem
 
+    @classmethod
+    @asynccontextmanager
+    async def _cpu_slot(cls):
+        """申请共享 CPU 计算槽，并维护可公开展示的进程内负载计数。"""
+        sem = cls._get_cpu_sem()
+        acquired = False
+        cls._cpu_waiting_tasks += 1
+        try:
+            await sem.acquire()
+            acquired = True
+            cls._cpu_waiting_tasks = max(0, cls._cpu_waiting_tasks - 1)
+            cls._cpu_active_tasks += 1
+            yield
+        finally:
+            if acquired:
+                cls._cpu_active_tasks = max(0, cls._cpu_active_tasks - 1)
+                sem.release()
+            else:
+                cls._cpu_waiting_tasks = max(0, cls._cpu_waiting_tasks - 1)
+
+    @classmethod
+    def get_load_snapshot(cls) -> dict[str, int]:
+        """返回当前进程的 CPU 计算任务负载快照。"""
+        capacity = cls._cpu_sem_limit or _resolve_cpu_sem_limit()
+        return {
+            'active': cls._cpu_active_tasks,
+            'waiting': cls._cpu_waiting_tasks,
+            'capacity': capacity,
+        }
+
     async def search_async(self, request: SearchRequest) -> SearchResponse:
         """search() 的并发安全异步封装：共享闸门串行化 + 线程池执行。
 
@@ -646,7 +681,7 @@ class DanbooruTagger:
         asyncio.to_thread(self.search)，以共享同一个 CPU 并发闸门。
         包含 60 秒超时，防止异常卡死导致信号量永久泄漏。
         """
-        async with self._get_cpu_sem():
+        async with self._cpu_slot():
             return await asyncio.wait_for(
                 asyncio.to_thread(self.search, request),
                 timeout=120.0,
@@ -661,7 +696,7 @@ class DanbooruTagger:
         target_categories: set[str] | None = None,
     ) -> list:
         """get_related() 的并发安全异步封装，共享同一个 CPU 闸门。"""
-        async with self._get_cpu_sem():
+        async with self._cpu_slot():
             return await asyncio.to_thread(
                 self.get_related, seed_tags, exclude, limit, show_nsfw, target_categories,
             )
@@ -672,7 +707,7 @@ class DanbooruTagger:
         show_nsfw: bool = True,
     ) -> list[dict]:
         """get_group_candidates() 的并发安全异步封装，共享同一个 CPU 闸门。"""
-        async with self._get_cpu_sem():
+        async with self._cpu_slot():
             return await asyncio.to_thread(
                 self.get_group_candidates, selected_tags, show_nsfw,
             )
@@ -684,7 +719,7 @@ class DanbooruTagger:
         min_cooc: int = 5,
     ) -> list:
         """search_artists_by_tags() 的并发安全异步封装。"""
-        async with self._get_cpu_sem():
+        async with self._cpu_slot():
             return await asyncio.to_thread(
                 self.search_artists_by_tags, tags, limit, min_cooc,
             )
@@ -698,7 +733,7 @@ class DanbooruTagger:
         target_categories: list[str] | None = None,
     ) -> tuple:
         """search_artists_pipeline() 的并发安全异步封装。"""
-        async with self._get_cpu_sem():
+        async with self._cpu_slot():
             return await asyncio.wait_for(
                 asyncio.to_thread(
                     self.search_artists_pipeline,

ui_nicegui.py

@@ -12,6 +12,7 @@ import sys
 sys.stdout.reconfigure(line_buffering=True)
 print("[UI] 脚本开始执行", flush=True)
 import asyncio
+import logging
 import os
 import re
 import time
@@ -88,7 +89,23 @@ from core.workspace import (
 from platform_utils import is_cloud, get_host_port, nsfw_allowed
 from mcp_server import mcp
 
-import logging
+
+# 仅统计当前进程中已建立 Socket.IO 连接的 UI 页面，不等同于唯一用户数。
+_ACTIVE_UI_CLIENT_IDS: set[str] = set()
+
+
+def _mark_ui_session_active(client_id: str) -> None:
+    _ACTIVE_UI_CLIENT_IDS.add(client_id)
+
+
+def _mark_ui_session_inactive(client_id: str) -> None:
+    _ACTIVE_UI_CLIENT_IDS.discard(client_id)
+
+
+def _get_active_ui_session_count() -> int:
+    return len(_ACTIVE_UI_CLIENT_IDS)
+
+
 logging.getLogger("httpx").setLevel(logging.WARNING)
 logging.getLogger("huggingface_hub").setLevel(logging.WARNING)
 logging.getLogger("mcp").setLevel(logging.WARNING)
@@ -380,6 +397,8 @@ class DanbooruSearchUI:
     def __init__(self):
         self.search_count_label = None
         self.service_status_container = None
+        self.service_status_timer = None
+        self._last_service_status_key = None
         self.current_search_interacted = True
         self._telemetry_search_started_at: float | None = None
         self._telemetry_selection_recorded = False
@@ -493,7 +512,7 @@ class DanbooruSearchUI:
         self._cat_checkboxes: dict[str, ui.checkbox] = {}
 
     def _update_footer_text(self):
-        if self.search_count_label is not None:
+        if self.search_count_label is not None and self._client_alive():
             try:
                 total = counter.get()
                 visits = counter.get_visits()
@@ -513,22 +532,47 @@ class DanbooruSearchUI:
                 pass
 
     def _update_service_status(self):
-        if self.service_status_container is None:
+        if self.service_status_container is None or not self._client_alive():
             return
+
+        ready = DanbooruTagger.is_ready()
+        online_sessions = _get_active_ui_session_count()
+        load = DanbooruTagger.get_load_snapshot()
+        active = load['active']
+        waiting = load['waiting']
+        capacity = load['capacity']
+        busy = ready and (waiting > 0 or active >= capacity)
+        status_key = (ready, busy, online_sessions, active, waiting, capacity)
+        if status_key == self._last_service_status_key:
+            return
+        self._last_service_status_key = status_key
+
         self.service_status_container.clear()
         with self.service_status_container:
-            if DanbooruTagger.is_ready():
-                with ui.row().classes(
-                    'w-full items-center gap-2 service-state-panel ready'
-                ):
-                    ui.icon('check_circle', size='18px', color='positive')
-                    ui.label('服务可用').classes('font-medium')
-            else:
+            if not ready:
                 with ui.row().classes(
                     'w-full items-center gap-2 service-state-panel loading'
                 ):
                     ui.spinner(size='18px', color='primary')
                     ui.label('引擎初始化中，请稍候…约需 5~10 分钟').classes('font-medium')
+            else:
+                with ui.row().classes(
+                    f'w-full items-center gap-2 service-state-panel {"busy" if busy else "ready"}'
+                ):
+                    ui.icon(
+                        'schedule' if busy else 'check_circle',
+                        size='18px',
+                        color='warning' if busy else 'positive',
+                    )
+                    parts = [
+                        '服务繁忙' if busy else '服务可用',
+                        f'{online_sessions} 个在线页面',
+                    ]
+                    if active > 0:
+                        parts.append(f'正在处理 {active} 个任务')
+                    if waiting > 0:
+                        parts.append(f'等待 {waiting} 个')
+                    ui.label(' · '.join(parts)).classes('font-medium')
 
     def _build_sponsor_dialog(self):
         with ui.dialog() as self.sponsor_dialog, ui.card().classes('w-full max-w-sm'):
@@ -872,6 +916,8 @@ class DanbooruSearchUI:
 
     def _save_config(self):
         """将当前控件状态序列化并写入 localStorage。"""
+        if not self._client_alive():
+            return
         cfg = self._collect_config_state()
         js = _json.dumps(cfg, ensure_ascii=False)
         ui.run_javascript(f"localStorage.setItem('{_CONFIG_LS_KEY}', {_json.dumps(js)});")
@@ -932,7 +978,7 @@ class DanbooruSearchUI:
     async def _restore_config(self):
         """从 localStorage 读取配置并恢复控件状态。"""
         try:
-            if getattr(ui.context.client, '_deleted', False):
+            if not self._client_alive():
                 return
             raw = await ui.run_javascript(
                 f"localStorage.getItem('{_CONFIG_LS_KEY}');",
@@ -1015,6 +1061,11 @@ class DanbooruSearchUI:
                     border: 1px solid #a7f3d0;
                     color: #047857;
                 }
+                .service-state-panel.busy {
+                    background: #fff7ed;
+                    border: 1px solid #fed7aa;
+                    color: #c2410c;
+                }
                 .query-insight-panel {
                     background: #f8fafc;
                     border: 1px solid #dbe4ee;
@@ -1342,6 +1393,7 @@ class DanbooruSearchUI:
             with ui.element('div').classes('w-full border-t border-slate-100 pt-3 mt-1'):
                 self.service_status_container = ui.column().classes('w-full gap-0')
                 self._update_service_status()
+                self.service_status_timer = ui.timer(1.0, self._update_service_status)
 
     # ── 工作区工具 ────────────────────────────────────────────────────────
 
@@ -2550,6 +2602,8 @@ class DanbooruSearchUI:
             pass  # 事件上下文已销毁（UI 重建中），数据仍在内存里，下次保存时会同步
 
     def _save_history(self):
+        if not self._client_alive():
+            return
         while True:
             try:
                 data = dump_collection(self.search_history, label='history')
@@ -2568,6 +2622,8 @@ class DanbooruSearchUI:
                 return
 
     def _save_favorites(self):
+        if not self._client_alive():
+            return False
         try:
             data = dump_collection(self.favorites, label='favorites')
             ui.run_javascript(
@@ -3024,7 +3080,11 @@ class DanbooruSearchUI:
 
     async def _hide_banner_when_ready(self):
         while not DanbooruTagger.is_ready():
+            if not self._client_alive():
+                return
             await asyncio.sleep(1)
+        if not self._client_alive():
+            return
         if self.init_banner:
             self.init_banner.set_visibility(False)
         self._update_service_status()
@@ -3035,9 +3095,11 @@ class DanbooruSearchUI:
 
     def _client_alive(self) -> bool:
         try:
-            _ = self.search_btn.client
-            return True
-        except RuntimeError:
+            client = self.client
+            if client is None and self.search_btn is not None:
+                client = self.search_btn.client
+            return client is not None and not bool(getattr(client, '_deleted', False))
+        except (AttributeError, RuntimeError):
             return False
 
     # ── 分词筛选 ──────────────────────────────────────────────────────────
@@ -4175,6 +4237,18 @@ class DanbooruSearchUI:
 
 @ui.page('/')
 async def main_page():
+    client = ui.context.client
+    client_id = client.id
+
+    def mark_connected(*_):
+        _mark_ui_session_active(client_id)
+
+    def mark_disconnected(*_):
+        _mark_ui_session_inactive(client_id)
+
+    client.on_connect(mark_connected)
+    client.on_disconnect(mark_disconnected)
+
     app_ui = DanbooruSearchUI()
     app_ui.build_page()