refactor: own shutdown lifecycle boundaries

app/agent/tools/base.py

@@ -3,7 +3,7 @@ import inspect
 import json
 import threading
 from abc import ABCMeta, abstractmethod
-from concurrent.futures import ThreadPoolExecutor
+from concurrent.futures import Future as ConcurrentFuture, ThreadPoolExecutor
 from contextvars import Context, copy_context
 from functools import partial
 from pathlib import Path
@@ -168,33 +168,130 @@ _blocking_semaphores = {
     for bucket, limit in _BLOCKING_BUCKET_LIMITS.items()
 }
 _blocking_executors: dict[str, ThreadPoolExecutor] = {}
-_blocking_executor_lock = threading.Lock()
+_blocking_retiring_executors: set[ThreadPoolExecutor] = set()
+_blocking_futures: dict[ConcurrentFuture[Any], ThreadPoolExecutor] = {}
+_blocking_executor_lock = threading.RLock()
+_blocking_executor_accepting = True
 
 
-def _get_blocking_executor(bucket: str) -> ThreadPoolExecutor:
-    """按桶懒加载线程池，避免在导入阶段创建过多 worker。"""
+def _discard_blocking_future(future: ConcurrentFuture[Any]) -> None:
+    """在同步调用到达终态后撤销 Future 与 retiring executor owner。"""
     with _blocking_executor_lock:
+        executor = _blocking_futures.pop(future, None)
+        if executor is None or executor not in _blocking_retiring_executors:
+            return
+        if executor not in _blocking_futures.values():
+            _blocking_retiring_executors.discard(executor)
+
+
+def _submit_blocking_call(
+    bucket: str,
+    bound_call: Callable[[], Any],
+) -> ConcurrentFuture[Any]:
+    """在提交门禁内原子取得 executor、提交调用并登记 Future owner。"""
+    context = copy_context()
+    with _blocking_executor_lock:
+        if not _blocking_executor_accepting:
+            raise RuntimeError("Agent 工具阻塞执行器正在关闭，不能再提交新任务")
         executor = _blocking_executors.get(bucket)
-        if executor:
-            return executor
-
-        limit = _BLOCKING_BUCKET_LIMITS[bucket]
-        executor = ThreadPoolExecutor(
-            max_workers=limit,
-            thread_name_prefix=f"agent-tool-{bucket}",
-        )
-        _blocking_executors[bucket] = executor
-        return executor
+        if executor is None:
+            limit = _BLOCKING_BUCKET_LIMITS[bucket]
+            executor = ThreadPoolExecutor(
+                max_workers=limit,
+                thread_name_prefix=f"agent-tool-{bucket}",
+            )
+            _blocking_executors[bucket] = executor
+        # 长期 worker 保持空底层上下文，每个任务只在自己的调用快照内运行。
+        future = Context().run(executor.submit, context.run, bound_call)
+        _blocking_futures[future] = executor
+        future.add_done_callback(_discard_blocking_future)
+        return future
 
 
-def shutdown_blocking_executors(*, wait: bool = True, cancel_futures: bool = False) -> None:
-    """关闭 Agent 工具阻塞线程池，释放长期运行进程或测试环境中的 worker。"""
+def _retire_blocking_executors(*, cancel_futures: bool) -> tuple[ThreadPoolExecutor, ...]:
+    """撤销活动 executor 的提交资格，并保留其运行 Future 对应的 owner。"""
     with _blocking_executor_lock:
-        executors = list(_blocking_executors.values())
+        executors = tuple(_blocking_executors.values())
         _blocking_executors.clear()
-
+        _blocking_retiring_executors.update(executors)
     for executor in executors:
-        executor.shutdown(wait=wait, cancel_futures=cancel_futures)
+        executor.shutdown(wait=False, cancel_futures=cancel_futures)
+    with _blocking_executor_lock:
+        owned_executors = set(_blocking_futures.values())
+        _blocking_retiring_executors.intersection_update(owned_executors)
+    return executors
+
+
+def begin_blocking_executor_shutdown(*, cancel_futures: bool = True) -> None:
+    """原子封住新阻塞工具提交，并请求取消尚未开始的同步调用。"""
+    global _blocking_executor_accepting
+    with _blocking_executor_lock:
+        _blocking_executor_accepting = False
+    _retire_blocking_executors(cancel_futures=cancel_futures)
+
+
+def reopen_blocking_executors() -> bool:
+    """仅在旧 Future 和 executor 全部收敛后重新开放测试生命周期。"""
+    global _blocking_executor_accepting
+    with _blocking_executor_lock:
+        if _blocking_futures or _blocking_retiring_executors:
+            return False
+        _blocking_executor_accepting = True
+        return True
+
+
+async def close_blocking_executors(
+    *,
+    timeout_seconds: float,
+    cancel_futures: bool = True,
+) -> bool:
+    """有限等待全部阻塞工具 Future，超时保留 Future 与 executor owner。"""
+    begin_blocking_executor_shutdown(cancel_futures=cancel_futures)
+    with _blocking_executor_lock:
+        futures = tuple(_blocking_futures)
+    wrapped_futures = tuple(asyncio.wrap_future(future) for future in futures)
+    if wrapped_futures:
+        done, pending = await asyncio.wait(
+            wrapped_futures,
+            timeout=max(0.0, timeout_seconds),
+        )
+        if done:
+            await asyncio.gather(*done, return_exceptions=True)
+        for pending_future in pending:
+            pending_future.add_done_callback(
+                lambda completed: completed.exception()
+                if not completed.cancelled()
+                else None
+            )
+
+    with _blocking_executor_lock:
+        unfinished = tuple(
+            future for future in _blocking_futures if not future.done()
+        )
+        retiring_count = len(_blocking_retiring_executors)
+    if unfinished:
+        logger.error(
+            "Agent 阻塞工具未在 %.1f 秒内收敛：futures=%d，executors=%d",
+            max(0.0, timeout_seconds),
+            len(unfinished),
+            retiring_count,
+        )
+        return False
+    return True
+
+
+def shutdown_blocking_executors(
+    *,
+    wait: bool = True,
+    cancel_futures: bool = False,
+) -> bool:
+    """同步清理测试 owner；非等待模式下保留尚未收敛的 executor 句柄。"""
+    executors = _retire_blocking_executors(cancel_futures=cancel_futures)
+    for executor in executors:
+        if wait:
+            executor.shutdown(wait=True, cancel_futures=cancel_futures)
+    with _blocking_executor_lock:
+        return not _blocking_futures and not _blocking_retiring_executors
 
 
 class ToolExecutionTimeoutError(TimeoutError):
@@ -226,13 +323,7 @@ async def run_agent_blocking(
 
     await semaphore.acquire()
     try:
-        context = copy_context()
-        # 长期 worker 保持空底层上下文，每个任务只在自己的调用快照内运行。
-        future = Context().run(
-            _get_blocking_executor(bucket_name).submit,
-            context.run,
-            bound_call,
-        )
+        future = _submit_blocking_call(bucket_name, bound_call)
     except Exception:
         semaphore.release()
         raise