fix(runtime): 收敛跨线程任务提交终态 (#6446)

This commit is contained in:
InfinityPacer
2026-08-25 06:43:37 +08:00
committed by GitHub
parent b277770080
commit 6c334f7b9b
8 changed files with 609 additions and 65 deletions
+35 -18
View File
@@ -1404,19 +1404,6 @@ class FailedRetryMixin:
redo_prompt = build_manual_redo_prompt(history) redo_prompt = build_manual_redo_prompt(history)
self.post_message(
Message(
channel=channel,
source=source,
userid=userid,
username=username,
title=f"已将整理记录 #{history_id} 交给智能助手处理",
text="处理完成后会在这里回复结果。",
link=self.runtime_config.history_url,
save_history=False,
)
)
async def _run_ai_takeover(): async def _run_ai_takeover():
final_output = "" final_output = ""
@@ -1425,6 +1412,18 @@ class FailedRetryMixin:
final_output = text_output or "" final_output = text_output or ""
try: try:
await self.async_post_message(
Message(
channel=channel,
source=source,
userid=userid,
username=username,
title=f"已将整理记录 #{history_id} 交给智能助手处理",
text="处理完成后会在这里回复结果。",
link=self.runtime_config.history_url,
save_history=False,
)
)
manager = get_running_agent_manager() manager = get_running_agent_manager()
if manager is None: if manager is None:
raise RuntimeError("智能助手服务未运行") raise RuntimeError("智能助手服务未运行")
@@ -1462,11 +1461,29 @@ class FailedRetryMixin:
) )
) )
get_task_registry().submit_threadsafe( try:
_run_ai_takeover(), registry = get_task_registry()
loop=global_vars.loop, loop = global_vars.loop
owner="chain.transfer.ai_takeover", registry.submit_threadsafe(
) _run_ai_takeover(),
loop=loop,
owner="chain.transfer.ai_takeover",
)
except RuntimeError as error:
logger.warning("智能助手整理任务提交失败:%s", error)
self.post_message(
Message(
channel=channel,
source=source,
userid=userid,
username=username,
title="智能助手整理失败",
text="系统正在关闭,无法提交处理任务,请稍后重试。",
link=self.runtime_config.history_url,
save_history=False,
)
)
return
def _re_transfer( def _re_transfer(
self, self,
+88 -24
View File
@@ -4,6 +4,7 @@ from __future__ import annotations
import asyncio import asyncio
import concurrent.futures import concurrent.futures
import threading
from collections.abc import Coroutine from collections.abc import Coroutine
from dataclasses import dataclass from dataclasses import dataclass
from functools import partial from functools import partial
@@ -19,6 +20,14 @@ class TaskRecord:
cancel_on_shutdown: bool cancel_on_shutdown: bool
@dataclass(slots=True)
class _ThreadsafeSubmission:
"""持有跨线程提交从排队到真实 Task 发布之间的生命周期。"""
coroutine: Coroutine[Any, Any, Any]
task: asyncio.Task[Any] | None = None
class TaskRegistry: class TaskRegistry:
"""管理由宿主创建的进程内后台任务,并提供统一取消与等待入口。""" """管理由宿主创建的进程内后台任务,并提供统一取消与等待入口。"""
@@ -27,6 +36,10 @@ class TaskRegistry:
self._records: dict[asyncio.Task[Any], TaskRecord] = {} self._records: dict[asyncio.Task[Any], TaskRecord] = {}
self._shutdown_cancel_requested: set[asyncio.Task[Any]] = set() self._shutdown_cancel_requested: set[asyncio.Task[Any]] = set()
self._shutdown_timeout_reported: set[asyncio.Task[Any]] = set() self._shutdown_timeout_reported: set[asyncio.Task[Any]] = set()
self._threadsafe_submissions: dict[
concurrent.futures.Future[Any], _ThreadsafeSubmission
] = {}
self._state_lock = threading.Lock()
self._accepting = True self._accepting = True
@property @property
@@ -77,37 +90,61 @@ class TaskRegistry:
owner: str, owner: str,
cancel_on_shutdown: bool = True, cancel_on_shutdown: bool = True,
) -> concurrent.futures.Future[Any]: ) -> concurrent.futures.Future[Any]:
"""从宿主线程提交协程,并在目标循环内原子登记 owner 后执行""" """从宿主线程提交协程,并持有排队阶段直至发布真实 Task"""
completion: concurrent.futures.Future[Any] = concurrent.futures.Future() completion: concurrent.futures.Future[Any] = concurrent.futures.Future()
task_holder: dict[str, asyncio.Task[Any]] = {} submission = _ThreadsafeSubmission(coroutine=coroutine)
def mirror_completion(task: asyncio.Task[Any]) -> None: def mirror_completion(task: asyncio.Task[Any]) -> None:
"""把登记任务的真实终态镜像给跨线程调用方。""" """把登记任务的真实终态镜像给跨线程调用方。"""
if completion.done(): if completion.done():
return return
if task.cancelled(): try:
completion.cancel() if task.cancelled():
completion.cancel()
return
exception = task.exception()
if exception is not None:
completion.set_exception(exception)
else:
completion.set_result(task.result())
except concurrent.futures.InvalidStateError:
# completion 可由调用线程同时取消,真实 Task 仍由 Registry 观察。
return return
exception = task.exception()
if exception is not None:
completion.set_exception(exception)
else:
completion.set_result(task.result())
def submit_on_loop() -> None: def submit_on_loop() -> None:
"""在目标循环内完成 accepting 检查、任务创建和 owner 登记""" """在目标循环内把 pending submission 原子移交给真实 Task"""
if completion.cancelled(): close_coroutine = False
task: asyncio.Task[Any] | None = None
error: Exception | None = None
with self._state_lock:
current = self._threadsafe_submissions.get(completion)
if current is not submission:
return
if completion.cancelled():
self._threadsafe_submissions.pop(completion, None)
close_coroutine = True
else:
try:
task = self.create(
coroutine,
owner=owner,
cancel_on_shutdown=cancel_on_shutdown,
)
except Exception as caught:
error = caught
close_coroutine = True
else:
submission.task = task
finally:
self._threadsafe_submissions.pop(completion, None)
if close_coroutine:
coroutine.close() coroutine.close()
return if error is not None:
try: try:
task = self.create(
coroutine,
owner=owner,
cancel_on_shutdown=cancel_on_shutdown,
)
except Exception as error:
if not completion.done():
completion.set_exception(error) completion.set_exception(error)
except concurrent.futures.InvalidStateError:
pass
loop.call_exception_handler( loop.call_exception_handler(
{ {
"message": "MoviePilot 跨线程后台任务提交失败", "message": "MoviePilot 跨线程后台任务提交失败",
@@ -116,7 +153,8 @@ class TaskRegistry:
} }
) )
return return
task_holder["task"] = task if task is None:
return
task.add_done_callback(mirror_completion) task.add_done_callback(mirror_completion)
if completion.cancelled() and not task.done(): if completion.cancelled() and not task.done():
task.cancel() task.cancel()
@@ -127,7 +165,17 @@ class TaskRegistry:
"""调用方取消 completion 时,把取消请求转交目标循环中的真实任务。""" """调用方取消 completion 时,把取消请求转交目标循环中的真实任务。"""
if not submitted.cancelled(): if not submitted.cancelled():
return return
task = task_holder.get("task") close_coroutine = False
with self._state_lock:
task = submission.task
if (
task is None
and self._threadsafe_submissions.get(submitted) is submission
):
self._threadsafe_submissions.pop(submitted, None)
close_coroutine = True
if close_coroutine:
coroutine.close()
if task is not None and not task.done(): if task is not None and not task.done():
try: try:
loop.call_soon_threadsafe(task.cancel) loop.call_soon_threadsafe(task.cancel)
@@ -135,10 +183,20 @@ class TaskRegistry:
pass pass
completion.add_done_callback(cancel_registered_task) completion.add_done_callback(cancel_registered_task)
with self._state_lock:
if not self._accepting:
coroutine.close()
raise RuntimeError("后台任务登记器正在关闭,不能再创建新任务")
self._threadsafe_submissions[completion] = submission
try: try:
loop.call_soon_threadsafe(submit_on_loop) loop.call_soon_threadsafe(submit_on_loop)
except RuntimeError: except RuntimeError:
coroutine.close() with self._state_lock:
close_coroutine = (
self._threadsafe_submissions.pop(completion, None) is submission
)
if close_coroutine:
coroutine.close()
raise raise
return completion return completion
@@ -182,7 +240,13 @@ class TaskRegistry:
async def shutdown(self, *, timeout_seconds: float = 10.0) -> bool: async def shutdown(self, *, timeout_seconds: float = 10.0) -> bool:
"""停止接收并有限等待存量任务,返回全部 owner 是否真实收敛。""" """停止接收并有限等待存量任务,返回全部 owner 是否真实收敛。"""
self._accepting = False with self._state_lock:
self._accepting = False
pending_submissions = tuple(self._threadsafe_submissions.items())
self._threadsafe_submissions.clear()
for completion, submission in pending_submissions:
submission.coroutine.close()
completion.cancel()
records = self.records records = self.records
tasks = [record.task for record in records] tasks = [record.task for record in records]
for record in records: for record in records:
+15
View File
@@ -224,6 +224,21 @@ preflight 会验证 Python 3.14、GIL 状态、`thread_inherit_context`、MovieP
依赖、语义、驱动、启动或样本完整性不成立。该工具只用于隔离的本地长 A/B,不接真实凭据、用户数据库、 依赖、语义、驱动、启动或样本完整性不成立。该工具只用于隔离的本地长 A/B,不接真实凭据、用户数据库、
媒体目录或外网,也不加入常规 CI。 媒体目录或外网,也不加入常规 CI。
## TaskRegistry 跨线程提交 A/B
`task_registry_ab.py` 验证目标事件循环尚未分发 callback 时执行 shutdownpending completion 与原始
coroutine 是否取得明确终态,同时采集跨线程提交最小协程的提交和完成耗时。分别在 Before/After revision
运行相同参数并保留两份 JSON,即可比较正确性与固定负载开销:
```bash
../.venv-test/bin/python scripts/perf/task_registry_ab.py \
--iterations 2000 \
--samples 7
```
该探针不访问数据库、配置或网络。吞吐结果用于识别可重复回退,不作为跨机器性能阈值;pending completion
取消且 coroutine 关闭属于正确性门禁。
### PostgreSQL 同步驱动三方案 ### PostgreSQL 同步驱动三方案
`postgresql_driver_ab.py` 在同一 PostgreSQL 容器中比较标准 V3/psycopg2、标准 `postgresql_driver_ab.py` 在同一 PostgreSQL 容器中比较标准 V3/psycopg2、标准
+123
View File
@@ -0,0 +1,123 @@
"""测量 TaskRegistry 跨线程 pending submission 的关停终态与提交吞吐。"""
from __future__ import annotations
import argparse
import asyncio
import inspect
import json
import statistics
import sys
import time
from pathlib import Path
PROJECT_ROOT = Path(__file__).resolve().parents[2]
if str(PROJECT_ROOT) not in sys.path:
sys.path.insert(0, str(PROJECT_ROOT))
from app.runtime.tasks import TaskRegistry
async def _value(value: int) -> int:
"""返回输入值,提供最小可完成协程。"""
return value
def pending_shutdown_probe() -> dict[str, object]:
"""在目标 loop 尚未分发 callback 时关闭 Registry,并报告提交终态。"""
registry = TaskRegistry()
target_loop = asyncio.new_event_loop()
coroutine = _value(1)
completion = registry.submit_threadsafe(
coroutine,
loop=target_loop,
owner="probe.pending",
)
shutdown_result = asyncio.run(registry.shutdown(timeout_seconds=0.001))
result = {
"shutdown_result": shutdown_result,
"completion_done": completion.done(),
"completion_cancelled": completion.cancelled(),
"coroutine_state": inspect.getcoroutinestate(coroutine),
}
completion.cancel()
coroutine.close()
target_loop.close()
return result
async def throughput_sample(iterations: int) -> dict[str, float]:
"""从工作线程提交一组最小协程,并测量提交和完整完成时间。"""
registry = TaskRegistry()
loop = asyncio.get_running_loop()
started = time.perf_counter()
def submit_all():
"""在同一宿主线程连续提交,保持各轮工作负载一致。"""
return [
registry.submit_threadsafe(
_value(index),
loop=loop,
owner="probe.throughput",
)
for index in range(iterations)
]
completions = await asyncio.to_thread(submit_all)
submitted = time.perf_counter()
values = await asyncio.gather(
*(asyncio.wrap_future(completion) for completion in completions)
)
finished = time.perf_counter()
assert sum(values) == iterations * (iterations - 1) // 2
assert await registry.shutdown(timeout_seconds=1.0) is True
return {
"submit_ms": (submitted - started) * 1000,
"total_ms": (finished - started) * 1000,
}
async def run_samples(iterations: int, samples: int) -> list[dict[str, float]]:
"""顺序执行多轮样本,避免并行样本互相争抢事件循环。"""
return [await throughput_sample(iterations) for _ in range(samples)]
def parse_args() -> argparse.Namespace:
"""解析探针负载规模。"""
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--iterations", type=int, default=2000)
parser.add_argument("--samples", type=int, default=7)
return parser.parse_args()
def main() -> None:
"""运行终态探针和吞吐样本并输出 JSON。"""
args = parse_args()
if args.iterations < 1 or args.samples < 1:
raise SystemExit("iterations 和 samples 必须大于 0")
samples = asyncio.run(run_samples(args.iterations, args.samples))
print(
json.dumps(
{
"pending": pending_shutdown_probe(),
"throughput": {
"iterations": args.iterations,
"samples": args.samples,
"submit_ms": [sample["submit_ms"] for sample in samples],
"total_ms": [sample["total_ms"] for sample in samples],
"submit_median_ms": statistics.median(
sample["submit_ms"] for sample in samples
),
"total_median_ms": statistics.median(
sample["total_ms"] for sample in samples
),
},
},
ensure_ascii=False,
indent=2,
)
)
if __name__ == "__main__":
main()
+20 -2
View File
@@ -3,7 +3,7 @@ from concurrent.futures import Future
from dataclasses import replace from dataclasses import replace
from unittest.mock import AsyncMock, Mock, patch from unittest.mock import AsyncMock, Mock, patch
from app.agent import MoviePilotAgent from app.agent.orchestrator import MoviePilotAgent
from app.agent.orchestrator import AgentManagerQueueFullError from app.agent.orchestrator import AgentManagerQueueFullError
from app.agent.tools.impl.ask_user_choice import ( from app.agent.tools.impl.ask_user_choice import (
AskUserChoiceTool, AskUserChoiceTool,
@@ -12,13 +12,14 @@ from app.agent.tools.impl.ask_user_choice import (
from app.agent.tools.impl.send_message import SendMessageTool from app.agent.tools.impl.send_message import SendMessageTool
from app.chain.message import MessageChain from app.chain.message import MessageChain
from app.runtime.config import global_vars, settings from app.runtime.config import global_vars, settings
from app.db import SessionFactory from app.db.session import SessionFactory
from app.db.oper.message import MessageOper from app.db.oper.message import MessageOper
from app.db.models.message import Message from app.db.models.message import Message
from app.application.messaging.agent import AgentInteractionOption, agent_interaction_manager from app.application.messaging.agent import AgentInteractionOption, agent_interaction_manager
from app.application.messaging.interaction import InteractionContext from app.application.messaging.interaction import InteractionContext
from app.application.messaging.media import media_interaction_manager from app.application.messaging.media import media_interaction_manager
from app.schemas.types import NotificationChannel, MessageType from app.schemas.types import NotificationChannel, MessageType
from app.runtime.tasks import TaskRegistry
def _clear_messages() -> None: def _clear_messages() -> None:
@@ -61,6 +62,23 @@ def test_agent_session_clear_uses_owned_threadsafe_submission():
} }
def test_agent_session_clear_handles_closed_task_registry():
"""宿主已停止接收任务时,同步消息清理链应记录拒绝而不是泄漏协程。"""
registry = TaskRegistry()
asyncio.run(registry.shutdown(timeout_seconds=0.01))
manager = Mock(clear_session=AsyncMock())
with patch.object(global_vars, "CURRENT_EVENT_LOOP", _running_loop_stub()), patch(
"app.chain.message.get_running_agent_manager", return_value=manager
), patch("app.chain.message.get_task_registry", return_value=registry), patch(
"app.chain.message.logger"
) as logger:
MessageChain._schedule_agent_session_clear("session-1", "10001")
logger.warning.assert_called_once()
assert "正在关闭" in logger.warning.call_args.args[0]
def test_remote_session_clear_reuses_owned_clear_scheduler(): def test_remote_session_clear_reuses_owned_clear_scheduler():
"""远程清理命令应复用唯一 Agent 会话清理入口。""" """远程清理命令应复用唯一 Agent 会话清理入口。"""
chain = MessageChain() chain = MessageChain()
+18
View File
@@ -1,9 +1,12 @@
"""服务端统计兼容入口的后台任务生命周期回归。""" """服务端统计兼容入口的后台任务生命周期回归。"""
import asyncio
from unittest.mock import Mock, patch from unittest.mock import Mock, patch
from app.adapters.external.server import MoviePilotServerHelper from app.adapters.external.server import MoviePilotServerHelper
from app.runtime.config import global_vars from app.runtime.config import global_vars
from app.runtime.tasks import TaskRegistry
def test_legacy_subscription_reports_use_owned_threadsafe_tasks() -> None: def test_legacy_subscription_reports_use_owned_threadsafe_tasks() -> None:
@@ -40,3 +43,18 @@ def test_legacy_subscription_report_rejects_without_runtime_loop() -> None:
"""宿主生命周期不可用时应拒绝提交,并保持布尔返回合同。""" """宿主生命周期不可用时应拒绝提交,并保持布尔返回合同。"""
with patch.object(global_vars, "CURRENT_EVENT_LOOP", None): with patch.object(global_vars, "CURRENT_EVENT_LOOP", None):
assert MoviePilotServerHelper.sub_done_async({"media_id": "1"}) is False assert MoviePilotServerHelper.sub_done_async({"media_id": "1"}) is False
def test_legacy_subscription_report_handles_closed_task_registry() -> None:
"""运行 loop 尚在但宿主已封口时,兼容入口应保持布尔失败合同。"""
registry = TaskRegistry()
asyncio.run(registry.shutdown(timeout_seconds=0.01))
loop = Mock(**{"is_running.return_value": True, "is_closed.return_value": False})
with patch.object(global_vars, "CURRENT_EVENT_LOOP", loop), patch(
"app.adapters.external.server.get_task_registry", return_value=registry
), patch("app.adapters.external.server.logger") as logger:
assert MoviePilotServerHelper.sub_done_async({"media_id": "1"}) is False
logger.warning.assert_called_once()
assert "正在关闭" in logger.warning.call_args.args[0]
+212 -14
View File
@@ -1,13 +1,39 @@
"""进程内后台任务登记与关停语义测试。""" """进程内后台任务登记与关停语义测试。"""
import asyncio import asyncio
import concurrent.futures
import inspect
import threading import threading
import weakref
from collections.abc import Coroutine
from typing import Any
import pytest import pytest
from app.runtime.tasks import TaskRegistry from app.runtime.tasks import TaskRegistry
class _CloseCountingCoroutine(Coroutine[Any, Any, None]):
"""记录 pending submission 的协程关闭所有权。"""
def __init__(self) -> None:
self.close_count = 0
def __await__(self):
return self
def send(self, _value):
raise StopIteration
def throw(self, typ, val=None, tb=None):
if val is None:
raise typ
raise val.with_traceback(tb)
def close(self) -> None:
self.close_count += 1
def test_task_registry_removes_completed_task() -> None: def test_task_registry_removes_completed_task() -> None:
"""正常完成的任务应自动退出登记表,避免长期持有请求对象。""" """正常完成的任务应自动退出登记表,避免长期持有请求对象。"""
@@ -124,37 +150,209 @@ def test_task_registry_owns_threadsafe_submission_until_shutdown() -> None:
def test_task_registry_rejects_threadsafe_submission_after_shutdown() -> None: def test_task_registry_rejects_threadsafe_submission_after_shutdown() -> None:
"""关停先赢得竞态时应关闭协程并通过 completion 报告拒绝原因""" """关停先赢得竞态时应关闭协程并同步拒绝提交"""
async def scenario() -> None: async def scenario() -> None:
registry = TaskRegistry() registry = TaskRegistry()
loop = asyncio.get_running_loop() loop = asyncio.get_running_loop()
reports: list[dict[str, object]] = []
previous_handler = loop.get_exception_handler()
loop.set_exception_handler(lambda _, context: reports.append(context))
async def late_worker() -> None: async def late_worker() -> None:
"""模拟关停完成后从宿主线程到达的晚任务。""" """模拟关停完成后从宿主线程到达的晚任务。"""
try: assert await registry.shutdown(timeout_seconds=1.0) is True
assert await registry.shutdown(timeout_seconds=1.0) is True coroutine = late_worker()
completion = await asyncio.to_thread( with pytest.raises(RuntimeError, match="正在关闭"):
await asyncio.to_thread(
registry.submit_threadsafe, registry.submit_threadsafe,
late_worker(), coroutine,
loop=loop, loop=loop,
owner="test.threadsafe-late", owner="test.threadsafe-late",
) )
with pytest.raises(RuntimeError, match="正在关闭"):
await asyncio.wrap_future(completion)
assert registry.records == () assert inspect.getcoroutinestate(coroutine) == inspect.CORO_CLOSED
assert reports[-1]["owner"] == "test.threadsafe-late" assert registry.records == ()
finally:
loop.set_exception_handler(previous_handler)
asyncio.run(scenario()) asyncio.run(scenario())
def test_task_registry_shutdown_closes_submission_before_loop_dispatch() -> None:
"""回调尚未执行时关停也必须关闭协程并结束 completion。"""
registry = TaskRegistry()
loop = asyncio.new_event_loop()
async def worker() -> None:
"""不应被迟到的 loop callback 启动。"""
coroutine = worker()
completion = registry.submit_threadsafe(
coroutine,
loop=loop,
owner="test.threadsafe-pending",
)
assert asyncio.run(registry.shutdown(timeout_seconds=0.01)) is True
assert completion.cancelled()
assert inspect.getcoroutinestate(coroutine) == inspect.CORO_CLOSED
loop.run_until_complete(asyncio.sleep(0))
assert registry.records == ()
loop.close()
def test_task_registry_caller_cancels_before_loop_dispatch() -> None:
"""调用方在 callback 前取消时应关闭协程,迟到 callback 不得创建 Task。"""
registry = TaskRegistry()
loop = asyncio.new_event_loop()
async def worker() -> None:
"""不应在 completion 取消后启动。"""
coroutine = worker()
completion = registry.submit_threadsafe(
coroutine,
loop=loop,
owner="test.threadsafe-cancel-before-dispatch",
)
assert completion.cancel()
assert inspect.getcoroutinestate(coroutine) == inspect.CORO_CLOSED
loop.run_until_complete(asyncio.sleep(0))
assert registry.records == ()
loop.close()
def test_task_registry_closes_submission_when_loop_rejects_callback() -> None:
"""目标 loop 已关闭时同步失败,并且不遗留未等待协程。"""
registry = TaskRegistry()
loop = asyncio.new_event_loop()
loop.close()
async def worker() -> None:
"""模拟无法进入目标 loop 的提交。"""
coroutine = worker()
with pytest.raises(RuntimeError, match="closed"):
registry.submit_threadsafe(
coroutine,
loop=loop,
owner="test.threadsafe-closed-loop",
)
assert inspect.getcoroutinestate(coroutine) == inspect.CORO_CLOSED
def test_task_registry_shutdown_owns_close_before_loop_rejection() -> None:
"""shutdown 已取走 pending 后,迟到的投递失败不得再次关闭同一协程。"""
registry = TaskRegistry()
dispatch_entered = threading.Event()
release_dispatch = threading.Event()
coroutine = _CloseCountingCoroutine()
class BlockingClosedLoop:
"""让 shutdown 稳定发生在投递抛错之前。"""
@staticmethod
def call_soon_threadsafe(_callback) -> None:
dispatch_entered.set()
release_dispatch.wait()
raise RuntimeError("loop closed")
with concurrent.futures.ThreadPoolExecutor(max_workers=1) as executor:
submission = executor.submit(
registry.submit_threadsafe,
coroutine,
loop=BlockingClosedLoop(),
owner="test.threadsafe-shutdown-before-rejection",
)
assert dispatch_entered.wait(1.0)
assert asyncio.run(registry.shutdown(timeout_seconds=0.01)) is True
release_dispatch.set()
with pytest.raises(RuntimeError, match="loop closed"):
submission.result(timeout=1.0)
assert coroutine.close_count == 1
def test_task_registry_shutdown_rejects_late_successful_dispatch() -> None:
"""shutdown 取走 pending 后,迟到的成功投递也不得发布真实 Task。"""
registry = TaskRegistry()
dispatch_entered = threading.Event()
release_dispatch = threading.Event()
callbacks = []
coroutine = _CloseCountingCoroutine()
class BlockingLoop:
"""让 shutdown 稳定发生在 loop 接受 callback 之前。"""
@staticmethod
def call_soon_threadsafe(callback) -> None:
dispatch_entered.set()
release_dispatch.wait()
callbacks.append(callback)
with concurrent.futures.ThreadPoolExecutor(max_workers=1) as executor:
submission = executor.submit(
registry.submit_threadsafe,
coroutine,
loop=BlockingLoop(),
owner="test.threadsafe-shutdown-before-dispatch",
)
assert dispatch_entered.wait(1.0)
assert asyncio.run(registry.shutdown(timeout_seconds=0.01)) is True
release_dispatch.set()
completion = submission.result(timeout=1.0)
assert completion.cancelled()
callbacks[0]()
assert coroutine.close_count == 1
assert registry.records == ()
def test_task_registry_completion_does_not_retain_itself_after_cancellation() -> None:
"""终态回调不得让已取消 completion 依赖循环 GC 才能释放。"""
registry = TaskRegistry()
loop = asyncio.new_event_loop()
completion = registry.submit_threadsafe(
_CloseCountingCoroutine(),
loop=loop,
owner="test.threadsafe-completion-release",
)
completion_ref = weakref.ref(completion)
assert completion.cancel()
loop.run_until_complete(asyncio.sleep(0))
del completion
assert completion_ref() is None
loop.close()
def test_task_registry_pending_non_cancellable_work_does_not_start_on_shutdown() -> None:
"""尚未发布成 Task 的 non-cancellable 工作仍应在关停时直接关闭。"""
registry = TaskRegistry()
loop = asyncio.new_event_loop()
async def worker() -> None:
"""未开始的同步兼容工作不应拖延关停。"""
coroutine = worker()
completion = registry.submit_threadsafe(
coroutine,
loop=loop,
owner="test.threadsafe-pending-non-cancellable",
cancel_on_shutdown=False,
)
assert asyncio.run(registry.shutdown(timeout_seconds=0.01)) is True
assert completion.cancelled()
assert inspect.getcoroutinestate(coroutine) == inspect.CORO_CLOSED
loop.run_until_complete(asyncio.sleep(0))
assert registry.records == ()
loop.close()
def test_task_registry_keeps_timed_out_sync_owner_until_real_completion() -> None: def test_task_registry_keeps_timed_out_sync_owner_until_real_completion() -> None:
"""同步线程超过关停预算后仍应保留 owner,不能把包装任务取消成伪完成。""" """同步线程超过关停预算后仍应保留 owner,不能把包装任务取消成伪完成。"""
+98 -7
View File
@@ -16,6 +16,7 @@ from app.chain.message import MessageChain
from app.chain.transfer import TransferChain from app.chain.transfer import TransferChain
from app.application.messaging.interaction import InteractionContext from app.application.messaging.interaction import InteractionContext
from app.runtime.config import global_vars, settings from app.runtime.config import global_vars, settings
from app.runtime.tasks import TaskRegistry
from app.schemas.types import NotificationChannel from app.schemas.types import NotificationChannel
@@ -128,10 +129,18 @@ class TestTransferFailedRetryButtons(unittest.TestCase):
errmsg="未识别到媒体信息", errmsg="未识别到媒体信息",
) )
def _close_pending_coro(coro, *args, **kwargs): async_messages = []
"""关闭被调度的协程:测试中事件循环未运行,不关闭会残留 never-awaited 警告。"""
coro.close()
def _run_pending_coro(coro, *args, **kwargs):
asyncio.run(coro)
async def _capture_message(message):
async_messages.append(message)
async def _finish_immediately(**kwargs):
kwargs["output_callback"]("ok")
manager = SimpleNamespace(run_background_prompt=_finish_immediately)
loop = Mock(**{"is_running.return_value": True, "is_closed.return_value": False}) loop = Mock(**{"is_running.return_value": True, "is_closed.return_value": False})
with patch.object(global_vars, "CURRENT_EVENT_LOOP", loop), patch.object( with patch.object(global_vars, "CURRENT_EVENT_LOOP", loop), patch.object(
settings, "AI_AGENT_ENABLE", True settings, "AI_AGENT_ENABLE", True
@@ -141,12 +150,14 @@ class TestTransferFailedRetryButtons(unittest.TestCase):
) as history_oper_cls, patch( ) as history_oper_cls, patch(
"app.chain._transfer.build_manual_redo_prompt", "app.chain._transfer.build_manual_redo_prompt",
return_value="retry transfer prompt", return_value="retry transfer prompt",
), patch(
"app.chain._transfer.get_running_agent_manager", return_value=manager
), patch("app.chain._transfer.get_task_registry") as get_registry: ), patch("app.chain._transfer.get_task_registry") as get_registry:
get_registry.return_value.submit_threadsafe.side_effect = ( get_registry.return_value.submit_threadsafe.side_effect = (
_close_pending_coro _run_pending_coro
) )
history_oper_cls.return_value.get.return_value = history history_oper_cls.return_value.get.return_value = history
with patch.object(chain, "post_message") as post_message: with patch.object(chain, "async_post_message", side_effect=_capture_message):
chain.handle_failed_transfer_callback( chain.handle_failed_transfer_callback(
callback_data="transfer_ai_retry_34", callback_data="transfer_ai_retry_34",
channel=NotificationChannel.Telegram, channel=NotificationChannel.Telegram,
@@ -160,10 +171,90 @@ class TestTransferFailedRetryButtons(unittest.TestCase):
get_registry.return_value.submit_threadsafe.call_args.kwargs["owner"], get_registry.return_value.submit_threadsafe.call_args.kwargs["owner"],
"chain.transfer.ai_takeover", "chain.transfer.ai_takeover",
) )
self.assertEqual(len(async_messages), 2)
self.assertEqual(
async_messages[0].title,
"已将整理记录 #34 交给智能助手处理",
)
self.assertEqual(async_messages[1].title, "智能助手整理完成")
def test_transfer_ai_retry_callback_reports_closed_task_registry(self):
"""宿主停止接收任务时,不得向用户报告智能助手已接管。"""
chain = TransferChain()
chain.runtime_config = replace(chain.runtime_config, ai_agent_enable=True)
history = SimpleNamespace(id=34)
registry = TaskRegistry()
asyncio.run(registry.shutdown(timeout_seconds=0.01))
loop = Mock(**{"is_running.return_value": True, "is_closed.return_value": False})
with patch.object(global_vars, "CURRENT_EVENT_LOOP", loop), patch.object(
settings, "AI_AGENT_ENABLE", True
), patch(
"app.chain._transfer.get_chain_transfer_history_port"
) as history_port, patch(
"app.chain._transfer.build_manual_redo_prompt",
return_value="retry transfer prompt",
), patch(
"app.chain._transfer.get_task_registry", return_value=registry
), patch(
"app.chain._transfer.logger"
) as logger, patch.object(
chain, "post_message"
) as post_message:
history_port.return_value.get.return_value = history
chain.handle_failed_transfer_callback(
callback_data="transfer_ai_retry_34",
channel=NotificationChannel.Telegram,
source="telegram-test",
userid="10001",
username="tester",
)
logger.warning.assert_called_once()
self.assertEqual(post_message.call_count, 1) self.assertEqual(post_message.call_count, 1)
self.assertEqual( self.assertEqual(
post_message.call_args_list[0].args[0].title, post_message.call_args.args[0].title,
"已将整理记录 #34 交给智能助手处理", "智能助手整理失败",
)
self.assertNotIn("已将", post_message.call_args.args[0].title)
def test_transfer_ai_retry_callback_reports_unavailable_event_loop(self):
"""主循环不可用时,应在创建后台协程前返回明确失败提示。"""
chain = TransferChain()
chain.runtime_config = replace(chain.runtime_config, ai_agent_enable=True)
history = SimpleNamespace(id=34)
with patch.object(global_vars, "CURRENT_EVENT_LOOP", None), patch.object(
settings, "AI_AGENT_ENABLE", True
), patch(
"app.chain._transfer.get_chain_transfer_history_port"
) as history_port, patch(
"app.chain._transfer.build_manual_redo_prompt",
return_value="retry transfer prompt",
), patch(
"app.chain._transfer.get_task_registry"
) as get_registry, patch(
"app.chain._transfer.logger"
) as logger, patch.object(
chain, "post_message"
) as post_message:
history_port.return_value.get.return_value = history
chain.handle_failed_transfer_callback(
callback_data="transfer_ai_retry_34",
channel=NotificationChannel.Telegram,
source="telegram-test",
userid="10001",
username="tester",
)
get_registry.return_value.submit_threadsafe.assert_not_called()
logger.warning.assert_called_once()
self.assertEqual(post_message.call_count, 1)
self.assertEqual(
post_message.call_args.args[0].title,
"智能助手整理失败",
) )
def test_transfer_ai_retry_callback_uses_successful_move_dest_as_source(self): def test_transfer_ai_retry_callback_uses_successful_move_dest_as_source(self):