From 14edd90120b1bd3a6a39a425d6b2ef50bb0c2457 Mon Sep 17 00:00:00 2001 From: YUZHEthefool <2804776511@qq.com> Date: Thu, 10 Sep 2026 23:51:37 +0800 Subject: [PATCH 1/3] feat(btw): deliver and finalize detached work requests Keep request-scoped WebChat delivery open until work completes and reclaim profile tasks before their runtime dependencies close. Related: #126 AI-Generated: true Generated-At: 2026-09-10T15:51:37Z --- astrbot/core/agent/btw/work_loop.py | 60 ++++- astrbot/core/agent/conversation_loop.py | 24 +- astrbot/core/core_lifecycle.py | 9 +- .../method/agent_sub_stages/internal.py | 17 +- astrbot/core/pipeline/process_stage/stage.py | 23 +- astrbot/core/pipeline/scheduler.py | 56 +++- docs/en/dev/astrbot-config.md | 2 + docs/zh/dev/astrbot-config.md | 2 + tests/unit/test_agent_internal_process.py | 12 +- tests/unit/test_btw_delivery.py | 252 ++++++++++++++++++ tests/unit/test_core_lifecycle.py | 25 +- 11 files changed, 465 insertions(+), 17 deletions(-) create mode 100644 tests/unit/test_btw_delivery.py diff --git a/astrbot/core/agent/btw/work_loop.py b/astrbot/core/agent/btw/work_loop.py index 7cc353f5ca..a504464f73 100644 --- a/astrbot/core/agent/btw/work_loop.py +++ b/astrbot/core/agent/btw/work_loop.py @@ -2,6 +2,7 @@ import asyncio from collections.abc import AsyncGenerator, Awaitable, Callable +from contextlib import aclosing from typing import Protocol from astrbot import logger @@ -50,6 +51,8 @@ def __init__( self._background_tasks: set[asyncio.Task] | None = None self._result_dispatcher: ResultDispatcher | None = None self._event_finalizer: EventFinalizer | None = None + self._tasks: dict[asyncio.Task, tuple[AstrMessageEvent, str]] = {} + self._closed = False def configure_detached_execution( self, @@ -92,6 +95,17 @@ async def submit(self, event: AstrMessageEvent) -> AsyncGenerator[None]: Falls back to inline execution when no runtime task registry is attached, which keeps the primitive usable in isolated tests. """ + if self._closed: + event.set_result( + MessageEventResult().message( + work_i18n.text( + work_i18n.resolve_event_locale(event), + "btw.work.status.cancelled", + ) + ) + ) + yield + return if ( self._background_tasks is None or self._result_dispatcher is None @@ -114,15 +128,43 @@ async def submit(self, event: AstrMessageEvent) -> AsyncGenerator[None]: ) yield + if self._closed: + await self.sessions.update_status(session.id, WorkSessionStatus.CANCELLED) + return + # The first yield returns only after the normal response stages deliver # the acknowledgement. Marking it here prevents the scheduler from # releasing event-owned temporary files before the worker needs them. event.set_extra("btw_detached_work", True) - create_tracked_task( + task = create_tracked_task( self._background_tasks, self._run_detached(event, session.id), name=f"btw_work:{session.id}", ) + self._tasks[task] = (event, session.id) + task.add_done_callback(lambda done: self._tasks.pop(done, None)) + + async def close(self) -> None: + """Cancel and finalize this profile's work, including unstarted tasks.""" + self._closed = True + tasks = dict(self._tasks) + for task in tasks: + task.cancel() + if tasks: + await asyncio.gather(*tasks, return_exceptions=True) + finalizers = [] + if self._event_finalizer is not None: + for event, session_id in tasks.values(): + if not event.get_extra("btw_detached_work_finished"): + await self.sessions.update_status( + session_id, WorkSessionStatus.CANCELLED + ) + finalizers.append(self._event_finalizer(event)) + if finalizers: + results = await asyncio.gather(*finalizers, return_exceptions=True) + for result in results: + if isinstance(result, BaseException): + raise result @staticmethod def _prepare_event(event: AstrMessageEvent, session_id: str) -> None: @@ -159,9 +201,18 @@ async def _execute( ) raise else: + failed = bool(event.get_extra("btw_work_failed")) + cancelled = bool(event.get_extra("agent_stop_requested")) await self.sessions.update_status( session_id, - WorkSessionStatus.COMPLETED, + WorkSessionStatus.FAILED + if failed + else ( + WorkSessionStatus.CANCELLED + if cancelled + else WorkSessionStatus.COMPLETED + ), + error="Work task failed." if failed else None, ) async def _run_detached(self, event: AstrMessageEvent, session_id: str) -> None: @@ -169,8 +220,9 @@ async def _run_detached(self, event: AstrMessageEvent, session_id: str) -> None: assert self._result_dispatcher is not None assert self._event_finalizer is not None try: - async for _ in self._execute(event, session_id): - await self._result_dispatcher(event) + async with aclosing(self._execute(event, session_id)) as execution: + async for _ in execution: + await self._result_dispatcher(event) except asyncio.CancelledError: raise except Exception as exc: diff --git a/astrbot/core/agent/conversation_loop.py b/astrbot/core/agent/conversation_loop.py index 6a1914ff84..c5d6a4eddf 100644 --- a/astrbot/core/agent/conversation_loop.py +++ b/astrbot/core/agent/conversation_loop.py @@ -1,6 +1,7 @@ """Opt-in conversation entry over the existing Agent request executor.""" -from collections.abc import AsyncGenerator +import asyncio +from collections.abc import AsyncGenerator, Awaitable, Callable from typing import TYPE_CHECKING from astrbot.core.agent.btw.types import is_work_loop_enabled @@ -43,6 +44,27 @@ async def initialize(self, ctx: PipelineContext) -> None: max_concurrent=concurrency if type(concurrency) is int else 2, ) + def configure_detached_work( + self, + *, + background_tasks: set[asyncio.Task], + result_dispatcher: Callable[[AstrMessageEvent], Awaitable[None]], + event_finalizer: Callable[[AstrMessageEvent], Awaitable[None]], + ) -> None: + """Attach the owning scheduler's delivery and cleanup services.""" + if self.work_loop is None: + raise RuntimeError("ConversationLoop is not initialized") + self.work_loop.configure_detached_execution( + background_tasks=background_tasks, + result_dispatcher=result_dispatcher, + event_finalizer=event_finalizer, + ) + + async def close(self) -> None: + """Stop work owned by this conversation entry.""" + if self.work_loop is not None: + await self.work_loop.close() + async def process(self, event: AstrMessageEvent) -> AsyncGenerator[None]: """Process one admitted conversation using the current Agent path.""" if ( diff --git a/astrbot/core/core_lifecycle.py b/astrbot/core/core_lifecycle.py index 89eb5fed87..06f2f31463 100644 --- a/astrbot/core/core_lifecycle.py +++ b/astrbot/core/core_lifecycle.py @@ -723,6 +723,8 @@ async def stop(self) -> None: if not self._cleanup_stack_closed: self._cleanup_stack_closed = True + for scheduler in self.pipeline_scheduler_mapping.values(): + self._register_cleanup("pipeline work tasks", scheduler.close) await self._cleanup_stack.aclose() self._initialized = False @@ -830,6 +832,9 @@ async def reload_pipeline_scheduler(self, conf_id: str) -> None: getattr(self.services, "authorization", None), ), ) + old_scheduler = self.pipeline_scheduler_mapping.get(conf_id) + if old_scheduler is not None: + await old_scheduler.close() await scheduler.initialize() self.pipeline_scheduler_mapping[conf_id] = scheduler manager = getattr(self, "turn_window_manager", None) @@ -838,4 +843,6 @@ async def reload_pipeline_scheduler(self, conf_id: str) -> None: async def remove_pipeline_scheduler(self, conf_id: str) -> None: """Remove the scheduler associated with a deleted configuration profile.""" - self.pipeline_scheduler_mapping.pop(conf_id, None) + scheduler = self.pipeline_scheduler_mapping.pop(conf_id, None) + if scheduler is not None: + await scheduler.close() diff --git a/astrbot/core/pipeline/process_stage/method/agent_sub_stages/internal.py b/astrbot/core/pipeline/process_stage/method/agent_sub_stages/internal.py index c692f1f07c..d3d258c44d 100644 --- a/astrbot/core/pipeline/process_stage/method/agent_sub_stages/internal.py +++ b/astrbot/core/pipeline/process_stage/method/agent_sub_stages/internal.py @@ -300,6 +300,7 @@ async def process(self, event: AstrMessageEvent) -> AsyncGenerator[None]: follow_up_consumed_marked = False follow_up_activated = False typing_requested = False + is_detached_work = bool(event.get_extra("btw_detached_work")) try: from astrbot.core.streaming_override import resolve_streaming_response @@ -356,7 +357,9 @@ async def process(self, event: AstrMessageEvent) -> AsyncGenerator[None]: logger.debug("ready to request llm provider") follow_up_capture = ( - self.ctx.execution_context.follow_up_coordinator.try_capture(event) + None + if is_detached_work + else self.ctx.execution_context.follow_up_coordinator.try_capture(event) ) if follow_up_capture: ( @@ -393,6 +396,9 @@ async def process(self, event: AstrMessageEvent) -> AsyncGenerator[None]: concurrent, lock_key, turn_cm, streaming_response = ( self._prepare_group_sender_concurrency(event, streaming_response) ) + work_lock = event.get_extra("btw_agent_lock_key") + if is_detached_work and isinstance(work_lock, str) and work_lock: + lock_key = work_lock async with ( turn_cm, @@ -409,6 +415,8 @@ async def process(self, event: AstrMessageEvent) -> AsyncGenerator[None]: streaming_response, ) if build_result is None: + if is_detached_work: + event.set_extra("btw_work_failed", True) return agent_runner = build_result.agent_runner @@ -469,8 +477,9 @@ async def process(self, event: AstrMessageEvent) -> AsyncGenerator[None]: ) else: runner_stop_callback = None - self._register_follow_up_runner(event, agent_runner, concurrent) - runner_registered = True + if not is_detached_work: + self._register_follow_up_runner(event, agent_runner, concurrent) + runner_registered = True event.trace.record( "astr_agent_prepare", system_prompt=req.system_prompt, @@ -550,6 +559,8 @@ async def process(self, event: AstrMessageEvent) -> AsyncGenerator[None]: ) except Exception as e: + if is_detached_work: + event.set_extra("btw_work_failed", True) logger.error( "Error occurred while processing agent: %s", safe_error("", e), diff --git a/astrbot/core/pipeline/process_stage/stage.py b/astrbot/core/pipeline/process_stage/stage.py index 5c6272e442..5d7b8cc27c 100644 --- a/astrbot/core/pipeline/process_stage/stage.py +++ b/astrbot/core/pipeline/process_stage/stage.py @@ -1,4 +1,5 @@ -from collections.abc import AsyncGenerator +import asyncio +from collections.abc import AsyncGenerator, Awaitable, Callable from astrbot.core.agent.conversation_loop import ConversationLoop from astrbot.core.agent.llm_types import ProviderRequest @@ -29,6 +30,26 @@ async def initialize(self, ctx: PipelineContext) -> None: self.star_request_sub_stage = StarRequestSubStage() await self.star_request_sub_stage.initialize(ctx) + def configure_detached_work( + self, + *, + background_tasks: set[asyncio.Task], + result_dispatcher: Callable[[AstrMessageEvent], Awaitable[None]], + event_finalizer: Callable[[AstrMessageEvent], Awaitable[None]], + ) -> None: + """Attach runtime services only to an enabled conversation loop.""" + if self.conversation_loop is not None: + self.conversation_loop.configure_detached_work( + background_tasks=background_tasks, + result_dispatcher=result_dispatcher, + event_finalizer=event_finalizer, + ) + + async def close(self) -> None: + """Reclaim work before this profile's scheduler is replaced.""" + if self.conversation_loop is not None: + await self.conversation_loop.close() + async def process( self, event: AstrMessageEvent, diff --git a/astrbot/core/pipeline/scheduler.py b/astrbot/core/pipeline/scheduler.py index 0cf797aca5..e45609ef6f 100644 --- a/astrbot/core/pipeline/scheduler.py +++ b/astrbot/core/pipeline/scheduler.py @@ -8,6 +8,7 @@ from .bootstrap import builtin_stage_classes from .context import PipelineContext +from .result_decorate.stage import ResultDecorateStage from .stage import Stage @@ -33,6 +34,50 @@ async def initialize(self) -> None: stage_instance = stage_cls() # 创建实例 await stage_instance.initialize(self.ctx) self.stages.append(stage_instance) + for stage in self.stages: + configure = getattr(stage, "configure_detached_work", None) + if callable(configure): + configure( + background_tasks=self.ctx.execution_context.background_tasks, + result_dispatcher=self.deliver_detached_result, + event_finalizer=self.finalize_detached_event, + ) + + async def close(self) -> None: + """Close profile-owned background work before replacing this scheduler.""" + for stage in reversed(self.stages): + close = getattr(stage, "close", None) + if callable(close): + await close() + + async def deliver_detached_result(self, event: AstrMessageEvent) -> None: + """Replay response decoration and delivery with onion ordering intact.""" + index = next( + ( + i + for i, stage in enumerate(self.stages) + if isinstance(stage, ResultDecorateStage) + ), + None, + ) + if index is None: + raise RuntimeError("ResultDecorateStage is not configured") + if not event.is_stopped(): + await self._process_stages(event, index) + + async def finalize_detached_event(self, event: AstrMessageEvent) -> None: + """Complete a retained request once, then release its resources.""" + if event.get_extra("btw_detached_work_finished"): + return + event.set_extra("btw_detached_work_finished", True) + try: + if event.requires_empty_completion and not event.get_extra( + "skip_empty_completion" + ): + await cast(_EmptyCompletionEvent, event).send(None) + finally: + event.cleanup_temporary_local_files() + self.ctx.execution_context.active_event_registry.unregister(event) async def _process_stages(self, event: AstrMessageEvent, from_stage=0) -> None: """依次执行各个阶段 @@ -91,8 +136,10 @@ async def execute(self, event: AstrMessageEvent) -> None: await self._process_stages(event) # 发送一个空消息, 以便于后续的处理 - if event.requires_empty_completion and not event.get_extra( - "skip_empty_completion" + if ( + event.requires_empty_completion + and not event.get_extra("skip_empty_completion") + and not event.get_extra("btw_detached_work") ): # Only adapters whose send implementation accepts ``None`` set this # flag. The base event contract deliberately remains message-only. @@ -112,5 +159,6 @@ async def execute(self, event: AstrMessageEvent) -> None: else: logger.debug("pipeline execution completed.") finally: - event.cleanup_temporary_local_files() - self.ctx.execution_context.active_event_registry.unregister(event) + if not event.get_extra("btw_detached_work"): + event.cleanup_temporary_local_files() + self.ctx.execution_context.active_event_registry.unregister(event) diff --git a/docs/en/dev/astrbot-config.md b/docs/en/dev/astrbot-config.md index a6567540e3..67d25eaf77 100644 --- a/docs/en/dev/astrbot-config.md +++ b/docs/en/dev/astrbot-config.md @@ -207,6 +207,8 @@ Automatic classifier candidates are evaluated separately. Enabling this entry do The work executor additionally requires `btw.work_loop.enabled`, also `false` by default. It reuses the Agent executor and records pending, running, completed, failed, and cancelled task states. `btw.work_loop.max_concurrent` limits active execution (default `2`); it does not impose a waiting-queue length limit. `btw.work_session.max_age_seconds` retains terminal states for `3600` seconds by default; active tasks do not expire, and expired terminal records are removed during the next session operation. Runtime-owned background services perform task execution and cleanup when attached by the scheduler. +Detached work acknowledges receipt before execution and returns results through the current response-decoration and delivery stages, including reply content checks. Inbound stages are not rerun. WebChat keeps the original request identifier open through the final result; acknowledgement does not end the request. Temporary event files remain available to the worker and are released on completion, failure, or cancellation. Replacing or removing a profile cancels its owned work; runtime shutdown also reclaims it. + ## WebUI and authentication Important `dashboard` defaults: diff --git a/docs/zh/dev/astrbot-config.md b/docs/zh/dev/astrbot-config.md index e3fb243319..622bb5c1ea 100644 --- a/docs/zh/dev/astrbot-config.md +++ b/docs/zh/dev/astrbot-config.md @@ -209,6 +209,8 @@ Alkaid [长期记忆](../use/long-term-memory) 当前没有对应的启停配置 工作执行器还需要开启 `btw.work_loop.enabled`，默认同样为 `false`。它复用 Agent 执行器并记录排队、运行、完成、失败、取消状态。`btw.work_loop.max_concurrent` 限制正在执行的任务数，默认 `2`，不限制等待队列长度。`btw.work_session.max_age_seconds` 默认保留终态记录 `3600` 秒；活动任务不会过期，终态过期记录在下次会话操作时清除。调度器接入后台服务后，由运行时拥有工作任务的执行和清理。 +后台工作在执行前确认接收，再通过当前回复装饰与发送阶段回送结果，包括回复内容检查；不重复运行入站阶段。WebChat 持续使用原请求标识，确认消息不会结束请求。事件临时文件保留到工作完成、失败或取消后再释放。配置档替换、删除以及运行时关闭会取消并回收其工作任务。 + ## WebUI 与认证 `dashboard` 的关键默认值： diff --git a/tests/unit/test_agent_internal_process.py b/tests/unit/test_agent_internal_process.py index 73d4038015..72aec84272 100644 --- a/tests/unit/test_agent_internal_process.py +++ b/tests/unit/test_agent_internal_process.py @@ -496,7 +496,8 @@ async def test_internal_process_stops_when_waiting_hook_blocks(monkeypatch): @pytest.mark.asyncio -async def test_internal_process_continues_when_send_typing_fails(monkeypatch): +@pytest.mark.parametrize("detached", [False, True]) +async def test_internal_process_continues_when_send_typing_fails(monkeypatch, detached): stage = internal.InternalAgentSubStage.__new__(internal.InternalAgentSubStage) stage.streaming_response = False stage.show_tool_use = True @@ -513,6 +514,7 @@ async def test_internal_process_continues_when_send_typing_fails(monkeypatch): extras={internal.LLM_ERROR_MESSAGE_EXTRA_KEY: "provider unavailable"}, ) event.send_typing.side_effect = RuntimeError("typing failed") + event.set_extra("btw_detached_work", detached) logger_warning = MagicMock() monkeypatch.setattr( @@ -537,6 +539,7 @@ async def test_internal_process_continues_when_send_typing_fails(monkeypatch): ) event.stop_typing.assert_awaited_once() logger_warning.assert_called() + assert bool(event.get_extra("btw_work_failed")) is detached @pytest.mark.asyncio @@ -568,7 +571,10 @@ async def test_internal_process_swallows_stop_typing_failures(monkeypatch): @pytest.mark.asyncio -async def test_internal_process_sends_error_for_blocked_provider_api_base(monkeypatch): +@pytest.mark.parametrize("detached", [False, True]) +async def test_internal_process_sends_error_for_blocked_provider_api_base( + monkeypatch, detached +): stage = internal.InternalAgentSubStage.__new__(internal.InternalAgentSubStage) stage.streaming_response = False stage.show_tool_use = True @@ -608,6 +614,7 @@ async def test_internal_process_sends_error_for_blocked_provider_api_base(monkey monkeypatch.setattr( internal, "build_main_agent", AsyncMock(return_value=build_result) ) + event.set_extra("btw_detached_work", detached) register_runner = MagicMock() stage.ctx.execution_context.follow_up_coordinator.register_active_runner = ( register_runner @@ -617,6 +624,7 @@ async def test_internal_process_sends_error_for_blocked_provider_api_base(monkey assert yielded == [] register_runner.assert_not_called() + assert bool(event.get_extra("btw_work_failed")) is detached event.send.assert_awaited_once() assert ( event.send.await_args.args[0].get_plain_text() diff --git a/tests/unit/test_btw_delivery.py b/tests/unit/test_btw_delivery.py new file mode 100644 index 0000000000..4d9d8b9282 --- /dev/null +++ b/tests/unit/test_btw_delivery.py @@ -0,0 +1,252 @@ +import asyncio +from types import SimpleNamespace + +import pytest + +from astrbot.core.agent.btw.types import WorkSessionStatus +from astrbot.core.agent.btw.work_loop import WorkLoop +from astrbot.core.agent.btw.work_sessions import WorkSessionManager +from astrbot.core.message.message_event_result import MessageChain, MessageEventResult +from astrbot.core.pipeline.result_decorate.stage import ResultDecorateStage +from astrbot.core.pipeline.scheduler import PipelineScheduler +from astrbot.core.pipeline.stage import Stage +from astrbot.core.utils.active_event_registry import ActiveEventRegistry +from astrbot.core.webchat.emitter import emit_webchat_response +from astrbot.core.webchat.queue_manager import WebChatQueueManager +from astrbot.core.webchat.run_coordinator import WebChatRunCoordinator + + +class WorkEvent: + requires_empty_completion = True + + def __init__(self, run, queues, attachments): + self.message_id = run.request_id + self.unified_msg_origin = "webchat:FriendMessage:shared" + self.message_str = "run work" + self.queues = queues + self.attachments = attachments + self.extras = {} + self.result = None + self.cleaned = 0 + self.trace = [] + + def get_extra(self, key, default=None): + return self.extras.get(key, default) + + def set_extra(self, key, value): + self.extras[key] = value + + def set_result(self, result): + self.result = result + + def is_stopped(self): + return False + + def get_platform_id(self): + return "webchat" + + def get_message_outline(self): + return self.message_str + + def cleanup_temporary_local_files(self): + self.cleaned += 1 + + async def send(self, message): + return await emit_webchat_response( + self.queues, self.message_id, message, attachments_dir=self.attachments + ) + + +class BlockingExecutor: + def __init__(self): + self.started = {} + self.release = {} + + async def process(self, event): + self.started[event.message_id].set() + await self.release[event.message_id].wait() + await event.send(MessageChain(type="agent_stats").message('{"calls": 1}')) + event.set_result(MessageEventResult().message("finished " + event.message_id)) + yield + + +class SubmitStage(Stage): + def __init__(self, work): + self.work = work + + async def initialize(self, ctx): + pass + + def configure_detached_work(self, **kwargs): + self.work.configure_detached_execution(**kwargs) + + async def process(self, event): + async for progress in self.work.submit(event): + yield progress + + async def close(self): + await self.work.close() + + +class DecorateStage(ResultDecorateStage): + async def initialize(self, ctx): + pass + + async def process(self, event): + if event.result is None: + return + event.trace.append("decorate-before") + yield + event.trace.append("decorate-after") + + +class SendStage(Stage): + async def initialize(self, ctx): + pass + + async def process(self, event): + if event.result is None: + return + event.trace.append("send") + await event.send(event.result) + event.result = None + + +async def setup_work(tmp_path): + queues = WebChatQueueManager() + coordinator = WebChatRunCoordinator(queues) + executor = BlockingExecutor() + sessions = WorkSessionManager() + work = WorkLoop(executor, sessions) + ctx = SimpleNamespace( + execution_context=SimpleNamespace( + active_event_registry=ActiveEventRegistry(), + background_tasks=set(), + ) + ) + scheduler = PipelineScheduler(ctx) + scheduler.stage_classes = [lambda: SubmitStage(work), DecorateStage, SendStage] + await scheduler.initialize() + events = [] + for request_id in ("first", "second"): + run = coordinator.create_run( + session_id="shared", username="test", request_id=request_id + ) + executor.started[request_id] = asyncio.Event() + executor.release[request_id] = asyncio.Event() + events.append(WorkEvent(run, queues, tmp_path)) + return scheduler, work, executor, queues, events + + +@pytest.mark.asyncio +async def test_background_webchat_keeps_each_request_until_its_final_result(tmp_path): + scheduler, work, executor, queues, events = await setup_work(tmp_path) + first, second = events + try: + await scheduler.execute(first) + await scheduler.execute(second) + await asyncio.wait_for(executor.started["first"].wait(), timeout=1) + await asyncio.wait_for(executor.started["second"].wait(), timeout=1) + for event in events: + ack = queues.back_queues[event.message_id].get_nowait() + assert ack["type"] == "plain" + assert ack["message_id"] == event.message_id + assert queues.back_queues[event.message_id].empty() + assert event.cleaned == 0 + + executor.release["first"].set() + first_task = next(t for t, (event, _) in work._tasks.items() if event is first) + await asyncio.wait_for(first_task, timeout=1) + messages = [queues.back_queues["first"].get_nowait() for _ in range(3)] + assert [message["type"] for message in messages] == ["plain", "plain", "end"] + assert messages[0]["chain_type"] == "agent_stats" + assert messages[1]["data"] == "finished first" + assert {message["message_id"] for message in messages} == {"first"} + assert queues.back_queues["second"].empty() + assert first.cleaned == 1 and second.cleaned == 0 + assert first.trace == ["decorate-before", "send", "decorate-after"] * 2 + await scheduler.finalize_detached_event(first) + assert first.cleaned == 1 + finally: + await scheduler.close() + assert second.cleaned == 1 + assert queues.back_queues["second"].get_nowait()["type"] == "end" + + +@pytest.mark.asyncio +async def test_closing_scheduler_cleans_work_cancelled_before_it_starts(tmp_path): + scheduler, work, _, queues, events = await setup_work(tmp_path) + event = events[0] + await scheduler.execute(event) + await scheduler.close() + session = await work.sessions.get_for_origin(event.unified_msg_origin) + assert session.status is WorkSessionStatus.CANCELLED + assert event.cleaned == 1 + assert not scheduler.ctx.execution_context.active_event_registry._events + assert [ + queues.back_queues[event.message_id].get_nowait()["type"] for _ in range(2) + ] == ["plain", "end"] + + +@pytest.mark.asyncio +async def test_finalizer_releases_resources_even_when_completion_delivery_fails( + tmp_path, +): + scheduler, _, _, _, events = await setup_work(tmp_path) + event = events[0] + scheduler.ctx.execution_context.active_event_registry.register(event) + + async def fail(message): + raise OSError("delivery unavailable") + + event.send = fail + with pytest.raises(OSError, match="delivery unavailable"): + await scheduler.finalize_detached_event(event) + assert event.cleaned == 1 + assert not scheduler.ctx.execution_context.active_event_registry._events + + +@pytest.mark.asyncio +async def test_work_delivery_failure_marks_failed_and_finishes_the_request(tmp_path): + scheduler, work, executor, _, events = await setup_work(tmp_path) + event = events[0] + await scheduler.execute(event) + await asyncio.wait_for(executor.started[event.message_id].wait(), timeout=1) + + async def fail_delivery(event): + raise OSError("cannot deliver") + + work._result_dispatcher = fail_delivery + task = next(iter(work._tasks)) + executor.release[event.message_id].set() + with pytest.raises(OSError, match="cannot deliver"): + await task + session = await work.sessions.get_for_origin(event.unified_msg_origin) + assert session.status is WorkSessionStatus.FAILED + assert session.error == "Work task failed." + assert event.cleaned == 1 + + +@pytest.mark.asyncio +async def test_close_during_acknowledgement_prevents_late_background_submission( + tmp_path, +): + scheduler, work, _, _, events = await setup_work(tmp_path) + event = events[0] + admission = work.submit(event) + await anext(admission) + await work.close() + assert [item async for item in admission] == [] + assert not work._tasks + assert not event.get_extra("btw_detached_work") + session = await work.sessions.get_for_origin(event.unified_msg_origin) + assert session.status is WorkSessionStatus.CANCELLED + + +@pytest.mark.asyncio +async def test_closed_work_rejects_submission_without_creating_a_session(tmp_path): + scheduler, work, _, _, events = await setup_work(tmp_path) + await work.close() + await scheduler.execute(events[0]) + assert await work.sessions.get_for_origin(events[0].unified_msg_origin) is None + assert not work._tasks diff --git a/tests/unit/test_core_lifecycle.py b/tests/unit/test_core_lifecycle.py index 3120a43691..3184e8f23a 100644 --- a/tests/unit/test_core_lifecycle.py +++ b/tests/unit/test_core_lifecycle.py @@ -1346,7 +1346,8 @@ async def test_reload_pipeline_scheduler_updates_existing( lifecycle.astrbot_config_mgr = mock_astrbot_config_mgr lifecycle.plugin_manager = mock_plugin_manager lifecycle.execution_context = MagicMock() - lifecycle.pipeline_scheduler_mapping = {} + old_scheduler = SimpleNamespace(close=AsyncMock()) + lifecycle.pipeline_scheduler_mapping = {"config1": old_scheduler} with ( patch( @@ -1360,6 +1361,7 @@ async def test_reload_pipeline_scheduler_updates_existing( # Verify scheduler was added to mapping assert "config1" in lifecycle.pipeline_scheduler_mapping + old_scheduler.close.assert_awaited_once() mock_new_scheduler.initialize.assert_awaited_once() @pytest.mark.asyncio @@ -1376,3 +1378,24 @@ async def test_reload_pipeline_scheduler_raises_for_missing_config( with pytest.raises(ValueError, match="配置文件 .* 不存在"): await lifecycle.reload_pipeline_scheduler("nonexistent") + + +@pytest.mark.asyncio +async def test_pipeline_work_closes_before_runtime_dependencies( + mock_log_broker, mock_db +): + lifecycle = AstrBotCoreLifecycle(mock_log_broker, mock_db) + order = [] + + async def close_work(): + order.append("work") + + async def close_transport(): + order.append("transport") + + lifecycle.pipeline_scheduler_mapping = { + "profile": SimpleNamespace(close=close_work), + } + lifecycle._register_cleanup("transport", close_transport) + await lifecycle.stop() + assert order == ["work", "transport"] From 28d7ca464c7e7921b7845539665433f61926698c Mon Sep 17 00:00:00 2001 From: BegoniaHe Date: Thu, 10 Sep 2026 21:20:57 +0200 Subject: [PATCH 2/3] fix(btw): restore scheduler close awaitable type AI-Generated: true Generated-At: 2026-09-10T19:20:50Z --- astrbot/core/pipeline/scheduler.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/astrbot/core/pipeline/scheduler.py b/astrbot/core/pipeline/scheduler.py index e45609ef6f..ebb43c5c18 100644 --- a/astrbot/core/pipeline/scheduler.py +++ b/astrbot/core/pipeline/scheduler.py @@ -48,7 +48,7 @@ async def close(self) -> None: for stage in reversed(self.stages): close = getattr(stage, "close", None) if callable(close): - await close() + await cast(Awaitable[None], close()) async def deliver_detached_result(self, event: AstrMessageEvent) -> None: """Replay response decoration and delivery with onion ordering intact.""" From 075a270541245358b8ecc3e8bfbcef9dcd38c45e Mon Sep 17 00:00:00 2001 From: BegoniaHe Date: Thu, 10 Sep 2026 21:24:50 +0200 Subject: [PATCH 3/3] fix(btw): reconcile detached delivery failures AI-Generated: true Generated-At: 2026-09-10T19:24:44Z --- astrbot/core/agent/btw/work_loop.py | 4 ++++ tests/unit/test_btw_delivery.py | 3 +-- 2 files changed, 5 insertions(+), 2 deletions(-) diff --git a/astrbot/core/agent/btw/work_loop.py b/astrbot/core/agent/btw/work_loop.py index a504464f73..31d46c58d6 100644 --- a/astrbot/core/agent/btw/work_loop.py +++ b/astrbot/core/agent/btw/work_loop.py @@ -224,8 +224,12 @@ async def _run_detached(self, event: AstrMessageEvent, session_id: str) -> None: async for _ in execution: await self._result_dispatcher(event) except asyncio.CancelledError: + await self.sessions.update_status(session_id, WorkSessionStatus.CANCELLED) raise except Exception as exc: + await self.sessions.update_status( + session_id, WorkSessionStatus.FAILED, error="Work task failed." + ) # The task registry logs unhandled exceptions with their traceback. # Consume executor failures here so provider details never reach it. logger.error("BTW work task failed: %s", safe_error("", exc)) diff --git a/tests/unit/test_btw_delivery.py b/tests/unit/test_btw_delivery.py index 4d9d8b9282..becf600cb5 100644 --- a/tests/unit/test_btw_delivery.py +++ b/tests/unit/test_btw_delivery.py @@ -219,8 +219,7 @@ async def fail_delivery(event): work._result_dispatcher = fail_delivery task = next(iter(work._tasks)) executor.release[event.message_id].set() - with pytest.raises(OSError, match="cannot deliver"): - await task + await task session = await work.sessions.get_for_origin(event.unified_msg_origin) assert session.status is WorkSessionStatus.FAILED assert session.error == "Work task failed."