From 5df3a477e224b7e9f8527521e8f6820133a2ff90 Mon Sep 17 00:00:00 2001 From: YUZHEthefool <2804776511@qq.com> Date: Fri, 11 Sep 2026 22:53:39 +0800 Subject: [PATCH] fix(btw): stop a third-party work run when its request is stopped /task stop for a third-party runner records the stop on the event, but the detached run never read it: it kept pulling responses from the runner and finished as completed. Nothing was cancelled, and the event's own stop flag was set the whole time. Honour the stop on all three paths. The third-party consumer stops draining the runner, the work loop stops dispatching delivered results and closes the execution so the runner is released, and a run still queued behind the semaphore never starts. A stopped run is a cancellation, not a runner failure, so the third-party handler reports nothing instead of the fallback error an unfinished stream would produce. This ends the local wait only: it does not revoke a task the remote service already accepted, which the /task stop docs now say. Fixes #157 AI-Generated: true Generated-At: 2026-09-11T14:53:39Z --- astrbot/core/agent/btw/work_loop.py | 31 ++- .../method/agent_sub_stages/third_party.py | 20 ++ docs/en/use/command.md | 2 +- docs/zh/use/command.md | 2 +- tests/unit/agent_sub_stage_support.py | 28 ++- tests/unit/test_btw_delivery.py | 181 ++++++++++++++++++ tests/unit/test_btw_work_loop.py | 33 ++-- tests/unit/test_conversation_loop.py | 7 + .../unit/test_third_party_agent_sub_stage.py | 3 + 9 files changed, 278 insertions(+), 29 deletions(-) diff --git a/astrbot/core/agent/btw/work_loop.py b/astrbot/core/agent/btw/work_loop.py index 1c51806351..683920b579 100644 --- a/astrbot/core/agent/btw/work_loop.py +++ b/astrbot/core/agent/btw/work_loop.py @@ -248,6 +248,14 @@ async def _execute( produced = False try: async with self._semaphore: + if event.is_stopped(): + # The request was stopped while this run waited its turn, so + # it never starts work the user already withdrew. + await self.sessions.update_status( + session_id, + WorkSessionStatus.CANCELLED, + ) + return await self.sessions.update_status( session_id, WorkSessionStatus.RUNNING, @@ -272,11 +280,15 @@ async def _execute( failed = bool(event.get_extra(WORK_FAILED_EXTRA)) or bool( event.get_extra(THIRD_PARTY_RUNNER_ERROR_EXTRA_KEY) ) - # An admitted stop request and a run that reported its own abort are - # both cancellations. ``run_agent`` clears ``agent_stop_requested`` - # when it reports the abort, so the stop flag alone misses that case. - cancelled = bool(event.get_extra("agent_stop_requested")) or bool( - event.get_extra("agent_user_aborted") + # A stopped event, an admitted stop request, and a run that reported + # its own abort are all cancellations. ``run_agent`` clears + # ``agent_stop_requested`` when it reports the abort, and a + # third-party stop only sets the event's own flag, so neither signal + # alone covers every way a work run is cancelled. + cancelled = ( + event.is_stopped() + or bool(event.get_extra("agent_stop_requested")) + or bool(event.get_extra("agent_user_aborted")) ) # An executor that produced nothing never reached an Agent: a run # the session turned away is the reported case. The generator @@ -301,6 +313,15 @@ async def _run_detached(self, event: AstrMessageEvent, session_id: str) -> None: try: async with aclosing(self._execute(event, session_id)) as execution: async for _ in execution: + if event.is_stopped(): + # Closing the execution releases the executor and the + # runner behind it. This ends the local run and its + # waiting only: whether the remote service stopped its + # own task is not something this side can claim. + await self.sessions.update_status( + session_id, WorkSessionStatus.CANCELLED + ) + return await self._result_dispatcher(event) await self._record_delivery_outcome(event, session_id) except asyncio.CancelledError: diff --git a/astrbot/core/pipeline/process_stage/method/agent_sub_stages/third_party.py b/astrbot/core/pipeline/process_stage/method/agent_sub_stages/third_party.py index 5a9b557974..ede8359b64 100644 --- a/astrbot/core/pipeline/process_stage/method/agent_sub_stages/third_party.py +++ b/astrbot/core/pipeline/process_stage/method/agent_sub_stages/third_party.py @@ -61,13 +61,21 @@ async def run_third_party_agent( max_step: int = 30, stream_to_general: bool = False, custom_error_message: str | None = None, + should_stop: Callable[[], bool] | None = None, ) -> AsyncGenerator[tuple[MessageChain, bool]]: """ 运行第三方 agent runner 并转换响应格式 类似于 run_agent 函数，但专门处理第三方 agent runner + + ``should_stop`` is polled as the stream is consumed, so a request that was + stopped stops draining this runner instead of waiting for it to finish. It + ends the local wait only; whether the remote service stops its own task is + not something this side can observe. """ try: async for resp in runner.step_until_done(max_step=max_step): # type: ignore[misc] + if should_stop is not None and should_stop(): + return if resp.type == "streaming_delta": if stream_to_general: continue @@ -241,6 +249,7 @@ async def _stream_runner_chain() -> AsyncGenerator[MessageChain]: max_step=max_step, stream_to_general=False, custom_error_message=custom_error_message, + should_stop=event.is_stopped, ): aggregator.add_chunk(chain, is_error) if is_error: @@ -258,6 +267,11 @@ async def _stream_runner_chain() -> AsyncGenerator[MessageChain]: ) yield + if event.is_stopped(): + # A stopped request did not fail. Report nothing instead of the + # fallback error an unfinished runner would otherwise produce. + return + if runner.done(): final_chain, is_runner_error = aggregator.finalize( runner.get_final_llm_resp() @@ -285,12 +299,18 @@ async def _handle_non_streaming_response( max_step=max_step, stream_to_general=stream_to_general, custom_error_message=custom_error_message, + should_stop=event.is_stopped, ): aggregator.add_chunk(chain, is_error) if is_error: event.set_extra(THIRD_PARTY_RUNNER_ERROR_EXTRA_KEY, True) yield + if event.is_stopped(): + # A stopped request did not fail. Report nothing instead of the + # fallback error an unfinished runner would otherwise produce. + return + final_chain, is_runner_error = aggregator.finalize(runner.get_final_llm_resp()) event.set_extra(THIRD_PARTY_RUNNER_ERROR_EXTRA_KEY, is_runner_error) result_content_type = ( diff --git a/docs/en/use/command.md b/docs/en/use/command.md index 62ddf86629..50af44b923 100644 --- a/docs/en/use/command.md +++ b/docs/en/use/command.md @@ -83,7 +83,7 @@ The user ID from `/session info` can be granted current-session `session_admin` - `/work `: Submit the remaining text to the BTW work loop without a `/chat` prefix or automatic classification. Requires `session.read`, with `btw.enabled` and `btw.work_loop.enabled` enabled on the current profile. The command identity is `builtin_commands:work`. Command quoting rules still apply. - `/work` or `/work status`: Show the latest task's text and status for the current profile and session: queued, running, completed, failed, cancelled, or delivery unconfirmed. `delivery unconfirmed` means the task itself finished but the platform never confirmed accepting its result. `status` queries only when it is the entire remainder, ignoring case; `/work status refactor` submits a task. Requires `session.read`. State is kept in memory, cleared on restart or profile reload/removal, and terminal tasks expire after `btw.work_session.max_age_seconds` (default 3600). -- `/task stop`: Stop running Agent or third-party Agent Runner tasks in the current session without deleting history. +- `/task stop`: Request that running Agent or third-party Agent Runner tasks in the current session stop, without deleting history. The local run stops consuming the runner and is recorded as cancelled. This ends the local wait only; a task already accepted by a third-party service is not revoked remotely. ### Providers and Models diff --git a/docs/zh/use/command.md b/docs/zh/use/command.md index 7d1c4b6dc2..e258db5c03 100644 --- a/docs/zh/use/command.md +++ b/docs/zh/use/command.md @@ -83,7 +83,7 @@ Orbit 不执行变量、命令、算术或波浪号展开，也不执行 glob、 - `/work <任务内容>`：将后面的文本显式提交给 BTW 工作循环，不需要 `/chat` 前缀或自动分类。要求 `session.read`，并在当前配置中启用 `btw.enabled` 和 `btw.work_loop.enabled`。指令标识为 `builtin_commands:work`；仍遵循指令引号规则。 - `/work` 或 `/work status`：查看当前配置与会话中最新任务的内容及状态：排队中、执行中、已完成、已失败、已取消或投递未确认。`投递未确认` 表示任务本身已结束，但平台没有确认接收其结果。仅当参数全部为 `status` 时查询，忽略大小写；`/work status 重构` 会提交任务。要求 `session.read`。状态保存在内存中，重启或配置重载、移除后清空；终态任务按 `btw.work_session.max_age_seconds` 过期，默认 3600 秒。 -- `/task stop`：停止当前会话中正在运行的 Agent 或第三方 Agent Runner 任务，不删除历史。 +- `/task stop`：请求停止当前会话中正在运行的 Agent 或第三方 Agent Runner 任务，不删除历史。本地运行会停止消费执行器并记为已取消；这里只结束本地等待，第三方服务已经接受的任务不会被远端撤销。 ### Provider 与模型 diff --git a/tests/unit/agent_sub_stage_support.py b/tests/unit/agent_sub_stage_support.py index 3a4e548a6d..c7e0c1d2a7 100644 --- a/tests/unit/agent_sub_stage_support.py +++ b/tests/unit/agent_sub_stage_support.py @@ -22,15 +22,22 @@ class FakeEvent: - def __init__(self, *, extras: dict | None = None): + def __init__(self, *, extras: dict | None = None, stopped: bool = False): self.unified_msg_origin = "webchat:FriendMessage:test-session" self._extras = extras or {} self.result_history: list[MessageEventResult] = [] self.temporary_local_files: list[str] = [] + self._stopped = stopped def get_extra(self, key: str): return self._extras.get(key) + def is_stopped(self) -> bool: + return self._stopped + + def stop_event(self) -> None: + self._stopped = True + def set_extra(self, key: str, value) -> None: self._extras[key] = value @@ -133,6 +140,25 @@ def done(self) -> bool: return self._done +class ThirdPartyResponseExecutor: + """Drive the real third-party response handler as a work-loop executor.""" + + def __init__(self, runner) -> None: + self.runner = runner + + async def process(self, event): + stage = third_party.ThirdPartyAgentSubStage.__new__( + third_party.ThirdPartyAgentSubStage + ) + async for _ in stage._handle_non_streaming_response( + runner=self.runner, + event=event, + stream_to_general=False, + custom_error_message=None, + ): + yield + + class FakeInternalRunner: def __init__( self, diff --git a/tests/unit/test_btw_delivery.py b/tests/unit/test_btw_delivery.py index 051c261314..c982829263 100644 --- a/tests/unit/test_btw_delivery.py +++ b/tests/unit/test_btw_delivery.py @@ -18,6 +18,10 @@ 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 +from tests.unit.agent_sub_stage_support import ( + FakeThirdPartyRunner, + ThirdPartyResponseExecutor, +) PROFILE_ID = "profile-1" @@ -487,3 +491,180 @@ async def test_unconfirmed_result_delivery_is_not_reported_as_completed(): session = await work.sessions.get_for_origin(event.unified_msg_origin) assert session is not None assert session.status is WorkSessionStatus.UNCONFIRMED + + +class StoppingThirdPartyRunner(FakeThirdPartyRunner): + """A third-party runner whose stream is stopped part-way through.""" + + def __init__(self, chunks, stop) -> None: + super().__init__(responses=chunks, final_resp=None) + self.chunks = chunks + self.stop = stop + self.consumed = 0 + + async def step_until_done(self, max_step: int = 30): + del max_step + for index, chunk in enumerate(self.chunks): + if index == 1: + # The user's stop lands while the runner is still streaming. + self.stop() + self.consumed += 1 + yield chunk + + +def _third_party_chunk(text: str, kind: str = "llm_result"): + return SimpleNamespace(type=kind, data={"chain": MessageChain().message(text)}) + + +def _work_event(tmp_path, request_id: str = "first"): + run = WebChatRunCoordinator(WebChatQueueManager()).create_run( + session_id="shared", username="test", request_id=request_id + ) + return WorkEvent(run, WebChatQueueManager(), tmp_path) + + +def _stoppable_work(runner, registry): + """Wire a work loop whose detached run finalizes through the scheduler.""" + work = WorkLoop(ThirdPartyResponseExecutor(runner), WorkSessionManager()) + dispatcher = AsyncMock() + execution_context = SimpleNamespace( + active_event_registry=registry, + background_tasks=set(), + ) + scheduler = PipelineScheduler(SimpleNamespace(execution_context=execution_context)) + work.configure_detached_execution( + background_tasks=execution_context.background_tasks, + result_dispatcher=dispatcher, + event_finalizer=scheduler.finalize_detached_event, + ) + return work, dispatcher + + +@pytest.mark.asyncio +async def test_stopped_third_party_work_stops_consuming_its_stream(tmp_path): + """A stop must end the local wait instead of draining the whole runner.""" + registry = ActiveEventRegistry() + event = _work_event(tmp_path) + registry.register(event) + chunks = [_third_party_chunk(text) for text in ("one", "two", "three")] + runner = StoppingThirdPartyRunner( + chunks, lambda: registry.stop_all(event.unified_msg_origin) + ) + work, dispatcher = _stoppable_work(runner, registry) + + session = await work.schedule(event) + await asyncio.wait_for(next(iter(work._tasks)), timeout=5) + + # The stop arrives as the second response is pulled, so the third is never + # consumed and nothing from the aborted runner is delivered. + assert runner.consumed < len(chunks) + dispatcher.assert_not_awaited() + assert event.result is None + assert session.status is WorkSessionStatus.CANCELLED + assert event.cleaned == 1 + assert not registry._events + + +@pytest.mark.asyncio +async def test_stopped_work_delivers_no_further_results(tmp_path): + """A stream that kept producing must not keep being delivered.""" + registry = ActiveEventRegistry() + event = _work_event(tmp_path) + registry.register(event) + chunks = [ + _third_party_chunk(text, "streaming_delta") for text in ("one", "two", "three") + ] + runner = FakeThirdPartyRunner(responses=chunks, final_resp=None) + work, dispatcher = _stoppable_work(runner, registry) + + async def stop_on_first_delivery(_event): + registry.stop_all(event.unified_msg_origin) + + dispatcher.side_effect = stop_on_first_delivery + + session = await work.schedule(event) + await asyncio.wait_for(next(iter(work._tasks)), timeout=5) + + assert dispatcher.await_count == 1 + assert session.status is WorkSessionStatus.CANCELLED + assert event.cleaned == 1 + assert not registry._events + + +class StopBetweenChunksExecutor: + """A local work run whose event is stopped while it is still producing.""" + + def __init__(self) -> None: + self.yielded = 0 + + async def process(self, event): + self.yielded += 1 + yield "first" + event.stop_event() + self.yielded += 1 + yield "second" + + +@pytest.mark.asyncio +async def test_stopped_run_never_delivers_the_chunks_after_the_stop(tmp_path): + """A local run keeps yielding; delivery must stop at the stop flag.""" + event = _work_event(tmp_path) + executor = StopBetweenChunksExecutor() + dispatcher = AsyncMock() + work = WorkLoop(executor, WorkSessionManager()) + work.configure_detached_execution( + background_tasks=set(), + result_dispatcher=dispatcher, + event_finalizer=AsyncMock(), + ) + + session = await work.schedule(event) + await asyncio.wait_for(next(iter(work._tasks)), timeout=5) + + assert executor.yielded == 2 + assert dispatcher.await_count == 1 + assert session.status is WorkSessionStatus.CANCELLED + + +class GatedExecutor: + """A work executor that runs one item at a time.""" + + def __init__(self) -> None: + self.started: list[str] = [] + self.release = asyncio.Event() + + async def process(self, event): + self.started.append(event.message_id) + await self.release.wait() + event.set_result(MessageEventResult().message("done")) + yield + + +@pytest.mark.asyncio +async def test_queued_work_never_starts_once_its_request_is_stopped(tmp_path): + """A stop withdraws a task that is still waiting its turn.""" + registry = ActiveEventRegistry() + executor = GatedExecutor() + work = WorkLoop(executor, WorkSessionManager(), max_concurrent=1) + work.configure_detached_execution( + background_tasks=set(), + result_dispatcher=AsyncMock(), + event_finalizer=AsyncMock(), + ) + running = _work_event(tmp_path, "running") + queued = _work_event(tmp_path, "queued") + registry.register(running) + registry.register(queued) + + await work.schedule(running) + queued_session = await work.schedule(queued) + await asyncio.sleep(0) + assert queued_session.status is WorkSessionStatus.PENDING + + registry.stop_all(queued.unified_msg_origin, exclude=running) + tasks = list(work._tasks) + executor.release.set() + await asyncio.wait_for(asyncio.gather(*tasks), timeout=5) + + assert executor.started == ["running"] + assert queued_session.status is WorkSessionStatus.CANCELLED diff --git a/tests/unit/test_btw_work_loop.py b/tests/unit/test_btw_work_loop.py index d1457d19e9..4a9d60505e 100644 --- a/tests/unit/test_btw_work_loop.py +++ b/tests/unit/test_btw_work_loop.py @@ -12,7 +12,10 @@ from astrbot.core.message.message_event_result import MessageChain from astrbot.core.pipeline.process_stage.method import agent_request from astrbot.core.pipeline.process_stage.method.agent_sub_stages import third_party -from tests.unit.agent_sub_stage_support import FakeThirdPartyRunner +from tests.unit.agent_sub_stage_support import ( + FakeThirdPartyRunner, + ThirdPartyResponseExecutor, +) from tests.unit.test_astr_agent_run_util import FakeEvent as RunnerEvent from tests.unit.test_astr_agent_run_util import FakeRunner @@ -23,6 +26,7 @@ def __init__(self, message: str) -> None: self.message_str = message self.extras = {} self.result = None + self._stopped = False def set_extra(self, key, value) -> None: self.extras[key] = value @@ -33,6 +37,12 @@ def get_extra(self, key): def set_result(self, value) -> None: self.result = value + def is_stopped(self) -> bool: + return self._stopped + + def stop_event(self) -> None: + self._stopped = True + class FailingExecutor: async def process(self, event): @@ -85,25 +95,6 @@ async def process(self, event): yield -class ThirdPartyExecutor: - """The real third-party response handler as the work loop's executor.""" - - def __init__(self, runner) -> None: - self.runner = runner - - async def process(self, event): - stage = third_party.ThirdPartyAgentSubStage.__new__( - third_party.ThirdPartyAgentSubStage - ) - async for _ in stage._handle_non_streaming_response( - runner=self.runner, - event=event, - stream_to_general=False, - custom_error_message=None, - ): - yield - - class AdmissionSkipExecutor: """The real Agent request stage for a run the session turns away.""" @@ -370,7 +361,7 @@ async def test_work_loop_records_third_party_runner_error_as_failed(): final_resp=None, ) sessions = WorkSessionManager() - work_loop = WorkLoop(ThirdPartyExecutor(runner), sessions) + work_loop = WorkLoop(ThirdPartyResponseExecutor(runner), sessions) _ = [item async for item in work_loop.process(event)] diff --git a/tests/unit/test_conversation_loop.py b/tests/unit/test_conversation_loop.py index a7321dcb13..4391a9687e 100644 --- a/tests/unit/test_conversation_loop.py +++ b/tests/unit/test_conversation_loop.py @@ -20,6 +20,7 @@ async def process(self, event): class FakeEvent: def __init__(self) -> None: self.extras = {} + self._stopped = False def set_extra(self, key, value) -> None: self.extras[key] = value @@ -27,6 +28,12 @@ def set_extra(self, key, value) -> None: def get_extra(self, key): return self.extras.get(key) + def is_stopped(self) -> bool: + return self._stopped + + def stop_event(self) -> None: + self._stopped = True + @pytest.mark.asyncio @pytest.mark.parametrize("btw", [{"enabled": True}, {"enabled": False}, {}, None]) diff --git a/tests/unit/test_third_party_agent_sub_stage.py b/tests/unit/test_third_party_agent_sub_stage.py index aea46debd8..487ebfaf3a 100644 --- a/tests/unit/test_third_party_agent_sub_stage.py +++ b/tests/unit/test_third_party_agent_sub_stage.py @@ -105,6 +105,9 @@ def __new__(cls): event.message_obj.message = [] event.platform_meta.support_streaming_message = True event.get_extra.return_value = None + # The stage polls the event's stop flag while it consumes the runner, so a + # loose mock must answer it instead of returning a truthy mock. + event.is_stopped.return_value = False results = [item async for item in stage.process(event)]