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)]