From b471c1e4fc40ffa64de996ea497af2a8976e75f7 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?C=E2=82=82=E2=82=82H=E2=82=82=E2=82=85NO=E2=82=86?= Date: Mon, 6 Jul 2026 01:22:15 +0800 Subject: [PATCH 1/3] =?UTF-8?q?fix:=20=E4=BF=AE=E5=A4=8D=20stop=20?= =?UTF-8?q?=E5=90=8E=E8=BF=9F=E5=88=B0=20LLM=20=E5=9B=9E=E5=A4=8D=E6=BC=8F?= =?UTF-8?q?=E5=8F=91?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../agent/runners/tool_loop_agent_runner.py | 183 +++- astrbot/core/astr_agent_hooks.py | 53 +- astrbot/core/astr_agent_run_util.py | 68 +- astrbot/core/pipeline/context_utils.py | 26 +- .../method/agent_sub_stages/internal.py | 16 +- tests/test_tool_loop_agent_runner.py | 891 +++++++++++++++++- 6 files changed, 1156 insertions(+), 81 deletions(-) diff --git a/astrbot/core/agent/runners/tool_loop_agent_runner.py b/astrbot/core/agent/runners/tool_loop_agent_runner.py index 88038473af..44123d664a 100644 --- a/astrbot/core/agent/runners/tool_loop_agent_runner.py +++ b/astrbot/core/agent/runners/tool_loop_agent_runner.py @@ -112,10 +112,6 @@ class ToolLoopAgentRunner(BaseAgentRunner[TContext]): EMPTY_OUTPUT_RETRY_ATTEMPTS = 3 EMPTY_OUTPUT_RETRY_WAIT_MIN_S = 1 EMPTY_OUTPUT_RETRY_WAIT_MAX_S = 4 - USER_INTERRUPTION_MESSAGE = ( - "[SYSTEM: User actively interrupted the response generation. " - "Partial output before interruption is preserved.]" - ) FOLLOW_UP_NOTICE_TEMPLATE = ( "\n\n[SYSTEM NOTICE] User sent follow-up messages while tool execution " "was in progress. Prioritize these follow-up instructions in your next " @@ -176,8 +172,12 @@ def _get_persona_custom_error_message(self) -> str | None: event = getattr(self.run_context.context, "event", None) return extract_persona_custom_error_message_from_event(event) - async def _complete_with_assistant_response(self, llm_resp: LLMResponse) -> None: + async def _complete_with_assistant_response(self, llm_resp: LLMResponse) -> bool: """Finalize the current step as a plain assistant response with no tool calls.""" + if self._is_stop_requested(): + await self._finalize_aborted_step(llm_resp) + return False + self.final_llm_resp = llm_resp self._transition_state(AgentState.DONE) self.stats.end_time = time.time() @@ -196,11 +196,24 @@ async def _complete_with_assistant_response(self, llm_resp: LLMResponse) -> None logger.warning("LLM returned empty assistant message with no tool calls.") self.run_context.messages.append(Message(role="assistant", content=parts)) + if self._is_stop_requested(): + if self._is_message_from_llm_response( + self.run_context.messages[-1], llm_resp + ): + self.run_context.messages.pop() + await self._finalize_aborted_step(llm_resp) + return False + try: await self.agent_hooks.on_agent_done(self.run_context, llm_resp) except Exception as e: logger.error(f"Error in on_agent_done hook: {e}", exc_info=True) + if self._is_stop_requested(): + self.discard_late_aborted_result() + self._resolve_unconsumed_follow_ups() + return False self._resolve_unconsumed_follow_ups() + return True @override async def reset( @@ -282,6 +295,7 @@ async def reset( self._follow_up_seq = 0 self._last_tool_name: str | None = None self._same_tool_streak = 0 + self._recorded_usage_ids: set[int] = set() # These two are used for tool schema mode handling # We now have two modes: @@ -323,6 +337,17 @@ async def reset( self.stats = AgentStats() self.stats.start_time = time.time() + def _record_llm_usage(self, llm_resp: LLMResponse) -> None: + if not llm_resp.usage: + return + usage_id = id(llm_resp.usage) + if usage_id in self._recorded_usage_ids: + return + self._recorded_usage_ids.add(usage_id) + self.stats.token_usage += llm_resp.usage + if self.req.conversation: + self.req.conversation.token_usage = llm_resp.usage.total + def _read_tool_hint(self) -> str: if self.read_tool is not None: return f"`{self.read_tool.name}`" @@ -488,6 +513,9 @@ async def _iter_llm_responses_with_fallback( last_err_response: LLMResponse | None = None for idx, candidate in enumerate(candidates): + if self._is_stop_requested(): + return + candidate_id = candidate.provider_config.get("id", "") is_last_candidate = idx == total_candidates - 1 if idx > 0: @@ -526,6 +554,8 @@ async def _iter_llm_responses_with_fallback( and not has_stream_output and (not is_last_candidate) ): + if self._is_stop_requested(): + return last_err_response = resp logger.warning( "Chat Model %s returns error response, trying fallback to next provider.", @@ -734,6 +764,10 @@ async def step(self): self._simple_print_message_role("[AftCompact]", self.run_context.messages) async for llm_response in self._iter_llm_responses_with_fallback(): + if self._is_stop_requested(): + llm_resp_result = llm_response + break + if llm_response.is_chunk: if self.stats.time_to_first_token == 0: self.stats.time_to_first_token = time.time() - self.stats.start_time @@ -747,11 +781,17 @@ async def step(self): ), ), ) + if self._is_stop_requested(): + llm_resp_result = llm_response + break if llm_response.result_chain: yield AgentResponse( type="streaming_delta", data=AgentResponseData(chain=llm_response.result_chain), ) + if self._is_stop_requested(): + llm_resp_result = llm_response + break elif llm_response.completion_text: yield AgentResponse( type="streaming_delta", @@ -759,22 +799,18 @@ async def step(self): chain=MessageChain().message(llm_response.completion_text), ), ) + if self._is_stop_requested(): + llm_resp_result = llm_response + break if self._is_stop_requested(): - llm_resp_result = LLMResponse( - role="assistant", - completion_text=self.USER_INTERRUPTION_MESSAGE, - reasoning_content=llm_response.reasoning_content, - reasoning_signature=llm_response.reasoning_signature, - ) + llm_resp_result = llm_response break continue llm_resp_result = llm_response if not llm_response.is_chunk and llm_response.usage: # only count the token usage of the final response for computation purpose - self.stats.token_usage += llm_response.usage - if self.req.conversation: - self.req.conversation.token_usage = llm_response.usage.total + self._record_llm_usage(llm_response) break # got final response if not llm_resp_result: @@ -809,7 +845,12 @@ async def step(self): return if not llm_resp.tools_call_name: - await self._complete_with_assistant_response(llm_resp) + if not await self._complete_with_assistant_response(llm_resp): + yield AgentResponse( + type="aborted", + data=AgentResponseData(chain=MessageChain(type="aborted")), + ) + return # 返回 LLM 结果 if llm_resp.reasoning_content: @@ -838,11 +879,21 @@ async def step(self): if llm_resp.tools_call_name: if self.tool_schema_mode == "skills_like": requery_resp, _ = await self._resolve_tool_exec(llm_resp) + if self._is_stop_requested(): + yield await self._finalize_aborted_step(requery_resp) + return if not requery_resp.tools_call_name: llm_resp = requery_resp logger.warning( "skills_like tool re-query returned no tool calls; fallback to assistant response." ) + if not await self._complete_with_assistant_response(llm_resp): + yield AgentResponse( + type="aborted", + data=AgentResponseData(chain=MessageChain(type="aborted")), + ) + return + if llm_resp.reasoning_content: yield AgentResponse( type="llm_result", @@ -864,8 +915,6 @@ async def step(self): chain=MessageChain().message(llm_resp.completion_text), ), ) - - await self._complete_with_assistant_response(llm_resp) return else: llm_resp.tools_call_name = requery_resp.tools_call_name @@ -1332,6 +1381,9 @@ async def _resolve_tool_exec( llm_resp = requery_resp self._sanitize_malformed_tool_calls(llm_resp) + if self._is_stop_requested(): + return llm_resp, subset + # If the re-query still returns no tool calls, and also does not have a meaningful assistant reply, # we consider it as a failure of the LLM to follow the tool-use instruction, # and we will retry once with a stronger instruction that explicitly requires the LLM to either call the tool or give an explanation. @@ -1370,7 +1422,24 @@ def request_stop(self) -> None: self._abort_signal.set() def _is_stop_requested(self) -> bool: - return self._abort_signal.is_set() + if self._abort_signal.is_set(): + return True + + event = getattr(self.run_context.context, "event", None) + if event is None: + return False + + is_stopped = getattr(event, "is_stopped", None) + if callable(is_stopped) and is_stopped(): + return True + + get_extra = getattr(event, "get_extra", None) + if callable(get_extra): + return bool(get_extra("agent_stop_requested")) or bool( + get_extra("agent_user_aborted") + ) + + return False def was_aborted(self) -> bool: return self._aborted @@ -1378,38 +1447,72 @@ def was_aborted(self) -> bool: def get_final_llm_resp(self) -> LLMResponse | None: return self.final_llm_resp + def _build_aborted_llm_response( + self, + llm_resp: LLMResponse | None = None, + ) -> LLMResponse: + """构造不会泄露迟到模型正文的空响应。""" + return LLMResponse( + role="assistant", + completion_text="", + id=getattr(llm_resp, "id", None) if llm_resp else None, + usage=getattr(llm_resp, "usage", None) if llm_resp else None, + ) + + def _is_message_from_llm_response( + self, + message: Message, + llm_resp: LLMResponse | None, + ) -> bool: + if llm_resp is None or message.role != "assistant" or message.tool_calls: + return False + content = message.content + if not isinstance(content, list): + return False + for part in content: + if isinstance(part, TextPart) and part.text != llm_resp.completion_text: + return False + if isinstance(part, ThinkPart) and part.think != ( + llm_resp.reasoning_content or "" + ): + return False + return True + + def discard_late_aborted_result(self) -> None: + """丢弃已完成但尚未对用户可见发送的迟到模型正文。""" + original_llm_resp = self.final_llm_resp + if original_llm_resp: + self._record_llm_usage(original_llm_resp) + self.final_llm_resp = self._build_aborted_llm_response(self.final_llm_resp) + self._aborted = True + event = getattr(self.run_context.context, "event", None) + if event is not None: + event.set_extra("agent_user_aborted", True) + event.set_extra("agent_stop_requested", False) + if self.run_context.messages and self._is_message_from_llm_response( + self.run_context.messages[-1], + original_llm_resp, + ): + self.run_context.messages.pop() + async def _finalize_aborted_step( self, llm_resp: LLMResponse | None = None, ) -> AgentResponse: logger.info("Agent execution was requested to stop by user.") - if llm_resp is None: - llm_resp = LLMResponse(role="assistant", completion_text="") - if llm_resp.role != "assistant": - llm_resp = LLMResponse( - role="assistant", - completion_text=self.USER_INTERRUPTION_MESSAGE, - ) - self.final_llm_resp = llm_resp + safe_llm_resp = self._build_aborted_llm_response(llm_resp) + self.final_llm_resp = safe_llm_resp + self._record_llm_usage(safe_llm_resp) self._aborted = True + event = getattr(self.run_context.context, "event", None) + if event is not None: + event.set_extra("agent_user_aborted", True) + event.set_extra("agent_stop_requested", False) self._transition_state(AgentState.DONE) self.stats.end_time = time.time() - parts = [] - if llm_resp.reasoning_content is not None or llm_resp.reasoning_signature: - parts.append( - ThinkPart( - think=llm_resp.reasoning_content or "", - encrypted=llm_resp.reasoning_signature, - ) - ) - if llm_resp.completion_text: - parts.append(TextPart(text=llm_resp.completion_text)) - if parts: - self.run_context.messages.append(Message(role="assistant", content=parts)) - try: - await self.agent_hooks.on_agent_done(self.run_context, llm_resp) + await self.agent_hooks.on_agent_done(self.run_context, safe_llm_resp) except Exception as e: logger.error(f"Error in on_agent_done hook: {e}", exc_info=True) diff --git a/astrbot/core/astr_agent_hooks.py b/astrbot/core/astr_agent_hooks.py index 3155257d71..b08238661e 100644 --- a/astrbot/core/astr_agent_hooks.py +++ b/astrbot/core/astr_agent_hooks.py @@ -7,9 +7,29 @@ from astrbot.core.agent.tool import FunctionTool from astrbot.core.astr_agent_context import AstrAgentContext from astrbot.core.pipeline.context_utils import call_event_hook +from astrbot.core.provider.entities import LLMResponse from astrbot.core.star.star_handler import EventType +def _event_requests_agent_stop(event) -> bool: + get_extra = getattr(event, "get_extra", None) + is_stopped = getattr(event, "is_stopped", None) + return ( + (bool(get_extra("agent_user_aborted")) if callable(get_extra) else False) + or (bool(get_extra("agent_stop_requested")) if callable(get_extra) else False) + or (bool(is_stopped()) if callable(is_stopped) else False) + ) + + +def _build_aborted_llm_response(llm_response) -> LLMResponse: + return LLMResponse( + role="assistant", + completion_text="", + id=getattr(llm_response, "id", None) if llm_response else None, + usage=getattr(llm_response, "usage", None) if llm_response else None, + ) + + class MainAgentHooks(BaseAgentRunHooks[AstrAgentContext]): async def on_agent_begin( self, run_context: ContextWrapper[AstrAgentContext] @@ -22,19 +42,34 @@ async def on_agent_begin( async def on_agent_done(self, run_context, llm_response) -> None: # 执行事件钩子 - if llm_response and llm_response.reasoning_content: + event = run_context.context.event + suppress_llm_response = _event_requests_agent_stop(event) + + if ( + not suppress_llm_response + and llm_response + and llm_response.reasoning_content + ): # we will use this in result_decorate stage to inject reasoning content to chain - run_context.context.event.set_extra( - "_llm_reasoning_content", llm_response.reasoning_content + set_extra = getattr(event, "set_extra", None) + if callable(set_extra): + set_extra("_llm_reasoning_content", llm_response.reasoning_content) + + if not suppress_llm_response: + await call_event_hook( + event, + EventType.OnLLMResponseEvent, + llm_response, ) + if _event_requests_agent_stop(event): + set_extra = getattr(event, "set_extra", None) + if callable(set_extra): + set_extra("_llm_reasoning_content", None) + llm_response = _build_aborted_llm_response(llm_response) + await call_event_hook( - run_context.context.event, - EventType.OnLLMResponseEvent, - llm_response, - ) - await call_event_hook( - run_context.context.event, + event, EventType.OnAgentDoneEvent, run_context, llm_response, diff --git a/astrbot/core/astr_agent_run_util.py b/astrbot/core/astr_agent_run_util.py index 465c5f6170..4896727119 100644 --- a/astrbot/core/astr_agent_run_util.py +++ b/astrbot/core/astr_agent_run_util.py @@ -25,7 +25,11 @@ def _should_stop_agent(astr_event) -> bool: - return astr_event.is_stopped() or bool(astr_event.get_extra("agent_stop_requested")) + return ( + astr_event.is_stopped() + or bool(astr_event.get_extra("agent_stop_requested")) + or bool(astr_event.get_extra("agent_user_aborted")) + ) def _truncate_tool_result(text: str, limit: int = 70) -> str: @@ -159,17 +163,8 @@ async def run_agent( agent_runner.request_stop() if resp.type == "aborted": - if can_buffer_llm_result: - merged_chain = _merge_buffered_llm_chains(buffered_llm_chains) - if merged_chain: - astr_event.set_result( - MessageEventResult( - chain=merged_chain.chain, - result_content_type=ResultContentType.LLM_RESULT, - ), - ) - yield merged_chain - astr_event.clear_result() + buffered_llm_chains.clear() + astr_event.clear_result() if not stop_watcher.done(): stop_watcher.cancel() try: @@ -272,16 +267,27 @@ async def run_agent( yield resp.data["chain"] # MessageChain if can_buffer_llm_result and agent_runner.done(): - merged_chain = _merge_buffered_llm_chains(buffered_llm_chains) - if merged_chain: - astr_event.set_result( - MessageEventResult( - chain=merged_chain.chain, - result_content_type=ResultContentType.LLM_RESULT, - ), - ) - yield merged_chain + if _should_stop_agent(astr_event): + buffered_llm_chains.clear() astr_event.clear_result() + discard_late_result = getattr( + agent_runner, + "discard_late_aborted_result", + None, + ) + if callable(discard_late_result): + discard_late_result() + else: + merged_chain = _merge_buffered_llm_chains(buffered_llm_chains) + if merged_chain: + astr_event.set_result( + MessageEventResult( + chain=merged_chain.chain, + result_content_type=ResultContentType.LLM_RESULT, + ), + ) + yield merged_chain + astr_event.clear_result() if not stop_watcher.done(): stop_watcher.cancel() @@ -439,6 +445,11 @@ async def run_live_agent( if queue_item is None: break + if agent_runner.was_aborted() or _should_stop_agent( + agent_runner.run_context.context.event + ): + continue + text = None if isinstance(queue_item, tuple): text, audio_data = queue_item @@ -550,8 +561,12 @@ async def _run_agent_feeder( # 更新 buffer 为剩余部分 buffer = temp_buffer + parts[-1] - # 处理剩余 buffer - if buffer.strip(): + # 处理剩余 buffer。若 stop 已到达，未送出的文本不能再进入 TTS。 + if ( + buffer.strip() + and not agent_runner.was_aborted() + and not _should_stop_agent(agent_runner.run_context.context.event) + ): await text_queue.put(buffer) except Exception as e: @@ -596,13 +611,20 @@ async def _simulated_stream_tts( text = await text_queue.get() if text is None: break + if _should_stop_agent(astr_event): + continue try: audio_path = await tts_provider.get_audio(text) + if _should_stop_agent(astr_event): + continue + if audio_path: with open(audio_path, "rb") as f: audio_data = f.read() + if _should_stop_agent(astr_event): + continue astr_event.track_temporary_local_file(audio_path) await audio_queue.put((text, audio_data)) except Exception as e: diff --git a/astrbot/core/pipeline/context_utils.py b/astrbot/core/pipeline/context_utils.py index 9402ce3e62..93104c7af9 100644 --- a/astrbot/core/pipeline/context_utils.py +++ b/astrbot/core/pipeline/context_utils.py @@ -9,6 +9,25 @@ from astrbot.core.star.star_handler import EventType, star_handlers_registry +def _event_requests_agent_stop(event: AstrMessageEvent) -> bool: + get_extra = getattr(event, "get_extra", None) + return ( + event.is_stopped() + or (bool(get_extra("agent_user_aborted")) if callable(get_extra) else False) + or (bool(get_extra("agent_stop_requested")) if callable(get_extra) else False) + ) + + +def _should_stop_hook_propagation( + event: AstrMessageEvent, hook_type: EventType +) -> bool: + if event.is_stopped(): + return True + return hook_type == EventType.OnLLMResponseEvent and _event_requests_agent_stop( + event + ) + + async def call_handler( event: AstrMessageEvent, handler: T.Callable[..., T.Awaitable[T.Any] | T.AsyncGenerator[T.Any, None]], @@ -90,6 +109,9 @@ async def call_event_hook( plugins_name=event.plugins_name, ) for handler in handlers: + if _should_stop_hook_propagation(event, hook_type): + return True + try: assert inspect.iscoroutinefunction(handler.handler) logger.debug( @@ -99,10 +121,10 @@ async def call_event_hook( except BaseException: logger.error(traceback.format_exc()) - if event.is_stopped(): + if _should_stop_hook_propagation(event, hook_type): logger.info( f"{star_map[handler.handler_module_path].name} - {handler.handler_name} 终止了事件传播。", ) return True - return event.is_stopped() + return _should_stop_hook_propagation(event, hook_type) 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 7b01dd2dc7..75f2bc79dc 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 @@ -320,6 +320,13 @@ async def process( ), ) yield + if agent_runner.was_aborted(): + event.set_result( + MessageEventResult( + chain=MessageChain().chain, + result_content_type=ResultContentType.STREAMING_FINISH, + ), + ) # 保存历史记录 if agent_runner.done() and ( @@ -351,7 +358,14 @@ async def process( ), ) yield - if agent_runner.done(): + if agent_runner.done() and agent_runner.was_aborted(): + event.set_result( + MessageEventResult( + chain=MessageChain().chain, + result_content_type=ResultContentType.STREAMING_FINISH, + ), + ) + elif agent_runner.done(): if final_llm_resp := agent_runner.get_final_llm_resp(): if final_llm_resp.completion_text: chain = ( diff --git a/tests/test_tool_loop_agent_runner.py b/tests/test_tool_loop_agent_runner.py index b4464680fb..3f40d9e267 100644 --- a/tests/test_tool_loop_agent_runner.py +++ b/tests/test_tool_loop_agent_runner.py @@ -16,12 +16,20 @@ from astrbot.core.agent.hooks import BaseAgentRunHooks from astrbot.core.agent.message import ImageURLPart, Message, TextPart from astrbot.core.agent.run_context import ContextWrapper +from astrbot.core.agent.runners.base import AgentState from astrbot.core.agent.runners.tool_loop_agent_runner import ToolLoopAgentRunner from astrbot.core.agent.tool import FunctionTool, ToolSet from astrbot.core.astr_agent_tool_exec import FunctionToolExecutor from astrbot.core.exceptions import EmptyModelOutputError -from astrbot.core.provider.entities import LLMResponse, ProviderRequest, TokenUsage -from astrbot.core.provider.provider import Provider +from astrbot.core.message.message_event_result import MessageChain +from astrbot.core.provider.entities import ( + LLMResponse, + ProviderMeta, + ProviderRequest, + ProviderType, + TokenUsage, +) +from astrbot.core.provider.provider import Provider, TTSProvider class MockProvider(Provider): @@ -214,6 +222,91 @@ async def text_chat_stream(self, **kwargs): ) +class MockLeakyAbortProvider(MockProvider): + def __init__(self, event=None): + super().__init__() + self.event = event + + async def text_chat(self, **kwargs) -> LLMResponse: + self.call_count += 1 + if self.event is not None: + self.event.set_extra("agent_stop_requested", True) + return LLMResponse( + role="assistant", + completion_text="late completion text", + result_chain=MessageChain().message("late result chain"), + tools_call_name=["late_tool"], + tools_call_args=[{"query": "late"}], + tools_call_ids=["call_late"], + tools_call_extra_content={"call_late": {"extra": "late"}}, + reasoning_content="late reasoning", + reasoning_signature="late signature", + raw_completion=object(), + id="late-id", + usage=TokenUsage(input_other=10, output=5), + ) + + +class MockDelayedTextProvider(MockProvider): + def __init__(self): + super().__init__() + self.started = asyncio.Event() + self.release = asyncio.Event() + + async def text_chat(self, **kwargs) -> LLMResponse: + self.call_count += 1 + self.started.set() + await self.release.wait() + return LLMResponse( + role="assistant", + completion_text="delayed visible text", + usage=TokenUsage(input_other=10, output=5), + ) + + +class MockStopAwareStreamProvider(MockProvider): + async def text_chat_stream(self, **kwargs): + abort_signal = kwargs.get("abort_signal") + if abort_signal is None: + return + await abort_signal.wait() + yield LLMResponse( + role="assistant", + completion_text="late chunk text", + reasoning_content="late chunk reasoning", + is_chunk=True, + ) + yield LLMResponse( + role="assistant", + completion_text="late final text", + is_chunk=False, + usage=TokenUsage(input_other=10, output=5), + ) + + +class MockDelayedTTSProvider(TTSProvider): + def __init__(self, audio_path: Path): + super().__init__({"type": "test_tts", "id": "test_tts"}, {}) + self.audio_path = audio_path + self.started = asyncio.Event() + self.release = asyncio.Event() + self.call_count = 0 + + async def get_audio(self, text: str) -> str: + self.call_count += 1 + self.started.set() + await self.release.wait() + return str(self.audio_path) + + def meta(self) -> ProviderMeta: + return ProviderMeta( + id="test_tts", + model=None, + type="test_tts", + provider_type=ProviderType.TEXT_TO_SPEECH, + ) + + class MockToolCallProvider(MockProvider): def __init__(self, tool_name: str, tool_args: dict[str, str] | None = None): super().__init__() @@ -326,6 +419,7 @@ def __init__(self): self.agent_done_called = False self.tool_start_called = False self.tool_end_called = False + self.agent_done_response = None async def on_agent_begin(self, run_context): self.agent_begin_called = True @@ -338,16 +432,120 @@ async def on_tool_end(self, run_context, tool, tool_args, tool_result): async def on_agent_done(self, run_context, llm_response): self.agent_done_called = True + self.agent_done_response = llm_response + + +class MockLateStopOnDoneHooks(MockHooks): + def __init__(self, event): + super().__init__() + self.event = event + + async def on_agent_done(self, run_context, llm_response): + await super().on_agent_done(run_context, llm_response) + self.event.set_extra("agent_stop_requested", True) + + +class HookPreStopRunner(ToolLoopAgentRunner): + def __init__(self, event): + super().__init__() + self.event = event + + def _is_stop_requested(self) -> bool: + if self.final_llm_resp is not None: + self.event.set_extra("agent_stop_requested", True) + return super()._is_stop_requested() + + +class LateStopAfterYieldRunner: + def __init__( + self, + event, + response_type: str, + text: str, + streaming: bool, + stop_after_yield: bool = True, + ): + self.run_context = ContextWrapper(context=MockAgentContext(event)) + self.response_type = response_type + self.text = text + self.streaming = streaming + self.stop_after_yield = stop_after_yield + self.stats = SimpleNamespace(to_dict=lambda: {}) + self._done = False + self._aborted = False + self.stop_requested = False + self.discard_late_result_called = False + + async def step(self): + yield SimpleNamespace( + type=self.response_type, + data={"chain": MessageChain().message(self.text)}, + ) + if self.stop_after_yield: + self.run_context.context.event.set_extra("agent_stop_requested", True) + self._done = True + + def done(self): + return self._done + + def request_stop(self): + self.stop_requested = True + + def was_aborted(self): + return self._aborted + + def discard_late_aborted_result(self): + self.discard_late_result_called = True + self._aborted = True + event = self.run_context.context.event + event.set_extra("agent_user_aborted", True) + event.set_extra("agent_stop_requested", False) class MockEvent: def __init__(self, umo: str, sender_id: str): self.unified_msg_origin = umo self._sender_id = sender_id + self._extras = {} + self.result = None + self.trace = SimpleNamespace(record=lambda *args, **kwargs: None) + self._stopped = False def get_sender_id(self): return self._sender_id + def set_extra(self, key, value): + self._extras[key] = value + + def get_extra(self, key=None, default=None): + if key is None: + return self._extras + return self._extras.get(key, default) + + def set_result(self, result): + self.result = result + + def clear_result(self): + self.result = None + + def is_stopped(self): + return self._stopped + + def stop(self): + self._stopped = True + + def track_temporary_local_file(self, _path): + return None + + def get_platform_name(self): + return "test" + + def get_platform_id(self): + return "test" + + async def send(self, *args, **kwargs): + return None + class MockAgentContext: def __init__(self, event): @@ -432,6 +630,10 @@ def _make_large_tool_result_text() -> str: return "x" * 100000 +async def _collect_async_iter(async_iter): + return [item async for item in async_iter] + + @pytest.mark.asyncio async def test_max_step_limit_functionality( runner, mock_provider, provider_request, mock_tool_executor, mock_hooks @@ -1080,7 +1282,7 @@ async def test_empty_output_retries_exhausted_then_uses_fallback_provider( @pytest.mark.asyncio -async def test_stop_signal_returns_aborted_and_persists_partial_message( +async def test_stop_signal_returns_aborted_and_discards_delayed_response( runner, provider_request, mock_tool_executor, mock_hooks ): provider = MockAbortableStreamProvider() @@ -1110,9 +1312,451 @@ async def test_stop_signal_returns_aborted_and_persists_partial_message( final_resp = runner.get_final_llm_resp() assert final_resp is not None assert final_resp.role == "assistant" - # When interrupted, the runner replaces completion_text with a system message - assert "interrupted" in final_resp.completion_text.lower() - assert runner.run_context.messages[-1].role == "assistant" + assert final_resp.completion_text == "" + assert final_resp.result_chain is None + assert final_resp.reasoning_content is None + assert final_resp.reasoning_signature is None + assert final_resp.tools_call_name == [] + assert final_resp.tools_call_args == [] + assert final_resp.tools_call_ids == [] + assert final_resp.tools_call_extra_content == {} + assert ( + not runner.run_context.messages + or runner.run_context.messages[-1].role != "assistant" + ) + assert mock_hooks.agent_done_response is final_resp + + +@pytest.mark.asyncio +async def test_aborted_final_response_sanitizes_all_model_output_fields( + runner, provider_request, mock_tool_executor, mock_hooks +): + event = MockEvent("test:FriendMessage:leaky_abort", "u1") + provider = MockLeakyAbortProvider(event) + + await runner.reset( + provider=provider, + request=provider_request, + run_context=ContextWrapper(context=MockAgentContext(event)), + tool_executor=mock_tool_executor, + agent_hooks=mock_hooks, + streaming=False, + ) + + responses = [] + async for response in runner.step(): + responses.append(response) + + assert [response.type for response in responses] == ["aborted"] + assert runner.was_aborted() is True + final_resp = runner.get_final_llm_resp() + assert final_resp is not None + assert final_resp.role == "assistant" + assert final_resp.completion_text == "" + assert final_resp.result_chain is None + assert final_resp.reasoning_content is None + assert final_resp.reasoning_signature is None + assert final_resp.raw_completion is None + assert final_resp.tools_call_args == [] + assert final_resp.tools_call_name == [] + assert final_resp.tools_call_ids == [] + assert final_resp.tools_call_extra_content == {} + assert final_resp.id == "late-id" + assert final_resp.usage is not None + assert final_resp.usage.total == 15 + assert runner.stats.token_usage.total == 15 + assert mock_hooks.agent_done_called is True + assert mock_hooks.agent_done_response is final_resp + + serialized_messages = repr(runner.run_context.messages) + assert "late completion text" not in serialized_messages + assert "late result chain" not in serialized_messages + assert "late reasoning" not in serialized_messages + assert "late_tool" not in serialized_messages + + +@pytest.mark.asyncio +async def test_runner_step_treats_agent_user_aborted_as_stop( + runner, provider_request, mock_tool_executor, mock_hooks +): + event = MockEvent("test:FriendMessage:runner_user_aborted", "u1") + event.set_extra("agent_user_aborted", True) + provider = MockLeakyAbortProvider() + + await runner.reset( + provider=provider, + request=provider_request, + run_context=ContextWrapper(context=MockAgentContext(event)), + tool_executor=mock_tool_executor, + agent_hooks=mock_hooks, + streaming=False, + ) + + responses = [response async for response in runner.step()] + + assert [response.type for response in responses] == ["aborted"] + assert provider.call_count == 0 + assert runner.was_aborted() is True + final_resp = runner.get_final_llm_resp() + assert final_resp is not None + assert final_resp.completion_text == "" + + +@pytest.mark.asyncio +async def test_run_agent_discards_buffered_llm_result_after_abort( + runner, provider_request, mock_tool_executor, mock_hooks +): + from astrbot.core.astr_agent_run_util import run_agent + + provider = MockDelayedTextProvider() + event = MockEvent("test:FriendMessage:buffer_abort", "u1") + + await runner.reset( + provider=provider, + request=provider_request, + run_context=ContextWrapper(context=MockAgentContext(event)), + tool_executor=mock_tool_executor, + agent_hooks=mock_hooks, + streaming=False, + ) + + agent_task = asyncio.create_task( + _collect_async_iter( + run_agent( + runner, + buffer_intermediate_messages=True, + ) + ) + ) + await asyncio.wait_for(provider.started.wait(), timeout=5) + event.set_extra("agent_stop_requested", True) + provider.release.set() + + chains = await asyncio.wait_for(agent_task, timeout=5) + + assert chains == [] + assert runner.was_aborted() is True + assert event.result is None + + +@pytest.mark.asyncio +async def test_run_agent_discards_buffered_result_when_stop_arrives_before_flush(): + from astrbot.core.astr_agent_run_util import run_agent + + event = MockEvent("test:FriendMessage:late_buffer_abort", "u1") + runner = LateStopAfterYieldRunner( + event, + response_type="llm_result", + text="buffered late text", + streaming=False, + ) + + chains = await _collect_async_iter( + run_agent( + cast(Any, runner), + buffer_intermediate_messages=True, + ) + ) + + assert chains == [] + assert event.result is None + assert runner.discard_late_result_called is True + assert runner.was_aborted() is True + + +@pytest.mark.asyncio +async def test_run_agent_treats_agent_user_aborted_as_stop(): + from astrbot.core.astr_agent_run_util import run_agent + + event = MockEvent("test:FriendMessage:user_aborted_abort", "u1") + event.set_extra("agent_user_aborted", True) + runner = LateStopAfterYieldRunner( + event, + response_type="llm_result", + text="aborted residual text", + streaming=False, + ) + + chains = await _collect_async_iter(run_agent(cast(Any, runner))) + + assert chains == [] + assert event.result is None + + +@pytest.mark.asyncio +async def test_stop_before_agent_done_hook_skips_llm_response_event( + provider_request, + mock_tool_executor, +): + from astrbot.core import astr_agent_hooks + from astrbot.core.astr_agent_hooks import MainAgentHooks + from astrbot.core.star.star_handler import EventType + + calls = [] + + async def fake_call_event_hook(event, hook_type, *args, **kwargs): + calls.append((hook_type, args)) + return False + + event = MockEvent("test:FriendMessage:hook_race_abort", "u1") + provider = MockProvider() + provider.should_call_tools = False + runner = HookPreStopRunner(event) + monkeypatch_context = pytest.MonkeyPatch() + monkeypatch_context.setattr( + astr_agent_hooks, + "call_event_hook", + fake_call_event_hook, + ) + + try: + await runner.reset( + provider=provider, + request=provider_request, + run_context=ContextWrapper(context=MockAgentContext(event)), + tool_executor=mock_tool_executor, + agent_hooks=MainAgentHooks(), + streaming=False, + ) + + responses = [response async for response in runner.step()] + finally: + monkeypatch_context.undo() + + assert [response.type for response in responses] == ["aborted"] + called_types = [hook_type for hook_type, _ in calls] + assert EventType.OnLLMResponseEvent not in called_types + assert EventType.OnAgentDoneEvent in called_types + assert event.get_extra("_llm_reasoning_content") is None + final_resp = runner.get_final_llm_resp() + assert final_resp is not None + assert final_resp.completion_text == "" + assert "这是我的最终回答" not in repr(runner.run_context.messages) + + +@pytest.mark.asyncio +async def test_stop_after_agent_done_hook_discards_final_assistant_message( + provider_request, + mock_tool_executor, +): + event = MockEvent("test:FriendMessage:post_hook_abort", "u1") + provider = MockProvider() + provider.should_call_tools = False + hooks = MockLateStopOnDoneHooks(event) + runner = ToolLoopAgentRunner() + + await runner.reset( + provider=provider, + request=provider_request, + run_context=ContextWrapper(context=MockAgentContext(event)), + tool_executor=mock_tool_executor, + agent_hooks=hooks, + streaming=False, + ) + + responses = [response async for response in runner.step()] + + assert [response.type for response in responses] == ["aborted"] + assert runner.was_aborted() is True + final_resp = runner.get_final_llm_resp() + assert final_resp is not None + assert final_resp.completion_text == "" + assert "这是我的最终回答" not in repr(runner.run_context.messages) + + +@pytest.mark.asyncio +async def test_live_agent_feeder_discards_residual_buffer_after_stop(): + from astrbot.core.astr_agent_run_util import _run_agent_feeder + + event = MockEvent("test:FriendMessage:live_late_buffer_abort", "u1") + runner = LateStopAfterYieldRunner( + event, + response_type="streaming_delta", + text="partial without punctuation", + streaming=True, + ) + text_queue = asyncio.Queue() + + await _run_agent_feeder( + cast(Any, runner), + text_queue, + max_step=30, + show_tool_use=True, + show_tool_call_result=False, + show_reasoning=False, + buffer_intermediate_messages=False, + ) + + assert await text_queue.get() is None + assert text_queue.empty() + + +@pytest.mark.asyncio +async def test_live_agent_discards_queued_tts_audio_after_stop(tmp_path): + from astrbot.core.astr_agent_run_util import run_live_agent + + audio_path = tmp_path / "late.wav" + audio_path.write_bytes(b"late-audio") + event = MockEvent("test:FriendMessage:live_tts_queued_abort", "u1") + runner = LateStopAfterYieldRunner( + event, + response_type="streaming_delta", + text="complete sentence!", + streaming=True, + stop_after_yield=False, + ) + tts_provider = MockDelayedTTSProvider(audio_path) + + live_task = asyncio.create_task( + _collect_async_iter(run_live_agent(cast(Any, runner), tts_provider)) + ) + await asyncio.wait_for(tts_provider.started.wait(), timeout=5) + event.set_extra("agent_stop_requested", True) + tts_provider.release.set() + + chains = await asyncio.wait_for(live_task, timeout=5) + + assert chains == [] + + +@pytest.mark.asyncio +async def test_stop_requested_before_stream_chunk_discards_chunk( + runner, provider_request, mock_tool_executor, mock_hooks +): + provider = MockStopAwareStreamProvider() + + await runner.reset( + provider=provider, + request=provider_request, + run_context=ContextWrapper(context=None), + tool_executor=mock_tool_executor, + agent_hooks=mock_hooks, + streaming=True, + ) + + runner.request_stop() + responses = [] + async for response in runner.step(): + responses.append(response) + + assert [response.type for response in responses] == ["aborted"] + assert runner.was_aborted() is True + final_resp = runner.get_final_llm_resp() + assert final_resp is not None + assert final_resp.completion_text == "" + assert final_resp.reasoning_content is None + + +@pytest.mark.asyncio +async def test_main_agent_hooks_skip_llm_response_hook_for_aborted_event(monkeypatch): + from astrbot.core import astr_agent_hooks + from astrbot.core.astr_agent_hooks import MainAgentHooks + from astrbot.core.star.star_handler import EventType + + calls = [] + + async def fake_call_event_hook(event, hook_type, *args, **kwargs): + calls.append((hook_type, args)) + return False + + monkeypatch.setattr(astr_agent_hooks, "call_event_hook", fake_call_event_hook) + + event = MockEvent("test:FriendMessage:hook_abort", "u1") + event.set_extra("agent_user_aborted", True) + response = LLMResponse( + role="assistant", + completion_text="", + reasoning_content="should not be stored", + ) + hooks = MainAgentHooks() + + await hooks.on_agent_done(ContextWrapper(context=MockAgentContext(event)), response) + + called_types = [hook_type for hook_type, _ in calls] + assert EventType.OnLLMResponseEvent not in called_types + assert EventType.OnAgentDoneEvent in called_types + assert event.get_extra("_llm_reasoning_content") is None + agent_done_call = next( + args for hook_type, args in calls if hook_type == EventType.OnAgentDoneEvent + ) + assert agent_done_call[1].completion_text == "" + assert agent_done_call[1].reasoning_content is None + + +@pytest.mark.asyncio +async def test_main_agent_hooks_skip_llm_response_hook_for_pending_stop(monkeypatch): + from astrbot.core import astr_agent_hooks + from astrbot.core.astr_agent_hooks import MainAgentHooks + from astrbot.core.star.star_handler import EventType + + calls = [] + + async def fake_call_event_hook(event, hook_type, *args, **kwargs): + calls.append((hook_type, args)) + return False + + monkeypatch.setattr(astr_agent_hooks, "call_event_hook", fake_call_event_hook) + + event = MockEvent("test:FriendMessage:hook_pending_stop", "u1") + event.set_extra("agent_stop_requested", True) + response = LLMResponse( + role="assistant", + completion_text="late hook text", + reasoning_content="late hook reasoning", + ) + hooks = MainAgentHooks() + + await hooks.on_agent_done(ContextWrapper(context=MockAgentContext(event)), response) + + called_types = [hook_type for hook_type, _ in calls] + assert EventType.OnLLMResponseEvent not in called_types + assert EventType.OnAgentDoneEvent in called_types + assert event.get_extra("_llm_reasoning_content") is None + agent_done_call = next( + args for hook_type, args in calls if hook_type == EventType.OnAgentDoneEvent + ) + assert agent_done_call[1].completion_text == "" + + +def test_llm_response_hook_propagation_stops_on_pending_agent_stop(): + from astrbot.core.pipeline.context_utils import _should_stop_hook_propagation + from astrbot.core.star.star_handler import EventType + + event = MockEvent("test:FriendMessage:hook_propagation_stop", "u1") + event.set_extra("agent_stop_requested", True) + + assert _should_stop_hook_propagation(event, EventType.OnLLMResponseEvent) is True + assert _should_stop_hook_propagation(event, EventType.OnAgentDoneEvent) is False + + +@pytest.mark.asyncio +async def test_main_agent_hooks_keep_normal_llm_response_hook(monkeypatch): + from astrbot.core import astr_agent_hooks + from astrbot.core.astr_agent_hooks import MainAgentHooks + from astrbot.core.star.star_handler import EventType + + calls = [] + + async def fake_call_event_hook(event, hook_type, *args, **kwargs): + calls.append((hook_type, args)) + return False + + monkeypatch.setattr(astr_agent_hooks, "call_event_hook", fake_call_event_hook) + + event = MockEvent("test:FriendMessage:hook_normal", "u1") + response = LLMResponse( + role="assistant", + completion_text="normal text", + reasoning_content="normal reasoning", + ) + hooks = MainAgentHooks() + + await hooks.on_agent_done(ContextWrapper(context=MockAgentContext(event)), response) + + called_types = [hook_type for hook_type, _ in calls] + assert EventType.OnLLMResponseEvent in called_types + assert EventType.OnAgentDoneEvent in called_types + assert event.get_extra("_llm_reasoning_content") == "normal reasoning" @pytest.mark.asyncio @@ -1364,6 +2008,234 @@ async def text_chat(self, **kwargs) -> LLMResponse: assert parts[0].text == "一张猫的照片" +@pytest.mark.asyncio +async def test_skills_like_requery_stop_discards_late_assistant_text(): + class SkillsLikeLateStopProvider(MockProvider): + def __init__(self, event): + super().__init__() + self.event = event + + async def text_chat(self, **kwargs) -> LLMResponse: + self.call_count += 1 + if self.call_count == 1: + return LLMResponse( + role="assistant", + completion_text="选择工具", + tools_call_name=["test_tool"], + tools_call_args=[{"query": "test"}], + tools_call_ids=["call_1"], + usage=TokenUsage(input_other=10, output=5), + ) + + self.event.set_extra("agent_stop_requested", True) + return LLMResponse( + role="assistant", + completion_text="skills_like late text", + usage=TokenUsage(input_other=10, output=5), + ) + + event = MockEvent(umo="test_umo", sender_id="test_sender") + provider = SkillsLikeLateStopProvider(event) + tool = FunctionTool( + name="test_tool", + description="测试", + parameters={"type": "object", "properties": {"query": {"type": "string"}}}, + handler=AsyncMock(), + ) + req = ProviderRequest( + prompt="调用工具", + func_tool=ToolSet(tools=[tool]), + contexts=[], + ) + runner = ToolLoopAgentRunner() + hooks = MockHooks() + + await runner.reset( + provider=provider, + request=req, + run_context=ContextWrapper(context=MockAgentContext(event)), + tool_executor=cast(Any, MockToolExecutor()), + agent_hooks=hooks, + tool_schema_mode="skills_like", + ) + + responses = [response async for response in runner.step()] + + assert [response.type for response in responses] == ["llm_result", "aborted"] + visible_text = "".join( + response.data["chain"].get_plain_text() + for response in responses + if response.type != "aborted" + ) + assert "skills_like late text" not in visible_text + assert runner.was_aborted() is True + final_resp = runner.get_final_llm_resp() + assert final_resp is not None + assert final_resp.completion_text == "" + assert hooks.agent_done_response is final_resp + assert "skills_like late text" not in repr(runner.run_context.messages) + + +@pytest.mark.asyncio +async def test_skills_like_requery_stop_skips_repair_request(): + class SkillsLikeRepairStopProvider(MockProvider): + def __init__(self, event): + super().__init__() + self.event = event + + async def text_chat(self, **kwargs) -> LLMResponse: + self.call_count += 1 + if self.call_count == 1: + return LLMResponse( + role="assistant", + completion_text="选择工具", + tools_call_name=["test_tool"], + tools_call_args=[{"query": "test"}], + tools_call_ids=["call_1"], + ) + + self.event.set_extra("agent_stop_requested", True) + return LLMResponse(role="assistant", completion_text="") + + event = MockEvent(umo="test_umo", sender_id="test_sender") + provider = SkillsLikeRepairStopProvider(event) + tool = FunctionTool( + name="test_tool", + description="测试", + parameters={"type": "object", "properties": {"query": {"type": "string"}}}, + handler=AsyncMock(), + ) + req = ProviderRequest( + prompt="调用工具", + func_tool=ToolSet(tools=[tool]), + contexts=[], + ) + runner = ToolLoopAgentRunner() + + await runner.reset( + provider=provider, + request=req, + run_context=ContextWrapper(context=MockAgentContext(event)), + tool_executor=cast(Any, MockToolExecutor()), + agent_hooks=MockHooks(), + tool_schema_mode="skills_like", + ) + + responses = [response async for response in runner.step()] + + assert [response.type for response in responses] == ["llm_result", "aborted"] + assert provider.call_count == 2 + + +@pytest.mark.asyncio +async def test_skills_like_requery_fallback_checks_stop_before_yielding_text(): + class StopAfterFallbackDoneHooks(MockHooks): + async def on_agent_done(self, run_context, llm_response): + await super().on_agent_done(run_context, llm_response) + run_context.context.event.set_extra("agent_stop_requested", True) + + class SkillsLikeFallbackProvider(MockProvider): + async def text_chat(self, **kwargs) -> LLMResponse: + self.call_count += 1 + if self.call_count == 1: + return LLMResponse( + role="assistant", + completion_text="选择工具", + tools_call_name=["test_tool"], + tools_call_args=[{"query": "test"}], + tools_call_ids=["call_1"], + usage=TokenUsage(input_other=10, output=5), + ) + + return LLMResponse( + role="assistant", + completion_text="skills_like fallback text", + reasoning_content="skills_like fallback reasoning", + usage=TokenUsage(input_other=10, output=5), + ) + + event = MockEvent(umo="test_umo", sender_id="test_sender") + provider = SkillsLikeFallbackProvider() + tool = FunctionTool( + name="test_tool", + description="测试", + parameters={"type": "object", "properties": {"query": {"type": "string"}}}, + handler=AsyncMock(), + ) + req = ProviderRequest( + prompt="调用工具", + func_tool=ToolSet(tools=[tool]), + contexts=[], + ) + runner = ToolLoopAgentRunner() + hooks = StopAfterFallbackDoneHooks() + + await runner.reset( + provider=provider, + request=req, + run_context=ContextWrapper(context=MockAgentContext(event)), + tool_executor=cast(Any, MockToolExecutor()), + agent_hooks=hooks, + tool_schema_mode="skills_like", + ) + + responses = [response async for response in runner.step()] + + assert [response.type for response in responses] == ["llm_result", "aborted"] + visible_text = "".join( + response.data["chain"].get_plain_text() + for response in responses + if response.type != "aborted" + ) + assert "skills_like fallback text" not in visible_text + assert "skills_like fallback reasoning" not in visible_text + assert runner.was_aborted() is True + final_resp = runner.get_final_llm_resp() + assert final_resp is not None + assert final_resp.completion_text == "" + assert "skills_like fallback text" not in repr(runner.run_context.messages) + assert "skills_like fallback reasoning" not in repr(runner.run_context.messages) + + +@pytest.mark.asyncio +async def test_fallback_provider_not_called_after_stop_on_primary_error_response( + provider_request, + mock_tool_executor, +): + class StopErrProvider(MockErrProvider): + def __init__(self, event): + super().__init__() + self.provider_config["id"] = "primary" + self.event = event + + async def text_chat(self, **kwargs) -> LLMResponse: + response = await super().text_chat(**kwargs) + self.event.set_extra("agent_stop_requested", True) + return response + + event = MockEvent("test:FriendMessage:fallback_stop", "u1") + primary = StopErrProvider(event) + fallback = MockProvider() + fallback.provider_config["id"] = "fallback" + fallback.should_call_tools = False + runner = ToolLoopAgentRunner() + + await runner.reset( + provider=primary, + request=provider_request, + run_context=ContextWrapper(context=MockAgentContext(event)), + tool_executor=mock_tool_executor, + agent_hooks=MockHooks(), + fallback_providers=[fallback], + ) + + responses = [response async for response in runner.step()] + + assert [response.type for response in responses] == ["aborted"] + assert primary.call_count == 1 + assert fallback.call_count == 0 + + @pytest.mark.asyncio async def test_follow_up_accepted_when_active_and_not_stopping( runner, mock_provider, provider_request, mock_tool_executor, mock_hooks @@ -1378,6 +2250,13 @@ async def test_follow_up_accepted_when_active_and_not_stopping( agent_hooks=mock_hooks, streaming=False, ) + runner._transition_state(AgentState.RUNNING) + + ticket = runner.follow_up(message_text="follow up while active") + + assert ticket is not None + assert ticket in runner._pending_follow_ups + assert len(runner._pending_follow_ups) == 1 @pytest.mark.asyncio From 53dc22e2f80d19944974395f2bb62ebce5bd5906 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?C=E2=82=82=E2=82=82H=E2=82=82=E2=82=85NO=E2=82=86?= Date: Mon, 6 Jul 2026 11:36:37 +0800 Subject: [PATCH 2/3] =?UTF-8?q?fix:=20=E5=A4=84=E7=90=86=20stop=20?= =?UTF-8?q?=E4=BF=AE=E5=A4=8D=E7=9A=84=20review=20=E9=A3=8E=E9=99=A9?= =?UTF-8?q?=E7=82=B9?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../agent/runners/tool_loop_agent_runner.py | 19 ++++++------ tests/test_tool_loop_agent_runner.py | 29 +++++++++++++++++++ 2 files changed, 39 insertions(+), 9 deletions(-) diff --git a/astrbot/core/agent/runners/tool_loop_agent_runner.py b/astrbot/core/agent/runners/tool_loop_agent_runner.py index 44123d664a..f9339aef3b 100644 --- a/astrbot/core/agent/runners/tool_loop_agent_runner.py +++ b/astrbot/core/agent/runners/tool_loop_agent_runner.py @@ -295,7 +295,7 @@ async def reset( self._follow_up_seq = 0 self._last_tool_name: str | None = None self._same_tool_streak = 0 - self._recorded_usage_ids: set[int] = set() + self._recorded_usages: list[T.Any] = [] # These two are used for tool schema mode handling # We now have two modes: @@ -340,10 +340,9 @@ async def reset( def _record_llm_usage(self, llm_resp: LLMResponse) -> None: if not llm_resp.usage: return - usage_id = id(llm_resp.usage) - if usage_id in self._recorded_usage_ids: + if any(usage is llm_resp.usage for usage in self._recorded_usages): return - self._recorded_usage_ids.add(usage_id) + self._recorded_usages.append(llm_resp.usage) self.stats.token_usage += llm_resp.usage if self.req.conversation: self.req.conversation.token_usage = llm_resp.usage.total @@ -1470,11 +1469,13 @@ def _is_message_from_llm_response( if not isinstance(content, list): return False for part in content: - if isinstance(part, TextPart) and part.text != llm_resp.completion_text: - return False - if isinstance(part, ThinkPart) and part.think != ( - llm_resp.reasoning_content or "" - ): + if isinstance(part, TextPart): + if part.text != llm_resp.completion_text: + return False + elif isinstance(part, ThinkPart): + if part.think != (llm_resp.reasoning_content or ""): + return False + else: return False return True diff --git a/tests/test_tool_loop_agent_runner.py b/tests/test_tool_loop_agent_runner.py index 3f40d9e267..48a736e152 100644 --- a/tests/test_tool_loop_agent_runner.py +++ b/tests/test_tool_loop_agent_runner.py @@ -626,6 +626,35 @@ def runner(): return ToolLoopAgentRunner() +def test_record_llm_usage_keeps_usage_reference_to_prevent_id_reuse(runner): + usage = TokenUsage(input_other=10, output=5) + runner.req = SimpleNamespace(conversation=SimpleNamespace(token_usage=0)) + runner.stats = SimpleNamespace(token_usage=TokenUsage()) + runner._recorded_usages = [] + + runner._record_llm_usage(LLMResponse(role="assistant", usage=usage)) + runner._record_llm_usage(LLMResponse(role="assistant", usage=usage)) + + assert runner.stats.token_usage.total == usage.total + assert runner.req.conversation.token_usage == usage.total + assert runner._recorded_usages == [usage] + + +def test_is_message_from_llm_response_rejects_extra_content_parts(runner): + llm_resp = LLMResponse(role="assistant", completion_text="late text") + message = Message( + role="assistant", + content=[ + TextPart(text="late text"), + ImageURLPart( + image_url=ImageURLPart.ImageURL(url="https://example.com/a.png") + ), + ], + ) + + assert runner._is_message_from_llm_response(message, llm_resp) is False + + def _make_large_tool_result_text() -> str: return "x" * 100000 From 88d376ffb01a14bac4640e671934cfad89bcd553 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?C=E2=82=82=E2=82=82H=E2=82=82=E2=82=85NO=E2=82=86?= Date: Sat, 11 Jul 2026 11:15:40 +0800 Subject: [PATCH 3/3] =?UTF-8?q?fix:=20=E5=AE=8C=E5=96=84=20stop=20?= =?UTF-8?q?=E5=85=A8=E9=93=BE=E8=B7=AF=E4=B8=AD=E6=96=AD=E4=B8=8E=E4=BA=A4?= =?UTF-8?q?=E4=BB=98=E4=B8=80=E8=87=B4=E6=80=A7?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 统一 Agent stop 状态判断和活动 runner 唤醒，阻断停止后的 Provider 重试、工具执行、装饰以及平台发送。 增加最终响应交付确认与历史快照回滚，完善流式、空流、hook、TTS 及临时文件的停止边界。 补齐 DingTalk、QQ Official 与多平台发送失败语义，并增加异步竞态、资源清理和正常路径回归测试。 --- .../builtin_commands/commands/conversation.py | 3 + .../agent/runners/tool_loop_agent_runner.py | 375 +++-- astrbot/core/agent/stop_policy.py | 31 + astrbot/core/astr_agent_hooks.py | 39 +- astrbot/core/astr_agent_run_util.py | 166 +- astrbot/core/pipeline/context_utils.py | 74 +- .../core/pipeline/process_stage/follow_up.py | 19 +- .../method/agent_sub_stages/internal.py | 120 +- astrbot/core/pipeline/respond/stage.py | 93 +- .../core/pipeline/result_decorate/stage.py | 43 +- astrbot/core/platform/astr_message_event.py | 28 + .../aiocqhttp/aiocqhttp_message_event.py | 21 +- .../sources/dingtalk/dingtalk_adapter.py | 279 +++- .../sources/dingtalk/dingtalk_event.py | 14 +- .../sources/discord/discord_platform_event.py | 48 +- .../core/platform/sources/lark/lark_event.py | 306 ++-- .../core/platform/sources/line/line_event.py | 153 +- .../platform/sources/mattermost/client.py | 48 +- .../sources/mattermost/mattermost_event.py | 19 +- .../sources/misskey/misskey_adapter.py | 37 +- .../platform/sources/misskey/misskey_event.py | 26 +- .../qqofficial/qqofficial_message_event.py | 243 ++- .../platform/sources/satori/satori_event.py | 33 +- .../platform/sources/slack/slack_event.py | 41 +- .../platform/sources/telegram/tg_event.py | 102 +- .../platform/sources/webchat/webchat_event.py | 171 +- .../platform/sources/wecom/wecom_event.py | 51 +- .../sources/wecom_ai_bot/wecomai_event.py | 74 +- .../sources/wecom_ai_bot/wecomai_webhook.py | 69 +- .../sources/weixin_oc/weixin_oc_adapter.py | 61 +- .../sources/weixin_oc/weixin_oc_event.py | 15 +- .../weixin_offacc_event.py | 25 + astrbot/core/provider/provider.py | 16 +- astrbot/core/provider/sources/genie_tts.py | 90 +- astrbot/core/utils/media_utils.py | 67 +- tests/test_dingtalk_adapter.py | 282 ++++ tests/test_discord_adapter.py | 45 + tests/test_lark_app_registration.py | 140 ++ tests/test_platform_stop_delivery.py | 610 ++++++++ tests/test_qqofficial_group_message_create.py | 261 +++- tests/test_telegram_adapter.py | 101 ++ tests/test_tool_loop_agent_runner.py | 1374 +++++++++++------ tests/unit/test_aiocqhttp_reply.py | 46 + tests/unit/test_astr_message_event.py | 35 + tests/unit/test_message_tools.py | 159 +- 45 files changed, 4923 insertions(+), 1130 deletions(-) create mode 100644 astrbot/core/agent/stop_policy.py create mode 100644 tests/test_platform_stop_delivery.py diff --git a/astrbot/builtin_stars/builtin_commands/commands/conversation.py b/astrbot/builtin_stars/builtin_commands/commands/conversation.py index 42282dc500..70fbd03aef 100644 --- a/astrbot/builtin_stars/builtin_commands/commands/conversation.py +++ b/astrbot/builtin_stars/builtin_commands/commands/conversation.py @@ -11,6 +11,7 @@ ) from astrbot.core.agent.runners.deerflow.deerflow_api_client import DeerFlowAPIClient from astrbot.core.db.po import ProviderStat +from astrbot.core.pipeline.process_stage.follow_up import request_active_runner_stop from astrbot.core.utils.active_event_registry import active_event_registry from .utils.rst_scene import RstScene @@ -206,6 +207,8 @@ async def stop(self, message: AstrMessageEvent) -> None: umo, exclude=message, ) + if request_active_runner_stop(umo) and stopped_count == 0: + stopped_count = 1 if stopped_count > 0: message.set_result( diff --git a/astrbot/core/agent/runners/tool_loop_agent_runner.py b/astrbot/core/agent/runners/tool_loop_agent_runner.py index f9339aef3b..75614a6a60 100644 --- a/astrbot/core/agent/runners/tool_loop_agent_runner.py +++ b/astrbot/core/agent/runners/tool_loop_agent_runner.py @@ -19,7 +19,7 @@ ) from tenacity import ( AsyncRetrying, - retry_if_exception_type, + retry_if_exception, stop_after_attempt, wait_exponential, ) @@ -60,6 +60,7 @@ ) from ..response import AgentResponseData, AgentStats from ..run_context import ContextWrapper, TContext +from ..stop_policy import event_requests_agent_stop from ..tool_executor import BaseFunctionToolExecutor from .base import AgentResponse, AgentState, BaseAgentRunner @@ -103,7 +104,12 @@ class _ToolExecutionInterrupted(Exception): """Raised when a running tool call is interrupted by a stop request.""" +class _AgentExecutionInterrupted(Exception): + """Raised when a pending agent operation is interrupted by a stop request.""" + + ToolExecutorResultT = T.TypeVar("ToolExecutorResultT") +AwaitResultT = T.TypeVar("AwaitResultT") class ToolLoopAgentRunner(BaseAgentRunner[TContext]): @@ -174,13 +180,11 @@ def _get_persona_custom_error_message(self) -> str | None: async def _complete_with_assistant_response(self, llm_resp: LLMResponse) -> bool: """Finalize the current step as a plain assistant response with no tool calls.""" - if self._is_stop_requested(): - await self._finalize_aborted_step(llm_resp) - return False - self.final_llm_resp = llm_resp + aborted_response_snapshot = self._build_aborted_llm_response(llm_resp) self._transition_state(AgentState.DONE) self.stats.end_time = time.time() + self._pre_final_response_messages = list(self.run_context.messages) parts = [] if llm_resp.reasoning_content is not None or llm_resp.reasoning_signature: @@ -196,25 +200,41 @@ async def _complete_with_assistant_response(self, llm_resp: LLMResponse) -> bool logger.warning("LLM returned empty assistant message with no tool calls.") self.run_context.messages.append(Message(role="assistant", content=parts)) - if self._is_stop_requested(): - if self._is_message_from_llm_response( - self.run_context.messages[-1], llm_resp - ): - self.run_context.messages.pop() - await self._finalize_aborted_step(llm_resp) - return False - try: await self.agent_hooks.on_agent_done(self.run_context, llm_resp) + except asyncio.CancelledError: + self.final_llm_resp = aborted_response_snapshot + raise except Exception as e: logger.error(f"Error in on_agent_done hook: {e}", exc_info=True) if self._is_stop_requested(): + self.final_llm_resp = aborted_response_snapshot self.discard_late_aborted_result() self._resolve_unconsumed_follow_ups() return False self._resolve_unconsumed_follow_ups() return True + async def _finalize_stop_after_response_yield( + self, + llm_resp: LLMResponse, + ) -> AgentResponse: + """Finalize a stop observed after publishing an LLM response. + + Args: + llm_resp: Response whose outward yield just resumed. + + Returns: + Safe aborted response for the caller. + """ + if llm_resp.tools_call_name: + return await self._finalize_aborted_step(llm_resp) + self.discard_late_aborted_result() + return AgentResponse( + type="aborted", + data=AgentResponseData(chain=MessageChain(type="aborted")), + ) + @override async def reset( self, @@ -295,7 +315,8 @@ async def reset( self._follow_up_seq = 0 self._last_tool_name: str | None = None self._same_tool_streak = 0 - self._recorded_usages: list[T.Any] = [] + self._pre_final_response_messages: list[Message] | None = None + self._final_response_delivery_committed = False # These two are used for tool schema mode handling # We now have two modes: @@ -337,16 +358,6 @@ async def reset( self.stats = AgentStats() self.stats.start_time = time.time() - def _record_llm_usage(self, llm_resp: LLMResponse) -> None: - if not llm_resp.usage: - return - if any(usage is llm_resp.usage for usage in self._recorded_usages): - return - self._recorded_usages.append(llm_resp.usage) - self.stats.token_usage += llm_resp.usage - if self.req.conversation: - self.req.conversation.token_usage = llm_resp.usage.total - def _read_tool_hint(self) -> str: if self.read_tool is not None: return f"`{self.read_tool.name}`" @@ -497,10 +508,19 @@ async def _iter_llm_responses( payload["model"] = self.req.model if self.streaming: stream = self.provider.text_chat_stream(**payload) - async for resp in stream: # type: ignore - yield resp + try: + while True: + try: + yield await self._await_or_abort(anext(stream)) # type: ignore + except StopAsyncIteration: + return + finally: + close_stream = getattr(stream, "aclose", None) + if close_stream is not None: + with suppress(RuntimeError, StopAsyncIteration): + await close_stream() else: - yield await self.provider.text_chat(**payload) + yield await self._await_or_abort(self.provider.text_chat(**payload)) async def _iter_llm_responses_with_fallback( self, @@ -511,6 +531,21 @@ async def _iter_llm_responses_with_fallback( last_exception: Exception | None = None last_err_response: LLMResponse | None = None + async def _interruptible_retry_sleep(delay: float) -> None: + deadline = asyncio.get_running_loop().time() + float(delay) + while not self._is_stop_requested(): + remaining = deadline - asyncio.get_running_loop().time() + if remaining <= 0: + return + try: + await asyncio.wait_for( + self._abort_signal.wait(), + timeout=min(remaining, 0.1), + ) + except asyncio.TimeoutError: + continue + return + for idx, candidate in enumerate(candidates): if self._is_stop_requested(): return @@ -526,17 +561,25 @@ async def _iter_llm_responses_with_fallback( self.provider = candidate try: retrying = AsyncRetrying( - retry=retry_if_exception_type(EmptyModelOutputError), + retry=retry_if_exception( + lambda exc: ( + isinstance(exc, EmptyModelOutputError) + and not self._is_stop_requested() + ) + ), stop=stop_after_attempt(self.EMPTY_OUTPUT_RETRY_ATTEMPTS), wait=wait_exponential( multiplier=1, min=self.EMPTY_OUTPUT_RETRY_WAIT_MIN_S, max=self.EMPTY_OUTPUT_RETRY_WAIT_MAX_S, ), + sleep=_interruptible_retry_sleep, reraise=True, ) async for attempt in retrying: + if self._is_stop_requested(): + return has_stream_output = False with attempt: try: @@ -568,12 +611,17 @@ async def _iter_llm_responses_with_fallback( if has_stream_output: return - except EmptyModelOutputError: + except EmptyModelOutputError as exc: if has_stream_output: logger.warning( "Chat Model %s returned empty output after streaming started; skipping empty-output retry.", candidate_id, ) + yield LLMResponse( + role="err", + completion_text=f"{type(exc).__name__}: {exc}", + ) + return else: logger.warning( "Chat Model %s returned empty output on attempt %s/%s.", @@ -583,6 +631,8 @@ async def _iter_llm_responses_with_fallback( ) raise except Exception as exc: # noqa: BLE001 + if self._is_stop_requested(): + return last_exception = exc logger.warning( "Chat Model %s request error: %s", @@ -749,6 +799,9 @@ async def step(self): await self.agent_hooks.on_agent_begin(self.run_context) except Exception as e: logger.error(f"Error in on_agent_begin hook: {e}", exc_info=True) + if self._is_stop_requested(): + yield await self._finalize_aborted_step() + return # 开始处理，转换到运行状态 self._transition_state(AgentState.RUNNING) @@ -757,12 +810,29 @@ async def step(self): # Process request-time context before sending it to the provider. token_usage = self.req.conversation.token_usage if self.req.conversation else 0 self._simple_print_message_role("[BefCompact]", self.run_context.messages) - self.run_context.messages = await self.request_context_manager.process( - self.run_context.messages, trusted_token_usage=token_usage - ) + try: + self.run_context.messages = await self._await_or_abort( + self.request_context_manager.process( + self.run_context.messages, + trusted_token_usage=token_usage, + ) + ) + except _AgentExecutionInterrupted: + yield await self._finalize_aborted_step() + return + if self._is_stop_requested(): + yield await self._finalize_aborted_step() + return self._simple_print_message_role("[AftCompact]", self.run_context.messages) async for llm_response in self._iter_llm_responses_with_fallback(): + if not llm_response.is_chunk and llm_response.usage: + # Count every completed provider response once, including one + # that arrives concurrently with a stop request. + self.stats.token_usage += llm_response.usage + if self.req.conversation: + self.req.conversation.token_usage = llm_response.usage.total + if self._is_stop_requested(): llm_resp_result = llm_response break @@ -801,15 +871,9 @@ async def step(self): if self._is_stop_requested(): llm_resp_result = llm_response break - if self._is_stop_requested(): - llm_resp_result = llm_response - break continue llm_resp_result = llm_response - if not llm_response.is_chunk and llm_response.usage: - # only count the token usage of the final response for computation purpose - self._record_llm_usage(llm_response) break # got final response if not llm_resp_result: @@ -861,11 +925,17 @@ async def step(self): ), ), ) + if self._is_stop_requested(): + yield await self._finalize_stop_after_response_yield(llm_resp) + return if llm_resp.result_chain: yield AgentResponse( type="llm_result", data=AgentResponseData(chain=llm_resp.result_chain), ) + if self._is_stop_requested(): + yield await self._finalize_stop_after_response_yield(llm_resp) + return elif llm_resp.completion_text: yield AgentResponse( type="llm_result", @@ -873,6 +943,9 @@ async def step(self): chain=MessageChain().message(llm_resp.completion_text), ), ) + if self._is_stop_requested(): + yield await self._finalize_stop_after_response_yield(llm_resp) + return # 如果有工具调用，还需处理工具调用 if llm_resp.tools_call_name: @@ -902,11 +975,21 @@ async def step(self): ), ), ) + if self._is_stop_requested(): + yield await self._finalize_stop_after_response_yield( + llm_resp + ) + return if llm_resp.result_chain: yield AgentResponse( type="llm_result", data=AgentResponseData(chain=llm_resp.result_chain), ) + if self._is_stop_requested(): + yield await self._finalize_stop_after_response_yield( + llm_resp + ) + return elif llm_resp.completion_text: yield AgentResponse( type="llm_result", @@ -914,6 +997,11 @@ async def step(self): chain=MessageChain().message(llm_resp.completion_text), ), ) + if self._is_stop_requested(): + yield await self._finalize_stop_after_response_yield( + llm_resp + ) + return return else: llm_resp.tools_call_name = requery_resp.tools_call_name @@ -1047,6 +1135,7 @@ async def _handle_function_tools( llm_response: LLMResponse, ) -> T.AsyncGenerator[_HandleFunctionToolsResult, None]: """处理函数工具调用。""" + self._raise_if_tool_execution_stopped() tool_call_result_blocks: list[ToolCallMessageSegment] = [] logger.info(f"Agent 使用工具: {llm_response.tools_call_name}") @@ -1082,6 +1171,7 @@ def _append_tool_call_result(tool_call_id: str, content: str) -> None: ], ) ) + self._raise_if_tool_execution_stopped() try: if not req.func_tool: return @@ -1148,6 +1238,7 @@ def _append_tool_call_result(tool_call_id: str, content: str) -> None: except Exception as e: logger.error(f"Error in on_tool_start hook: {e}", exc_info=True) + self._raise_if_tool_execution_stopped() executor = self.tool_executor.execute( tool=func_tool, run_context=self.run_context, @@ -1188,6 +1279,7 @@ def _append_tool_call_result(tool_call_id: str, content: str) -> None: yield _HandleFunctionToolsResult.from_cached_image( cached_img ) + self._raise_if_tool_execution_stopped() elif isinstance(content_item, EmbeddedResource): resource = content_item.resource if isinstance(resource, TextResourceContents): @@ -1214,6 +1306,7 @@ def _append_tool_call_result(tool_call_id: str, content: str) -> None: yield _HandleFunctionToolsResult.from_cached_image( cached_img ) + self._raise_if_tool_execution_stopped() else: result_parts.append( "The tool has returned a data type that is not supported." @@ -1261,6 +1354,7 @@ def _append_tool_call_result(tool_call_id: str, content: str) -> None: ), ) + self._raise_if_tool_execution_stopped() try: await self.agent_hooks.on_tool_end( self.run_context, @@ -1270,6 +1364,7 @@ def _append_tool_call_result(tool_call_id: str, content: str) -> None: ) except Exception as e: logger.error(f"Error in on_tool_end hook: {e}", exc_info=True) + self._raise_if_tool_execution_stopped() except Exception as e: if isinstance(e, _ToolExecutionInterrupted): raise @@ -1298,6 +1393,7 @@ def _append_tool_call_result(tool_call_id: str, content: str) -> None: ], ) ) + self._raise_if_tool_execution_stopped() logger.info(f"Tool `{func_tool_name}` Result: {tool_result_content}") # 处理函数调用响应 @@ -1305,6 +1401,12 @@ def _append_tool_call_result(tool_call_id: str, content: str) -> None: yield _HandleFunctionToolsResult.from_tool_call_result_blocks( tool_call_result_blocks ) + self._raise_if_tool_execution_stopped() + + def _raise_if_tool_execution_stopped(self) -> None: + """Prevent further tool-side effects after a stop request.""" + if self._is_stop_requested(): + raise _ToolExecutionInterrupted def _build_tool_requery_context( self, @@ -1349,6 +1451,8 @@ async def _resolve_tool_exec( llm_resp: LLMResponse, ) -> tuple[LLMResponse, ToolSet | None]: """Used in 'skills_like' tool schema mode to re-query LLM with param-only tool schemas.""" + if self._is_stop_requested(): + return llm_resp, self.req.func_tool tool_names = llm_resp.tools_call_name if not tool_names: return llm_resp, self.req.func_tool @@ -1366,16 +1470,21 @@ async def _resolve_tool_exec( ) if param_subset.tools and tool_names: contexts = self._build_tool_requery_context(tool_names) - requery_resp = await self.provider.text_chat( - contexts=self._sanitize_contexts_for_provider(contexts), - func_tool=param_subset, - model=self.req.model, - session_id=self.req.session_id, - extra_user_content_parts=self.req.extra_user_content_parts, - # tool_choice="required", - abort_signal=self._abort_signal, - request_max_retries=self.request_max_retries, - ) + try: + requery_resp = await self._await_or_abort( + self.provider.text_chat( + contexts=self._sanitize_contexts_for_provider(contexts), + func_tool=param_subset, + model=self.req.model, + session_id=self.req.session_id, + extra_user_content_parts=self.req.extra_user_content_parts, + # tool_choice="required", + abort_signal=self._abort_signal, + request_max_retries=self.request_max_retries, + ) + ) + except _AgentExecutionInterrupted: + return llm_resp, subset if requery_resp: llm_resp = requery_resp self._sanitize_malformed_tool_calls(llm_resp) @@ -1397,16 +1506,25 @@ async def _resolve_tool_exec( tool_names, extra_instruction=self.SKILLS_LIKE_REQUERY_REPAIR_INSTRUCTION, ) - repair_resp = await self.provider.text_chat( - contexts=self._sanitize_contexts_for_provider(repair_contexts), - func_tool=param_subset, - model=self.req.model, - session_id=self.req.session_id, - extra_user_content_parts=self.req.extra_user_content_parts, - # tool_choice="required", - abort_signal=self._abort_signal, - request_max_retries=self.request_max_retries, - ) + try: + repair_resp = await self._await_or_abort( + self.provider.text_chat( + contexts=self._sanitize_contexts_for_provider( + repair_contexts + ), + func_tool=param_subset, + model=self.req.model, + session_id=self.req.session_id, + extra_user_content_parts=( + self.req.extra_user_content_parts + ), + # tool_choice="required", + abort_signal=self._abort_signal, + request_max_retries=self.request_max_retries, + ) + ) + except _AgentExecutionInterrupted: + return llm_resp, subset if repair_resp: llm_resp = repair_resp self._sanitize_malformed_tool_calls(llm_resp) @@ -1420,25 +1538,66 @@ def done(self) -> bool: def request_stop(self) -> None: self._abort_signal.set() + async def _await_or_abort( + self, + awaitable: T.Awaitable[AwaitResultT], + ) -> AwaitResultT: + """Await an owned operation while monitoring the current stop state. + + Args: + awaitable: Operation owned by the current agent step. + + Returns: + Result produced by the operation. + + Raises: + _AgentExecutionInterrupted: The current event requested an agent stop. + asyncio.CancelledError: The outer task was cancelled. + """ + + async def _wait_for_stop_request() -> None: + while not self._is_stop_requested(): + try: + await asyncio.wait_for(self._abort_signal.wait(), timeout=0.1) + except asyncio.TimeoutError: + continue + + operation_task = asyncio.ensure_future(awaitable) + stop_task = asyncio.create_task(_wait_for_stop_request()) + try: + if not self._is_stop_requested(): + await asyncio.wait( + {operation_task, stop_task}, + return_when=asyncio.FIRST_COMPLETED, + ) + if stop_task.done() or self._is_stop_requested(): + if operation_task.done() and not operation_task.cancelled(): + exception = operation_task.exception() + if exception is None: + return operation_task.result() + if not operation_task.done(): + operation_task.cancel() + await asyncio.gather(operation_task, return_exceptions=True) + raise _AgentExecutionInterrupted + return operation_task.result() + except asyncio.CancelledError: + if not operation_task.done(): + operation_task.cancel() + await asyncio.gather(operation_task, return_exceptions=True) + raise + finally: + if not stop_task.done(): + stop_task.cancel() + await asyncio.gather(stop_task, return_exceptions=True) + def _is_stop_requested(self) -> bool: + if self._final_response_delivery_committed and self.done(): + return False if self._abort_signal.is_set(): return True event = getattr(self.run_context.context, "event", None) - if event is None: - return False - - is_stopped = getattr(event, "is_stopped", None) - if callable(is_stopped) and is_stopped(): - return True - - get_extra = getattr(event, "get_extra", None) - if callable(get_extra): - return bool(get_extra("agent_stop_requested")) or bool( - get_extra("agent_user_aborted") - ) - - return False + return event_requests_agent_stop(event) def was_aborted(self) -> bool: return self._aborted @@ -1450,65 +1609,57 @@ def _build_aborted_llm_response( self, llm_resp: LLMResponse | None = None, ) -> LLMResponse: - """构造不会泄露迟到模型正文的空响应。""" + """Build an empty response that cannot leak late model output.""" return LLMResponse( role="assistant", completion_text="", id=getattr(llm_resp, "id", None) if llm_resp else None, - usage=getattr(llm_resp, "usage", None) if llm_resp else None, + usage=(replace(llm_resp.usage) if llm_resp and llm_resp.usage else None), ) - def _is_message_from_llm_response( - self, - message: Message, - llm_resp: LLMResponse | None, - ) -> bool: - if llm_resp is None or message.role != "assistant" or message.tool_calls: - return False - content = message.content - if not isinstance(content, list): - return False - for part in content: - if isinstance(part, TextPart): - if part.text != llm_resp.completion_text: - return False - elif isinstance(part, ThinkPart): - if part.think != (llm_resp.reasoning_content or ""): - return False - else: - return False - return True - def discard_late_aborted_result(self) -> None: - """丢弃已完成但尚未对用户可见发送的迟到模型正文。""" - original_llm_resp = self.final_llm_resp - if original_llm_resp: - self._record_llm_usage(original_llm_resp) + """Discard late model output that was not delivered to the user.""" + if self._final_response_delivery_committed: + return self.final_llm_resp = self._build_aborted_llm_response(self.final_llm_resp) self._aborted = True event = getattr(self.run_context.context, "event", None) if event is not None: event.set_extra("agent_user_aborted", True) event.set_extra("agent_stop_requested", False) - if self.run_context.messages and self._is_message_from_llm_response( - self.run_context.messages[-1], - original_llm_resp, - ): - self.run_context.messages.pop() + if self._pre_final_response_messages is not None: + self.run_context.messages[:] = self._pre_final_response_messages + self._pre_final_response_messages = None + + def commit_final_response_delivery( + self, + *, + delivery_confirmed: bool = False, + ) -> bool: + """Commit final-response history after downstream delivery succeeds. + + Args: + delivery_confirmed: Whether delivery completed before a later stop. + + Returns: + Whether the delivery was committed before a stop request arrived. + """ + if not delivery_confirmed and self._is_stop_requested(): + self.discard_late_aborted_result() + return False + self._final_response_delivery_committed = True + self._pre_final_response_messages = None + return True async def _finalize_aborted_step( self, llm_resp: LLMResponse | None = None, ) -> AgentResponse: logger.info("Agent execution was requested to stop by user.") - safe_llm_resp = self._build_aborted_llm_response(llm_resp) - self.final_llm_resp = safe_llm_resp - self._record_llm_usage(safe_llm_resp) - self._aborted = True - event = getattr(self.run_context.context, "event", None) - if event is not None: - event.set_extra("agent_user_aborted", True) - event.set_extra("agent_stop_requested", False) + self.final_llm_resp = llm_resp + self.discard_late_aborted_result() + safe_llm_resp = self.final_llm_resp + assert safe_llm_resp is not None self._transition_state(AgentState.DONE) self.stats.end_time = time.time() diff --git a/astrbot/core/agent/stop_policy.py b/astrbot/core/agent/stop_policy.py new file mode 100644 index 0000000000..3840ef67f2 --- /dev/null +++ b/astrbot/core/agent/stop_policy.py @@ -0,0 +1,31 @@ +from typing import Any + +AGENT_OUTPUT_DELIVERY_CONFIRMED_KEY = "_agent_output_delivery_confirmed" + + +class AgentOutputStopped(Exception): + """Unwind streaming adapters without triggering broad error handlers.""" + + +def event_requests_agent_stop(event: Any) -> bool: + """Return whether an event requests the current agent run to stop. + + Args: + event: Event-like object that may expose stop state and extras. + + Returns: + Whether hard stop, soft stop, or user-aborted state is active. + """ + if event is None: + return False + + is_stopped = getattr(event, "is_stopped", None) + if callable(is_stopped) and is_stopped(): + return True + + get_extra = getattr(event, "get_extra", None) + if not callable(get_extra): + return False + return bool(get_extra("agent_stop_requested")) or bool( + get_extra("agent_user_aborted") + ) diff --git a/astrbot/core/astr_agent_hooks.py b/astrbot/core/astr_agent_hooks.py index b08238661e..a3f528b007 100644 --- a/astrbot/core/astr_agent_hooks.py +++ b/astrbot/core/astr_agent_hooks.py @@ -4,32 +4,13 @@ from astrbot.core.agent.hooks import BaseAgentRunHooks from astrbot.core.agent.run_context import ContextWrapper +from astrbot.core.agent.stop_policy import event_requests_agent_stop from astrbot.core.agent.tool import FunctionTool from astrbot.core.astr_agent_context import AstrAgentContext -from astrbot.core.pipeline.context_utils import call_event_hook -from astrbot.core.provider.entities import LLMResponse +from astrbot.core.pipeline.context_utils import call_agent_done_hook, call_event_hook from astrbot.core.star.star_handler import EventType -def _event_requests_agent_stop(event) -> bool: - get_extra = getattr(event, "get_extra", None) - is_stopped = getattr(event, "is_stopped", None) - return ( - (bool(get_extra("agent_user_aborted")) if callable(get_extra) else False) - or (bool(get_extra("agent_stop_requested")) if callable(get_extra) else False) - or (bool(is_stopped()) if callable(is_stopped) else False) - ) - - -def _build_aborted_llm_response(llm_response) -> LLMResponse: - return LLMResponse( - role="assistant", - completion_text="", - id=getattr(llm_response, "id", None) if llm_response else None, - usage=getattr(llm_response, "usage", None) if llm_response else None, - ) - - class MainAgentHooks(BaseAgentRunHooks[AstrAgentContext]): async def on_agent_begin( self, run_context: ContextWrapper[AstrAgentContext] @@ -43,7 +24,7 @@ async def on_agent_begin( async def on_agent_done(self, run_context, llm_response) -> None: # 执行事件钩子 event = run_context.context.event - suppress_llm_response = _event_requests_agent_stop(event) + suppress_llm_response = event_requests_agent_stop(event) if ( not suppress_llm_response @@ -51,9 +32,7 @@ async def on_agent_done(self, run_context, llm_response) -> None: and llm_response.reasoning_content ): # we will use this in result_decorate stage to inject reasoning content to chain - set_extra = getattr(event, "set_extra", None) - if callable(set_extra): - set_extra("_llm_reasoning_content", llm_response.reasoning_content) + event.set_extra("_llm_reasoning_content", llm_response.reasoning_content) if not suppress_llm_response: await call_event_hook( @@ -62,15 +41,11 @@ async def on_agent_done(self, run_context, llm_response) -> None: llm_response, ) - if _event_requests_agent_stop(event): - set_extra = getattr(event, "set_extra", None) - if callable(set_extra): - set_extra("_llm_reasoning_content", None) - llm_response = _build_aborted_llm_response(llm_response) + if event_requests_agent_stop(event): + event.set_extra("_llm_reasoning_content", None) - await call_event_hook( + await call_agent_done_hook( event, - EventType.OnAgentDoneEvent, run_context, llm_response, ) diff --git a/astrbot/core/astr_agent_run_util.py b/astrbot/core/astr_agent_run_util.py index 4896727119..b1b69af918 100644 --- a/astrbot/core/astr_agent_run_util.py +++ b/astrbot/core/astr_agent_run_util.py @@ -8,6 +8,7 @@ from astrbot.core import logger from astrbot.core.agent.message import Message from astrbot.core.agent.runners.tool_loop_agent_runner import ToolLoopAgentRunner +from astrbot.core.agent.stop_policy import AgentOutputStopped, event_requests_agent_stop from astrbot.core.astr_agent_context import AstrAgentContext from astrbot.core.message.components import BaseMessageComponent, Json, Plain from astrbot.core.message.message_event_result import ( @@ -24,14 +25,6 @@ AgentRunner = ToolLoopAgentRunner[AstrAgentContext] -def _should_stop_agent(astr_event) -> bool: - return ( - astr_event.is_stopped() - or bool(astr_event.get_extra("agent_stop_requested")) - or bool(astr_event.get_extra("agent_user_aborted")) - ) - - def _truncate_tool_result(text: str, limit: int = 70) -> str: if limit <= 0: return "" @@ -135,10 +128,20 @@ async def run_agent( stream_to_general, agent_runner, ) + + def _request_stop_if_needed() -> bool: + """Synchronize event stop state into the runner abort signal.""" + if not event_requests_agent_stop(astr_event): + return False + agent_runner.request_stop() + return True + while step_idx < max_step + 1: step_idx += 1 - if step_idx == max_step + 1: + stop_requested = _request_stop_if_needed() + + if step_idx == max_step + 1 and not stop_requested: logger.warning( f"Agent reached max steps ({max_step}), forcing a final response." ) @@ -159,23 +162,14 @@ async def run_agent( ) try: async for resp in agent_runner.step(): - if _should_stop_agent(astr_event): - agent_runner.request_stop() + stop_requested = _request_stop_if_needed() if resp.type == "aborted": buffered_llm_chains.clear() astr_event.clear_result() - if not stop_watcher.done(): - stop_watcher.cancel() - try: - await stop_watcher - except asyncio.CancelledError: - pass - astr_event.set_extra("agent_user_aborted", True) - astr_event.set_extra("agent_stop_requested", False) return - if _should_stop_agent(astr_event): + if stop_requested: continue if resp.type == "tool_call_result": @@ -191,9 +185,11 @@ async def run_agent( if msg_chain.type == "tool_direct_result": # tool_direct_result 用于标记 llm tool 需要直接发送给用户的内容 await astr_event.send(msg_chain) + _request_stop_if_needed() continue if astr_event.get_platform_id() == "webchat": await astr_event.send(msg_chain) + _request_stop_if_needed() elif show_tool_use and show_tool_call_result: status_msg = _build_tool_result_status_message( msg_chain, tool_name_by_call_id @@ -201,6 +197,7 @@ async def run_agent( await astr_event.send( MessageChain(type="tool_call").message(status_msg) ) + _request_stop_if_needed() # 对于其他情况，暂时先不处理 continue elif resp.type == "tool_call": @@ -212,6 +209,8 @@ async def run_agent( # 需要分段才能保证消息顺序正确。 # 若 show_tool_use 为 False，不会有独立消息插入，无需分段。 yield MessageChain(chain=[], type="break") + if _request_stop_if_needed(): + continue tool_info = _extract_chain_json_data(resp.data["chain"]) astr_event.trace.record( @@ -222,6 +221,7 @@ async def run_agent( if astr_event.get_platform_name() == "webchat": await astr_event.send(resp.data["chain"]) + _request_stop_if_needed() elif show_tool_use: if show_tool_call_result and isinstance(tool_info, dict): # Delay tool status notification until tool_call_result. @@ -230,6 +230,7 @@ async def run_agent( _build_tool_call_status_message(tool_info) ) await astr_event.send(chain) + _request_stop_if_needed() continue elif resp.type == "llm_result": chain = resp.data["chain"] @@ -259,24 +260,20 @@ async def run_agent( ) yield resp.data["chain"] astr_event.clear_result() + _request_stop_if_needed() elif resp.type == "streaming_delta": chain = resp.data["chain"] if chain.type == "reasoning" and not show_reasoning: # display the reasoning content only when configured continue yield resp.data["chain"] # MessageChain + _request_stop_if_needed() if can_buffer_llm_result and agent_runner.done(): - if _should_stop_agent(astr_event): + if _request_stop_if_needed(): buffered_llm_chains.clear() astr_event.clear_result() - discard_late_result = getattr( - agent_runner, - "discard_late_aborted_result", - None, - ) - if callable(discard_late_result): - discard_late_result() + agent_runner.discard_late_aborted_result() else: merged_chain = _merge_buffered_llm_chains(buffered_llm_chains) if merged_chain: @@ -288,16 +285,15 @@ async def run_agent( ) yield merged_chain astr_event.clear_result() + if _request_stop_if_needed(): + agent_runner.discard_late_aborted_result() - if not stop_watcher.done(): - stop_watcher.cancel() - try: - await stop_watcher - except asyncio.CancelledError: - pass if agent_runner.done(): # send agent stats to webchat - if astr_event.get_platform_name() == "webchat": + if ( + astr_event.get_platform_name() == "webchat" + and not event_requests_agent_stop(astr_event) + ): await astr_event.send( MessageChain( type="agent_stats", @@ -307,13 +303,14 @@ async def run_agent( break + except (GeneratorExit, asyncio.CancelledError): + agent_runner.request_stop() + if agent_runner.done(): + agent_runner.discard_late_aborted_result() + else: + await agent_runner._finalize_aborted_step() + raise except Exception as e: - if "stop_watcher" in locals() and not stop_watcher.done(): - stop_watcher.cancel() - try: - await stop_watcher - except asyncio.CancelledError: - pass logger.error(traceback.format_exc()) custom_error_message = extract_persona_custom_error_message_from_event( @@ -332,6 +329,14 @@ async def run_agent( role="err", completion_text=err_msg, ) + if event_requests_agent_stop(astr_event): + agent_runner.request_stop() + astr_event.clear_result() + if agent_runner.done(): + agent_runner.discard_late_aborted_result() + else: + await agent_runner._finalize_aborted_step(error_llm_response) + return try: await agent_runner.agent_hooks.on_agent_done( agent_runner.run_context, error_llm_response @@ -344,14 +349,21 @@ async def run_agent( else: astr_event.set_result(MessageEventResult().message(err_msg)) return + finally: + if not stop_watcher.done(): + stop_watcher.cancel() + try: + await stop_watcher + except asyncio.CancelledError: + pass async def _watch_agent_stop_signal(agent_runner: AgentRunner, astr_event) -> None: - while not agent_runner.done(): - if _should_stop_agent(astr_event): + while True: + if event_requests_agent_stop(astr_event): agent_runner.request_stop() return - await asyncio.sleep(0.5) + await asyncio.sleep(0.1) async def run_live_agent( @@ -425,7 +437,12 @@ async def run_live_agent( # 2. 启动 TTS 任务：负责从 text_queue 读取文本并生成音频到 audio_queue if support_stream: tts_task = asyncio.create_task( - _safe_tts_stream_wrapper(tts_provider, text_queue, audio_queue) + _safe_tts_stream_wrapper( + tts_provider, + text_queue, + audio_queue, + agent_runner.run_context.context.event, + ) ) else: tts_task = asyncio.create_task( @@ -445,7 +462,7 @@ async def run_live_agent( if queue_item is None: break - if agent_runner.was_aborted() or _should_stop_agent( + if agent_runner.was_aborted() or event_requests_agent_stop( agent_runner.run_context.context.event ): continue @@ -479,16 +496,17 @@ async def run_live_agent( feeder_task.cancel() if not tts_task.done(): tts_task.cancel() - - # 确保队列被消费 - pass + await asyncio.gather(feeder_task, tts_task, return_exceptions=True) tts_end_time = time.time() # 发送 TTS 统计信息 try: astr_event = agent_runner.run_context.context.event - if astr_event.get_platform_name() == "webchat": + if ( + astr_event.get_platform_name() == "webchat" + and not event_requests_agent_stop(astr_event) + ): tts_duration = tts_end_time - tts_start_time await astr_event.send( MessageChain( @@ -520,6 +538,7 @@ async def _run_agent_feeder( ) -> None: """运行 Agent 并将文本输出分句放入队列""" buffer = "" + astr_event = agent_runner.run_context.context.event try: async for chain in run_agent( agent_runner, @@ -532,6 +551,8 @@ async def _run_agent_feeder( ): if chain is None: continue + if event_requests_agent_stop(astr_event): + continue # 提取文本 text = chain.get_plain_text() @@ -553,7 +574,9 @@ async def _run_agent_feeder( temp_buffer += full_sentence if len(temp_buffer) >= 10: - if temp_buffer.strip(): + if temp_buffer.strip() and not event_requests_agent_stop( + astr_event + ): logger.info(f"[Live Agent Feeder] 分句: {temp_buffer}") await text_queue.put(temp_buffer) temp_buffer = "" @@ -561,11 +584,11 @@ async def _run_agent_feeder( # 更新 buffer 为剩余部分 buffer = temp_buffer + parts[-1] - # 处理剩余 buffer。若 stop 已到达，未送出的文本不能再进入 TTS。 + # Do not enqueue residual text after a stop request. if ( buffer.strip() and not agent_runner.was_aborted() - and not _should_stop_agent(agent_runner.run_context.context.event) + and not event_requests_agent_stop(astr_event) ): await text_queue.put(buffer) @@ -580,10 +603,34 @@ async def _safe_tts_stream_wrapper( tts_provider: TTSProvider, text_queue: asyncio.Queue[str | None], audio_queue: "asyncio.Queue[bytes | tuple[str, bytes] | None]", + astr_event: Any, ) -> None: """包装原生流式 TTS 确保异常处理和队列关闭""" + + class _StopAwareTextQueue: + """Expose the source queue while blocking new reads after stop.""" + + def __init__(self) -> None: + self.astr_event = astr_event + + async def get(self) -> str | None: + if event_requests_agent_stop(astr_event): + raise AgentOutputStopped + item = await text_queue.get() + if event_requests_agent_stop(astr_event): + raise AgentOutputStopped + return item + + def __getattr__(self, name: str) -> Any: + return getattr(text_queue, name) + try: - await tts_provider.get_audio_stream(text_queue, audio_queue) + await tts_provider.get_audio_stream( # type: ignore[arg-type] + _StopAwareTextQueue(), + audio_queue, + ) + except AgentOutputStopped: + pass except Exception as e: logger.error(f"[Live TTS Stream] Error: {e}", exc_info=True) finally: @@ -611,21 +658,18 @@ async def _simulated_stream_tts( text = await text_queue.get() if text is None: break - if _should_stop_agent(astr_event): + if event_requests_agent_stop(astr_event): continue try: audio_path = await tts_provider.get_audio(text) - if _should_stop_agent(astr_event): - continue - if audio_path: + astr_event.track_temporary_local_file(audio_path) + if event_requests_agent_stop(astr_event): + continue with open(audio_path, "rb") as f: audio_data = f.read() - if _should_stop_agent(astr_event): - continue - astr_event.track_temporary_local_file(audio_path) await audio_queue.put((text, audio_data)) except Exception as e: logger.error( diff --git a/astrbot/core/pipeline/context_utils.py b/astrbot/core/pipeline/context_utils.py index 93104c7af9..acdacf7efe 100644 --- a/astrbot/core/pipeline/context_utils.py +++ b/astrbot/core/pipeline/context_utils.py @@ -1,31 +1,79 @@ +import asyncio import inspect import traceback import typing as T +from dataclasses import replace from astrbot import logger +from astrbot.core.agent.stop_policy import event_requests_agent_stop from astrbot.core.message.message_event_result import CommandResult, MessageEventResult from astrbot.core.platform.astr_message_event import AstrMessageEvent +from astrbot.core.provider.entities import LLMResponse from astrbot.core.star.star import star_map from astrbot.core.star.star_handler import EventType, star_handlers_registry -def _event_requests_agent_stop(event: AstrMessageEvent) -> bool: - get_extra = getattr(event, "get_extra", None) +def _should_stop_hook_propagation( + event: AstrMessageEvent, hook_type: EventType +) -> bool: return ( - event.is_stopped() - or (bool(get_extra("agent_user_aborted")) if callable(get_extra) else False) - or (bool(get_extra("agent_stop_requested")) if callable(get_extra) else False) + event_requests_agent_stop(event) + if hook_type + in { + EventType.OnLLMResponseEvent, + EventType.OnUsingLLMToolEvent, + EventType.OnLLMToolRespondEvent, + } + else event.is_stopped() ) -def _should_stop_hook_propagation( - event: AstrMessageEvent, hook_type: EventType -) -> bool: - if event.is_stopped(): - return True - return hook_type == EventType.OnLLMResponseEvent and _event_requests_agent_stop( - event +async def call_agent_done_hook( + event: AstrMessageEvent, + run_context: T.Any, + llm_response: LLMResponse | None, +) -> None: + """Dispatch all agent-done handlers with per-handler safe responses. + + Args: + event: Current message event. + run_context: Agent run context passed to lifecycle handlers. + llm_response: Original final model response, if one exists. + """ + handlers = star_handlers_registry.get_handlers_by_event_type( + EventType.OnAgentDoneEvent, + plugins_name=getattr(event, "plugins_name", None), ) + safe_response_id = getattr(llm_response, "id", None) if llm_response else None + safe_response_usage = ( + replace(llm_response.usage) if llm_response and llm_response.usage else None + ) + stop_seen = False + for handler in handlers: + response_for_handler = llm_response + stop_seen = stop_seen or event_requests_agent_stop(event) + if stop_seen: + response_for_handler = LLMResponse( + role="assistant", + completion_text="", + id=safe_response_id, + usage=(replace(safe_response_usage) if safe_response_usage else None), + ) + + try: + assert inspect.iscoroutinefunction(handler.handler) + logger.debug( + f"hook({EventType.OnAgentDoneEvent.name}) -> " + f"{star_map[handler.handler_module_path].name} - {handler.handler_name}", + ) + await handler.handler(event, run_context, response_for_handler) + except asyncio.CancelledError: + raise + except BaseException: + logger.error(traceback.format_exc()) + stop_seen = stop_seen or event_requests_agent_stop(event) + if stop_seen: + event.set_extra("agent_stop_requested", True) async def call_handler( @@ -118,6 +166,8 @@ async def call_event_hook( f"hook({hook_type.name}) -> {star_map[handler.handler_module_path].name} - {handler.handler_name}", ) await handler.handler(event, *args, **kwargs) + except asyncio.CancelledError: + raise except BaseException: logger.error(traceback.format_exc()) diff --git a/astrbot/core/pipeline/process_stage/follow_up.py b/astrbot/core/pipeline/process_stage/follow_up.py index 79ec16a85b..249d185de7 100644 --- a/astrbot/core/pipeline/process_stage/follow_up.py +++ b/astrbot/core/pipeline/process_stage/follow_up.py @@ -5,6 +5,7 @@ from astrbot import logger from astrbot.core.agent.runners.tool_loop_agent_runner import FollowUpTicket +from astrbot.core.agent.stop_policy import event_requests_agent_stop from astrbot.core.astr_agent_run_util import AgentRunner from astrbot.core.platform.astr_message_event import AstrMessageEvent @@ -43,6 +44,22 @@ def unregister_active_runner(umo: str, runner: AgentRunner) -> None: _ACTIVE_AGENT_RUNNERS.pop(umo, None) +def request_active_runner_stop(umo: str) -> bool: + """Wake the active runner for a session immediately. + + Args: + umo: Unified message origin of the active session. + + Returns: + Whether an active runner was found and signalled. + """ + runner = _ACTIVE_AGENT_RUNNERS.get(umo) + if runner is None: + return False + runner.request_stop() + return True + + def _get_follow_up_order_state(umo: str) -> dict[str, object]: state = _FOLLOW_UP_ORDER_STATE.get(umo) if state is None: @@ -172,7 +189,7 @@ def try_capture_follow_up(event: AstrMessageEvent) -> FollowUpCapture | None: if not active_sender_id or active_sender_id != sender_id: return None - if runner_event.get_extra("agent_stop_requested"): + if event_requests_agent_stop(runner_event): return None ticket = runner.follow_up(message_text=_event_follow_up_text(event)) 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 75f2bc79dc..44d33ac3e9 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 @@ -13,6 +13,10 @@ dump_messages_with_checkpoints, ) from astrbot.core.agent.response import AgentStats +from astrbot.core.agent.stop_policy import ( + AGENT_OUTPUT_DELIVERY_CONFIRMED_KEY, + event_requests_agent_stop, +) from astrbot.core.astr_main_agent import ( LLM_ERROR_MESSAGE_EXTRA_KEY, MainAgentBuildConfig, @@ -51,6 +55,39 @@ class InternalAgentSubStage(Stage): + @staticmethod + def _confirm_agent_output_delivery( + event: AstrMessageEvent, + agent_runner: AgentRunner, + ) -> bool: + """Commit delivered output or discard it when stop wins the boundary. + + Args: + event: Current message event. + agent_runner: Runner that owns the pending final-response snapshot. + + Returns: + Whether downstream delivery completed and processing may continue. + """ + delivery_confirmed = bool( + event.get_extra(AGENT_OUTPUT_DELIVERY_CONFIRMED_KEY, False) + ) + if event_requests_agent_stop(event): + agent_runner.request_stop() + event.clear_result() + if agent_runner.done(): + if delivery_confirmed: + agent_runner.commit_final_response_delivery(delivery_confirmed=True) + else: + agent_runner.discard_late_aborted_result() + return False + if agent_runner.done(): + if not delivery_confirmed: + return False + return agent_runner.commit_final_response_delivery(delivery_confirmed=True) + event.set_extra(AGENT_OUTPUT_DELIVERY_CONFIRMED_KEY, False) + return True + async def initialize(self, ctx: PipelineContext) -> None: self.ctx = ctx conf = ctx.astrbot_config @@ -157,7 +194,8 @@ async def initialize(self, ctx: PipelineContext) -> None: async def _send_llm_error_message( self, event: AstrMessageEvent, message: object ) -> None: - await event.send(MessageChain().message(str(message))) + if not event_requests_agent_stop(event): + await event.send(MessageChain().message(str(message))) async def process( self, event: AstrMessageEvent, provider_wake_prefix: str @@ -167,6 +205,8 @@ async def process( follow_up_activated = False typing_requested = False try: + if event_requests_agent_stop(event): + return streaming_response = self.streaming_response if (enable_streaming := event.get_extra("enable_streaming")) is not None: streaming_response = bool(enable_streaming) @@ -205,12 +245,19 @@ async def process( ) return + if event_requests_agent_stop(event): + return + try: typing_requested = True await event.send_typing() except Exception: logger.warning("send_typing failed", exc_info=True) - if await call_event_hook(event, EventType.OnWaitingLLMRequestEvent): + if event_requests_agent_stop(event): + return + if await call_event_hook( + event, EventType.OnWaitingLLMRequestEvent + ) or event_requests_agent_stop(event): return async with session_lock_manager.acquire_lock(event.unified_msg_origin): @@ -218,6 +265,8 @@ async def process( agent_runner: AgentRunner | None = None runner_registered = False try: + if event_requests_agent_stop(event): + return build_cfg = replace( self.main_agent_cfg, provider_wake_prefix=provider_wake_prefix, @@ -246,6 +295,12 @@ async def process( provider = build_result.provider reset_coro = build_result.reset_coro + if event_requests_agent_stop(event): + if reset_coro: + reset_coro.close() + agent_runner.request_stop() + return + api_base = provider.provider_config.get("api_base", "") for host in decoded_blocked: if host in api_base: @@ -262,15 +317,22 @@ async def process( and not event.platform_meta.support_streaming_message ) - if await call_event_hook(event, EventType.OnLLMRequestEvent, req): + if await call_event_hook( + event, EventType.OnLLMRequestEvent, req + ) or event_requests_agent_stop(event): if reset_coro: reset_coro.close() + agent_runner.request_stop() return # apply reset if reset_coro: await reset_coro + if event_requests_agent_stop(event): + agent_runner.request_stop() + return + register_active_runner(event.unified_msg_origin, agent_runner) runner_registered = True action_type = event.get_extra("action_type") @@ -304,6 +366,7 @@ async def process( ) # 使用 run_live_agent，总是使用流式响应 + event.set_extra(AGENT_OUTPUT_DELIVERY_CONFIRMED_KEY, False) event.set_result( MessageEventResult() .set_result_content_type(ResultContentType.STREAMING_RESULT) @@ -320,29 +383,14 @@ async def process( ), ) yield - if agent_runner.was_aborted(): - event.set_result( - MessageEventResult( - chain=MessageChain().chain, - result_content_type=ResultContentType.STREAMING_FINISH, - ), - ) - - # 保存历史记录 - if agent_runner.done() and ( - not event.is_stopped() or agent_runner.was_aborted() - ): - await self._save_to_history( - event, - req, - agent_runner.get_final_llm_resp(), - agent_runner.run_context.messages, - agent_runner.stats, - user_aborted=agent_runner.was_aborted(), - ) + self._confirm_agent_output_delivery( + event, + agent_runner, + ) elif streaming_response and not stream_to_general: # 流式响应 + event.set_extra(AGENT_OUTPUT_DELIVERY_CONFIRMED_KEY, False) event.set_result( MessageEventResult() .set_result_content_type(ResultContentType.STREAMING_RESULT) @@ -358,14 +406,11 @@ async def process( ), ) yield - if agent_runner.done() and agent_runner.was_aborted(): - event.set_result( - MessageEventResult( - chain=MessageChain().chain, - result_content_type=ResultContentType.STREAMING_FINISH, - ), - ) - elif agent_runner.done(): + delivery_confirmed = self._confirm_agent_output_delivery( + event, + agent_runner, + ) + if delivery_confirmed and agent_runner.done(): if final_llm_resp := agent_runner.get_final_llm_resp(): if final_llm_resp.completion_text: chain = ( @@ -393,8 +438,16 @@ async def process( show_reasoning=self.show_reasoning, buffer_intermediate_messages=self.buffer_intermediate_messages, ): + event.set_extra( + AGENT_OUTPUT_DELIVERY_CONFIRMED_KEY, + False, + ) yield + self._confirm_agent_output_delivery(event, agent_runner) + delivery_confirmed = bool( + event.get_extra(AGENT_OUTPUT_DELIVERY_CONFIRMED_KEY, False) + ) final_resp = agent_runner.get_final_llm_resp() event.trace.record( @@ -413,7 +466,7 @@ async def process( ) # 检查事件是否被停止，如果被停止则不保存历史记录 - if not event.is_stopped() or agent_runner.was_aborted(): + if agent_runner.was_aborted() or delivery_confirmed: await self._save_to_history( event, req, @@ -436,6 +489,9 @@ async def process( except Exception as e: logger.error(f"Error occurred while processing agent: {e}") + if event_requests_agent_stop(event): + event.clear_result() + return custom_error_message = extract_persona_custom_error_message_from_event( event ) diff --git a/astrbot/core/pipeline/respond/stage.py b/astrbot/core/pipeline/respond/stage.py index 888fb65a4b..03b5354a5f 100644 --- a/astrbot/core/pipeline/respond/stage.py +++ b/astrbot/core/pipeline/respond/stage.py @@ -5,6 +5,11 @@ import astrbot.core.message.components as Comp from astrbot.core import logger +from astrbot.core.agent.stop_policy import ( + AGENT_OUTPUT_DELIVERY_CONFIRMED_KEY, + AgentOutputStopped, + event_requests_agent_stop, +) from astrbot.core.message.components import BaseMessageComponent, ComponentType from astrbot.core.message.message_event_result import MessageChain, ResultContentType from astrbot.core.platform.astr_message_event import AstrMessageEvent @@ -15,6 +20,13 @@ from ..stage import Stage, register_stage +def _discard_result_if_stopped(event: AstrMessageEvent) -> bool: + if not event_requests_agent_stop(event): + return False + event.clear_result() + return True + + @register_stage class RespondStage(Stage): # 组件类型到其非空判断函数的映射 @@ -170,6 +182,8 @@ async def process( self, event: AstrMessageEvent, ) -> None | AsyncGenerator[None, None]: + if _discard_result_if_stopped(event): + return result = event.get_result() if result is None: return @@ -201,6 +215,7 @@ async def process( logger.info( "send_message_to_user already delivered the same text in this session, skip respond stage to avoid duplicate reply.", ) + event.set_extra(AGENT_OUTPUT_DELIVERY_CONFIRMED_KEY, True) return logger.info( @@ -220,8 +235,52 @@ async def process( == "realtime_segmenting" ) logger.info(f"应用流式输出({event.get_platform_id()})") - await event.send_streaming(result.async_stream, realtime_segmenting) + source_stream = result.async_stream + source_exhausted = False + source_has_payload = False + stream_delivery_succeeded = False + + async def _guarded_stream(): + nonlocal source_exhausted, source_has_payload + async for chain in source_stream: + if event_requests_agent_stop(event): + raise AgentOutputStopped + if any( + not isinstance(component, Comp.Plain) + or bool(component.text.strip()) + for component in chain.chain + ): + source_has_payload = True + yield chain + if event_requests_agent_stop(event): + raise AgentOutputStopped + if event_requests_agent_stop(event): + raise AgentOutputStopped + source_exhausted = True + + guarded_stream = _guarded_stream() + try: + await event.send_streaming(guarded_stream, realtime_segmenting) + except AgentOutputStopped: + pass + except Exception as e: + logger.error("发送流式消息失败: %s", e, exc_info=True) + else: + if source_exhausted and source_has_payload: + stream_delivery_succeeded = True + event.set_extra(AGENT_OUTPUT_DELIVERY_CONFIRMED_KEY, True) + finally: + await guarded_stream.aclose() + close_source = getattr(source_stream, "aclose", None) + if close_source is not None: + await close_source() + if not stream_delivery_succeeded: + event.clear_result() + _discard_result_if_stopped(event) return + + sent_any = False + send_failed = False if len(result.chain) > 0: # 检查路径映射 if mappings := self.platform_settings.get("path_mapping", []): @@ -266,18 +325,28 @@ async def process( return for comp in result.chain: i = await self._calc_comp_interval(comp) + if _discard_result_if_stopped(event): + return await asyncio.sleep(i) + if _discard_result_if_stopped(event): + return try: if comp.type in need_separately: await event.send(result.derive([comp])) else: await event.send(result.derive([*header_comps, comp])) header_comps.clear() + except AgentOutputStopped: + event.clear_result() + return except Exception as e: + send_failed = True logger.error( f"发送消息链失败: chain = {MessageChain([comp])}, error = {e}", exc_info=True, ) + else: + sent_any = True else: if all( comp.type in {ComponentType.Reply, ComponentType.At} @@ -294,24 +363,46 @@ async def process( modify_raw_chain=True, ) for comp in sep_comps: + if _discard_result_if_stopped(event): + return chain = result.derive([comp]) try: await event.send(chain) + except AgentOutputStopped: + event.clear_result() + return except Exception as e: + send_failed = True logger.error( f"发送消息链失败: chain = {chain}, error = {e}", exc_info=True, ) + else: + sent_any = True chain = result.derive(result.chain) if result.chain and len(result.chain) > 0: + if _discard_result_if_stopped(event): + return try: await event.send(chain) + except AgentOutputStopped: + event.clear_result() + return except Exception as e: + send_failed = True logger.error( f"发送消息链失败: chain = {chain}, error = {e}", exc_info=True, ) + else: + sent_any = True + if not sent_any or send_failed: + event.clear_result() + return + event.set_extra(AGENT_OUTPUT_DELIVERY_CONFIRMED_KEY, True) + if _discard_result_if_stopped(event): + return if await call_event_hook(event, EventType.OnAfterMessageSentEvent): return diff --git a/astrbot/core/pipeline/result_decorate/stage.py b/astrbot/core/pipeline/result_decorate/stage.py index 5956c8b5ef..c85e0269cd 100644 --- a/astrbot/core/pipeline/result_decorate/stage.py +++ b/astrbot/core/pipeline/result_decorate/stage.py @@ -1,3 +1,4 @@ +import asyncio import random import re import time @@ -5,6 +6,7 @@ from collections.abc import AsyncGenerator from astrbot.core import file_token_service, html_renderer, logger +from astrbot.core.agent.stop_policy import event_requests_agent_stop from astrbot.core.message.components import At, Image, Json, Node, Plain, Record, Reply from astrbot.core.message.message_event_result import ResultContentType from astrbot.core.pipeline.content_safety_check.stage import ContentSafetyCheckStage @@ -18,6 +20,13 @@ from ..stage import Stage, register_stage, registered_stages +def _discard_result_if_stopped(event: AstrMessageEvent) -> bool: + if not event_requests_agent_stop(event): + return False + event.clear_result() + return True + + @register_stage class ResultDecorateStage(Stage): async def initialize(self, ctx: PipelineContext) -> None: @@ -127,6 +136,8 @@ async def process( self, event: AstrMessageEvent, ) -> None | AsyncGenerator[None, None]: + if _discard_result_if_stopped(event): + return result = event.get_result() if result is None or not result.chain: return @@ -154,6 +165,10 @@ async def process( check_text=text, ): yield + if _discard_result_if_stopped(event): + return + if _discard_result_if_stopped(event): + return # 发送消息前事件钩子 handlers = star_handlers_registry.get_handlers_by_event_type( @@ -175,10 +190,12 @@ async def process( logger.debug( f"hook(on_decorating_result) -> {star_map[handler.handler_module_path].name} - {handler.handler_name} 将消息结果清空。", ) + except asyncio.CancelledError: + raise except BaseException: logger.error(traceback.format_exc()) - if event.is_stopped(): + if _discard_result_if_stopped(event): logger.info( f"{star_map[handler.handler_module_path].name} - {handler.handler_name} 终止了事件传播。", ) @@ -264,10 +281,8 @@ async def process( and random.random() <= self.tts_trigger_probability and tts_provider ) - if should_tts and not tts_provider: - logger.warning( - f"会话 {event.unified_msg_origin} 未配置文本转语音模型。", - ) + if _discard_result_if_stopped(event): + return if ( not should_tts @@ -297,10 +312,16 @@ async def process( new_chain = [] for comp in result.chain: if isinstance(comp, Plain) and len(comp.text) > 1: + if _discard_result_if_stopped(event): + return try: logger.info(f"TTS 请求: {comp.text}") audio_path = await tts_provider.get_audio(comp.text) logger.info(f"TTS 结果: {audio_path}") + if audio_path: + event.track_temporary_local_file(audio_path) + if _discard_result_if_stopped(event): + return if not audio_path: logger.error( f"由于 TTS 音频文件未找到，消息段转语音失败: {comp.text}", @@ -308,8 +329,6 @@ async def process( new_chain.append(comp) continue - event.track_temporary_local_file(audio_path) - use_file_service = self.ctx.astrbot_config[ "provider_tts_settings" ]["use_file_service"] @@ -325,6 +344,8 @@ async def process( token = await file_token_service.register_file( audio_path, ) + if _discard_result_if_stopped(event): + return url = f"{callback_api_base}/api/file/{token}" logger.debug(f"已注册：{url}") @@ -338,6 +359,8 @@ async def process( if dual_output: new_chain.append(comp) except Exception: + if _discard_result_if_stopped(event): + return logger.error(traceback.format_exc()) logger.error("TTS 失败，使用文本发送。") new_chain.append(comp) @@ -365,6 +388,10 @@ async def process( use_network=self.t2i_use_network, template_name=self.t2i_active_template, ) + if _discard_result_if_stopped(event): + return + except asyncio.CancelledError: + raise except BaseException: logger.error("文本转图片失败，使用文本发送。") return @@ -380,6 +407,8 @@ async def process( and self.ctx.astrbot_config["callback_api_base"] ): token = await file_token_service.register_file(url) + if _discard_result_if_stopped(event): + return url = f"{self.ctx.astrbot_config['callback_api_base']}/api/file/{token}" logger.debug(f"已注册：{url}") result.chain = [Image.fromURL(url)] diff --git a/astrbot/core/platform/astr_message_event.py b/astrbot/core/platform/astr_message_event.py index c82332770b..de7c319c77 100644 --- a/astrbot/core/platform/astr_message_event.py +++ b/astrbot/core/platform/astr_message_event.py @@ -9,6 +9,10 @@ from typing import Any from astrbot import logger +from astrbot.core.agent.stop_policy import ( + AgentOutputStopped, + event_requests_agent_stop, +) from astrbot.core.agent.tool import ToolSet from astrbot.core.db.po import Conversation from astrbot.core.message.components import ( @@ -266,16 +270,40 @@ def is_admin(self) -> bool: async def process_buffer(self, buffer: str, pattern: re.Pattern) -> str: """将消息缓冲区中的文本按指定正则表达式分割后发送至消息平台，作为不支持流式输出平台的Fallback。""" while True: + if event_requests_agent_stop(self): + raise AgentOutputStopped match = re.search(pattern, buffer) if not match: break matched_text = match.group().strip() if matched_text: await self.send(MessageChain([Plain(matched_text)])) + if event_requests_agent_stop(self): + raise AgentOutputStopped await asyncio.sleep(1.5) # 限速 buffer = buffer[match.end() :] return buffer + async def send_streaming_fallback_component( + self, + component: BaseMessageComponent, + delay: float = 1.5, + ) -> None: + """Send one fallback component with stop-aware rate limiting. + + Args: + component: Non-text component to send. + delay: Rate-limit delay after delivery. + """ + if event_requests_agent_stop(self): + raise AgentOutputStopped + await self.send(MessageChain([component])) + if event_requests_agent_stop(self): + raise AgentOutputStopped + await asyncio.sleep(delay) + if event_requests_agent_stop(self): + raise AgentOutputStopped + async def send_streaming( self, generator: AsyncGenerator[MessageChain, None], diff --git a/astrbot/core/platform/sources/aiocqhttp/aiocqhttp_message_event.py b/astrbot/core/platform/sources/aiocqhttp/aiocqhttp_message_event.py index 91a7444f38..16823022cd 100644 --- a/astrbot/core/platform/sources/aiocqhttp/aiocqhttp_message_event.py +++ b/astrbot/core/platform/sources/aiocqhttp/aiocqhttp_message_event.py @@ -17,6 +17,7 @@ Video, ) from astrbot.api.platform import Group, MessageMember +from astrbot.core.agent.stop_policy import AgentOutputStopped, event_requests_agent_stop class AiocqhttpMessageEvent(AstrMessageEvent): @@ -130,6 +131,7 @@ async def send_message( event: Event | None = None, is_group: bool = False, session_id: str | None = None, + stop_event: AstrMessageEvent | None = None, ) -> None: """发送消息至 QQ 协议端（aiocqhttp）。 @@ -139,19 +141,26 @@ async def send_message( event (Event | None, optional): aiocqhttp 事件对象. is_group (bool, optional): 是否为群消息. session_id (str | None, optional): 会话 ID（群号或 QQ 号 + stop_event: AstrBot event used to guard platform writes. """ + if event_requests_agent_stop(stop_event): + raise AgentOutputStopped # 转发消息、文件消息不能和普通消息混在一起发送 send_one_by_one = any( isinstance(seg, Node | Nodes | File) for seg in message_chain.chain ) if not send_one_by_one: ret = await cls._parse_onebot_json(message_chain) + if event_requests_agent_stop(stop_event): + raise AgentOutputStopped if not ret: return await cls._dispatch_send(bot, event, is_group, session_id, ret) return for seg in message_chain.chain: + if event_requests_agent_stop(stop_event): + raise AgentOutputStopped if isinstance(seg, Node | Nodes): # 合并转发消息 if isinstance(seg, Node): @@ -159,6 +168,8 @@ async def send_message( seg = nodes payload = await seg.to_dict() + if event_requests_agent_stop(stop_event): + raise AgentOutputStopped if is_group: payload["group_id"] = session_id @@ -172,9 +183,13 @@ async def send_message( await bot.call_action("send_private_forward_msg", **payload) elif isinstance(seg, File): d = await cls._from_segment_to_dict(seg) + if event_requests_agent_stop(stop_event): + raise AgentOutputStopped await cls._dispatch_send(bot, event, is_group, session_id, [d]) else: messages = await cls._parse_onebot_json(MessageChain([seg])) + if event_requests_agent_stop(stop_event): + raise AgentOutputStopped if not messages: continue await cls._dispatch_send(bot, event, is_group, session_id, messages) @@ -182,6 +197,8 @@ async def send_message( async def send(self, message: MessageChain) -> None: """发送消息""" + if event_requests_agent_stop(self): + raise AgentOutputStopped event = getattr(self.message_obj, "raw_message", None) is_group = bool(self.get_group_id()) @@ -193,6 +210,7 @@ async def send(self, message: MessageChain) -> None: event=event, # 不强制要求一定是 Event is_group=is_group, session_id=session_id, + stop_event=self, ) await super().send(message) @@ -225,8 +243,7 @@ async def send_streaming( if any(p in buffer for p in "。？！~…"): buffer = await self.process_buffer(buffer, pattern) else: - await self.send(MessageChain(chain=[comp])) - await asyncio.sleep(1.5) # 限速 + await self.send_streaming_fallback_component(comp) buffer = buffer.strip() if buffer: diff --git a/astrbot/core/platform/sources/dingtalk/dingtalk_adapter.py b/astrbot/core/platform/sources/dingtalk/dingtalk_adapter.py index 9a07608c3c..f1b562b459 100644 --- a/astrbot/core/platform/sources/dingtalk/dingtalk_adapter.py +++ b/astrbot/core/platform/sources/dingtalk/dingtalk_adapter.py @@ -12,7 +12,7 @@ from dingtalk_stream import AckMessage from astrbot import logger -from astrbot.api.event import MessageChain +from astrbot.api.event import AstrMessageEvent, MessageChain from astrbot.api.message_components import At, File, Image, Plain, Record, Video from astrbot.api.platform import ( AstrBotMessage, @@ -22,6 +22,7 @@ PlatformMetadata, ) from astrbot.core import sp +from astrbot.core.agent.stop_policy import AgentOutputStopped, event_requests_agent_stop from astrbot.core.platform.astr_message_event import MessageSesion from astrbot.core.utils.astrbot_path import get_astrbot_temp_path from astrbot.core.utils.io import download_file @@ -112,7 +113,10 @@ async def send_by_session( self, session: MessageSesion, message_chain: MessageChain, + stop_event: AstrMessageEvent | None = None, ) -> None: + if event_requests_agent_stop(stop_event): + raise AgentOutputStopped robot_code = self.client_id if session.message_type == MessageType.GROUP_MESSAGE: @@ -121,18 +125,21 @@ async def send_by_session( open_conversation_id=open_conversation_id, robot_code=robot_code, message_chain=message_chain, + stop_event=stop_event, ) else: staff_id = await self._get_sender_staff_id(session) + if event_requests_agent_stop(stop_event): + raise AgentOutputStopped if not staff_id: - logger.warning( - "钉钉私聊会话缺少 staff_id 映射，回退使用 session_id 作为 userId 发送", + raise RuntimeError( + "DingTalk private session is missing a staff_id mapping." ) - staff_id = session.session_id await self.send_message_chain_to_user( staff_id=staff_id, robot_code=robot_code, message_chain=message_chain, + stop_event=stop_event, ) await super().send_by_session(session, message_chain) @@ -386,30 +393,59 @@ async def download_ding_file( await download_file(download_url, str(f_path)) return str(f_path) - async def get_access_token(self) -> str: + async def get_access_token( + self, + stop_event: AstrMessageEvent | None = None, + ) -> str: try: + if event_requests_agent_stop(stop_event): + raise AgentOutputStopped access_token = await asyncio.get_running_loop().run_in_executor( None, self.client_.get_access_token, ) + if event_requests_agent_stop(stop_event): + raise AgentOutputStopped if access_token: return access_token + except AgentOutputStopped: + raise except Exception as e: logger.warning(f"通过 dingtalk_stream 获取 access_token 失败: {e}") + if event_requests_agent_stop(stop_event): + raise AgentOutputStopped payload = {"appKey": self.client_id, "appSecret": self.client_secret} async with aiohttp.ClientSession() as session: + if event_requests_agent_stop(stop_event): + raise AgentOutputStopped async with session.post( "https://api.dingtalk.com/v1.0/oauth2/accessToken", json=payload, ) as resp: - if resp.status != 200: - logger.error( - f"获取钉钉机器人 access_token 失败: {resp.status}, {await resp.text()}", + if event_requests_agent_stop(stop_event): + raise AgentOutputStopped + if not 200 <= resp.status < 300: + response_text = await resp.text() + if event_requests_agent_stop(stop_event): + raise AgentOutputStopped + raise RuntimeError( + "DingTalk access token request failed: " + f"{resp.status}, {response_text}" ) - return "" data = await resp.json() - return cast(str, data.get("data", {}).get("accessToken", "")) + if event_requests_agent_stop(stop_event): + raise AgentOutputStopped + access_token = cast( + str, + data.get("accessToken") + or data.get("data", {}).get("accessToken", ""), + ) + if not access_token: + raise RuntimeError( + "DingTalk access token response did not include a token." + ) + return access_token async def _get_sender_staff_id(self, session: MessageSesion) -> str: try: @@ -430,11 +466,13 @@ async def _send_group_message( robot_code: str, msg_key: str, msg_param: dict, + stop_event: AstrMessageEvent | None = None, ) -> None: - access_token = await self.get_access_token() + access_token = await self.get_access_token(stop_event) + if event_requests_agent_stop(stop_event): + raise AgentOutputStopped if not access_token: - logger.error("钉钉群消息发送失败: access_token 为空") - return + raise RuntimeError("DingTalk group message access token is empty.") payload = { "msgKey": msg_key, @@ -446,16 +484,36 @@ async def _send_group_message( "Content-Type": "application/json", "x-acs-dingtalk-access-token": access_token, } + if event_requests_agent_stop(stop_event): + raise AgentOutputStopped async with aiohttp.ClientSession() as session: + if event_requests_agent_stop(stop_event): + raise AgentOutputStopped async with session.post( "https://api.dingtalk.com/v1.0/robot/groupMessages/send", headers=headers, json=payload, ) as resp: - if resp.status != 200: - logger.error( - f"钉钉群消息发送失败: {resp.status}, {await resp.text()}", + if event_requests_agent_stop(stop_event): + raise AgentOutputStopped + response_text = await resp.text() + if event_requests_agent_stop(stop_event): + raise AgentOutputStopped + if not 200 <= resp.status < 300: + raise RuntimeError( + "DingTalk group message delivery failed: " + f"{resp.status}, {response_text}" ) + try: + response_data = json.loads(response_text) if response_text else {} + except json.JSONDecodeError: + response_data = {} + if isinstance(response_data, dict): + error_code = response_data.get("errcode", response_data.get("code")) + if error_code not in (None, "", 0, "0"): + raise RuntimeError( + f"DingTalk group message delivery failed: {response_text}" + ) async def _send_private_message( self, @@ -463,11 +521,13 @@ async def _send_private_message( robot_code: str, msg_key: str, msg_param: dict, + stop_event: AstrMessageEvent | None = None, ) -> None: - access_token = await self.get_access_token() + access_token = await self.get_access_token(stop_event) + if event_requests_agent_stop(stop_event): + raise AgentOutputStopped if not access_token: - logger.error("钉钉私聊消息发送失败: access_token 为空") - return + raise RuntimeError("DingTalk private message access token is empty.") payload = { "robotCode": robot_code, @@ -479,16 +539,36 @@ async def _send_private_message( "Content-Type": "application/json", "x-acs-dingtalk-access-token": access_token, } + if event_requests_agent_stop(stop_event): + raise AgentOutputStopped async with aiohttp.ClientSession() as session: + if event_requests_agent_stop(stop_event): + raise AgentOutputStopped async with session.post( "https://api.dingtalk.com/v1.0/robot/oToMessages/batchSend", headers=headers, json=payload, ) as resp: - if resp.status != 200: - logger.error( - f"钉钉私聊消息发送失败: {resp.status}, {await resp.text()}", + if event_requests_agent_stop(stop_event): + raise AgentOutputStopped + response_text = await resp.text() + if event_requests_agent_stop(stop_event): + raise AgentOutputStopped + if not 200 <= resp.status < 300: + raise RuntimeError( + "DingTalk private message delivery failed: " + f"{resp.status}, {response_text}" ) + try: + response_data = json.loads(response_text) if response_text else {} + except json.JSONDecodeError: + response_data = {} + if isinstance(response_data, dict): + error_code = response_data.get("errcode", response_data.get("code")) + if error_code not in (None, "", 0, "0"): + raise RuntimeError( + f"DingTalk private message delivery failed: {response_text}" + ) def _safe_remove_file(self, file_path: str | None) -> None: if not file_path: @@ -514,12 +594,18 @@ async def _prepare_voice_for_dingtalk(self, input_path: str) -> tuple[str, bool] converted = await convert_audio_format(input_path, "amr") return converted, converted != input_path - async def upload_media(self, file_path: str, media_type: str) -> str: + async def upload_media( + self, + file_path: str, + media_type: str, + stop_event: AstrMessageEvent | None = None, + ) -> str: media_file_path = Path(file_path) - access_token = await self.get_access_token() + access_token = await self.get_access_token(stop_event) + if event_requests_agent_stop(stop_event): + raise AgentOutputStopped if not access_token: - logger.error("钉钉媒体上传失败: access_token 为空") - return "" + raise RuntimeError("DingTalk media upload access token is empty.") form = aiohttp.FormData() form.add_field( @@ -528,25 +614,41 @@ async def upload_media(self, file_path: str, media_type: str) -> str: filename=media_file_path.name, content_type="application/octet-stream", ) + if event_requests_agent_stop(stop_event): + raise AgentOutputStopped async with aiohttp.ClientSession() as session: + if event_requests_agent_stop(stop_event): + raise AgentOutputStopped async with session.post( f"https://oapi.dingtalk.com/media/upload?access_token={access_token}&type={media_type}", data=form, ) as resp: - if resp.status != 200: - logger.error( - f"钉钉媒体上传失败: {resp.status}, {await resp.text()}" + if event_requests_agent_stop(stop_event): + raise AgentOutputStopped + if not 200 <= resp.status < 300: + raise RuntimeError( + "DingTalk media upload failed: " + f"{resp.status}, {await resp.text()}" ) - return "" data = await resp.json() + if event_requests_agent_stop(stop_event): + raise AgentOutputStopped if data.get("errcode") != 0: - logger.error(f"钉钉媒体上传失败: {data}") - return "" - return cast(str, data.get("media_id", "")) + raise RuntimeError(f"DingTalk media upload failed: {data}") + media_id = cast(str, data.get("media_id", "")) + if not media_id: + raise RuntimeError("DingTalk media upload returned no media id.") + return media_id - async def upload_image(self, image: Image) -> str: + async def upload_image( + self, + image: Image, + stop_event: AstrMessageEvent | None = None, + ) -> str: image_file_path = await image.convert_to_file_path() - return await self.upload_media(image_file_path, "image") + if event_requests_agent_stop(stop_event): + raise AgentOutputStopped + return await self.upload_media(image_file_path, "image", stop_event) async def _send_message_chain( self, @@ -555,14 +657,21 @@ async def _send_message_chain( robot_code: str, message_chain: MessageChain, at_str: str = "", + stop_event: AstrMessageEvent | None = None, ) -> None: + sent_any = False + async def send_message(msg_key: str, msg_param: dict) -> None: + nonlocal sent_any + if event_requests_agent_stop(stop_event): + raise AgentOutputStopped if target_type == "group": await self._send_group_message( open_conversation_id=target_id, robot_code=robot_code, msg_key=msg_key, msg_param=msg_param, + stop_event=stop_event, ) else: await self._send_private_message( @@ -570,9 +679,13 @@ async def send_message(msg_key: str, msg_param: dict) -> None: robot_code=robot_code, msg_key=msg_key, msg_param=msg_param, + stop_event=stop_event, ) + sent_any = True for segment in message_chain.chain: + if event_requests_agent_stop(stop_event): + raise AgentOutputStopped if isinstance(segment, Plain): text = segment.text.strip() if not text and not at_str: @@ -589,9 +702,11 @@ async def send_message(msg_key: str, msg_param: dict) -> None: if photo_url.startswith(("http://", "https://")): pass else: - photo_url = await self.upload_image(segment) + photo_url = await self.upload_image(segment, stop_event) + if event_requests_agent_stop(stop_event): + raise AgentOutputStopped if not photo_url: - continue + raise RuntimeError("DingTalk image has no deliverable URL.") await send_message( msg_key="sampleImageMsg", msg_param={"photoURL": photo_url}, @@ -600,14 +715,28 @@ async def send_message(msg_key: str, msg_param: dict) -> None: converted_audio = None try: audio_path = await segment.convert_to_file_path() + if event_requests_agent_stop(stop_event): + raise AgentOutputStopped ( audio_path, converted_audio, ) = await self._prepare_voice_for_dingtalk(audio_path) - media_id = await self.upload_media(audio_path, "voice") + if event_requests_agent_stop(stop_event): + raise AgentOutputStopped + media_id = await self.upload_media( + audio_path, + "voice", + stop_event, + ) + if event_requests_agent_stop(stop_event): + raise AgentOutputStopped if not media_id: - continue + raise RuntimeError( + "DingTalk voice upload returned no media id." + ) duration_ms = await get_media_duration(audio_path) + if event_requests_agent_stop(stop_event): + raise AgentOutputStopped await send_message( msg_key="sampleAudio", msg_param={ @@ -615,9 +744,11 @@ async def send_message(msg_key: str, msg_param: dict) -> None: "duration": str(duration_ms or 1000), }, ) + except AgentOutputStopped: + raise except Exception as e: logger.warning(f"钉钉语音发送失败: {e}") - continue + raise finally: if converted_audio: self._safe_remove_file(audio_path) @@ -626,16 +757,38 @@ async def send_message(msg_key: str, msg_param: dict) -> None: cover_path = None try: source_video_path = await segment.convert_to_file_path() + if event_requests_agent_stop(stop_event): + raise AgentOutputStopped video_path = source_video_path if not video_path.lower().endswith(".mp4"): video_path = await convert_video_format(video_path, "mp4") + if event_requests_agent_stop(stop_event): + raise AgentOutputStopped converted_video = video_path != source_video_path cover_path = await extract_video_cover(video_path) - video_media_id = await self.upload_media(video_path, "file") - pic_media_id = await self.upload_media(cover_path, "image") + if event_requests_agent_stop(stop_event): + raise AgentOutputStopped + video_media_id = await self.upload_media( + video_path, + "file", + stop_event, + ) + if event_requests_agent_stop(stop_event): + raise AgentOutputStopped + pic_media_id = await self.upload_media( + cover_path, + "image", + stop_event, + ) + if event_requests_agent_stop(stop_event): + raise AgentOutputStopped if not video_media_id or not pic_media_id: - continue + raise RuntimeError( + "DingTalk video upload returned incomplete media ids." + ) duration_ms = await get_media_duration(video_path) + if event_requests_agent_stop(stop_event): + raise AgentOutputStopped duration_sec = max(1, int((duration_ms or 1000) / 1000)) await send_message( msg_key="sampleVideo", @@ -646,9 +799,11 @@ async def send_message(msg_key: str, msg_param: dict) -> None: "picMediaId": pic_media_id, }, ) + except AgentOutputStopped: + raise except Exception as e: logger.warning(f"钉钉视频发送失败: {e}") - continue + raise finally: self._safe_remove_file(cover_path) if converted_video: @@ -656,12 +811,19 @@ async def send_message(msg_key: str, msg_param: dict) -> None: elif isinstance(segment, File): try: file_path = await segment.get_file() + if event_requests_agent_stop(stop_event): + raise AgentOutputStopped if not file_path: - logger.warning("钉钉文件发送失败: 无法解析文件路径") - continue - media_id = await self.upload_media(file_path, "file") + raise RuntimeError("DingTalk file path could not be resolved.") + media_id = await self.upload_media( + file_path, + "file", + stop_event, + ) + if event_requests_agent_stop(stop_event): + raise AgentOutputStopped if not media_id: - continue + raise RuntimeError("DingTalk file upload returned no media id.") file_name = segment.name or Path(file_path).name file_type = Path(file_name).suffix.lstrip(".") await send_message( @@ -672,9 +834,14 @@ async def send_message(msg_key: str, msg_param: dict) -> None: "fileType": file_type, }, ) + except AgentOutputStopped: + raise except Exception as e: logger.warning(f"钉钉文件发送失败: {e}") - continue + raise + + if not sent_any: + raise RuntimeError("DingTalk message chain produced no platform delivery.") async def send_message_chain_to_group( self, @@ -682,6 +849,7 @@ async def send_message_chain_to_group( robot_code: str, message_chain: MessageChain, at_str: str = "", + stop_event: AstrMessageEvent | None = None, ) -> None: await self._send_message_chain( target_type="group", @@ -689,6 +857,7 @@ async def send_message_chain_to_group( robot_code=robot_code, message_chain=message_chain, at_str=at_str, + stop_event=stop_event, ) async def send_message_chain_to_user( @@ -697,6 +866,7 @@ async def send_message_chain_to_user( robot_code: str, message_chain: MessageChain, at_str: str = "", + stop_event: AstrMessageEvent | None = None, ) -> None: await self._send_message_chain( target_type="user", @@ -704,13 +874,17 @@ async def send_message_chain_to_user( robot_code=robot_code, message_chain=message_chain, at_str=at_str, + stop_event=stop_event, ) async def send_message_chain_with_incoming( self, incoming_message: dingtalk_stream.ChatbotMessage, message_chain: MessageChain, + stop_event: AstrMessageEvent | None = None, ) -> None: + if event_requests_agent_stop(stop_event): + raise AgentOutputStopped robot_code = self.client_id # at_list: list[str] = [] @@ -734,6 +908,7 @@ async def send_message_chain_with_incoming( open_conversation_id=cast(str, incoming_message.conversation_id), robot_code=robot_code, message_chain=message_chain, + stop_event=stop_event, # at_str=at_str, ) else: @@ -743,13 +918,15 @@ async def send_message_chain_with_incoming( session_id=normalized_sender_id, ) staff_id = sender_staff_id or await self._get_sender_staff_id(session) + if event_requests_agent_stop(stop_event): + raise AgentOutputStopped if not staff_id: - logger.error("钉钉私聊回复失败: 缺少 sender_staff_id") - return + raise RuntimeError("DingTalk private reply is missing sender_staff_id.") await self.send_message_chain_to_user( staff_id=staff_id, robot_code=robot_code, message_chain=message_chain, + stop_event=stop_event, # at_str=at_str, ) diff --git a/astrbot/core/platform/sources/dingtalk/dingtalk_event.py b/astrbot/core/platform/sources/dingtalk/dingtalk_event.py index 3331c51476..5a43154752 100644 --- a/astrbot/core/platform/sources/dingtalk/dingtalk_event.py +++ b/astrbot/core/platform/sources/dingtalk/dingtalk_event.py @@ -1,7 +1,7 @@ from typing import Any -from astrbot import logger from astrbot.api.event import AstrMessageEvent, MessageChain +from astrbot.core.agent.stop_policy import AgentOutputStopped, event_requests_agent_stop class DingtalkMessageEvent(AstrMessageEvent): @@ -19,12 +19,14 @@ def __init__( self.adapter = adapter async def send(self, message: MessageChain) -> None: + if event_requests_agent_stop(self): + raise AgentOutputStopped if not self.adapter: - logger.error("钉钉消息发送失败: 缺少 adapter") - return + raise RuntimeError("DingTalk message adapter is unavailable.") await self.adapter.send_message_chain_with_incoming( incoming_message=self.message_obj.raw_message, message_chain=message, + stop_event=self, ) await super().send(message) @@ -32,12 +34,16 @@ async def send_streaming(self, generator, use_fallback: bool = False): # 钉钉统一回退为缓冲发送：最终发送仍使用新的 HTTP 消息接口。 buffer = None async for chain in generator: + if event_requests_agent_stop(self): + raise AgentOutputStopped if not buffer: buffer = chain else: buffer.chain.extend(chain.chain) if not buffer: - return None + raise RuntimeError("DingTalk streaming message produced no delivery.") buffer.squash_plain() + if event_requests_agent_stop(self): + raise AgentOutputStopped await self.send(buffer) return await super().send_streaming(generator, use_fallback) diff --git a/astrbot/core/platform/sources/discord/discord_platform_event.py b/astrbot/core/platform/sources/discord/discord_platform_event.py index ff085dafd6..314393d658 100644 --- a/astrbot/core/platform/sources/discord/discord_platform_event.py +++ b/astrbot/core/platform/sources/discord/discord_platform_event.py @@ -18,6 +18,7 @@ Reply, ) from astrbot.api.platform import AstrBotMessage, At, PlatformMetadata +from astrbot.core.agent.stop_policy import AgentOutputStopped, event_requests_agent_stop from astrbot.core.utils.media_utils import ( MEDIA_MIME_EXTENSIONS, MediaResolver, @@ -52,6 +53,8 @@ def __init__( async def send(self, message: MessageChain) -> None: """发送消息到Discord平台""" + if event_requests_agent_stop(self): + raise AgentOutputStopped # 解析消息链为 Discord 所需的对象 try: ( @@ -60,10 +63,15 @@ async def send(self, message: MessageChain) -> None: view, embeds, reference_message_id, - ) = await self._parse_to_discord(message) + ) = await self._parse_to_discord(message, self) + except AgentOutputStopped: + raise except Exception as e: logger.error(f"[Discord] 解析消息链时失败: {e}", exc_info=True) - return + raise RuntimeError("Discord message conversion failed.") from e + + if event_requests_agent_stop(self): + raise AgentOutputStopped kwargs = {} if content: @@ -77,27 +85,34 @@ async def send(self, message: MessageChain) -> None: if reference_message_id and not self.interaction_followup_webhook: kwargs["reference"] = self.client.get_message(int(reference_message_id)) if not kwargs: - logger.debug("[Discord] 尝试发送空消息，已忽略。") - return + raise RuntimeError("Discord message conversion produced no content.") # 根据上下文执行发送/回复操作 try: # -- 斜杠指令/交互上下文 -- if self.interaction_followup_webhook: + if event_requests_agent_stop(self): + raise AgentOutputStopped await self.interaction_followup_webhook.send(**kwargs) # -- 常规消息上下文 -- else: channel = await self._get_channel() + if event_requests_agent_stop(self): + raise AgentOutputStopped if not channel: - return + raise RuntimeError("Discord message channel is unavailable.") if not isinstance(channel, discord.abc.Messageable): - logger.error(f"[Discord] 频道 {channel.id} 不是可发送消息的类型") - return + raise RuntimeError( + f"Discord channel {channel.id} is not messageable." + ) await channel.send(**kwargs) + except AgentOutputStopped: + raise except Exception as e: logger.error(f"[Discord] 发送消息时发生未知错误: {e}", exc_info=True) + raise RuntimeError("Discord message delivery failed.") from e await super().send(message) @@ -132,6 +147,7 @@ async def _get_channel( async def _parse_to_discord( self, message: MessageChain, + stop_event: AstrMessageEvent | None = None, ) -> tuple[ str, list[discord.File], @@ -146,6 +162,8 @@ async def _parse_to_discord( embeds = [] reference_message_id = None for i in message.chain: # 遍历消息链 + if event_requests_agent_stop(stop_event): + raise AgentOutputStopped if isinstance(i, Plain): # 如果是文字类型的 content_parts.append(i.text) elif isinstance(i, Reply): @@ -175,6 +193,8 @@ async def _parse_to_discord( file_content, media_type="image", ).to_base64_data(strict=True) + if event_requests_agent_stop(stop_event): + raise AgentOutputStopped if not image_data: logger.warning( "[Discord] 图片解析失败: %s", @@ -190,6 +210,8 @@ async def _parse_to_discord( ) ) + except AgentOutputStopped: + raise except Exception: # 使用 getattr 来安全地访问 i.file，以防 i 本身就是问题 file_info = getattr(i, "file", "未知") @@ -214,6 +236,8 @@ async def _parse_to_discord( strict=True, target_format="wav", ) + if event_requests_agent_stop(stop_event): + raise AgentOutputStopped if not audio_data: logger.warning( "[Discord] 语音解析失败: %s", @@ -227,6 +251,8 @@ async def _parse_to_discord( filename="audio.wav", ) ) + except AgentOutputStopped: + raise except Exception: audio_ref = getattr(i, "file", "未知") logger.error( @@ -237,10 +263,16 @@ async def _parse_to_discord( elif isinstance(i, File): try: file_path_str = await i.get_file() + if event_requests_agent_stop(stop_event): + raise AgentOutputStopped if file_path_str: path = Path(file_path_str) if await asyncio.to_thread(path.exists): + if event_requests_agent_stop(stop_event): + raise AgentOutputStopped file_bytes = await asyncio.to_thread(path.read_bytes) + if event_requests_agent_stop(stop_event): + raise AgentOutputStopped files.append( discord.File(BytesIO(file_bytes), filename=i.name), ) @@ -250,6 +282,8 @@ async def _parse_to_discord( ) else: logger.warning(f"[Discord] 获取文件失败: {i.name}") + except AgentOutputStopped: + raise except Exception as e: logger.warning(f"[Discord] 处理文件失败: {i.name}, 错误: {e}") elif isinstance(i, DiscordEmbed): diff --git a/astrbot/core/platform/sources/lark/lark_event.py b/astrbot/core/platform/sources/lark/lark_event.py index ccc90455f2..0dd8fd032a 100644 --- a/astrbot/core/platform/sources/lark/lark_event.py +++ b/astrbot/core/platform/sources/lark/lark_event.py @@ -28,6 +28,7 @@ from astrbot.api.event import AstrMessageEvent, MessageChain from astrbot.api.message_components import At, File, Json, Plain, Record, Video from astrbot.api.message_components import Image as AstrBotImage +from astrbot.core.agent.stop_policy import AgentOutputStopped, event_requests_agent_stop from astrbot.core.utils.media_utils import ( MediaResolver, convert_audio_to_opus, @@ -188,24 +189,31 @@ async def _upload_lark_file( return None @staticmethod - async def _convert_to_lark(message: MessageChain, lark_client: lark.Client) -> list: + async def _convert_to_lark( + message: MessageChain, + lark_client: lark.Client, + stop_event: AstrMessageEvent | None = None, + ) -> list: ret = [] _stage = [] for comp in message.chain: + if event_requests_agent_stop(stop_event): + raise AgentOutputStopped if isinstance(comp, Plain): _stage.append({"tag": "md", "text": comp.text}) elif isinstance(comp, At): _stage.append({"tag": "at", "user_id": comp.qq, "style": []}) elif isinstance(comp, AstrBotImage): if not comp.file: - logger.error("[Lark] 图片路径为空，无法上传") - continue + raise RuntimeError("Lark image path is empty.") try: async with MediaResolver( comp.file, media_type="image", ).as_path() as image: + if event_requests_agent_stop(stop_event): + raise AgentOutputStopped with image.open("rb") as image_file: request = ( CreateImageRequest.builder() @@ -219,23 +227,27 @@ async def _convert_to_lark(message: MessageChain, lark_client: lark.Client) -> l ) if lark_client.im is None: - logger.error( - "[Lark] API Client im 模块未初始化，无法上传图片" + raise RuntimeError( + "Lark IM client is unavailable for image upload." ) - continue response = await lark_client.im.v1.image.acreate(request) + except AgentOutputStopped: + raise except Exception as e: logger.error(f"[Lark] 无法打开或上传图片文件: {e}") - continue + raise RuntimeError("Lark image upload failed.") from e + + if event_requests_agent_stop(stop_event): + raise AgentOutputStopped if not response.success(): - logger.error(f"无法上传飞书图片({response.code}): {response.msg}") - continue + raise RuntimeError( + f"Lark image upload failed ({response.code}): {response.msg}" + ) if response.data is None: - logger.error("[Lark] 上传图片成功但未返回数据(data is None)") - continue + raise RuntimeError("Lark image upload returned no data.") image_key = response.data.image_key logger.debug(image_key) @@ -348,7 +360,10 @@ async def _send_interactive_card( reply_message_id: str | None = None, receive_id: str | None = None, receive_id_type: str | None = None, + stop_event: AstrMessageEvent | None = None, ) -> bool: + if event_requests_agent_stop(stop_event): + raise AgentOutputStopped if lark_client.cardkit is None: logger.error("[Lark] API Client cardkit 模块未初始化，无法发送卡片") return False @@ -368,6 +383,9 @@ async def _send_interactive_card( logger.error(f"[Lark] 创建卡片失败: {e}") return False + if event_requests_agent_stop(stop_event): + raise AgentOutputStopped + if not response.success(): logger.error(f"[Lark] 创建卡片失败({response.code}): {response.msg}") return False @@ -396,6 +414,7 @@ async def _send_collapsible_reasoning_panel( reply_message_id: str | None = None, receive_id: str | None = None, receive_id_type: str | None = None, + stop_event: AstrMessageEvent | None = None, ) -> bool: if not reasoning_content: return True @@ -409,6 +428,7 @@ async def _send_collapsible_reasoning_panel( reply_message_id=reply_message_id, receive_id=receive_id, receive_id_type=receive_id_type, + stop_event=stop_event, ) @staticmethod @@ -418,6 +438,7 @@ async def send_message_chain( reply_message_id: str | None = None, receive_id: str | None = None, receive_id_type: str | None = None, + stop_event: AstrMessageEvent | None = None, ) -> None: """通用的消息链发送方法 @@ -429,8 +450,9 @@ async def send_message_chain( receive_id_type: 接收者ID类型，如 'open_id', 'chat_id'（用于主动发送） """ if lark_client.im is None: - logger.error("[Lark] API Client im 模块未初始化") - return + raise RuntimeError("Lark IM client is unavailable.") + if event_requests_agent_stop(stop_event): + raise AgentOutputStopped # 分离文件、音频、视频组件和其他组件 file_components: list[File] = [] @@ -467,8 +489,11 @@ async def send_message_chain( reply_message_id=reply_message_id, receive_id=receive_id, receive_id_type=receive_id_type, + stop_event=stop_event, ): return + if event_requests_agent_stop(stop_event): + raise AgentOutputStopped # 先发送非文件内容（如果有） if other_components: @@ -486,7 +511,10 @@ async def _flush_buffer() -> None: res = await LarkMessageEvent._convert_to_lark( pending_chain, lark_client, + stop_event, ) + if event_requests_agent_stop(stop_event): + raise AgentOutputStopped if res: # 只在有内容时发送 wrapped = { "zh_cn": { @@ -494,7 +522,7 @@ async def _flush_buffer() -> None: "content": res, }, } - await LarkMessageEvent._send_im_message( + sent = await LarkMessageEvent._send_im_message( lark_client, content=json.dumps(wrapped), msg_type="post", @@ -502,9 +530,13 @@ async def _flush_buffer() -> None: receive_id=receive_id, receive_id_type=receive_id_type, ) + if not sent: + raise RuntimeError("Lark message delivery failed.") # 维持组件顺序：遇到折叠面板标记先 flush 当前普通内容并发送卡片 for comp in other_components: + if event_requests_agent_stop(stop_event): + raise AgentOutputStopped if isinstance(comp, Json) and isinstance(comp.data, dict): comp_type = comp.data.get("type") if comp_type == "lark_collapsible_panel_reasoning": @@ -520,6 +552,7 @@ async def _flush_buffer() -> None: reply_message_id=reply_message_id, receive_id=receive_id, receive_id_type=receive_id_type, + stop_event=stop_event, ) if not success: buffered_components.append( @@ -534,18 +567,39 @@ async def _flush_buffer() -> None: # 发送附件 for file_comp in file_components: + if event_requests_agent_stop(stop_event): + raise AgentOutputStopped await LarkMessageEvent._send_file_message( - file_comp, lark_client, reply_message_id, receive_id, receive_id_type + file_comp, + lark_client, + reply_message_id, + receive_id, + receive_id_type, + stop_event, ) for audio_comp in audio_components: + if event_requests_agent_stop(stop_event): + raise AgentOutputStopped await LarkMessageEvent._send_audio_message( - audio_comp, lark_client, reply_message_id, receive_id, receive_id_type + audio_comp, + lark_client, + reply_message_id, + receive_id, + receive_id_type, + stop_event, ) for media_comp in media_components: + if event_requests_agent_stop(stop_event): + raise AgentOutputStopped await LarkMessageEvent._send_media_message( - media_comp, lark_client, reply_message_id, receive_id, receive_id_type + media_comp, + lark_client, + reply_message_id, + receive_id, + receive_id_type, + stop_event, ) async def send(self, message: MessageChain) -> None: @@ -554,6 +608,7 @@ async def send(self, message: MessageChain) -> None: message, self.bot, reply_message_id=self.message_obj.message_id, + stop_event=self, ) await super().send(message) @@ -564,6 +619,7 @@ async def _send_file_message( reply_message_id: str | None = None, receive_id: str | None = None, receive_id_type: str | None = None, + stop_event: AstrMessageEvent | None = None, ) -> None: """发送文件消息 @@ -574,15 +630,19 @@ async def _send_file_message( receive_id: 接收者ID（用于主动发送） receive_id_type: 接收者ID类型（用于主动发送） """ + if event_requests_agent_stop(stop_event): + raise AgentOutputStopped file_path = file_comp.file or "" file_key = await LarkMessageEvent._upload_lark_file( lark_client, path=file_path, file_type="stream" ) + if event_requests_agent_stop(stop_event): + raise AgentOutputStopped if not file_key: - return + raise RuntimeError("Lark file upload failed.") content = json.dumps({"file_key": file_key}) - await LarkMessageEvent._send_im_message( + sent = await LarkMessageEvent._send_im_message( lark_client, content=content, msg_type="file", @@ -590,6 +650,8 @@ async def _send_file_message( receive_id=receive_id, receive_id_type=receive_id_type, ) + if not sent: + raise RuntimeError("Lark file message delivery failed.") @staticmethod async def _send_audio_message( @@ -598,6 +660,7 @@ async def _send_audio_message( reply_message_id: str | None = None, receive_id: str | None = None, receive_id_type: str | None = None, + stop_event: AstrMessageEvent | None = None, ) -> None: """发送音频消息 @@ -608,61 +671,70 @@ async def _send_audio_message( receive_id: 接收者ID（用于主动发送） receive_id_type: 接收者ID类型（用于主动发送） """ + if event_requests_agent_stop(stop_event): + raise AgentOutputStopped # 获取音频文件路径 try: original_audio_path = await audio_comp.convert_to_file_path() except Exception as e: logger.error(f"[Lark] 无法获取音频文件路径: {e}") - return + raise RuntimeError("Lark audio path conversion failed.") from e + + if event_requests_agent_stop(stop_event): + raise AgentOutputStopped if not original_audio_path or not os.path.exists(original_audio_path): - logger.error(f"[Lark] 音频文件不存在: {original_audio_path}") - return + raise RuntimeError(f"Lark audio file does not exist: {original_audio_path}") # 转换为opus格式 converted_audio_path = None try: - audio_path = await convert_audio_to_opus(original_audio_path) - # 如果转换后路径与原路径不同，说明生成了新文件 - if audio_path != original_audio_path: - converted_audio_path = audio_path - else: - audio_path = original_audio_path - except Exception as e: - logger.error(f"[Lark] 音频格式转换失败，将尝试直接上传: {e}") - # 如果转换失败，继续尝试直接上传原始文件 - audio_path = original_audio_path - - # 获取音频时长 - duration = await get_media_duration(audio_path) - - # 上传音频文件 - file_key = await LarkMessageEvent._upload_lark_file( - lark_client, - path=audio_path, - file_type="opus", - duration=duration, - ) - - # 清理转换后的临时音频文件 - if converted_audio_path and os.path.exists(converted_audio_path): try: - os.remove(converted_audio_path) - logger.debug(f"[Lark] 已删除转换后的音频文件: {converted_audio_path}") + audio_path = await convert_audio_to_opus(original_audio_path) + # 如果转换后路径与原路径不同，说明生成了新文件 + if audio_path != original_audio_path: + converted_audio_path = audio_path + else: + audio_path = original_audio_path except Exception as e: - logger.warning(f"[Lark] 删除转换后的音频文件失败: {e}") - - if not file_key: - return + logger.error(f"[Lark] 音频格式转换失败，将尝试直接上传: {e}") + audio_path = original_audio_path - await LarkMessageEvent._send_im_message( - lark_client, - content=json.dumps({"file_key": file_key}), - msg_type="audio", - reply_message_id=reply_message_id, - receive_id=receive_id, - receive_id_type=receive_id_type, - ) + if event_requests_agent_stop(stop_event): + raise AgentOutputStopped + duration = await get_media_duration(audio_path) + if event_requests_agent_stop(stop_event): + raise AgentOutputStopped + file_key = await LarkMessageEvent._upload_lark_file( + lark_client, + path=audio_path, + file_type="opus", + duration=duration, + ) + if event_requests_agent_stop(stop_event): + raise AgentOutputStopped + if not file_key: + raise RuntimeError("Lark audio upload failed.") + + sent = await LarkMessageEvent._send_im_message( + lark_client, + content=json.dumps({"file_key": file_key}), + msg_type="audio", + reply_message_id=reply_message_id, + receive_id=receive_id, + receive_id_type=receive_id_type, + ) + if not sent: + raise RuntimeError("Lark audio message delivery failed.") + finally: + if converted_audio_path and os.path.exists(converted_audio_path): + try: + os.remove(converted_audio_path) + logger.debug( + f"[Lark] 已删除转换后的音频文件: {converted_audio_path}" + ) + except Exception as e: + logger.warning(f"[Lark] 删除转换后的音频文件失败: {e}") @staticmethod async def _send_media_message( @@ -671,6 +743,7 @@ async def _send_media_message( reply_message_id: str | None = None, receive_id: str | None = None, receive_id_type: str | None = None, + stop_event: AstrMessageEvent | None = None, ) -> None: """发送视频消息 @@ -681,61 +754,70 @@ async def _send_media_message( receive_id: 接收者ID（用于主动发送） receive_id_type: 接收者ID类型（用于主动发送） """ + if event_requests_agent_stop(stop_event): + raise AgentOutputStopped # 获取视频文件路径 try: original_video_path = await media_comp.convert_to_file_path() except Exception as e: logger.error(f"[Lark] 无法获取视频文件路径: {e}") - return + raise RuntimeError("Lark video path conversion failed.") from e + + if event_requests_agent_stop(stop_event): + raise AgentOutputStopped if not original_video_path or not os.path.exists(original_video_path): - logger.error(f"[Lark] 视频文件不存在: {original_video_path}") - return + raise RuntimeError(f"Lark video file does not exist: {original_video_path}") # 转换为mp4格式 converted_video_path = None try: - video_path = await convert_video_format(original_video_path, "mp4") - # 如果转换后路径与原路径不同，说明生成了新文件 - if video_path != original_video_path: - converted_video_path = video_path - else: - video_path = original_video_path - except Exception as e: - logger.error(f"[Lark] 视频格式转换失败，将尝试直接上传: {e}") - # 如果转换失败，继续尝试直接上传原始文件 - video_path = original_video_path - - # 获取视频时长 - duration = await get_media_duration(video_path) - - # 上传视频文件 - file_key = await LarkMessageEvent._upload_lark_file( - lark_client, - path=video_path, - file_type="mp4", - duration=duration, - ) - - # 清理转换后的临时视频文件 - if converted_video_path and os.path.exists(converted_video_path): try: - os.remove(converted_video_path) - logger.debug(f"[Lark] 已删除转换后的视频文件: {converted_video_path}") + video_path = await convert_video_format(original_video_path, "mp4") + # 如果转换后路径与原路径不同，说明生成了新文件 + if video_path != original_video_path: + converted_video_path = video_path + else: + video_path = original_video_path except Exception as e: - logger.warning(f"[Lark] 删除转换后的视频文件失败: {e}") - - if not file_key: - return + logger.error(f"[Lark] 视频格式转换失败，将尝试直接上传: {e}") + video_path = original_video_path - await LarkMessageEvent._send_im_message( - lark_client, - content=json.dumps({"file_key": file_key}), - msg_type="media", - reply_message_id=reply_message_id, - receive_id=receive_id, - receive_id_type=receive_id_type, - ) + if event_requests_agent_stop(stop_event): + raise AgentOutputStopped + duration = await get_media_duration(video_path) + if event_requests_agent_stop(stop_event): + raise AgentOutputStopped + file_key = await LarkMessageEvent._upload_lark_file( + lark_client, + path=video_path, + file_type="mp4", + duration=duration, + ) + if event_requests_agent_stop(stop_event): + raise AgentOutputStopped + if not file_key: + raise RuntimeError("Lark video upload failed.") + + sent = await LarkMessageEvent._send_im_message( + lark_client, + content=json.dumps({"file_key": file_key}), + msg_type="media", + reply_message_id=reply_message_id, + receive_id=receive_id, + receive_id_type=receive_id_type, + ) + if not sent: + raise RuntimeError("Lark video message delivery failed.") + finally: + if converted_video_path and os.path.exists(converted_video_path): + try: + os.remove(converted_video_path) + logger.debug( + f"[Lark] 已删除转换后的视频文件: {converted_video_path}" + ) + except Exception as e: + logger.warning(f"[Lark] 删除转换后的视频文件失败: {e}") async def react(self, emoji: str) -> None: if self.bot.im is None: @@ -959,10 +1041,16 @@ async def _sender_loop() -> None: while not done: await text_changed.wait() text_changed.clear() + if done: + return + if event_requests_agent_stop(self): + raise AgentOutputStopped snapshot = delta if snapshot and snapshot != last_sent and card_id: sequence += 1 ok = await self._update_streaming_text(card_id, snapshot, sequence) + if event_requests_agent_stop(self): + raise AgentOutputStopped if ok: last_sent = snapshot if delta != snapshot: @@ -971,6 +1059,8 @@ async def _sender_loop() -> None: async def _consume_rest_and_fallback(gen, initial_text: str) -> None: """Card creation failed; consume remaining chunks and send non-streaming.""" nonlocal fallback_used + if event_requests_agent_stop(self): + raise AgentOutputStopped fallback_used = True buffer = MessageChain().message(initial_text) if initial_text else None async for chain in gen: @@ -981,6 +1071,8 @@ async def _consume_rest_and_fallback(gen, initial_text: str) -> None: else: buffer.chain.extend(chain.chain) if buffer: + if event_requests_agent_stop(self): + raise AgentOutputStopped buffer.squash_plain() await self.send(buffer) asyncio.create_task( @@ -992,10 +1084,16 @@ async def _flush_and_close_card() -> None: """补发最终文本并关闭当前卡片的流式模式。""" if not card_id: return + if event_requests_agent_stop(self): + raise AgentOutputStopped nonlocal sequence if delta and delta != last_sent: sequence += 1 - await self._update_streaming_text(card_id, delta, sequence) + updated = await self._update_streaming_text(card_id, delta, sequence) + if not updated: + raise RuntimeError("Lark final streaming update failed.") + if event_requests_agent_stop(self): + return sequence += 1 await self._close_streaming_mode(card_id, sequence) @@ -1029,6 +1127,8 @@ async def _flush_and_close_card() -> None: # Lazy card creation on first text token if card_id is None: card_id = await self._create_streaming_card() + if event_requests_agent_stop(self): + raise AgentOutputStopped if not card_id: logger.warning( "[Lark] 无法创建流式卡片，回退到非流式发送" @@ -1040,6 +1140,8 @@ async def _flush_and_close_card() -> None: card_id, reply_message_id=self.message_obj.message_id, ) + if event_requests_agent_stop(self): + raise AgentOutputStopped if not sent: logger.error( "[Lark] 发送流式卡片消息失败，回退到非流式发送" diff --git a/astrbot/core/platform/sources/line/line_event.py b/astrbot/core/platform/sources/line/line_event.py index 8b82ad1820..6a7f842156 100644 --- a/astrbot/core/platform/sources/line/line_event.py +++ b/astrbot/core/platform/sources/line/line_event.py @@ -16,6 +16,7 @@ Record, Video, ) +from astrbot.core.agent.stop_policy import AgentOutputStopped, event_requests_agent_stop from astrbot.core.utils.astrbot_path import get_astrbot_temp_path from astrbot.core.utils.media_utils import get_media_duration @@ -37,7 +38,10 @@ def __init__( @staticmethod async def _component_to_message_object( segment: BaseMessageComponent, + stop_event: AstrMessageEvent | None = None, ) -> dict | None: + if event_requests_agent_stop(stop_event): + raise AgentOutputStopped if isinstance(segment, Plain): text = segment.text.strip() if not text: @@ -51,7 +55,7 @@ async def _component_to_message_object( return {"type": "text", "text": f"@{name}"[:5000]} if isinstance(segment, Image): - image_url = await LineMessageEvent._resolve_image_url(segment) + image_url = await LineMessageEvent._resolve_image_url(segment, stop_event) if not image_url: return None return { @@ -61,10 +65,13 @@ async def _component_to_message_object( } if isinstance(segment, Record): - audio_url = await LineMessageEvent._resolve_record_url(segment) + audio_url = await LineMessageEvent._resolve_record_url(segment, stop_event) if not audio_url: return None - duration = await LineMessageEvent._resolve_record_duration(segment) + duration = await LineMessageEvent._resolve_record_duration( + segment, + stop_event, + ) return { "type": "audio", "originalContentUrl": audio_url, @@ -72,10 +79,13 @@ async def _component_to_message_object( } if isinstance(segment, Video): - video_url = await LineMessageEvent._resolve_video_url(segment) + video_url = await LineMessageEvent._resolve_video_url(segment, stop_event) if not video_url: return None - preview_url = await LineMessageEvent._resolve_video_preview_url(segment) + preview_url = await LineMessageEvent._resolve_video_preview_url( + segment, + stop_event, + ) if not preview_url: return None return { @@ -85,11 +95,11 @@ async def _component_to_message_object( } if isinstance(segment, File): - file_url = await LineMessageEvent._resolve_file_url(segment) + file_url = await LineMessageEvent._resolve_file_url(segment, stop_event) if not file_url: return None file_name = str(segment.name or "").strip() or "file.bin" - file_size = await LineMessageEvent._resolve_file_size(segment) + file_size = await LineMessageEvent._resolve_file_size(segment, stop_event) if file_size <= 0: return None return { @@ -102,51 +112,87 @@ async def _component_to_message_object( return None @staticmethod - async def _resolve_image_url(segment: Image) -> str: + async def _resolve_image_url( + segment: Image, + stop_event: AstrMessageEvent | None = None, + ) -> str: candidate = (segment.url or segment.file or "").strip() if candidate.startswith("https://"): return candidate try: - return await segment.register_to_file_service() + result = await segment.register_to_file_service() + if event_requests_agent_stop(stop_event): + raise AgentOutputStopped + return result + except AgentOutputStopped: + raise except Exception as e: logger.debug("[LINE] resolve image url failed: %s", e) return "" @staticmethod - async def _resolve_record_url(segment: Record) -> str: + async def _resolve_record_url( + segment: Record, + stop_event: AstrMessageEvent | None = None, + ) -> str: candidate = (segment.url or segment.file or "").strip() if candidate.startswith("https://"): return candidate try: - return await segment.register_to_file_service() + result = await segment.register_to_file_service() + if event_requests_agent_stop(stop_event): + raise AgentOutputStopped + return result + except AgentOutputStopped: + raise except Exception as e: logger.debug("[LINE] resolve record url failed: %s", e) return "" @staticmethod - async def _resolve_record_duration(segment: Record) -> int: + async def _resolve_record_duration( + segment: Record, + stop_event: AstrMessageEvent | None = None, + ) -> int: try: file_path = await segment.convert_to_file_path() + if event_requests_agent_stop(stop_event): + raise AgentOutputStopped duration_ms = await get_media_duration(file_path) + if event_requests_agent_stop(stop_event): + raise AgentOutputStopped if isinstance(duration_ms, int) and duration_ms > 0: return duration_ms + except AgentOutputStopped: + raise except Exception as e: logger.debug("[LINE] resolve record duration failed: %s", e) return 1000 @staticmethod - async def _resolve_video_url(segment: Video) -> str: + async def _resolve_video_url( + segment: Video, + stop_event: AstrMessageEvent | None = None, + ) -> str: candidate = (segment.file or "").strip() if candidate.startswith("https://"): return candidate try: - return await segment.register_to_file_service() + result = await segment.register_to_file_service() + if event_requests_agent_stop(stop_event): + raise AgentOutputStopped + return result + except AgentOutputStopped: + raise except Exception as e: logger.debug("[LINE] resolve video url failed: %s", e) return "" @staticmethod - async def _resolve_video_preview_url(segment: Video) -> str: + async def _resolve_video_preview_url( + segment: Video, + stop_event: AstrMessageEvent | None = None, + ) -> str: cover_candidate = (segment.cover or "").strip() if cover_candidate.startswith("https://"): return cover_candidate @@ -154,15 +200,27 @@ async def _resolve_video_preview_url(segment: Video) -> str: if cover_candidate: try: cover_seg = Image(file=cover_candidate) - return await cover_seg.register_to_file_service() + result = await cover_seg.register_to_file_service() + if event_requests_agent_stop(stop_event): + raise AgentOutputStopped + return result + except AgentOutputStopped: + raise except Exception as e: logger.debug("[LINE] resolve video cover failed: %s", e) try: video_path = await segment.convert_to_file_path() + if event_requests_agent_stop(stop_event): + raise AgentOutputStopped temp_dir = Path(get_astrbot_temp_path()) temp_dir.mkdir(parents=True, exist_ok=True) thumb_path = temp_dir / f"line_video_preview_{uuid.uuid4().hex}.jpg" + if stop_event is not None: + stop_event.track_temporary_local_file(str(thumb_path)) + + if event_requests_agent_stop(stop_event): + raise AgentOutputStopped process = await asyncio.create_subprocess_exec( "ffmpeg", @@ -178,40 +236,68 @@ async def _resolve_video_preview_url(segment: Video) -> str: stderr=asyncio.subprocess.PIPE, ) await process.communicate() + if event_requests_agent_stop(stop_event): + raise AgentOutputStopped if process.returncode != 0 or not thumb_path.exists(): return "" cover_seg = Image.fromFileSystem(str(thumb_path)) - return await cover_seg.register_to_file_service() + result = await cover_seg.register_to_file_service() + if event_requests_agent_stop(stop_event): + raise AgentOutputStopped + return result + except AgentOutputStopped: + raise except Exception as e: logger.debug("[LINE] generate video preview failed: %s", e) return "" @staticmethod - async def _resolve_file_url(segment: File) -> str: + async def _resolve_file_url( + segment: File, + stop_event: AstrMessageEvent | None = None, + ) -> str: if segment.url and segment.url.startswith("https://"): return segment.url try: - return await segment.register_to_file_service() + result = await segment.register_to_file_service() + if event_requests_agent_stop(stop_event): + raise AgentOutputStopped + return result + except AgentOutputStopped: + raise except Exception as e: logger.debug("[LINE] resolve file url failed: %s", e) return "" @staticmethod - async def _resolve_file_size(segment: File) -> int: + async def _resolve_file_size( + segment: File, + stop_event: AstrMessageEvent | None = None, + ) -> int: try: file_path = await segment.get_file(allow_return_url=False) + if event_requests_agent_stop(stop_event): + raise AgentOutputStopped if file_path and os.path.exists(file_path): return int(os.path.getsize(file_path)) + except AgentOutputStopped: + raise except Exception as e: logger.debug("[LINE] resolve file size failed: %s", e) return 0 @classmethod - async def build_line_messages(cls, message_chain: MessageChain) -> list[dict]: + async def build_line_messages( + cls, + message_chain: MessageChain, + stop_event: AstrMessageEvent | None = None, + ) -> list[dict]: messages: list[dict] = [] for segment in message_chain.chain: - obj = await cls._component_to_message_object(segment) + if event_requests_agent_stop(stop_event): + raise AgentOutputStopped + obj = await cls._component_to_message_object(segment, stop_event) if obj: messages.append(obj) @@ -226,9 +312,13 @@ async def build_line_messages(cls, message_chain: MessageChain) -> list[dict]: return messages async def send(self, message: MessageChain) -> None: - messages = await self.build_line_messages(message) + if event_requests_agent_stop(self): + raise AgentOutputStopped + messages = await self.build_line_messages(message, self) if not messages: - return + raise RuntimeError( + "LINE message conversion produced no deliverable content." + ) raw = self.message_obj.raw_message reply_token = "" @@ -240,9 +330,13 @@ async def send(self, message: MessageChain) -> None: sent = await self.line_api.reply_message(reply_token, messages) if not sent: + if event_requests_agent_stop(self): + raise AgentOutputStopped target_id = self.get_group_id() or self.get_sender_id() if target_id: - await self.line_api.push_message(target_id, messages) + sent = await self.line_api.push_message(target_id, messages) + if not sent: + raise RuntimeError("LINE message delivery failed.") await super().send(message) @@ -254,6 +348,8 @@ async def send_streaming( if not use_fallback: buffer = None async for chain in generator: + if event_requests_agent_stop(self): + raise AgentOutputStopped if not buffer: buffer = chain else: @@ -268,15 +364,18 @@ async def send_streaming( pattern = re.compile(r"[^。？！~…]+[。？！~…]+") async for chain in generator: + if event_requests_agent_stop(self): + raise AgentOutputStopped if isinstance(chain, MessageChain): for comp in chain.chain: + if event_requests_agent_stop(self): + raise AgentOutputStopped if isinstance(comp, Plain): buffer += comp.text if any(p in buffer for p in "。？！~…"): buffer = await self.process_buffer(buffer, pattern) else: - await self.send(MessageChain(chain=[comp])) - await asyncio.sleep(1.5) + await self.send_streaming_fallback_component(comp) if buffer.strip(): await self.send(MessageChain([Plain(buffer)])) diff --git a/astrbot/core/platform/sources/mattermost/client.py b/astrbot/core/platform/sources/mattermost/client.py index c35b893873..52afece288 100644 --- a/astrbot/core/platform/sources/mattermost/client.py +++ b/astrbot/core/platform/sources/mattermost/client.py @@ -7,8 +7,9 @@ import aiohttp from astrbot.api import logger -from astrbot.api.event import MessageChain +from astrbot.api.event import AstrMessageEvent, MessageChain from astrbot.api.message_components import At, File, Image, Plain, Record, Reply, Video +from astrbot.core.agent.stop_policy import AgentOutputStopped, event_requests_agent_stop from astrbot.core.utils.astrbot_path import get_astrbot_temp_path from astrbot.core.utils.media_utils import MediaResolver, detect_image_mime_type_async @@ -149,12 +150,15 @@ async def send_message_chain( self, channel_id: str, message_chain: MessageChain, + stop_event: AstrMessageEvent | None = None, ) -> dict[str, Any]: text_parts: list[str] = [] file_ids: list[str] = [] root_id: str | None = None for segment in message_chain.chain: + if event_requests_agent_stop(stop_event): + raise AgentOutputStopped if isinstance(segment, Plain): text_parts.append(segment.text) elif isinstance(segment, At): @@ -166,8 +170,12 @@ async def send_message_chain( root_id = str(segment.id) elif isinstance(segment, Image): path = await segment.convert_to_file_path() + if event_requests_agent_stop(stop_event): + raise AgentOutputStopped file_path = Path(path) file_bytes = await asyncio.to_thread(file_path.read_bytes) + if event_requests_agent_stop(stop_event): + raise AgentOutputStopped mime_type = ( await detect_image_mime_type_async( file_bytes, @@ -175,14 +183,17 @@ async def send_message_chain( ) or mimetypes.guess_type(file_path.name)[0] ) - file_ids.append( - await self.upload_file( - channel_id, - file_bytes, - file_path.name, - mime_type or "image/jpeg", - ) + if event_requests_agent_stop(stop_event): + raise AgentOutputStopped + file_id = await self.upload_file( + channel_id, + file_bytes, + file_path.name, + mime_type or "image/jpeg", ) + if event_requests_agent_stop(stop_event): + raise AgentOutputStopped + file_ids.append(file_id) elif isinstance(segment, (File, Record, Video)): if isinstance(segment, File): path = await segment.get_file() @@ -190,22 +201,29 @@ async def send_message_chain( else: path = await segment.convert_to_file_path() filename = Path(path).name + if event_requests_agent_stop(stop_event): + raise AgentOutputStopped file_path = Path(path) file_bytes = await asyncio.to_thread(file_path.read_bytes) - file_ids.append( - await self.upload_file( - channel_id, - file_bytes, - filename, - mimetypes.guess_type(filename)[0] or "application/octet-stream", - ) + if event_requests_agent_stop(stop_event): + raise AgentOutputStopped + file_id = await self.upload_file( + channel_id, + file_bytes, + filename, + mimetypes.guess_type(filename)[0] or "application/octet-stream", ) + if event_requests_agent_stop(stop_event): + raise AgentOutputStopped + file_ids.append(file_id) else: logger.debug( "Mattermost send_message_chain skipped unsupported segment: %s", segment.type, ) + if event_requests_agent_stop(stop_event): + raise AgentOutputStopped return await self.create_post( channel_id, "".join(text_parts).strip(), diff --git a/astrbot/core/platform/sources/mattermost/mattermost_event.py b/astrbot/core/platform/sources/mattermost/mattermost_event.py index 5faaf71345..b73bbc8585 100644 --- a/astrbot/core/platform/sources/mattermost/mattermost_event.py +++ b/astrbot/core/platform/sources/mattermost/mattermost_event.py @@ -1,10 +1,10 @@ -import asyncio import re from collections.abc import AsyncGenerator from astrbot.api.event import AstrMessageEvent, MessageChain from astrbot.api.message_components import Plain from astrbot.api.platform import Group, MessageMember +from astrbot.core.agent.stop_policy import AgentOutputStopped, event_requests_agent_stop from .client import MattermostClient @@ -26,7 +26,13 @@ def __init__( self.track_temporary_local_file(path) async def send(self, message: MessageChain) -> None: - await self.client.send_message_chain(self.get_session_id(), message) + if event_requests_agent_stop(self): + raise AgentOutputStopped + await self.client.send_message_chain( + self.get_session_id(), + message, + stop_event=self, + ) await super().send(message) async def send_streaming( @@ -39,6 +45,8 @@ async def send_streaming( if not use_fallback: message_buffer: MessageChain | None = None async for chain in generator: + if event_requests_agent_stop(self): + raise AgentOutputStopped if not message_buffer: message_buffer = chain else: @@ -52,8 +60,12 @@ async def send_streaming( text_buffer = "" async for chain in generator: + if event_requests_agent_stop(self): + raise AgentOutputStopped if isinstance(chain, MessageChain): for comp in chain.chain: + if event_requests_agent_stop(self): + raise AgentOutputStopped if isinstance(comp, Plain): text_buffer += comp.text if any(p in text_buffer for p in "。？！~…"): @@ -62,8 +74,7 @@ async def send_streaming( self._FALLBACK_SENTENCE_PATTERN, ) else: - await self.send(MessageChain(chain=[comp])) - await asyncio.sleep(1.5) + await self.send_streaming_fallback_component(comp) if text_buffer.strip(): await self.send(MessageChain([Plain(text_buffer)])) diff --git a/astrbot/core/platform/sources/misskey/misskey_adapter.py b/astrbot/core/platform/sources/misskey/misskey_adapter.py index c488007442..b243e3b1d4 100644 --- a/astrbot/core/platform/sources/misskey/misskey_adapter.py +++ b/astrbot/core/platform/sources/misskey/misskey_adapter.py @@ -5,13 +5,14 @@ import astrbot.api.message_components as Comp from astrbot.api import logger -from astrbot.api.event import MessageChain +from astrbot.api.event import AstrMessageEvent, MessageChain from astrbot.api.platform import ( AstrBotMessage, Platform, PlatformMetadata, register_platform_adapter, ) +from astrbot.core.agent.stop_policy import AgentOutputStopped, event_requests_agent_stop from astrbot.core.platform.astr_message_event import MessageSession from .misskey_api import MisskeyAPI @@ -372,7 +373,10 @@ async def send_by_session( self, session: MessageSession, message_chain: MessageChain, + stop_event: AstrMessageEvent | None = None, ) -> None: + if event_requests_agent_stop(stop_event): + raise AgentOutputStopped if not self.api: logger.error("[Misskey] API 客户端未初始化") return await super().send_by_session(session, message_chain) @@ -422,6 +426,8 @@ async def send_by_session( fallback_urls: list[str] = [] if not self.enable_file_upload: + if event_requests_agent_stop(stop_event): + raise AgentOutputStopped return await self._send_text_only_message( session_id, text, @@ -449,6 +455,8 @@ async def _upload_comp(comp) -> object | None: local_path = None try: async with sem: + if event_requests_agent_stop(stop_event): + raise AgentOutputStopped if not self.api: return None @@ -456,6 +464,8 @@ async def _upload_comp(comp) -> object | None: url_candidate, local_path = await resolve_component_url_or_path( comp, ) + if event_requests_agent_stop(stop_event): + raise AgentOutputStopped if not url_candidate and not local_path: return None @@ -473,26 +483,38 @@ async def _upload_comp(comp) -> object | None: preferred_name, folder_id=self.upload_folder, ) + if event_requests_agent_stop(stop_event): + raise AgentOutputStopped if isinstance(result, dict) and result.get("id"): return str(result["id"]) # 本地文件上传 if local_path: + if event_requests_agent_stop(stop_event): + raise AgentOutputStopped file_id = await upload_local_with_retries( self.api, str(local_path), preferred_name, self.upload_folder, ) + if event_requests_agent_stop(stop_event): + raise AgentOutputStopped if file_id: return file_id # 所有上传都失败，尝试获取 URL 作为回退 if hasattr(comp, "register_to_file_service"): try: + if event_requests_agent_stop(stop_event): + raise AgentOutputStopped url = await comp.register_to_file_service() + if event_requests_agent_stop(stop_event): + raise AgentOutputStopped if url: return {"fallback_url": url} + except AgentOutputStopped: + raise except Exception: pass @@ -540,6 +562,8 @@ async def _upload_comp(comp) -> object | None: try: results = await asyncio.gather(*upload_tasks) if upload_tasks else [] + if event_requests_agent_stop(stop_event): + raise AgentOutputStopped for r in results: if not r: continue @@ -554,6 +578,8 @@ async def _upload_comp(comp) -> object | None: file_ids.append(fid_str) except Exception: pass + except AgentOutputStopped: + raise except Exception: logger.debug("[Misskey] 并发上传过程中出现异常，继续发送文本") @@ -567,6 +593,8 @@ async def _upload_comp(comp) -> object | None: payload: dict[str, Any] = {"toRoomId": room_id, "text": text} if file_ids: payload["fileIds"] = file_ids + if event_requests_agent_stop(stop_event): + raise AgentOutputStopped await self.api.send_room_message(payload) elif session_id: from .misskey_utils import ( @@ -587,6 +615,8 @@ async def _upload_comp(comp) -> object | None: logger.warning( f"[Misskey] 聊天消息只支持单个文件，忽略其余 {len(file_ids) - 1} 个文件", ) + if event_requests_agent_stop(stop_event): + raise AgentOutputStopped await self.api.send_message(payload) else: # 回退到发帖逻辑 @@ -616,6 +646,8 @@ async def _upload_comp(comp) -> object | None: # 从缓存中获取原消息ID作为reply_id reply_id = user_info_for_reply.get("reply_to_note_id") + if event_requests_agent_stop(stop_event): + raise AgentOutputStopped await self.api.create_note( text=text, visibility=visibility, @@ -629,8 +661,11 @@ async def _upload_comp(comp) -> object | None: channel_id=fields["channel_id"], ) + except AgentOutputStopped: + raise except Exception as e: logger.error(f"[Misskey] 发送消息失败: {e}") + raise return await super().send_by_session(session, message_chain) diff --git a/astrbot/core/platform/sources/misskey/misskey_event.py b/astrbot/core/platform/sources/misskey/misskey_event.py index 068f7e7a28..fd8e932c72 100644 --- a/astrbot/core/platform/sources/misskey/misskey_event.py +++ b/astrbot/core/platform/sources/misskey/misskey_event.py @@ -1,4 +1,3 @@ -import asyncio import re from collections.abc import AsyncGenerator @@ -6,6 +5,7 @@ from astrbot.api.event import AstrMessageEvent, MessageChain from astrbot.api.message_components import Plain from astrbot.api.platform import AstrBotMessage, PlatformMetadata +from astrbot.core.agent.stop_policy import AgentOutputStopped, event_requests_agent_stop from .misskey_utils import ( add_at_mention_if_needed, @@ -42,6 +42,8 @@ def _is_system_command(self, message_str: str) -> bool: async def send(self, message: MessageChain) -> None: """发送消息，使用适配器的完整上传和发送逻辑""" + if event_requests_agent_stop(self): + raise AgentOutputStopped try: logger.debug( f"[MisskeyEvent] send 方法被调用，消息链包含 {len(message.chain)} 个组件", @@ -72,7 +74,7 @@ async def send(self, message: MessageChain) -> None: # 调用适配器的 send_by_session 方法 if hasattr(self.client, "send_by_session"): logger.debug("[MisskeyEvent] 调用适配器的 send_by_session 方法") - await self.client.send_by_session(session, message) + await self.client.send_by_session(session, message, stop_event=self) else: # 回退到原来的简化发送逻辑 content, has_at = serialize_message_chain(message.chain) @@ -99,15 +101,21 @@ async def send(self, message: MessageChain) -> None: if hasattr(self.client, "send_message") and is_valid_user_session_id( self.session_id, ): + if event_requests_agent_stop(self): + raise AgentOutputStopped user_id = extract_user_id_from_session_id(self.session_id) await self.client.send_message(user_id, content) elif hasattr( self.client, "send_room_message", ) and is_valid_room_session_id(self.session_id): + if event_requests_agent_stop(self): + raise AgentOutputStopped room_id = extract_room_id_from_session_id(self.session_id) await self.client.send_room_message(room_id, content) elif original_message_id and hasattr(self.client, "create_note"): + if event_requests_agent_stop(self): + raise AgentOutputStopped visibility, visible_user_ids = resolve_visibility_from_raw_message( raw_message, ) @@ -118,13 +126,18 @@ async def send(self, message: MessageChain) -> None: visible_user_ids=visible_user_ids, ) elif hasattr(self.client, "create_note"): + if event_requests_agent_stop(self): + raise AgentOutputStopped logger.debug("[MisskeyEvent] 创建新帖子") await self.client.create_note(content) await super().send(message) + except AgentOutputStopped: + raise except Exception as e: logger.error(f"[MisskeyEvent] 发送失败: {e}") + raise async def send_streaming( self, @@ -134,6 +147,8 @@ async def send_streaming( if not use_fallback: buffer = None async for chain in generator: + if event_requests_agent_stop(self): + raise AgentOutputStopped if not buffer: buffer = chain else: @@ -148,15 +163,18 @@ async def send_streaming( pattern = re.compile(r"[^。？！~…]+[。？！~…]+") async for chain in generator: + if event_requests_agent_stop(self): + raise AgentOutputStopped if isinstance(chain, MessageChain): for comp in chain.chain: + if event_requests_agent_stop(self): + raise AgentOutputStopped if isinstance(comp, Plain): buffer += comp.text if any(p in buffer for p in "。？！~…"): buffer = await self.process_buffer(buffer, pattern) else: - await self.send(MessageChain(chain=[comp])) - await asyncio.sleep(1.5) # 限速 + await self.send_streaming_fallback_component(comp) if buffer.strip(): await self.send(MessageChain([Plain(buffer)])) diff --git a/astrbot/core/platform/sources/qqofficial/qqofficial_message_event.py b/astrbot/core/platform/sources/qqofficial/qqofficial_message_event.py index 1e7d4f96ef..bda891406e 100644 --- a/astrbot/core/platform/sources/qqofficial/qqofficial_message_event.py +++ b/astrbot/core/platform/sources/qqofficial/qqofficial_message_event.py @@ -3,6 +3,7 @@ import logging import os import random +from collections.abc import Callable from typing import cast import aiofiles @@ -18,7 +19,7 @@ from tenacity import ( before_sleep_log, retry, - retry_if_exception_type, + retry_if_exception, stop_after_attempt, wait_exponential, ) @@ -27,6 +28,7 @@ from astrbot.api.event import AstrMessageEvent, MessageChain from astrbot.api.message_components import File, Image, Plain, Record, Video from astrbot.api.platform import AstrBotMessage, PlatformMetadata +from astrbot.core.agent.stop_policy import AgentOutputStopped, event_requests_agent_stop from astrbot.core.utils.media_utils import MediaResolver, file_uri_to_path, is_file_uri @@ -54,20 +56,43 @@ def _patch_qq_botpy_formdata() -> None: _patch_qq_botpy_formdata() -def _qqofficial_retry(max_attempts: int = 5): +def _qqofficial_retry( + max_attempts: int = 5, + stop_requested: Callable[[], bool] | None = None, +): """Retry decorator for QQ Official API transient errors (HTTP 500/504)""" + + stop_requested = stop_requested or (lambda: False) + + async def _interruptible_sleep(delay: float) -> None: + deadline = asyncio.get_running_loop().time() + float(delay) + while True: + if stop_requested(): + raise AgentOutputStopped + remaining = deadline - asyncio.get_running_loop().time() + if remaining <= 0: + return + await asyncio.sleep(min(remaining, 0.1)) + return retry( - retry=retry_if_exception_type( - ( - botpy.errors.ServerError, - botpy.errors.SequenceNumberError, - OSError, - asyncio.TimeoutError, - APIReturnNoneError, + retry=retry_if_exception( + lambda exc: ( + isinstance( + exc, + ( + botpy.errors.ServerError, + botpy.errors.SequenceNumberError, + OSError, + asyncio.TimeoutError, + APIReturnNoneError, + ), + ) + and not stop_requested() ) ), stop=stop_after_attempt(max_attempts), wait=wait_exponential(multiplier=2, min=2, max=30), + sleep=_interruptible_sleep, before_sleep=before_sleep_log(logger, logging.WARNING), reraise=True, ) @@ -103,11 +128,15 @@ def __init__( self.send_buffer = None async def send(self, message: MessageChain) -> None: + if event_requests_agent_stop(self): + raise AgentOutputStopped self.send_buffer = message await self._post_send() async def send_streaming(self, generator, use_fallback: bool = False): """流式输出仅支持消息列表私聊（C2C），其他消息源退化为普通发送""" + if event_requests_agent_stop(self): + raise AgentOutputStopped # 先标记事件层“已执行发送操作”，避免异常路径遗漏 await super().send_streaming(generator, use_fallback) # QQ C2C 流式协议：开始/中间分片使用 state=1，结束分片使用 state=10 @@ -115,11 +144,14 @@ async def send_streaming(self, generator, use_fallback: bool = False): last_edit_time = 0 # 上次发送分片的时间 throttle_interval = 1 # 分片间最短间隔 (秒) ret = None + stream_active = False source = ( self.message_obj.raw_message ) # 提前获取，避免 generator 为空时 NameError try: async for chain in generator: + if event_requests_agent_stop(self): + raise AgentOutputStopped source = self.message_obj.raw_message if not isinstance(source, botpy.message.C2CMessage): @@ -134,12 +166,16 @@ async def send_streaming(self, generator, use_fallback: bool = False): # tool_call break 信号：工具开始执行，先把已有 buffer 以 state=10 结束当前流式段 if chain.type == "break": - if self.send_buffer: + if self.send_buffer or stream_active: + if not self.send_buffer: + self.send_buffer = MessageChain().message("\n") stream_payload["state"] = 10 ret = await self._post_send(stream=stream_payload) - ret_id = self._extract_response_message_id(ret) - if ret_id is not None: - stream_payload["id"] = ret_id + if ret is None: + raise RuntimeError( + "QQ Official streaming message delivery failed." + ) + stream_active = False # 重置 stream_payload，为下一段流式做准备 stream_payload = { "state": 1, @@ -150,6 +186,34 @@ async def send_streaming(self, generator, use_fallback: bool = False): last_edit_time = 0 continue + if any( + isinstance(component, Image | Record | Video | File) + for component in chain.chain + ): + if self.send_buffer or stream_active: + if not self.send_buffer: + self.send_buffer = MessageChain().message("\n") + stream_payload["state"] = 10 + ret = await self._post_send(stream=stream_payload) + if ret is None: + raise RuntimeError( + "QQ Official streaming message delivery failed." + ) + stream_active = False + + self.send_buffer = chain + ret = await self._post_send() + if ret is None: + raise RuntimeError("QQ Official message delivery failed.") + stream_payload = { + "state": 1, + "id": None, + "index": 0, + "reset": False, + } + last_edit_time = 0 + continue + # 累积内容 if not self.send_buffer: self.send_buffer = chain @@ -163,25 +227,48 @@ async def send_streaming(self, generator, use_fallback: bool = False): message.Message, await self._post_send(stream=stream_payload), ) + if ret is None: + raise RuntimeError( + "QQ Official streaming message delivery failed." + ) stream_payload["index"] += 1 ret_id = self._extract_response_message_id(ret) - if ret_id is not None: - stream_payload["id"] = ret_id + if ret_id is None: + raise RuntimeError( + "QQ Official streaming response did not include a message id." + ) + stream_payload["id"] = ret_id + stream_active = True last_edit_time = asyncio.get_running_loop().time() self.send_buffer = None # 清空已发送的分片，避免下次重复发送旧内容 if isinstance(source, botpy.message.C2CMessage): # 结束流式对话，发送 buffer 中剩余内容 - stream_payload["state"] = 10 - ret = await self._post_send(stream=stream_payload) + if self.send_buffer or stream_active: + if not self.send_buffer: + self.send_buffer = MessageChain().message("\n") + stream_payload["state"] = 10 + ret = await self._post_send(stream=stream_payload) + if ret is None: + raise RuntimeError( + "QQ Official streaming message delivery failed." + ) else: ret = await self._post_send() + if ret is None: + raise RuntimeError( + "QQ Official streaming message produced no platform delivery." + ) + except AgentOutputStopped: + self.send_buffer = None + raise except Exception as e: logger.error(f"发送流式消息时出错: {e}", exc_info=True) # 避免累计内容在异常后被整包重复发送：仅清理缓存，不做非流式整包兜底 # 如需兜底，应该只发送未发送 delta（后续可继续优化） self.send_buffer = None + raise return None @@ -230,27 +317,35 @@ def _split_message_chain_by_media(message: MessageChain) -> list[MessageChain]: return chunks async def _post_send(self, stream: dict | None = None): - if not self.send_buffer: - return None - - message_chains = self._split_message_chain_by_media(self.send_buffer) - stream_for_chain = stream if len(message_chains) == 1 else None - - ret = None - for message_chain in message_chains: - ret = await self._post_send_one(message_chain, stream_for_chain) - - self.send_buffer = None - - return ret + try: + if event_requests_agent_stop(self): + raise AgentOutputStopped + if self.send_buffer is None or not self.send_buffer.chain: + raise RuntimeError("QQ Official message buffer is empty.") + + message_chains = self._split_message_chain_by_media(self.send_buffer) + stream_for_chain = stream if len(message_chains) == 1 else None + + ret = None + for message_chain in message_chains: + if event_requests_agent_stop(self): + raise AgentOutputStopped + ret = await self._post_send_one(message_chain, stream_for_chain) + if ret is None: + raise RuntimeError( + "QQ Official message delivery returned no response." + ) + return ret + finally: + self.send_buffer = None async def _post_send_one( self, message_to_send: MessageChain, stream: dict | None = None, ): - if not message_to_send: - return None + if not message_to_send.chain: + raise RuntimeError("QQ Official message chain is empty.") source = self.message_obj.raw_message @@ -261,8 +356,9 @@ async def _post_send_one( | botpy.message.DirectMessage | botpy.message.C2CMessage, ): - logger.warning(f"[QQOfficial] 不支持的消息源类型: {type(source)}") - return None + raise RuntimeError( + f"QQ Official does not support source type: {type(source)}" + ) ( plain_text, @@ -275,6 +371,8 @@ async def _post_send_one( ) = await QQOfficialMessageEvent._parse_to_qqofficial(message_to_send) if record_file_path: self.track_temporary_local_file(record_file_path) + if event_requests_agent_stop(self): + raise AgentOutputStopped # C2C 流式仅用于文本分片，富媒体时降级为普通发送，避免平台侧流式校验报错。 if stream and ( @@ -291,7 +389,7 @@ async def _post_send_one( and not video_file_source and not file_source ): - return None + raise RuntimeError("QQ Official message has no deliverable content.") # QQ C2C 流式 API 说明： # - 开始/中间分片（state=1）：增量追加内容，不需要 \n（加了会导致强制换行） @@ -327,8 +425,9 @@ async def _post_send_one( match source: case botpy.message.GroupMessage(): if not source.group_openid: - logger.error("[QQOfficial] GroupMessage 缺少 group_openid") - return None + raise RuntimeError( + "QQ Official group message is missing group_openid." + ) if image_base64: media = await self.upload_group_and_c2c_image( @@ -483,7 +582,12 @@ async def _post_send_one( ) case _: - pass + raise RuntimeError( + f"QQ Official does not support source type: {type(source)}" + ) + + if ret is None: + raise RuntimeError("QQ Official message delivery returned no response.") await super().send(message_to_send) @@ -496,13 +600,19 @@ async def _send_with_markdown_fallback( plain_text: str, stream: dict | None = None, ): + if event_requests_agent_stop(self): + raise AgentOutputStopped try: return await send_func(payload) except _QQOFFICIAL_SEND_API_ERRORS as err: + if event_requests_agent_stop(self): + raise AgentOutputStopped from None logger.info("[QQOfficial] 回复消息失败: %s, 尝试使用主动发送接口。", err) if payload.get("msg_id"): fallback_payload = payload.copy() try: + if event_requests_agent_stop(self): + raise AgentOutputStopped ret = await send_func(fallback_payload) logger.info("[QQOfficial] 使用主动发送接口发送成功。") return ret @@ -531,6 +641,8 @@ async def _send_with_markdown_fallback( logger.warning( "[QQOfficial] 流式 markdown 分片换行校验失败，已修正后重试一次。" ) + if event_requests_agent_stop(self): + raise AgentOutputStopped return await send_func(retry_payload) if ( @@ -552,6 +664,8 @@ async def _send_with_markdown_fallback( fallback_content = cast(str, fallback_payload.get("content") or "") if fallback_content and not fallback_content.endswith("\n"): fallback_payload["content"] = fallback_content + "\n" + if event_requests_agent_stop(self): + raise AgentOutputStopped return await send_func(fallback_payload) async def upload_group_and_c2c_image( @@ -566,8 +680,10 @@ async def upload_group_and_c2c_image( "srv_send_msg": False, } - @_qqofficial_retry() + @_qqofficial_retry(stop_requested=lambda: event_requests_agent_stop(self)) async def _do_upload(): + if event_requests_agent_stop(self): + raise AgentOutputStopped if "openid" in kwargs: payload["openid"] = kwargs["openid"] route = Route( @@ -645,8 +761,10 @@ async def upload_group_and_c2c_media( else: return None - @_qqofficial_retry() + @_qqofficial_retry(stop_requested=lambda: event_requests_agent_stop(self)) async def _do_upload(): + if event_requests_agent_stop(self): + raise AgentOutputStopped result = await self.bot.api._http.request(route, json=payload) if result is None: err_msg = "上传文件API返回None，触发重试" @@ -656,24 +774,29 @@ async def _do_upload(): try: result = await _do_upload() - if result: - if not isinstance(result, dict): - logger.error(f"上传文件响应格式错误: {result}") - return None - - return Media( - file_uuid=result["file_uuid"], - file_info=result["file_info"], - ttl=result.get("ttl", 0), + if not isinstance(result, dict): + raise RuntimeError( + f"QQ Official media API returned an invalid response: {result}" ) + + return Media( + file_uuid=result["file_uuid"], + file_info=result["file_info"], + ttl=result.get("ttl", 0), + ) + except AgentOutputStopped: + raise except APIReturnNoneError: logger.warning(f"上传文件API返回None，共尝试5次后放弃: {file_source}") + raise RuntimeError( + "QQ Official media upload returned no response." + ) from None except (botpy.errors.ServerError, botpy.errors.SequenceNumberError): logger.error(f"上传媒体文件失败，共尝试5次后放弃: {file_source}") + raise except Exception as e: logger.error(f"上传请求错误: {e}") - - return None + raise async def post_c2c_message( self, @@ -705,8 +828,13 @@ async def post_c2c_message( retry_times = 3 - @_qqofficial_retry(retry_times) + @_qqofficial_retry( + retry_times, + stop_requested=lambda: event_requests_agent_stop(self), + ) async def _do_request(): + if event_requests_agent_stop(self): + raise AgentOutputStopped result = await self.bot.api._http.request(route, json=payload) if result is None: err_msg = "发送消息API返回None，触发重试" @@ -717,14 +845,19 @@ async def _do_request(): try: result = await _do_request() except APIReturnNoneError: + if event_requests_agent_stop(self): + raise AgentOutputStopped from None logger.warning( f"[QQOfficial] post_c2c_message: 发送消息失败，API 返回 None，共尝试{retry_times}次后放弃" ) - return None + raise RuntimeError( + "QQ Official message API returned no response." + ) from None if not isinstance(result, dict): - logger.error(f"[QQOfficial] post_c2c_message: 响应不是 dict: {result}") - return None + raise RuntimeError( + f"QQ Official message API returned an invalid response: {result}" + ) return message.Message(**result) diff --git a/astrbot/core/platform/sources/satori/satori_event.py b/astrbot/core/platform/sources/satori/satori_event.py index 7e2e92eb40..e0d3581196 100644 --- a/astrbot/core/platform/sources/satori/satori_event.py +++ b/astrbot/core/platform/sources/satori/satori_event.py @@ -15,6 +15,7 @@ Video, ) from astrbot.api.platform import AstrBotMessage, PlatformMetadata +from astrbot.core.agent.stop_policy import AgentOutputStopped, event_requests_agent_stop from astrbot.core.utils.media_utils import resolve_media_ref_to_base64_data if TYPE_CHECKING: @@ -147,7 +148,11 @@ async def send(self, message: MessageChain) -> None: content_parts = [] for component in message.chain: + if event_requests_agent_stop(self): + raise AgentOutputStopped component_content = await self._convert_component_to_satori(component) + if event_requests_agent_stop(self): + raise AgentOutputStopped if component_content: content_parts.append(component_content) @@ -155,12 +160,16 @@ async def send(self, message: MessageChain) -> None: if isinstance(component, Node): # 单个转发节点 node_content = await self._convert_node_to_satori(component) + if event_requests_agent_stop(self): + raise AgentOutputStopped if node_content: content_parts.append(node_content) elif isinstance(component, Nodes): # 合并转发消息 node_content = await self._convert_nodes_to_satori(component) + if event_requests_agent_stop(self): + raise AgentOutputStopped if node_content: content_parts.append(node_content) @@ -168,6 +177,8 @@ async def send(self, message: MessageChain) -> None: channel_id = self.session_id data = {"channel_id": channel_id, "content": content} + if event_requests_agent_stop(self): + raise AgentOutputStopped result = await self.adapter.send_http_request( "POST", "/message.create", @@ -176,9 +187,12 @@ async def send(self, message: MessageChain) -> None: user_id, ) if not result: - logger.error("Satori 消息发送失败") + raise RuntimeError("Satori message delivery failed.") + except AgentOutputStopped: + raise except Exception as e: logger.error(f"Satori 消息发送异常: {e}") + raise await super().send(message) @@ -197,18 +211,27 @@ async def send_streaming(self, generator, use_fallback: bool = False): continue for component in chain.chain: + if event_requests_agent_stop(self): + raise AgentOutputStopped if isinstance(component, Plain): content_parts.append(component.text) elif isinstance(component, Image): if content_parts: + if event_requests_agent_stop(self): + content_parts = [] + continue content = "".join(content_parts) temp_chain = MessageChain([Plain(text=content)]) await self.send(temp_chain) content_parts = [] + if event_requests_agent_stop(self): + continue try: image_data_url = await self._image_to_data_url( component ) + if event_requests_agent_stop(self): + raise AgentOutputStopped if image_data_url: img_chain = MessageChain( [ @@ -218,18 +241,26 @@ async def send_streaming(self, generator, use_fallback: bool = False): ], ) await self.send(img_chain) + except AgentOutputStopped: + raise except Exception as e: logger.error(f"图片转换为base64失败: {e}") + raise else: content_parts.append(str(component)) if content_parts: + if event_requests_agent_stop(self): + raise AgentOutputStopped content = "".join(content_parts) temp_chain = MessageChain([Plain(text=content)]) await self.send(temp_chain) + except AgentOutputStopped: + raise except Exception as e: logger.error(f"Satori 流式消息发送异常: {e}") + raise return await super().send_streaming(generator, use_fallback) diff --git a/astrbot/core/platform/sources/slack/slack_event.py b/astrbot/core/platform/sources/slack/slack_event.py index 5fb26d22f0..b04a369da8 100644 --- a/astrbot/core/platform/sources/slack/slack_event.py +++ b/astrbot/core/platform/sources/slack/slack_event.py @@ -1,4 +1,3 @@ -import asyncio import re from collections.abc import AsyncGenerator, Iterable from pathlib import Path @@ -15,6 +14,7 @@ Plain, ) from astrbot.api.platform import Group, MessageMember +from astrbot.core.agent.stop_policy import AgentOutputStopped, event_requests_agent_stop class SlackMessageEvent(AstrMessageEvent): @@ -33,8 +33,11 @@ def __init__( async def _from_segment_to_slack_block( segment: BaseMessageComponent, web_client: AsyncWebClient, + stop_event: AstrMessageEvent | None = None, ) -> dict | None: """将消息段转换为 Slack 块格式""" + if event_requests_agent_stop(stop_event): + raise AgentOutputStopped if isinstance(segment, Plain): return {"type": "section", "text": {"type": "mrkdwn", "text": segment.text}} if isinstance(segment, Image): @@ -47,6 +50,8 @@ async def _from_segment_to_slack_block( "alt_text": "图片", } path = await segment.convert_to_file_path() + if event_requests_agent_stop(stop_event): + raise AgentOutputStopped response = await web_client.files_upload_v2( file=path, filename=Path(path).name, @@ -92,12 +97,15 @@ async def _from_segment_to_slack_block( async def _parse_slack_blocks( message_chain: MessageChain, web_client: AsyncWebClient, + stop_event: AstrMessageEvent | None = None, ): """解析成 Slack 块格式""" blocks = [] text_content = "" for segment in message_chain.chain: + if event_requests_agent_stop(stop_event): + raise AgentOutputStopped if isinstance(segment, Plain): text_content += segment.text else: @@ -115,7 +123,10 @@ async def _parse_slack_blocks( block = await SlackMessageEvent._from_segment_to_slack_block( segment, web_client, + stop_event, ) + if event_requests_agent_stop(stop_event): + raise AgentOutputStopped if block: blocks.append(block) @@ -128,27 +139,34 @@ async def _parse_slack_blocks( return blocks, "" if blocks else text_content async def send(self, message: MessageChain) -> None: + if event_requests_agent_stop(self): + raise AgentOutputStopped blocks, text = await SlackMessageEvent._parse_slack_blocks( message, self.web_client, + self, ) try: if self.get_group_id(): # 发送到频道 - await self.web_client.chat_postMessage( + response = await self.web_client.chat_postMessage( channel=self.get_group_id(), text=text, blocks=blocks or None, ) else: # 发送私信 - await self.web_client.chat_postMessage( + response = await self.web_client.chat_postMessage( channel=self.get_sender_id(), text=text, blocks=blocks or None, ) + if not response.get("ok"): + raise RuntimeError("Slack message delivery failed.") except Exception: + if event_requests_agent_stop(self): + raise AgentOutputStopped from None # 如果块发送失败，尝试只发送文本 parts = [] for segment in message.chain: @@ -160,16 +178,20 @@ async def send(self, message: MessageChain) -> None: parts.append(" [图片] ") fallback_text = "".join(parts) + if event_requests_agent_stop(self): + raise AgentOutputStopped if self.get_group_id(): - await self.web_client.chat_postMessage( + response = await self.web_client.chat_postMessage( channel=self.get_group_id(), text=fallback_text, ) else: - await self.web_client.chat_postMessage( + response = await self.web_client.chat_postMessage( channel=self.get_sender_id(), text=fallback_text, ) + if not response.get("ok"): + raise RuntimeError("Slack fallback message delivery failed.") await super().send(message) @@ -181,6 +203,8 @@ async def send_streaming( if not use_fallback: buffer = None async for chain in generator: + if event_requests_agent_stop(self): + raise AgentOutputStopped if not buffer: buffer = chain else: @@ -195,15 +219,18 @@ async def send_streaming( pattern = re.compile(r"[^。？！~…]+[。？！~…]+") async for chain in generator: + if event_requests_agent_stop(self): + raise AgentOutputStopped if isinstance(chain, MessageChain): for comp in chain.chain: + if event_requests_agent_stop(self): + raise AgentOutputStopped if isinstance(comp, Plain): buffer += comp.text if any(p in buffer for p in "。？！~…"): buffer = await self.process_buffer(buffer, pattern) else: - await self.send(MessageChain(chain=[comp])) - await asyncio.sleep(1.5) # 限速 + await self.send_streaming_fallback_component(comp) if buffer.strip(): await self.send(MessageChain([Plain(buffer)])) diff --git a/astrbot/core/platform/sources/telegram/tg_event.py b/astrbot/core/platform/sources/telegram/tg_event.py index 8445a8ea1e..a93c99303e 100644 --- a/astrbot/core/platform/sources/telegram/tg_event.py +++ b/astrbot/core/platform/sources/telegram/tg_event.py @@ -22,6 +22,7 @@ Video, ) from astrbot.api.platform import AstrBotMessage, MessageType, PlatformMetadata +from astrbot.core.agent.stop_policy import AgentOutputStopped, event_requests_agent_stop from astrbot.core.utils.metrics import Metric @@ -111,9 +112,12 @@ async def _send_text_chunks( client: ExtBot, text: str, payload: dict[str, Any], + stop_event: AstrMessageEvent | None = None, ) -> None: """按 Telegram 限制切分文本后逐段发送。""" for chunk in cls._split_message(text): + if event_requests_agent_stop(stop_event): + raise AgentOutputStopped try: markdown_text = telegramify_markdown.markdownify( chunk, @@ -127,6 +131,8 @@ async def _send_text_chunks( logger.warning( f"Failed to convert message to Markdown，using normal text: {e!s}" ) + if event_requests_agent_stop(stop_event): + raise AgentOutputStopped await client.send_message(text=chunk, **cast(Any, payload)) @classmethod @@ -163,6 +169,7 @@ async def _send_media_with_action( *, user_name: str, message_thread_id: str | None = None, + stop_event: AstrMessageEvent | None = None, **payload: Any, ) -> None: """发送媒体时显示 upload action，发送完成后恢复 typing""" @@ -172,10 +179,14 @@ async def _send_media_with_action( await cls._send_chat_action( client, user_name, upload_action, effective_thread_id ) + if event_requests_agent_stop(stop_event): + raise AgentOutputStopped send_payload = dict(payload) if effective_thread_id and "message_thread_id" not in send_payload: send_payload["message_thread_id"] = effective_thread_id await send_coro(**send_payload) + if event_requests_agent_stop(stop_event): + return await cls._send_chat_action( client, user_name, ChatAction.TYPING, effective_thread_id ) @@ -191,6 +202,7 @@ async def _send_voice_with_fallback( user_name: str = "", message_thread_id: str | None = None, use_media_action: bool = False, + stop_event: AstrMessageEvent | None = None, ) -> None: """Send a voice message, falling back to a document if the user's privacy settings forbid voice messages (``BadRequest`` with @@ -199,6 +211,8 @@ async def _send_voice_with_fallback( When *use_media_action* is ``True`` the helper wraps the send calls with ``_send_media_with_action`` (used by the streaming path). """ + if event_requests_agent_stop(stop_event): + raise AgentOutputStopped try: if use_media_action: media_payload = dict(payload) @@ -209,6 +223,7 @@ async def _send_voice_with_fallback( ChatAction.UPLOAD_VOICE, client.send_voice, user_name=user_name, + stop_event=stop_event, voice=path, **cast(Any, media_payload), ) @@ -219,6 +234,8 @@ async def _send_voice_with_fallback( # distinguish the voice-privacy case via the API error message. if "Voice_messages_forbidden" not in e.message: raise + if event_requests_agent_stop(stop_event): + raise AgentOutputStopped logger.warning( "User privacy settings prevent receiving voice messages, falling back to sending an audio file. " "To enable voice messages, go to Telegram Settings → Privacy and Security → Voice Messages → set to 'Everyone'." @@ -232,6 +249,7 @@ async def _send_voice_with_fallback( ChatAction.UPLOAD_DOCUMENT, client.send_document, user_name=user_name, + stop_event=stop_event, document=path, caption=caption, **cast(Any, media_payload), @@ -271,6 +289,7 @@ async def send_with_client( client: ExtBot, message: MessageChain, user_name: str, + stop_event: AstrMessageEvent | None = None, ) -> None: image_path = None @@ -293,8 +312,12 @@ async def send_with_client( # 根据消息链确定合适的 chat action 并发送 action = cls._get_chat_action_for_chain(message.chain) await cls._send_chat_action(client, user_name, action, message_thread_id) + if event_requests_agent_stop(stop_event): + raise AgentOutputStopped for i in message.chain: + if event_requests_agent_stop(stop_event): + raise AgentOutputStopped payload = { "chat_id": user_name, } @@ -307,9 +330,11 @@ async def send_with_client( if at_user_id and not at_flag: i.text = f"@{at_user_id} {i.text}" at_flag = True - await cls._send_text_chunks(client, i.text, payload) + await cls._send_text_chunks(client, i.text, payload, stop_event) elif isinstance(i, Image): image_path = await i.convert_to_file_path() + if event_requests_agent_stop(stop_event): + raise AgentOutputStopped if _is_gif(image_path): send_coro = client.send_animation media_kwarg = {"animation": image_path} @@ -319,21 +344,28 @@ async def send_with_client( await send_coro(**media_kwarg, **cast(Any, payload)) elif isinstance(i, File): path = await i.get_file() + if event_requests_agent_stop(stop_event): + raise AgentOutputStopped name = i.name or os.path.basename(path) await client.send_document( document=path, filename=name, **cast(Any, payload) ) elif isinstance(i, Record): path = await i.convert_to_file_path() + if event_requests_agent_stop(stop_event): + raise AgentOutputStopped await cls._send_voice_with_fallback( client, path, payload, caption=i.text or None, use_media_action=False, + stop_event=stop_event, ) elif isinstance(i, Video): path = await i.convert_to_file_path() + if event_requests_agent_stop(stop_event): + raise AgentOutputStopped await client.send_video( video=path, caption=getattr(i, "text", None) or None, @@ -342,9 +374,19 @@ async def send_with_client( async def send(self, message: MessageChain) -> None: if self.get_message_type() == MessageType.GROUP_MESSAGE: - await self.send_with_client(self.client, message, self.message_obj.group_id) + await self.send_with_client( + self.client, + message, + self.message_obj.group_id, + self, + ) else: - await self.send_with_client(self.client, message, self.get_sender_id()) + await self.send_with_client( + self.client, + message, + self.get_sender_id(), + self, + ) await super().send(message) async def react(self, emoji: str | None, big: bool = False) -> None: @@ -398,7 +440,7 @@ async def _send_message_draft( message_thread_id: 可选，目标消息线程 ID parse_mode: 可选，消息文本的解析模式 """ - if not text or not text.strip(): + if not text or not text.strip() or event_requests_agent_stop(self): return kwargs: dict[str, Any] = {} @@ -430,10 +472,14 @@ async def _process_chain_items( ) -> None: """处理 MessageChain 中的各类组件，文本通过 on_text 回调追加，媒体直接发送。""" for i in chain.chain: + if event_requests_agent_stop(self): + raise AgentOutputStopped if isinstance(i, Plain): on_text(i.text) elif isinstance(i, Image): image_path = await i.convert_to_file_path() + if event_requests_agent_stop(self): + raise AgentOutputStopped if _is_gif(image_path): action = ChatAction.UPLOAD_VIDEO send_coro = self.client.send_animation @@ -447,23 +493,29 @@ async def _process_chain_items( action, send_coro, user_name=user_name, + stop_event=self, **media_kwarg, **cast(Any, payload), ) elif isinstance(i, File): path = await i.get_file() + if event_requests_agent_stop(self): + raise AgentOutputStopped name = i.name or os.path.basename(path) await self._send_media_with_action( self.client, ChatAction.UPLOAD_DOCUMENT, self.client.send_document, user_name=user_name, + stop_event=self, document=path, filename=name, **cast(Any, payload), ) elif isinstance(i, Record): path = await i.convert_to_file_path() + if event_requests_agent_stop(self): + raise AgentOutputStopped await self._send_voice_with_fallback( self.client, path, @@ -472,14 +524,18 @@ async def _process_chain_items( user_name=user_name, message_thread_id=message_thread_id, use_media_action=True, + stop_event=self, ) elif isinstance(i, Video): path = await i.convert_to_file_path() + if event_requests_agent_stop(self): + raise AgentOutputStopped await self._send_media_with_action( self.client, ChatAction.UPLOAD_VIDEO, self.client.send_video, user_name=user_name, + stop_event=self, video=path, **cast(Any, payload), ) @@ -488,7 +544,7 @@ async def _process_chain_items( async def _send_final_segment(self, delta: str, payload: dict[str, Any]) -> None: """将累积文本作为 MarkdownV2 真实消息发送，失败时回退到纯文本。""" - await self._send_text_chunks(self.client, delta, payload) + await self._send_text_chunks(self.client, delta, payload, self) async def send_streaming(self, generator, use_fallback: bool = False): message_thread_id = None @@ -553,6 +609,8 @@ async def _draft_sender_loop() -> None: while not done: await text_changed.wait() text_changed.clear() + if done or event_requests_agent_stop(self): + return # 发送最新的缓冲区内容（MarkdownV2 渲染，与真实消息一致） if delta and delta != last_sent_text: draft_text = delta[: self.MAX_MESSAGE_LENGTH] @@ -593,6 +651,8 @@ def _append_text(t: str) -> None: try: async for chain in generator: + if event_requests_agent_stop(self): + continue if not isinstance(chain, MessageChain): continue @@ -606,6 +666,8 @@ def _append_text(t: str) -> None: "\u23f3", message_thread_id, ) + if event_requests_agent_stop(self): + raise AgentOutputStopped await self._send_final_segment(delta, payload) delta = "" last_sent_text = "" @@ -622,12 +684,16 @@ def _append_text(t: str) -> None: # 流式结束：用 emoji 清空 draft，然后发真实消息持久化 if delta: + if event_requests_agent_stop(self): + raise AgentOutputStopped await self._send_message_draft( user_name, draft_id, "\u23f3", message_thread_id, ) + if event_requests_agent_stop(self): + raise AgentOutputStopped await self._send_final_segment(delta, payload) async def _send_streaming_edit( @@ -647,7 +713,10 @@ async def _send_streaming_edit( chat_action_interval = 0.5 # chat action 的节流间隔 (秒) # 发送初始 typing 状态 - await self._ensure_typing(user_name, message_thread_id) + if not event_requests_agent_stop(self): + await self._ensure_typing(user_name, message_thread_id) + if event_requests_agent_stop(self): + return last_chat_action_time = asyncio.get_running_loop().time() def _append_text(t: str) -> None: @@ -655,6 +724,8 @@ def _append_text(t: str) -> None: delta += t async for chain in generator: + if event_requests_agent_stop(self): + continue if not isinstance(chain, MessageChain): continue @@ -669,6 +740,7 @@ def _append_text(t: str) -> None: ) except Exception as e: logger.warning(f"编辑消息失败(streaming-break): {e!s}") + raise message_id = None delta = "" continue @@ -676,6 +748,8 @@ def _append_text(t: str) -> None: await self._process_chain_items( chain, payload, user_name, message_thread_id, _append_text ) + if event_requests_agent_stop(self): + continue # 编辑或发送消息 if message_id and len(delta) <= self.MAX_MESSAGE_LENGTH: @@ -686,6 +760,8 @@ def _append_text(t: str) -> None: current_time = asyncio.get_running_loop().time() if current_time - last_chat_action_time >= chat_action_interval: await self._ensure_typing(user_name, message_thread_id) + if event_requests_agent_stop(self): + continue last_chat_action_time = current_time try: await self.client.edit_message_text( @@ -701,6 +777,8 @@ def _append_text(t: str) -> None: current_time = asyncio.get_running_loop().time() if current_time - last_chat_action_time >= chat_action_interval: await self._ensure_typing(user_name, message_thread_id) + if event_requests_agent_stop(self): + continue last_chat_action_time = current_time try: msg = await self.client.send_message( @@ -709,11 +787,16 @@ def _append_text(t: str) -> None: current_content = delta except Exception as e: logger.warning(f"发送消息失败(streaming): {e!s}") + raise message_id = msg.message_id last_edit_time = asyncio.get_running_loop().time() try: - if delta and current_content != delta: + if ( + delta + and current_content != delta + and not event_requests_agent_stop(self) + ): try: markdown_text = telegramify_markdown.markdownify( delta, @@ -726,10 +809,15 @@ def _append_text(t: str) -> None: ) except Exception as e: logger.warning(f"Markdown转换失败，使用普通文本: {e!s}") + if event_requests_agent_stop(self): + raise AgentOutputStopped from e await self.client.edit_message_text( text=delta, chat_id=payload["chat_id"], message_id=message_id, ) + except AgentOutputStopped: + raise except Exception as e: logger.warning(f"编辑消息失败(streaming): {e!s}") + raise diff --git a/astrbot/core/platform/sources/webchat/webchat_event.py b/astrbot/core/platform/sources/webchat/webchat_event.py index e440b3d38e..329b3ed8a5 100644 --- a/astrbot/core/platform/sources/webchat/webchat_event.py +++ b/astrbot/core/platform/sources/webchat/webchat_event.py @@ -9,6 +9,7 @@ from astrbot.api import logger from astrbot.api.event import AstrMessageEvent, MessageChain from astrbot.api.message_components import File, Image, Json, Plain, Record +from astrbot.core.agent.stop_policy import AgentOutputStopped, event_requests_agent_stop from astrbot.core.utils.astrbot_path import get_astrbot_data_path from astrbot.core.utils.media_utils import ( MEDIA_MIME_EXTENSIONS, @@ -41,6 +42,7 @@ async def _send( session_id: str, streaming: bool = False, emit_complete: bool = False, + stop_event: AstrMessageEvent | None = None, ) -> str | None: request_id = str(message_id) conversation_id = _extract_conversation_id(session_id) @@ -58,9 +60,13 @@ async def _send( }, # end means this request is finished ) return + if event_requests_agent_stop(stop_event): + raise AgentOutputStopped data = "" for comp in message.chain: + if event_requests_agent_stop(stop_event): + raise AgentOutputStopped if isinstance(comp, Plain): data = comp.text await web_chat_back_queue.put( @@ -85,43 +91,113 @@ async def _send( elif isinstance(comp, Image): # save image to local image_base64 = await comp.convert_to_base64() + if event_requests_agent_stop(stop_event): + raise AgentOutputStopped image_bytes = base64.b64decode(image_base64) mime_type = await detect_image_mime_type_async( image_bytes, default_mime_type=None, ) + if event_requests_agent_stop(stop_event): + raise AgentOutputStopped suffix = MEDIA_MIME_EXTENSIONS.get(mime_type or "", ".jpg") filename = f"{str(uuid.uuid4())}{suffix}" - path = os.path.join(attachments_dir, filename) - await asyncio.to_thread(Path(path).write_bytes, image_bytes) - data = f"[IMAGE]{filename}" - await web_chat_back_queue.put( - { - "type": "image", - "data": data, - "streaming": streaming, - "message_id": message_id, - }, - ) + path = Path(attachments_dir) / filename + published = False + try: + write_task = asyncio.create_task( + asyncio.to_thread(path.write_bytes, image_bytes) + ) + try: + await asyncio.shield(write_task) + except asyncio.CancelledError as exc: + while not write_task.done(): + try: + await asyncio.shield(write_task) + except asyncio.CancelledError: + continue + if not write_task.cancelled(): + try: + write_task.result() + except Exception: + pass + raise exc + if event_requests_agent_stop(stop_event): + raise AgentOutputStopped + data = f"[IMAGE]{filename}" + await web_chat_back_queue.put( + { + "type": "image", + "data": data, + "streaming": streaming, + "message_id": message_id, + }, + ) + published = True + finally: + if not published: + try: + path.unlink(missing_ok=True) + except OSError as exc: + logger.warning( + "Failed to clean unpublished WebChat image %s: %s", + path, + exc, + ) elif isinstance(comp, Record): # save record to local filename = f"{str(uuid.uuid4())}.wav" - path = os.path.join(attachments_dir, filename) + path = Path(attachments_dir) / filename record_base64 = await comp.convert_to_base64() + if event_requests_agent_stop(stop_event): + raise AgentOutputStopped record_bytes = base64.b64decode(record_base64) - await asyncio.to_thread(Path(path).write_bytes, record_bytes) - data = f"[RECORD]{filename}" - await web_chat_back_queue.put( - { - "type": "record", - "data": data, - "streaming": streaming, - "message_id": message_id, - }, - ) + published = False + try: + write_task = asyncio.create_task( + asyncio.to_thread(path.write_bytes, record_bytes) + ) + try: + await asyncio.shield(write_task) + except asyncio.CancelledError as exc: + while not write_task.done(): + try: + await asyncio.shield(write_task) + except asyncio.CancelledError: + continue + if not write_task.cancelled(): + try: + write_task.result() + except Exception: + pass + raise exc + if event_requests_agent_stop(stop_event): + raise AgentOutputStopped + data = f"[RECORD]{filename}" + await web_chat_back_queue.put( + { + "type": "record", + "data": data, + "streaming": streaming, + "message_id": message_id, + }, + ) + published = True + finally: + if not published: + try: + path.unlink(missing_ok=True) + except OSError as exc: + logger.warning( + "Failed to clean unpublished WebChat record %s: %s", + path, + exc, + ) elif isinstance(comp, File): # save file to local file_path = await comp.get_file() + if event_requests_agent_stop(stop_event): + raise AgentOutputStopped raw_original_name = comp.name or os.path.basename(file_path) original_name = ( PurePosixPath(str(raw_original_name).replace("\\", "/")) @@ -132,21 +208,38 @@ async def _send( original_name = os.path.basename(file_path) or "file" ext = os.path.splitext(original_name)[1] or "" filename = f"{uuid.uuid4()!s}{ext}" - dest_path = os.path.join(attachments_dir, filename) - shutil.copy2(file_path, dest_path) - data = f"[FILE]{filename}|{original_name}" - await web_chat_back_queue.put( - { - "type": "file", - "data": data, - "streaming": streaming, - "message_id": message_id, - }, - ) + dest_path = Path(attachments_dir) / filename + published = False + try: + shutil.copy2(file_path, dest_path) + if event_requests_agent_stop(stop_event): + raise AgentOutputStopped + data = f"[FILE]{filename}|{original_name}" + await web_chat_back_queue.put( + { + "type": "file", + "data": data, + "streaming": streaming, + "message_id": message_id, + }, + ) + published = True + finally: + if not published: + try: + dest_path.unlink(missing_ok=True) + except OSError as exc: + logger.warning( + "Failed to clean unpublished WebChat file %s: %s", + dest_path, + exc, + ) else: logger.debug(f"webchat 忽略: {comp.type}") if emit_complete: + if event_requests_agent_stop(stop_event): + raise AgentOutputStopped await web_chat_back_queue.put( { "type": "complete", @@ -161,7 +254,12 @@ async def _send( async def send(self, message: MessageChain | None) -> None: message_id = self.message_obj.message_id - await WebChatMessageEvent._send(message_id, message, session_id=self.session_id) + await WebChatMessageEvent._send( + message_id, + message, + session_id=self.session_id, + stop_event=self, + ) await super().send(MessageChain([])) async def send_streaming(self, generator, use_fallback: bool = False) -> None: @@ -196,6 +294,8 @@ async def send_streaming(self, generator, use_fallback: bool = False) -> None: if text: payload["text"] = text + if event_requests_agent_stop(self): + raise AgentOutputStopped await web_chat_back_queue.put(payload) continue @@ -216,6 +316,7 @@ async def send_streaming(self, generator, use_fallback: bool = False) -> None: message=chain, session_id=self.session_id, streaming=True, + stop_event=self, ) if not r: continue @@ -224,6 +325,8 @@ async def send_streaming(self, generator, use_fallback: bool = False) -> None: else: final_data += r + if event_requests_agent_stop(self): + raise AgentOutputStopped await web_chat_back_queue.put( { "type": "complete", # complete means we return the final result diff --git a/astrbot/core/platform/sources/wecom/wecom_event.py b/astrbot/core/platform/sources/wecom/wecom_event.py index 265b41014f..b449834e2f 100644 --- a/astrbot/core/platform/sources/wecom/wecom_event.py +++ b/astrbot/core/platform/sources/wecom/wecom_event.py @@ -8,6 +8,7 @@ from astrbot.api.event import AstrMessageEvent, MessageChain from astrbot.api.message_components import File, Image, Plain, Record, Video from astrbot.api.platform import AstrBotMessage, PlatformMetadata +from astrbot.core.agent.stop_policy import AgentOutputStopped, event_requests_agent_stop from astrbot.core.utils.media_utils import convert_audio_to_amr from .wecom_kf_message import WeChatKFMessage @@ -80,6 +81,8 @@ async def split_plain(self, plain: str) -> list[str]: return result async def send(self, message: MessageChain) -> None: + if event_requests_agent_stop(self): + raise AgentOutputStopped message_obj = self.message_obj is_wechat_kf = hasattr(self.client, "kf_message") @@ -92,10 +95,18 @@ async def send(self, message: MessageChain) -> None: user_id = self.get_sender_id() for comp in message.chain: + if event_requests_agent_stop(self): + raise AgentOutputStopped if isinstance(comp, Plain): # Split long text messages if needed plain_chunks = await self.split_plain(comp.text) - for chunk in plain_chunks: + if event_requests_agent_stop(self): + raise AgentOutputStopped + for index, chunk in enumerate(plain_chunks): + if index: + await asyncio.sleep(0.5) + if event_requests_agent_stop(self): + raise AgentOutputStopped try: kf_message_api.send_text(user_id, self.get_self_id(), chunk) except WeChatClientException as e: @@ -104,14 +115,17 @@ async def send(self, message: MessageChain) -> None: logger.warning( f"kf API error 40096 for user {user_id}, falling back to regular message API" ) + if event_requests_agent_stop(self): + raise AgentOutputStopped self.client.message.send_text( self.get_self_id(), user_id, chunk ) else: raise - await asyncio.sleep(0.5) # Avoid sending too fast elif isinstance(comp, Image): img_path = await comp.convert_to_file_path() + if event_requests_agent_stop(self): + raise AgentOutputStopped with open(img_path, "rb") as f: try: @@ -130,9 +144,13 @@ async def send(self, message: MessageChain) -> None: ) elif isinstance(comp, Record): record_path = await comp.convert_to_file_path() + if event_requests_agent_stop(self): + raise AgentOutputStopped record_path_amr = await convert_audio_to_amr(record_path) try: + if event_requests_agent_stop(self): + raise AgentOutputStopped with open(record_path_amr, "rb") as f: try: response = self.client.media.upload("voice", f) @@ -160,6 +178,8 @@ async def send(self, message: MessageChain) -> None: logger.warning(f"删除临时音频文件失败: {e}") elif isinstance(comp, File): file_path = await comp.get_file() + if event_requests_agent_stop(self): + raise AgentOutputStopped with open(file_path, "rb") as f: try: @@ -178,6 +198,8 @@ async def send(self, message: MessageChain) -> None: ) elif isinstance(comp, Video): video_path = await comp.convert_to_file_path() + if event_requests_agent_stop(self): + raise AgentOutputStopped with open(video_path, "rb") as f: try: @@ -199,18 +221,27 @@ async def send(self, message: MessageChain) -> None: else: # 企业微信应用 for comp in message.chain: + if event_requests_agent_stop(self): + raise AgentOutputStopped if isinstance(comp, Plain): # Split long text messages if needed plain_chunks = await self.split_plain(comp.text) - for chunk in plain_chunks: + if event_requests_agent_stop(self): + raise AgentOutputStopped + for index, chunk in enumerate(plain_chunks): + if index: + await asyncio.sleep(0.5) + if event_requests_agent_stop(self): + raise AgentOutputStopped self.client.message.send_text( message_obj.self_id, message_obj.session_id, chunk, ) - await asyncio.sleep(0.5) # Avoid sending too fast elif isinstance(comp, Image): img_path = await comp.convert_to_file_path() + if event_requests_agent_stop(self): + raise AgentOutputStopped with open(img_path, "rb") as f: try: @@ -229,9 +260,13 @@ async def send(self, message: MessageChain) -> None: ) elif isinstance(comp, Record): record_path = await comp.convert_to_file_path() + if event_requests_agent_stop(self): + raise AgentOutputStopped record_path_amr = await convert_audio_to_amr(record_path) try: + if event_requests_agent_stop(self): + raise AgentOutputStopped with open(record_path_amr, "rb") as f: try: response = self.client.media.upload("voice", f) @@ -259,6 +294,8 @@ async def send(self, message: MessageChain) -> None: logger.warning(f"删除临时音频文件失败: {e}") elif isinstance(comp, File): file_path = await comp.get_file() + if event_requests_agent_stop(self): + raise AgentOutputStopped with open(file_path, "rb") as f: try: @@ -277,6 +314,8 @@ async def send(self, message: MessageChain) -> None: ) elif isinstance(comp, Video): video_path = await comp.convert_to_file_path() + if event_requests_agent_stop(self): + raise AgentOutputStopped with open(video_path, "rb") as f: try: @@ -301,6 +340,8 @@ async def send(self, message: MessageChain) -> None: async def send_streaming(self, generator, use_fallback: bool = False): buffer = None async for chain in generator: + if event_requests_agent_stop(self): + raise AgentOutputStopped if not buffer: buffer = chain else: @@ -308,5 +349,7 @@ async def send_streaming(self, generator, use_fallback: bool = False): if not buffer: return None buffer.squash_plain() + if event_requests_agent_stop(self): + raise AgentOutputStopped await self.send(buffer) return await super().send_streaming(generator, use_fallback) diff --git a/astrbot/core/platform/sources/wecom_ai_bot/wecomai_event.py b/astrbot/core/platform/sources/wecom_ai_bot/wecomai_event.py index 74d120f5f3..30be6e9053 100644 --- a/astrbot/core/platform/sources/wecom_ai_bot/wecomai_event.py +++ b/astrbot/core/platform/sources/wecom_ai_bot/wecomai_event.py @@ -6,6 +6,7 @@ from astrbot.api import logger from astrbot.api.event import AstrMessageEvent, MessageChain from astrbot.api.message_components import At, Image, Plain +from astrbot.core.agent.stop_policy import AgentOutputStopped, event_requests_agent_stop from .wecomai_api import WecomAIBotAPIClient from .wecomai_queue_mgr import WecomAIQueueMgr @@ -64,9 +65,12 @@ async def _send( queue_mgr: WecomAIQueueMgr, streaming: bool = False, suppress_unsupported_log: bool = False, + stop_event: AstrMessageEvent | None = None, ): back_queue = queue_mgr.get_or_create_back_queue(stream_id) + if event_requests_agent_stop(stop_event): + raise AgentOutputStopped if not message_chain: await back_queue.put( { @@ -79,6 +83,8 @@ async def _send( data = "" for comp in message_chain.chain: + if event_requests_agent_stop(stop_event): + raise AgentOutputStopped if isinstance(comp, At): data = f"@{comp.name} " await back_queue.put( @@ -103,6 +109,8 @@ async def _send( # 处理图片消息 try: image_base64 = await comp.convert_to_base64() + if event_requests_agent_stop(stop_event): + raise AgentOutputStopped if image_base64: await back_queue.put( { @@ -114,8 +122,11 @@ async def _send( ) else: logger.warning("图片数据为空，跳过") + except AgentOutputStopped: + raise except Exception as e: logger.error("处理图片消息失败: %s", e) + raise else: if not suppress_unsupported_log: logger.warning( @@ -151,6 +162,8 @@ async def send(self, message: MessageChain | None) -> None: """发送消息""" if message is None: return + if event_requests_agent_stop(self): + raise AgentOutputStopped raw = self.message_obj.raw_message assert isinstance(raw, dict), ( "wecom_ai_bot platform event raw_message should be a dict" @@ -169,7 +182,7 @@ async def send(self, message: MessageChain | None) -> None: and req_id ): if self.only_use_webhook_url_to_send and self.webhook_client and message: - await self.webhook_client.send_message_chain(message) + await self.webhook_client.send_message_chain(message, stop_event=self) await super().send(MessageChain([])) return @@ -177,10 +190,13 @@ async def send(self, message: MessageChain | None) -> None: await self.webhook_client.send_message_chain( message, unsupported_only=True, + stop_event=self, ) content = self._extract_plain_text_from_chain(message) - await self.long_connection_sender( + if event_requests_agent_stop(self): + raise AgentOutputStopped + sent = await self.long_connection_sender( req_id, { "msgtype": "stream", @@ -191,11 +207,13 @@ async def send(self, message: MessageChain | None) -> None: }, }, ) + if not sent: + raise RuntimeError("WeCom long-connection message delivery failed.") await super().send(MessageChain([])) return if self.only_use_webhook_url_to_send and self.webhook_client and message: - await self.webhook_client.send_message_chain(message) + await self.webhook_client.send_message_chain(message, stop_event=self) await self._mark_stream_complete(stream_id) await super().send(MessageChain([])) return @@ -204,6 +222,7 @@ async def send(self, message: MessageChain | None) -> None: await self.webhook_client.send_message_chain( message, unsupported_only=True, + stop_event=self, ) await WecomAIBotMessageEvent._send( @@ -211,6 +230,7 @@ async def send(self, message: MessageChain | None) -> None: stream_id, self.queue_mgr, suppress_unsupported_log=self.webhook_client is not None, + stop_event=self, ) await super().send(MessageChain([])) @@ -240,7 +260,14 @@ async def send_streaming(self, generator, use_fallback=False) -> None: async for chain in generator: merged_chain.chain.extend(chain.chain) merged_chain.squash_plain() - await self.webhook_client.send_message_chain(merged_chain) + if event_requests_agent_stop(self): + raise AgentOutputStopped + await self.webhook_client.send_message_chain( + merged_chain, + stop_event=self, + ) + if event_requests_agent_stop(self): + return await self.long_connection_sender( req_id, { @@ -262,7 +289,10 @@ async def send_streaming(self, generator, use_fallback=False) -> None: await self.webhook_client.send_message_chain( chain, unsupported_only=True, + stop_event=self, ) + if event_requests_agent_stop(self): + raise AgentOutputStopped chain.squash_plain() # 流式输出不 strip，保留换行等格式字符 @@ -273,7 +303,9 @@ async def send_streaming(self, generator, use_fallback=False) -> None: increment_plain += chunk_text now = asyncio.get_running_loop().time() if now - last_stream_update_time >= self.STREAM_FLUSH_INTERVAL: - await self.long_connection_sender( + if event_requests_agent_stop(self): + raise AgentOutputStopped + sent = await self.long_connection_sender( req_id, { "msgtype": "stream", @@ -284,9 +316,15 @@ async def send_streaming(self, generator, use_fallback=False) -> None: }, }, ) + if not sent: + raise RuntimeError( + "WeCom long-connection message delivery failed." + ) last_stream_update_time = now - await self.long_connection_sender( + if event_requests_agent_stop(self): + raise AgentOutputStopped + sent = await self.long_connection_sender( req_id, { "msgtype": "stream", @@ -297,6 +335,8 @@ async def send_streaming(self, generator, use_fallback=False) -> None: }, }, ) + if not sent: + raise RuntimeError("WeCom long-connection message delivery failed.") await super().send_streaming(generator, use_fallback) return @@ -305,7 +345,14 @@ async def send_streaming(self, generator, use_fallback=False) -> None: async for chain in generator: merged_chain.chain.extend(chain.chain) merged_chain.squash_plain() - await self.webhook_client.send_message_chain(merged_chain) + if event_requests_agent_stop(self): + raise AgentOutputStopped + await self.webhook_client.send_message_chain( + merged_chain, + stop_event=self, + ) + if event_requests_agent_stop(self): + return await self._mark_stream_complete(stream_id) await super().send_streaming(generator, use_fallback) return @@ -317,6 +364,8 @@ async def send_streaming(self, generator, use_fallback=False) -> None: async def enqueue_stream_plain(text: str) -> None: if not text: return + if event_requests_agent_stop(self): + raise AgentOutputStopped await back_queue.put( { "type": "plain", @@ -329,12 +378,18 @@ async def enqueue_stream_plain(text: str) -> None: async for chain in generator: if self.webhook_client: await self.webhook_client.send_message_chain( - chain, unsupported_only=True + chain, + unsupported_only=True, + stop_event=self, ) + if event_requests_agent_stop(self): + raise AgentOutputStopped if chain.type == "break" and final_data: if increment_plain: await enqueue_stream_plain(increment_plain) + if event_requests_agent_stop(self): + raise AgentOutputStopped # 分割符 await back_queue.put( { @@ -367,10 +422,13 @@ async def enqueue_stream_plain(text: str) -> None: queue_mgr=self.queue_mgr, streaming=True, suppress_unsupported_log=self.webhook_client is not None, + stop_event=self, ) await enqueue_stream_plain(increment_plain) + if event_requests_agent_stop(self): + raise AgentOutputStopped await back_queue.put( { "type": "complete", # complete means we return the final result diff --git a/astrbot/core/platform/sources/wecom_ai_bot/wecomai_webhook.py b/astrbot/core/platform/sources/wecom_ai_bot/wecomai_webhook.py index 6f42f264b9..daed7822db 100644 --- a/astrbot/core/platform/sources/wecom_ai_bot/wecomai_webhook.py +++ b/astrbot/core/platform/sources/wecom_ai_bot/wecomai_webhook.py @@ -12,8 +12,9 @@ import aiohttp from astrbot.api import logger -from astrbot.api.event import MessageChain +from astrbot.api.event import AstrMessageEvent, MessageChain from astrbot.api.message_components import At, File, Image, Plain, Record, Video +from astrbot.core.agent.stop_policy import AgentOutputStopped, event_requests_agent_stop from astrbot.core.utils.media_utils import convert_audio_format @@ -78,8 +79,14 @@ async def send_payload(self, payload: dict[str, Any]) -> None: ) logger.debug("企业微信消息推送成功: %s", payload.get("msgtype", "unknown")) - async def send_markdown_v2(self, content: str) -> None: + async def send_markdown_v2( + self, + content: str, + stop_event: AstrMessageEvent | None = None, + ) -> None: for chunk in self._split_markdown_v2_content(content): + if event_requests_agent_stop(stop_event): + raise AgentOutputStopped await self.send_payload( { "msgtype": "markdown_v2", @@ -87,9 +94,15 @@ async def send_markdown_v2(self, content: str) -> None: } ) - async def send_image_base64(self, image_base64: str) -> None: + async def send_image_base64( + self, + image_base64: str, + stop_event: AstrMessageEvent | None = None, + ) -> None: image_bytes = base64.b64decode(image_base64) md5 = hashlib.md5(image_bytes).hexdigest() + if event_requests_agent_stop(stop_event): + raise AgentOutputStopped await self.send_payload( { "msgtype": "image", @@ -101,8 +114,13 @@ async def send_image_base64(self, image_base64: str) -> None: ) async def upload_media( - self, file_path: Path, media_type: Literal["file", "voice"] + self, + file_path: Path, + media_type: Literal["file", "voice"], + stop_event: AstrMessageEvent | None = None, ) -> str: + if event_requests_agent_stop(stop_event): + raise AgentOutputStopped if not file_path.exists() or not file_path.is_file(): raise WecomAIBotWebhookError(f"文件不存在: {file_path}") @@ -138,8 +156,14 @@ async def upload_media( raise WecomAIBotWebhookError("上传媒体失败: 返回缺少 media_id") return str(media_id) - async def send_file(self, file_path: Path) -> None: - media_id = await self.upload_media(file_path, "file") + async def send_file( + self, + file_path: Path, + stop_event: AstrMessageEvent | None = None, + ) -> None: + media_id = await self.upload_media(file_path, "file", stop_event) + if event_requests_agent_stop(stop_event): + raise AgentOutputStopped await self.send_payload( { "msgtype": "file", @@ -147,8 +171,14 @@ async def send_file(self, file_path: Path) -> None: } ) - async def send_voice(self, file_path: Path) -> None: - media_id = await self.upload_media(file_path, "voice") + async def send_voice( + self, + file_path: Path, + stop_event: AstrMessageEvent | None = None, + ) -> None: + media_id = await self.upload_media(file_path, "voice", stop_event) + if event_requests_agent_stop(stop_event): + raise AgentOutputStopped await self.send_payload( { "msgtype": "voice", @@ -164,16 +194,19 @@ async def send_message_chain( self, message_chain: MessageChain, unsupported_only: bool = False, + stop_event: AstrMessageEvent | None = None, ) -> None: async def flush_markdown_buffer(parts: list[str]) -> None: content = "".join(parts).strip() parts.clear() if content: - await self.send_markdown_v2(content) + await self.send_markdown_v2(content, stop_event) markdown_buffer: list[str] = [] for component in message_chain.chain: + if event_requests_agent_stop(stop_event): + raise AgentOutputStopped if unsupported_only and self.is_stream_supported_component(component): continue if isinstance(component, Plain): @@ -184,21 +217,29 @@ async def flush_markdown_buffer(parts: list[str]) -> None: elif isinstance(component, Image): await flush_markdown_buffer(markdown_buffer) image_base64 = await component.convert_to_base64() - await self.send_image_base64(image_base64) + if event_requests_agent_stop(stop_event): + raise AgentOutputStopped + await self.send_image_base64(image_base64, stop_event) elif isinstance(component, File): await flush_markdown_buffer(markdown_buffer) file_path = await component.get_file() + if event_requests_agent_stop(stop_event): + raise AgentOutputStopped if not file_path: logger.warning("文件消息缺少有效文件路径，已跳过: %s", component) continue - await self.send_file(Path(file_path)) + await self.send_file(Path(file_path), stop_event) elif isinstance(component, Video): await flush_markdown_buffer(markdown_buffer) video_path = await component.convert_to_file_path() - await self.send_file(Path(video_path)) + if event_requests_agent_stop(stop_event): + raise AgentOutputStopped + await self.send_file(Path(video_path), stop_event) elif isinstance(component, Record): await flush_markdown_buffer(markdown_buffer) source_voice_path = Path(await component.convert_to_file_path()) + if event_requests_agent_stop(stop_event): + raise AgentOutputStopped target_voice_path = source_voice_path converted = False if source_voice_path.suffix.lower() != ".amr": @@ -207,7 +248,9 @@ async def flush_markdown_buffer(parts: list[str]) -> None: ) converted = target_voice_path != source_voice_path try: - await self.send_voice(target_voice_path) + if event_requests_agent_stop(stop_event): + raise AgentOutputStopped + await self.send_voice(target_voice_path, stop_event) finally: if converted and target_voice_path.exists(): try: diff --git a/astrbot/core/platform/sources/weixin_oc/weixin_oc_adapter.py b/astrbot/core/platform/sources/weixin_oc/weixin_oc_adapter.py index 5bbf24a72a..7bef6bdb99 100644 --- a/astrbot/core/platform/sources/weixin_oc/weixin_oc_adapter.py +++ b/astrbot/core/platform/sources/weixin_oc/weixin_oc_adapter.py @@ -16,7 +16,7 @@ import qrcode as qrcode_lib from astrbot import logger -from astrbot.api.event import MessageChain +from astrbot.api.event import AstrMessageEvent, MessageChain from astrbot.api.message_components import File, Image, Plain, Record, Reply, Video from astrbot.api.platform import ( AstrBotMessage, @@ -27,6 +27,7 @@ register_platform_adapter, ) from astrbot.core import astrbot_config +from astrbot.core.agent.stop_policy import AgentOutputStopped, event_requests_agent_stop from astrbot.core.platform.astr_message_event import MessageSesion from astrbot.core.utils.astrbot_path import get_astrbot_temp_path from astrbot.core.utils.media_utils import ( @@ -642,7 +643,10 @@ async def _prepare_media_item( upload_media_type: int, item_type: int, file_name: str, + stop_event: AstrMessageEvent | None = None, ) -> dict[str, Any]: + if event_requests_agent_stop(stop_event): + raise AgentOutputStopped raw_bytes = media_path.read_bytes() raw_size = len(raw_bytes) raw_md5 = hashlib.md5(raw_bytes).hexdigest() @@ -669,6 +673,8 @@ async def _prepare_media_item( token_required=True, timeout_ms=self.api_timeout_ms, ) + if event_requests_agent_stop(stop_event): + raise AgentOutputStopped logger.debug( "weixin_oc(%s): getuploadurl response user=%s media_type=%s raw_size=%s raw_md5=%s filekey=%s file=%s upload_param_len=%s", self.meta().id, @@ -690,6 +696,8 @@ async def _prepare_media_item( aes_key_hex, media_path, ) + if event_requests_agent_stop(stop_event): + raise AgentOutputStopped logger.debug( "weixin_oc(%s): prepared media item type=%s file=%s user=%s mid_size=%s upload_param_len=%s query_len=%s", self.meta().id, @@ -843,8 +851,12 @@ async def _resolve_inbound_media_component( return None async def _resolve_media_file_path( - self, segment: Image | Video | File + self, + segment: Image | Video | File, + stop_event: AstrMessageEvent | None = None, ) -> Path | None: + if event_requests_agent_stop(stop_event): + raise AgentOutputStopped try: if isinstance(segment, File): path = await segment.get_file() @@ -856,6 +868,9 @@ async def _resolve_media_file_path( logger.warning("weixin_oc(%s): media resolve failed: %s", self.meta().id, e) return None + if event_requests_agent_stop(stop_event): + raise AgentOutputStopped + if not path: return None media_path = Path(path) @@ -870,7 +885,10 @@ async def _send_items_to_session( *, cache_components: list[Any] | None = None, cache_message_str: str | None = None, + stop_event: AstrMessageEvent | None = None, ) -> bool: + if event_requests_agent_stop(stop_event): + raise AgentOutputStopped if not self.token: logger.warning("weixin_oc(%s): missing token, skip send", self.meta().id) return False @@ -888,6 +906,8 @@ async def _send_items_to_session( user_id, ) return False + if event_requests_agent_stop(stop_event): + raise AgentOutputStopped payload = await self.client.request_json( "POST", "ilink/bot/sendmessage", @@ -989,13 +1009,16 @@ async def _send_media_segment( user_id: str, segment: Image | Video | File, text: str | None = None, + stop_event: AstrMessageEvent | None = None, ) -> bool: + if event_requests_agent_stop(stop_event): + raise AgentOutputStopped if not self.token: logger.warning( "weixin_oc(%s): missing token, skip media send", self.meta().id ) return False - media_path = await self._resolve_media_file_path(segment) + media_path = await self._resolve_media_file_path(segment, stop_event) if media_path is None: logger.warning( "weixin_oc(%s): skip media segment, media file not resolvable", @@ -1024,7 +1047,10 @@ async def _send_media_segment( upload_media_type, item_type, file_name, + stop_event, ) + except AgentOutputStopped: + raise except Exception as e: logger.error( "weixin_oc(%s): prepare media failed: %s", @@ -1035,12 +1061,19 @@ async def _send_media_segment( return False if text: - await self._send_items_to_session( + if event_requests_agent_stop(stop_event): + raise AgentOutputStopped + text_sent = await self._send_items_to_session( user_id, [self._build_plain_text_item(text)], cache_components=[Plain(text)], cache_message_str=text, + stop_event=stop_event, ) + if not text_sent: + return False + if event_requests_agent_stop(stop_event): + raise AgentOutputStopped return await self._send_items_to_session( user_id, [media_item], @@ -1049,6 +1082,7 @@ async def _send_media_segment( [media_item], include_ref_text=False, ), + stop_event=stop_event, ) async def _start_login_session(self) -> OpenClawLoginSession: @@ -1619,7 +1653,11 @@ def _message_chain_to_text(self, message_chain: MessageChain) -> str: return text.strip() async def _send_to_session( - self, user_id: str, text: str, _components: list[Any] | None = None + self, + user_id: str, + text: str, + _components: list[Any] | None = None, + stop_event: AstrMessageEvent | None = None, ) -> bool: if not text: text = self._message_chain_to_text(MessageChain(_components or [])) @@ -1632,18 +1670,24 @@ async def _send_to_session( return await self._send_items_to_session( user_id, [self._build_plain_text_item(text)], + stop_event=stop_event, ) async def send_by_session( self, session: MessageSesion, message_chain: MessageChain, + stop_event: AstrMessageEvent | None = None, ) -> None: + if event_requests_agent_stop(stop_event): + raise AgentOutputStopped target_user = session.session_id pending_text = "" has_supported_segment = False failed_segments = 0 for segment in message_chain.chain: + if event_requests_agent_stop(stop_event): + raise AgentOutputStopped if isinstance(segment, Plain): pending_text += segment.text continue @@ -1654,6 +1698,7 @@ async def send_by_session( target_user, segment, text=pending_text.strip() or None, + stop_event=stop_event, ) if not sent: failed_segments += 1 @@ -1668,7 +1713,11 @@ async def send_by_session( if pending_text: has_supported_segment = True - sent = await self._send_to_session(target_user, pending_text.strip()) + sent = await self._send_to_session( + target_user, + pending_text.strip(), + stop_event=stop_event, + ) if not sent: failed_segments += 1 diff --git a/astrbot/core/platform/sources/weixin_oc/weixin_oc_event.py b/astrbot/core/platform/sources/weixin_oc/weixin_oc_event.py index 84a19a9e7b..e8208f2b83 100644 --- a/astrbot/core/platform/sources/weixin_oc/weixin_oc_event.py +++ b/astrbot/core/platform/sources/weixin_oc/weixin_oc_event.py @@ -1,6 +1,5 @@ from __future__ import annotations -import asyncio import uuid from typing import TYPE_CHECKING @@ -14,6 +13,7 @@ Record, Video, ) +from astrbot.core.agent.stop_policy import AgentOutputStopped, event_requests_agent_stop if TYPE_CHECKING: # pragma: no cover - typing helper from .weixin_oc_adapter import WeixinOCAdapter @@ -60,9 +60,11 @@ def _build_plain_text(message: MessageChain) -> str: ) async def send(self, message: MessageChain) -> None: + if event_requests_agent_stop(self): + raise AgentOutputStopped if not message.chain: return - await self.platform.send_by_session(self.session, message) + await self.platform.send_by_session(self.session, message, stop_event=self) await super().send(message) async def send_typing(self) -> None: @@ -81,6 +83,8 @@ async def send_streaming(self, generator, use_fallback: bool = False): if not use_fallback: buffer = None async for chain in generator: + if event_requests_agent_stop(self): + raise AgentOutputStopped if not buffer: buffer = chain else: @@ -92,12 +96,15 @@ async def send_streaming(self, generator, use_fallback: bool = False): buffer = "" async for chain in generator: + if event_requests_agent_stop(self): + raise AgentOutputStopped if not isinstance(chain, MessageChain): continue for component in chain.chain: + if event_requests_agent_stop(self): + raise AgentOutputStopped if not isinstance(component, Plain): - await self.send(MessageChain(chain=[component])) - await asyncio.sleep(1.2) + await self.send_streaming_fallback_component(component, delay=1.2) continue buffer += component.text if buffer.strip(): diff --git a/astrbot/core/platform/sources/weixin_official_account/weixin_offacc_event.py b/astrbot/core/platform/sources/weixin_official_account/weixin_offacc_event.py index ae536593c5..f70e19ffd0 100644 --- a/astrbot/core/platform/sources/weixin_official_account/weixin_offacc_event.py +++ b/astrbot/core/platform/sources/weixin_official_account/weixin_offacc_event.py @@ -9,6 +9,7 @@ from astrbot.api.event import AstrMessageEvent, MessageChain from astrbot.api.message_components import Image, Plain, Record from astrbot.api.platform import AstrBotMessage, PlatformMetadata +from astrbot.core.agent.stop_policy import AgentOutputStopped, event_requests_agent_stop from astrbot.core.utils.media_utils import convert_audio_to_amr @@ -81,25 +82,37 @@ async def split_plain(self, plain: str, max_length: int = 1024) -> list[str]: return result async def send(self, message: MessageChain) -> None: + if event_requests_agent_stop(self): + raise AgentOutputStopped message_obj = self.message_obj active_send_mode = cast(dict, message_obj.raw_message).get( "active_send_mode", False ) for comp in message.chain: + if event_requests_agent_stop(self): + raise AgentOutputStopped if isinstance(comp, Plain): # Split long text messages if needed plain_chunks = await self.split_plain(comp.text) + if event_requests_agent_stop(self): + raise AgentOutputStopped if active_send_mode: for chunk in plain_chunks: + if event_requests_agent_stop(self): + raise AgentOutputStopped self.client.message.send_text(message_obj.sender.user_id, chunk) else: # disable passive sending, just store the chunks in logger.debug( f"split plain into {len(plain_chunks)} chunks for passive reply. Message not sent." ) + if event_requests_agent_stop(self): + raise AgentOutputStopped self.message_out["cached_xml"] = plain_chunks elif isinstance(comp, Image): img_path = await comp.convert_to_file_path() + if event_requests_agent_stop(self): + raise AgentOutputStopped with open(img_path, "rb") as f: try: @@ -111,6 +124,8 @@ async def send(self, message: MessageChain) -> None: ) return logger.debug(f"微信公众平台上传图片返回: {response}") + if event_requests_agent_stop(self): + raise AgentOutputStopped if active_send_mode: self.client.message.send_image( @@ -129,9 +144,13 @@ async def send(self, message: MessageChain) -> None: elif isinstance(comp, Record): record_path = await comp.convert_to_file_path() + if event_requests_agent_stop(self): + raise AgentOutputStopped record_path_amr = await convert_audio_to_amr(record_path) try: + if event_requests_agent_stop(self): + raise AgentOutputStopped with open(record_path_amr, "rb") as f: try: response = self.client.media.upload("voice", f) @@ -144,6 +163,8 @@ async def send(self, message: MessageChain) -> None: ) return logger.info(f"微信公众平台上传语音返回: {response}") + if event_requests_agent_stop(self): + raise AgentOutputStopped if active_send_mode: self.client.message.send_voice( @@ -178,6 +199,8 @@ async def send(self, message: MessageChain) -> None: async def send_streaming(self, generator, use_fallback: bool = False): buffer = None async for chain in generator: + if event_requests_agent_stop(self): + raise AgentOutputStopped if not buffer: buffer = chain else: @@ -185,5 +208,7 @@ async def send_streaming(self, generator, use_fallback: bool = False): if not buffer: return None buffer.squash_plain() + if event_requests_agent_stop(self): + raise AgentOutputStopped await self.send(buffer) return await super().send_streaming(generator, use_fallback) diff --git a/astrbot/core/provider/provider.py b/astrbot/core/provider/provider.py index 0cc9f1ca1c..8a5531f7be 100644 --- a/astrbot/core/provider/provider.py +++ b/astrbot/core/provider/provider.py @@ -5,6 +5,7 @@ from typing import Literal, TypeAlias, Union from astrbot.core.agent.message import ContentPart, Message, is_checkpoint_message +from astrbot.core.agent.stop_policy import event_requests_agent_stop from astrbot.core.agent.tool import ToolSet from astrbot.core.provider.entities import ( LLMResponse, @@ -279,14 +280,19 @@ async def get_audio_stream( if text_part is None: # 输入结束，处理累积的文本 - if accumulated_text: + astr_event = getattr(text_queue, "astr_event", None) + if accumulated_text and not event_requests_agent_stop(astr_event): try: # 调用原有的 get_audio 方法获取音频文件路径 audio_path = await self.get_audio(accumulated_text) - # 读取音频文件内容 - with open(audio_path, "rb") as f: - audio_data = f.read() - await audio_queue.put((accumulated_text, audio_data)) + if audio_path and astr_event is not None: + astr_event.track_temporary_local_file(audio_path) + if audio_path and not event_requests_agent_stop(astr_event): + # 读取音频文件内容 + with open(audio_path, "rb") as f: + audio_data = f.read() + if not event_requests_agent_stop(astr_event): + await audio_queue.put((accumulated_text, audio_data)) except Exception: # 出错时也要发送 None 结束标记 pass diff --git a/astrbot/core/provider/sources/genie_tts.py b/astrbot/core/provider/sources/genie_tts.py index b76bf6b465..8196e44361 100644 --- a/astrbot/core/provider/sources/genie_tts.py +++ b/astrbot/core/provider/sources/genie_tts.py @@ -1,8 +1,11 @@ import asyncio import os +import threading import uuid +from contextlib import suppress from astrbot.core import logger +from astrbot.core.agent.stop_policy import event_requests_agent_stop from astrbot.core.provider.entities import ProviderType from astrbot.core.provider.provider import TTSProvider from astrbot.core.provider.register import register_provider_adapter @@ -60,24 +63,41 @@ async def get_audio(self, text: str) -> str: path = os.path.join(temp_dir, filename) loop = asyncio.get_running_loop() + cleanup_requested = threading.Event() def _generate(save_path: str) -> None: assert genie is not None - genie.tts( - character_name=self.character_name, - text=text, - save_path=save_path, - ) + try: + genie.tts( + character_name=self.character_name, + text=text, + save_path=save_path, + ) + finally: + if cleanup_requested.is_set(): + with suppress(OSError): + os.remove(save_path) try: - await loop.run_in_executor(None, _generate, path) + generation_future = loop.run_in_executor(None, _generate, path) + await asyncio.shield(generation_future) if os.path.exists(path): return path raise RuntimeError("Genie TTS did not save to file.") + except asyncio.CancelledError: + cleanup_requested.set() + generation_future.add_done_callback( + lambda future: future.exception() if not future.cancelled() else None + ) + with suppress(OSError): + os.remove(path) + raise except Exception as e: + with suppress(OSError): + os.remove(path) raise RuntimeError(f"Genie TTS generation failed: {e}") async def get_audio_stream( @@ -86,6 +106,7 @@ async def get_audio_stream( audio_queue: "asyncio.Queue[bytes | tuple[str, bytes] | None]", ) -> None: loop = asyncio.get_running_loop() + astr_event = getattr(text_queue, "astr_event", None) while True: text = await text_queue.get() @@ -98,29 +119,54 @@ async def get_audio_stream( os.makedirs(temp_dir, exist_ok=True) filename = f"genie_tts_{uuid.uuid4()}.wav" path = os.path.join(temp_dir, filename) + cleanup_requested = threading.Event() def _generate(save_path: str, t: str) -> None: assert genie is not None - genie.tts( - character_name=self.character_name, - text=t, - save_path=save_path, + try: + genie.tts( + character_name=self.character_name, + text=t, + save_path=save_path, + ) + finally: + if cleanup_requested.is_set(): + with suppress(OSError): + os.remove(save_path) + + if astr_event is not None: + astr_event.track_temporary_local_file(path) + generation_future = loop.run_in_executor(None, _generate, path, text) + try: + await asyncio.shield(generation_future) + except asyncio.CancelledError: + # The worker may outlive this coroutine, so let it remove any + # file created after cancellation without delaying shutdown. + cleanup_requested.set() + generation_future.add_done_callback( + lambda future: ( + future.exception() if not future.cancelled() else None + ) ) - - await loop.run_in_executor(None, _generate, path, text) + with suppress(OSError): + os.remove(path) + raise if os.path.exists(path): - with open(path, "rb") as f: - audio_data = f.read() - - # Put (text, bytes) into queue so frontend can display text - await audio_queue.put((text, audio_data)) - - # Clean up try: - os.remove(path) - except OSError: - pass + if event_requests_agent_stop(astr_event): + continue + with open(path, "rb") as f: + audio_data = f.read() + if not event_requests_agent_stop(astr_event): + # Put (text, bytes) into queue so frontend can display text + await audio_queue.put((text, audio_data)) + finally: + # The event cleanup remains a fallback if immediate removal fails. + try: + os.remove(path) + except OSError: + pass else: logger.error(f"Genie TTS failed to generate audio for: {text}") diff --git a/astrbot/core/utils/media_utils.py b/astrbot/core/utils/media_utils.py index a3990775ac..3746e009da 100644 --- a/astrbot/core/utils/media_utils.py +++ b/astrbot/core/utils/media_utils.py @@ -1130,10 +1130,25 @@ async def convert_audio_format( if audio_path.lower().endswith(f".{output_format}"): return audio_path + replace_output_path = None + replace_target_path = None if output_path is None: temp_dir = Path(get_astrbot_temp_path()) temp_dir.mkdir(parents=True, exist_ok=True) output_path = str(temp_dir / f"media_audio_{uuid.uuid4().hex}.{output_format}") + elif os.path.lexists(output_path): + replace_output_path = output_path + requested_path = Path(output_path) + replace_target_path = ( + requested_path.resolve(strict=False) + if requested_path.is_symlink() + else requested_path + ) + output_path = str( + replace_target_path.with_name( + f".{requested_path.stem}.{uuid.uuid4().hex}{requested_path.suffix}" + ) + ) args = ["ffmpeg", "-y", "-i", audio_path] if output_format == "amr": @@ -1161,6 +1176,16 @@ async def convert_audio_format( args.extend(["-acodec", "libopus", "-ac", "1", "-ar", "16000"]) args.append(output_path) + def _remove_incomplete_output() -> None: + """Remove the converter-owned output after an interrupted run.""" + if not output_path or not os.path.exists(output_path): + return + try: + os.remove(output_path) + except OSError as exc: + logger.warning("Failed to remove incomplete audio output file: %s", exc) + + process = None try: process = await asyncio.create_subprocess_exec( *args, @@ -1169,24 +1194,46 @@ async def convert_audio_format( ) _, stderr = await process.communicate() if process.returncode != 0: - if output_path and os.path.exists(output_path): - try: - os.remove(output_path) - except OSError as e: - logger.warning( - "Failed to remove failed audio output file: %s", - e, - ) error_msg = stderr.decode() if stderr else "unknown error" raise Exception(f"ffmpeg conversion failed: {error_msg}") + if replace_output_path is not None: + os.replace(output_path, replace_target_path) + output_path = replace_output_path logger.debug( "Audio converted successfully: %s -> %s", audio_path, output_path, ) return output_path - except FileNotFoundError: - raise Exception("ffmpeg not found") + except asyncio.CancelledError as cancel_error: + if process is not None and getattr(process, "returncode", None) is None: + kill = getattr(process, "kill", None) + if callable(kill): + try: + kill() + except ProcessLookupError: + pass + wait = getattr(process, "wait", None) + if callable(wait): + wait_task = asyncio.create_task(wait()) + while not wait_task.done(): + try: + await asyncio.shield(wait_task) + except asyncio.CancelledError: + continue + if not wait_task.cancelled(): + try: + wait_task.result() + except Exception: + pass + _remove_incomplete_output() + raise cancel_error + except FileNotFoundError as exc: + _remove_incomplete_output() + raise Exception("ffmpeg not found") from exc + except Exception: + _remove_incomplete_output() + raise async def convert_audio_to_amr(audio_path: str, output_path: str | None = None) -> str: diff --git a/tests/test_dingtalk_adapter.py b/tests/test_dingtalk_adapter.py index aa1e638c8c..dee3a0b0bd 100644 --- a/tests/test_dingtalk_adapter.py +++ b/tests/test_dingtalk_adapter.py @@ -1,8 +1,14 @@ import asyncio import threading +from types import SimpleNamespace +from unittest.mock import AsyncMock import pytest +from astrbot.api.event import MessageChain +from astrbot.api.message_components import File, Record, Video +from astrbot.api.platform import MessageType +from astrbot.core.agent.stop_policy import AgentOutputStopped from astrbot.core.platform.sources.dingtalk import dingtalk_adapter from astrbot.core.platform.sources.dingtalk.dingtalk_adapter import ( DINGTALK_RECONNECT_INITIAL_DELAY, @@ -10,6 +16,9 @@ DingtalkPlatformAdapter, _dingtalk_reconnect_delay, ) +from astrbot.core.platform.sources.dingtalk.dingtalk_event import ( + DingtalkMessageEvent, +) def test_dingtalk_reconnect_delay_uses_exponential_backoff(): @@ -30,6 +39,279 @@ def test_dingtalk_reconnect_delay_is_capped(): assert _dingtalk_reconnect_delay(20) == DINGTALK_RECONNECT_MAX_DELAY +@pytest.mark.asyncio +async def test_dingtalk_stop_after_staff_lookup_blocks_message_write(): + adapter = DingtalkPlatformAdapter.__new__(DingtalkPlatformAdapter) + adapter.client_id = "bot" + adapter._id_to_sid = lambda value: value + adapter.meta = lambda: SimpleNamespace(id="dingtalk") + lookup_started = asyncio.Event() + release_lookup = asyncio.Event() + + async def get_staff_id(_session): + lookup_started.set() + await release_lookup.wait() + return "staff" + + adapter._get_sender_staff_id = get_staff_id + adapter.send_message_chain_to_user = AsyncMock() + extras = {} + stop_event = SimpleNamespace( + is_stopped=lambda: False, + get_extra=lambda key, default=None: extras.get(key, default), + set_extra=lambda key, value: extras.__setitem__(key, value), + ) + incoming = SimpleNamespace( + sender_id="user", + sender_staff_id="", + conversation_type="1", + ) + + task = asyncio.create_task( + adapter.send_message_chain_with_incoming( + incoming, + MessageChain().message("answer"), + stop_event=stop_event, + ) + ) + await asyncio.wait_for(lookup_started.wait(), timeout=1) + stop_event.set_extra("agent_stop_requested", True) + release_lookup.set() + + with pytest.raises(AgentOutputStopped): + await task + adapter.send_message_chain_to_user.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_dingtalk_private_session_requires_staff_id_mapping(): + adapter = DingtalkPlatformAdapter.__new__(DingtalkPlatformAdapter) + adapter.client_id = "bot" + adapter._get_sender_staff_id = AsyncMock(return_value="") + adapter.send_message_chain_to_user = AsyncMock() + session = SimpleNamespace( + message_type=MessageType.FRIEND_MESSAGE, + session_id="opaque-session", + ) + + with pytest.raises(RuntimeError, match="missing a staff_id mapping"): + await adapter.send_by_session(session, MessageChain().message("answer")) + + adapter.send_message_chain_to_user.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_dingtalk_empty_chain_is_not_reported_as_delivered(): + adapter = DingtalkPlatformAdapter.__new__(DingtalkPlatformAdapter) + + with pytest.raises(RuntimeError, match="no platform delivery"): + await adapter._send_message_chain( + "group", + "conversation", + "robot", + MessageChain(), + ) + + +@pytest.mark.asyncio +async def test_dingtalk_empty_stream_is_not_reported_as_delivered(): + event = object.__new__(DingtalkMessageEvent) + event._extras = {} + event._force_stopped = False + event._result = None + + async def empty_source(): + if False: + yield MessageChain().message("never") + + with pytest.raises(RuntimeError, match="produced no delivery"): + await event.send_streaming(empty_source()) + + +@pytest.mark.asyncio +async def test_dingtalk_stop_during_token_lookup_blocks_http_write(monkeypatch): + adapter = DingtalkPlatformAdapter.__new__(DingtalkPlatformAdapter) + token_started = asyncio.Event() + release_token = asyncio.Event() + post_calls = 0 + extras = {} + stop_event = SimpleNamespace( + is_stopped=lambda: False, + get_extra=lambda key, default=None: extras.get(key, default), + ) + + async def get_access_token(_stop_event=None): + token_started.set() + await release_token.wait() + return "token" + + class Session: + async def __aenter__(self): + return self + + async def __aexit__(self, *_args): + return False + + def post(self, *_args, **_kwargs): + nonlocal post_calls + post_calls += 1 + raise AssertionError("DingTalk POST must not start after stop") + + adapter.get_access_token = get_access_token + monkeypatch.setattr(dingtalk_adapter.aiohttp, "ClientSession", Session) + task = asyncio.create_task( + adapter._send_message_chain( + "group", + "conversation", + "robot", + MessageChain().message("late"), + stop_event=stop_event, + ) + ) + await asyncio.wait_for(token_started.wait(), timeout=1) + extras["agent_stop_requested"] = True + release_token.set() + + with pytest.raises(AgentOutputStopped): + await task + assert post_calls == 0 + + +@pytest.mark.parametrize( + ("status", "response_text"), + [ + (500, "rejected"), + (200, '{"errcode": 123, "errmsg": "rejected"}'), + ], +) +@pytest.mark.asyncio +async def test_dingtalk_rejected_message_raises_delivery_error( + monkeypatch, + status, + response_text, +): + adapter = DingtalkPlatformAdapter.__new__(DingtalkPlatformAdapter) + adapter.get_access_token = AsyncMock(return_value="token") + + class Response: + async def __aenter__(self): + return self + + async def __aexit__(self, *_args): + return False + + async def text(self): + return response_text + + Response.status = status + + class Session: + async def __aenter__(self): + return self + + async def __aexit__(self, *_args): + return False + + def post(self, *_args, **_kwargs): + return Response() + + monkeypatch.setattr(dingtalk_adapter.aiohttp, "ClientSession", Session) + + with pytest.raises(RuntimeError, match="delivery failed"): + await adapter._send_message_chain( + "group", + "conversation", + "robot", + MessageChain().message("required"), + ) + + +@pytest.mark.asyncio +async def test_dingtalk_successful_message_is_still_accepted(monkeypatch): + adapter = DingtalkPlatformAdapter.__new__(DingtalkPlatformAdapter) + adapter.get_access_token = AsyncMock(return_value="token") + + class Response: + status = 200 + + async def __aenter__(self): + return self + + async def __aexit__(self, *_args): + return False + + async def text(self): + return '{"processQueryKey": "query"}' + + class Session: + async def __aenter__(self): + return self + + async def __aexit__(self, *_args): + return False + + def post(self, *_args, **_kwargs): + return Response() + + monkeypatch.setattr(dingtalk_adapter.aiohttp, "ClientSession", Session) + + await adapter._send_message_chain( + "group", + "conversation", + "robot", + MessageChain().message("delivered"), + ) + + +@pytest.mark.parametrize("component_type", ["record", "video", "file"]) +@pytest.mark.asyncio +async def test_dingtalk_required_media_failure_is_not_swallowed( + monkeypatch, + component_type, +): + adapter = DingtalkPlatformAdapter.__new__(DingtalkPlatformAdapter) + adapter.upload_media = AsyncMock(side_effect=RuntimeError("upload failed")) + adapter._safe_remove_file = lambda _path: None + + if component_type == "record": + component = Record(file="source.wav") + monkeypatch.setattr( + Record, + "convert_to_file_path", + AsyncMock(return_value="/tmp/source.wav"), + ) + adapter._prepare_voice_for_dingtalk = AsyncMock( + return_value=("/tmp/source.ogg", False) + ) + elif component_type == "video": + component = Video(file="source.mp4") + monkeypatch.setattr( + Video, + "convert_to_file_path", + AsyncMock(return_value="/tmp/source.mp4"), + ) + monkeypatch.setattr( + dingtalk_adapter, + "extract_video_cover", + AsyncMock(return_value="/tmp/cover.jpg"), + ) + else: + component = File(name="source.txt", file="source.txt") + monkeypatch.setattr( + File, + "get_file", + AsyncMock(return_value="/tmp/source.txt"), + ) + + with pytest.raises(RuntimeError, match="upload failed"): + await adapter._send_message_chain( + "group", + "conversation", + "robot", + MessageChain([component]), + ) + + @pytest.mark.asyncio async def test_dingtalk_reconnect_delay_wakes_on_terminate(monkeypatch): class ObservedEvent: diff --git a/tests/test_discord_adapter.py b/tests/test_discord_adapter.py index 0d0e281a8d..982276bd56 100644 --- a/tests/test_discord_adapter.py +++ b/tests/test_discord_adapter.py @@ -1,10 +1,13 @@ +import asyncio import base64 from io import BytesIO from types import SimpleNamespace +from unittest.mock import AsyncMock import pytest from astrbot.api.message_components import Image, Record +from astrbot.core.agent.stop_policy import AgentOutputStopped from astrbot.core.message.message_event_result import MessageChain from astrbot.core.platform.sources.discord import ( discord_platform_adapter, @@ -130,3 +133,45 @@ def __init__(self, fp: BytesIO, filename: str) -> None: assert view is None assert embeds == [] assert reference_message_id is None + + +@pytest.mark.asyncio +async def test_discord_stop_after_conversion_blocks_webhook_write(): + event = DiscordPlatformEvent.__new__(DiscordPlatformEvent) + event._extras = {} + event._force_stopped = False + event._result = None + parse_started = asyncio.Event() + release_parse = asyncio.Event() + + async def parse(_message, _stop_event=None): + parse_started.set() + await release_parse.wait() + return "answer", [], None, [], None + + event._parse_to_discord = parse + event.interaction_followup_webhook = SimpleNamespace(send=AsyncMock()) + + task = asyncio.create_task(event.send(MessageChain().message("answer"))) + await asyncio.wait_for(parse_started.wait(), timeout=1) + event.set_extra("agent_stop_requested", True) + release_parse.set() + + with pytest.raises(AgentOutputStopped): + await task + event.interaction_followup_webhook.send.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_discord_platform_failure_is_reported_to_respond_stage(): + event = DiscordPlatformEvent.__new__(DiscordPlatformEvent) + event._extras = {} + event._force_stopped = False + event._result = None + event._parse_to_discord = AsyncMock(return_value=("answer", [], None, [], None)) + event.interaction_followup_webhook = SimpleNamespace( + send=AsyncMock(side_effect=RuntimeError("network failed")) + ) + + with pytest.raises(RuntimeError, match="delivery failed"): + await event.send(MessageChain().message("answer")) diff --git a/tests/test_lark_app_registration.py b/tests/test_lark_app_registration.py index b7e9b0bad5..0e452f6c5a 100644 --- a/tests/test_lark_app_registration.py +++ b/tests/test_lark_app_registration.py @@ -1,9 +1,19 @@ +import asyncio +from types import SimpleNamespace +from unittest.mock import AsyncMock + +import pytest + +from astrbot.api.event import MessageChain +from astrbot.api.message_components import Json, Plain +from astrbot.core.agent.stop_policy import AgentOutputStopped from astrbot.core.platform.sources.lark.app_registration import ( DEFAULT_FEISHU_OPEN_DOMAIN, DEFAULT_LARK_OPEN_DOMAIN, _registration_data, resolve_app_registration_endpoints, ) +from astrbot.core.platform.sources.lark.lark_event import LarkMessageEvent def test_resolve_app_registration_endpoints_uses_feishu_accounts_domain(): @@ -30,3 +40,133 @@ def test_registration_data_accepts_wrapped_and_plain_payloads(): assert _registration_data(wrapped) == {"device_code": "device"} assert _registration_data(plain) == {"device_code": "device"} + + +@pytest.mark.asyncio +async def test_lark_file_upload_stop_blocks_message_send(monkeypatch): + extras = {} + event = SimpleNamespace( + is_stopped=lambda: False, + get_extra=lambda key, default=None: extras.get(key, default), + ) + + async def upload_then_stop(*_args, **_kwargs): + extras["agent_stop_requested"] = True + return "uploaded-file-key" + + send_message = AsyncMock() + monkeypatch.setattr(LarkMessageEvent, "_upload_lark_file", upload_then_stop) + monkeypatch.setattr(LarkMessageEvent, "_send_im_message", send_message) + + with pytest.raises(AgentOutputStopped): + await LarkMessageEvent._send_file_message( + SimpleNamespace(file="file.txt"), + SimpleNamespace(), + stop_event=event, + ) + + send_message.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_lark_card_creation_stop_blocks_message_send(monkeypatch): + extras = {} + card_started = asyncio.Event() + release_card = asyncio.Event() + event = SimpleNamespace( + is_stopped=lambda: False, + get_extra=lambda key, default=None: extras.get(key, default), + ) + + class CardResult: + data = SimpleNamespace(card_id="card-id") + + @staticmethod + def success(): + return True + + async def create_card(_request): + card_started.set() + await release_card.wait() + return CardResult() + + send_message = AsyncMock(return_value=True) + monkeypatch.setattr(LarkMessageEvent, "_send_im_message", send_message) + client = SimpleNamespace( + im=object(), + cardkit=SimpleNamespace( + v1=SimpleNamespace(card=SimpleNamespace(acreate=create_card)) + ), + ) + task = asyncio.create_task( + LarkMessageEvent.send_message_chain( + MessageChain( + [ + Json( + data={ + "type": "lark_collapsible_panel_reasoning", + "content": "secret", + } + ) + ] + ), + client, + stop_event=event, + ) + ) + await asyncio.wait_for(card_started.wait(), timeout=1) + extras["agent_stop_requested"] = True + release_card.set() + + with pytest.raises(AgentOutputStopped): + await task + send_message.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_lark_audio_path_stop_blocks_conversion(monkeypatch, tmp_path): + extras = {} + event = SimpleNamespace( + is_stopped=lambda: False, + get_extra=lambda key, default=None: extras.get(key, default), + ) + audio_path = tmp_path / "audio.wav" + audio_path.write_bytes(b"audio") + + async def resolve_then_stop(): + extras["agent_stop_requested"] = True + return str(audio_path) + + convert_audio = AsyncMock() + monkeypatch.setattr( + "astrbot.core.platform.sources.lark.lark_event.convert_audio_to_opus", + convert_audio, + ) + + with pytest.raises(AgentOutputStopped): + await LarkMessageEvent._send_audio_message( + SimpleNamespace(convert_to_file_path=resolve_then_stop), + SimpleNamespace(), + stop_event=event, + ) + convert_audio.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_lark_rejected_post_is_reported_as_delivery_failure(monkeypatch): + monkeypatch.setattr( + LarkMessageEvent, + "_convert_to_lark", + AsyncMock(return_value=[[{"tag": "md", "text": "answer"}]]), + ) + monkeypatch.setattr( + LarkMessageEvent, + "_send_im_message", + AsyncMock(return_value=False), + ) + + with pytest.raises(RuntimeError, match="delivery failed"): + await LarkMessageEvent.send_message_chain( + MessageChain([Plain("answer")]), + SimpleNamespace(im=object()), + ) diff --git a/tests/test_platform_stop_delivery.py b/tests/test_platform_stop_delivery.py new file mode 100644 index 0000000000..2f0418ce68 --- /dev/null +++ b/tests/test_platform_stop_delivery.py @@ -0,0 +1,610 @@ +import asyncio +import threading +from pathlib import Path +from types import SimpleNamespace +from unittest.mock import AsyncMock, MagicMock + +import pytest + +from astrbot.api.event import MessageChain +from astrbot.api.message_components import File, Image, Record +from astrbot.core.agent.stop_policy import AgentOutputStopped +from astrbot.core.platform.sources.satori.satori_event import SatoriPlatformEvent +from astrbot.core.platform.sources.webchat import webchat_event +from astrbot.core.platform.sources.webchat.webchat_event import WebChatMessageEvent +from astrbot.core.platform.sources.wecom import wecom_event +from astrbot.core.platform.sources.wecom.wecom_event import WecomPlatformEvent +from astrbot.core.platform.sources.wecom_ai_bot.wecomai_event import ( + WecomAIBotMessageEvent, +) +from astrbot.core.platform.sources.wecom_ai_bot.wecomai_webhook import ( + WecomAIBotWebhookClient, +) +from astrbot.core.platform.sources.weixin_oc.weixin_oc_adapter import WeixinOCAdapter +from astrbot.core.platform.sources.weixin_official_account.weixin_offacc_event import ( + WeixinOfficialAccountPlatformEvent, +) +from astrbot.core.utils import media_utils + + +class _StopEvent: + def __init__(self): + self.extras = {} + + def is_stopped(self): + return False + + def get_extra(self, key, default=None): + return self.extras.get(key, default) + + def set_extra(self, key, value): + self.extras[key] = value + + +@pytest.mark.asyncio +async def test_wecom_webhook_stop_blocks_later_markdown_chunks(): + event = _StopEvent() + webhook = object.__new__(WecomAIBotWebhookClient) + writes = 0 + + async def send_payload(_payload): + nonlocal writes + writes += 1 + event.set_extra("agent_stop_requested", True) + + webhook.send_payload = send_payload + + with pytest.raises(AgentOutputStopped): + await webhook.send_message_chain( + MessageChain().message("x" * 4097), + stop_event=event, + ) + assert writes == 1 + + +@pytest.mark.asyncio +async def test_wecom_partial_webhook_stop_blocks_required_long_reply(): + event = object.__new__(WecomAIBotMessageEvent) + event._extras = {} + event._force_stopped = False + event._result = None + event.session = SimpleNamespace(session_id="stream") + event.message_obj = SimpleNamespace(raw_message={"stream_id": "stream"}) + event.queue_mgr = SimpleNamespace( + get_pending_response=lambda _stream_id: { + "callback_params": { + "connection_mode": "long_connection", + "req_id": "request", + } + } + ) + event.only_use_webhook_url_to_send = False + event.long_connection_sender = AsyncMock(return_value=True) + + async def webhook_send(*_args, **_kwargs): + event.set_extra("agent_stop_requested", True) + + event.webhook_client = SimpleNamespace(send_message_chain=webhook_send) + + with pytest.raises(AgentOutputStopped): + await event.send(MessageChain().message("required reply")) + event.long_connection_sender.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_satori_streaming_stop_is_not_swallowed(): + event = object.__new__(SatoriPlatformEvent) + event._extras = {} + event._force_stopped = False + event._result = None + event.send = AsyncMock(side_effect=AgentOutputStopped) + + async def source(): + yield MessageChain().message("late") + + with pytest.raises(AgentOutputStopped): + await event.send_streaming(source()) + + +@pytest.mark.asyncio +async def test_satori_rejected_message_is_reported_as_delivery_failure(): + event = object.__new__(SatoriPlatformEvent) + event._extras = {} + event._force_stopped = False + event._result = None + event.session = SimpleNamespace(session_id="channel") + event.adapter = SimpleNamespace( + logins=[{"platform": "test", "user": {"id": "bot"}}], + send_http_request=AsyncMock(return_value=None), + ) + + with pytest.raises(RuntimeError, match="delivery failed"): + await event.send(MessageChain().message("answer")) + + +@pytest.mark.asyncio +async def test_webchat_stop_after_conversion_is_not_reported_as_success(monkeypatch): + event = _StopEvent() + queue = asyncio.Queue() + + async def convert_to_base64(_image): + event.set_extra("agent_stop_requested", True) + return "aGVsbG8=" + + monkeypatch.setattr(Image, "convert_to_base64", convert_to_base64) + monkeypatch.setattr( + "astrbot.core.platform.sources.webchat.webchat_event." + "webchat_queue_mgr.get_or_create_back_queue", + lambda *_args, **_kwargs: queue, + ) + + with pytest.raises(AgentOutputStopped): + await WebChatMessageEvent._send( + "message", + MessageChain([Image(file="ignored")]), + "webchat!user!conversation", + stop_event=event, + ) + assert queue.empty() + + +@pytest.mark.parametrize("component_type", ["image", "record", "file"]) +@pytest.mark.parametrize("stop_after_write", [False, True]) +@pytest.mark.asyncio +async def test_webchat_attachment_ownership_after_write( + monkeypatch, + tmp_path, + component_type, + stop_after_write, +): + event = _StopEvent() + queue = asyncio.Queue() + attachments = tmp_path / "attachments" + attachments.mkdir() + monkeypatch.setattr(webchat_event, "attachments_dir", str(attachments)) + monkeypatch.setattr( + webchat_event.webchat_queue_mgr, + "get_or_create_back_queue", + lambda *_args, **_kwargs: queue, + ) + + if component_type == "image": + component = Image(file="ignored") + monkeypatch.setattr( + Image, + "convert_to_base64", + AsyncMock(return_value="aGVsbG8="), + ) + elif component_type == "record": + component = Record(file="ignored") + monkeypatch.setattr( + Record, + "convert_to_base64", + AsyncMock(return_value="YXVkaW8="), + ) + else: + source = tmp_path / "source.txt" + source.write_text("file", encoding="utf-8") + component = File(file=str(source), name="source.txt") + monkeypatch.setattr( + File, + "get_file", + AsyncMock(return_value=str(source)), + ) + + if stop_after_write and component_type in {"image", "record"}: + + async def write_then_stop(func, *args, **kwargs): + result = func(*args, **kwargs) + event.set_extra("agent_stop_requested", True) + return result + + monkeypatch.setattr(webchat_event.asyncio, "to_thread", write_then_stop) + elif stop_after_write: + original_copy = webchat_event.shutil.copy2 + + def copy_then_stop(*args, **kwargs): + result = original_copy(*args, **kwargs) + event.set_extra("agent_stop_requested", True) + return result + + monkeypatch.setattr(webchat_event.shutil, "copy2", copy_then_stop) + + send = WebChatMessageEvent._send( + "message", + MessageChain([component]), + "webchat!user!conversation", + stop_event=event, + ) + if stop_after_write: + with pytest.raises(AgentOutputStopped): + await send + assert queue.empty() + assert list(attachments.iterdir()) == [] + else: + await send + assert queue.qsize() == 1 + assert len(list(attachments.iterdir())) == 1 + + +@pytest.mark.parametrize("component_type", ["image", "record"]) +@pytest.mark.asyncio +async def test_webchat_cancellation_waits_for_background_write_cleanup( + monkeypatch, + tmp_path, + component_type, +): + event = _StopEvent() + queue = asyncio.Queue() + attachments = tmp_path / "attachments" + attachments.mkdir() + monkeypatch.setattr(webchat_event, "attachments_dir", str(attachments)) + monkeypatch.setattr( + webchat_event.webchat_queue_mgr, + "get_or_create_back_queue", + lambda *_args, **_kwargs: queue, + ) + + if component_type == "image": + component = Image(file="ignored") + monkeypatch.setattr( + Image, + "convert_to_base64", + AsyncMock(return_value="aGVsbG8="), + ) + else: + component = Record(file="ignored") + monkeypatch.setattr( + Record, + "convert_to_base64", + AsyncMock(return_value="YXVkaW8="), + ) + + write_started = threading.Event() + release_write = threading.Event() + + async def delayed_to_thread(func, *args, **kwargs): + def run(): + result = func(*args, **kwargs) + write_started.set() + release_write.wait() + return result + + return await asyncio.get_running_loop().run_in_executor(None, run) + + monkeypatch.setattr(webchat_event.asyncio, "to_thread", delayed_to_thread) + task = asyncio.create_task( + WebChatMessageEvent._send( + "message", + MessageChain([component]), + "webchat!user!conversation", + stop_event=event, + ) + ) + await asyncio.wait_for( + asyncio.get_running_loop().run_in_executor(None, write_started.wait), + timeout=1, + ) + task.cancel() + await asyncio.sleep(0) + release_write.set() + + with pytest.raises(asyncio.CancelledError): + await task + + assert queue.empty() + assert list(attachments.iterdir()) == [] + + +@pytest.mark.parametrize("event_type", ["wecom_kf", "wecom_app", "weixin"]) +@pytest.mark.asyncio +async def test_amr_conversion_stop_still_cleans_temporary_file( + monkeypatch, + tmp_path, + event_type, +): + converted_path = tmp_path / f"{event_type}.amr" + converted_path.write_bytes(b"audio") + + async def convert_to_file_path(_record): + return str(tmp_path / "source.wav") + + monkeypatch.setattr(Record, "convert_to_file_path", convert_to_file_path) + + if event_type.startswith("wecom"): + event = object.__new__(WecomPlatformEvent) + event._extras = {} + event._force_stopped = False + event._result = None + event.message_obj = SimpleNamespace(self_id="bot", session_id="user") + event.get_sender_id = lambda: "user" + event.get_self_id = lambda: "bot" + if event_type == "wecom_kf": + + class FakeKFMessage: + pass + + monkeypatch.setattr(wecom_event, "WeChatKFMessage", FakeKFMessage) + event.client = SimpleNamespace( + kf_message=FakeKFMessage(), + message=SimpleNamespace(), + media=SimpleNamespace(), + ) + else: + event.client = SimpleNamespace( + message=SimpleNamespace(), + media=SimpleNamespace(), + ) + module = wecom_event + else: + from astrbot.core.platform.sources.weixin_official_account import ( + weixin_offacc_event, + ) + + event = object.__new__(WeixinOfficialAccountPlatformEvent) + event._extras = {} + event._force_stopped = False + event._result = None + event.message_out = {} + event.message_obj = SimpleNamespace( + raw_message={"active_send_mode": True}, + sender=SimpleNamespace(user_id="user"), + ) + event.client = SimpleNamespace( + message=SimpleNamespace(), + media=SimpleNamespace(), + ) + module = weixin_offacc_event + + async def convert_audio(_path): + event.set_extra("agent_stop_requested", True) + return str(converted_path) + + monkeypatch.setattr(module, "convert_audio_to_amr", convert_audio) + + with pytest.raises(AgentOutputStopped): + await event.send(MessageChain([Record(file="source.wav")])) + + assert not converted_path.exists() + + +@pytest.mark.parametrize("output_kind", ["new", "existing", "dangling_symlink"]) +@pytest.mark.asyncio +async def test_cancelled_audio_conversion_reaps_process_and_preserves_ownership( + monkeypatch, + tmp_path, + output_kind, +): + output_path = tmp_path / "cancelled.amr" + target_path = None + if output_kind == "existing": + output_path.write_bytes(b"original audio") + elif output_kind == "dangling_symlink": + target_path = tmp_path / "missing-target.amr" + output_path.symlink_to(target_path.name) + output_created = asyncio.Event() + wait_started = asyncio.Event() + release_wait = asyncio.Event() + + class Process: + returncode = None + killed = False + waited = False + output_path = None + + async def communicate(self): + self.output_path.write_bytes(b"partial audio") + output_created.set() + await asyncio.Event().wait() + + def kill(self): + self.killed = True + + async def wait(self): + self.waited = True + wait_started.set() + await release_wait.wait() + self.returncode = -9 + return self.returncode + + process = Process() + + async def create_subprocess_exec(*args, **_kwargs): + process.output_path = Path(args[-1]) + return process + + monkeypatch.setattr( + media_utils.asyncio, + "create_subprocess_exec", + create_subprocess_exec, + ) + task = asyncio.create_task( + media_utils.convert_audio_format( + "source.wav", + "amr", + str(output_path), + ) + ) + await asyncio.wait_for(output_created.wait(), timeout=1) + task.cancel() + await asyncio.wait_for(wait_started.wait(), timeout=1) + task.cancel() + await asyncio.sleep(0) + release_wait.set() + + with pytest.raises(asyncio.CancelledError): + await task + + assert process.killed is True + assert process.waited is True + assert not process.output_path.exists() + if output_kind == "existing": + assert output_path.read_bytes() == b"original audio" + elif output_kind == "dangling_symlink": + assert output_path.is_symlink() + assert not target_path.exists() + else: + assert not output_path.exists() + + +@pytest.mark.parametrize("failure", [FileNotFoundError("ffmpeg"), OSError("spawn")]) +@pytest.mark.asyncio +async def test_audio_spawn_failure_preserves_preexisting_explicit_output( + monkeypatch, + tmp_path, + failure, +): + output_path = tmp_path / "existing.amr" + output_path.write_bytes(b"original audio") + + async def create_subprocess_exec(*_args, **_kwargs): + raise failure + + monkeypatch.setattr( + media_utils.asyncio, + "create_subprocess_exec", + create_subprocess_exec, + ) + + with pytest.raises(Exception): + await media_utils.convert_audio_format( + "source.wav", + "amr", + str(output_path), + ) + + assert output_path.read_bytes() == b"original audio" + + +@pytest.mark.parametrize("output_kind", ["new", "existing", "symlink"]) +@pytest.mark.asyncio +async def test_successful_audio_conversion_keeps_explicit_output( + monkeypatch, + tmp_path, + output_kind, +): + output_path = tmp_path / "converted.amr" + target_path = None + if output_kind == "existing": + output_path.write_bytes(b"original audio") + elif output_kind == "symlink": + target_path = tmp_path / "converted-target.amr" + target_path.write_bytes(b"original audio") + output_path.symlink_to(target_path.name) + + class Process: + returncode = 0 + + def __init__(self, converter_output_path): + self.output_path = converter_output_path + + async def communicate(self): + self.output_path.write_bytes(b"converted audio") + return b"", b"" + + async def create_subprocess_exec(*args, **_kwargs): + return Process(Path(args[-1])) + + monkeypatch.setattr( + media_utils.asyncio, + "create_subprocess_exec", + create_subprocess_exec, + ) + + result = await media_utils.convert_audio_format( + "source.wav", + "amr", + str(output_path), + ) + + assert result == str(output_path) + assert output_path.read_bytes() == b"converted audio" + if output_kind == "symlink": + assert output_path.is_symlink() + assert target_path.read_bytes() == b"converted audio" + + +@pytest.mark.asyncio +async def test_weixin_oc_text_failure_is_not_hidden_by_media_success(): + adapter = object.__new__(WeixinOCAdapter) + adapter.token = "token" + calls = [] + + async def resolve_path(*_args, **_kwargs): + return Path("/tmp/ignored.jpg") + + async def prepare_item(*_args, **_kwargs): + return {"type": adapter.IMAGE_ITEM_TYPE, "image_item": {}} + + async def send_items(_user_id, item_list, **_kwargs): + calls.append(item_list) + return False + + adapter._resolve_media_file_path = resolve_path + adapter._prepare_media_item = prepare_item + adapter._send_items_to_session = send_items + + sent = await adapter._send_media_segment( + "user", + Image(file="ignored"), + text="required text", + ) + + assert sent is False + assert len(calls) == 1 + assert calls[0][0]["type"] == 1 + + +@pytest.mark.asyncio +async def test_wecom_stop_between_plain_chunks_blocks_later_writes(monkeypatch): + event = object.__new__(WecomPlatformEvent) + event._extras = {} + event._force_stopped = False + event._result = None + event.message_obj = SimpleNamespace(self_id="bot", session_id="user") + writes = [] + + def send_text(*_args): + writes.append(_args[-1]) + event.set_extra("agent_stop_requested", True) + + event.client = SimpleNamespace( + message=SimpleNamespace(send_text=send_text), + ) + monkeypatch.setattr(wecom_event.asyncio, "sleep", AsyncMock()) + + with pytest.raises(AgentOutputStopped): + await event.send(MessageChain().message("x" * 2050)) + assert len(writes) == 1 + + +@pytest.mark.asyncio +async def test_weixin_official_stop_after_conversion_blocks_platform_write( + monkeypatch, +): + event = object.__new__(WeixinOfficialAccountPlatformEvent) + event._extras = {} + event._force_stopped = False + event._result = None + event.message_out = {} + event.message_obj = SimpleNamespace( + raw_message={"active_send_mode": True}, + sender=SimpleNamespace(user_id="user"), + ) + upload = MagicMock(return_value={"media_id": "media"}) + send_image = MagicMock() + event.client = SimpleNamespace( + media=SimpleNamespace(upload=upload), + message=SimpleNamespace(send_image=send_image), + ) + + async def convert_to_file_path(_image): + event.set_extra("agent_stop_requested", True) + return "/tmp/ignored.jpg" + + monkeypatch.setattr(Image, "convert_to_file_path", convert_to_file_path) + + with pytest.raises(AgentOutputStopped): + await event.send(MessageChain([Image(file="ignored")])) + upload.assert_not_called() + send_image.assert_not_called() diff --git a/tests/test_qqofficial_group_message_create.py b/tests/test_qqofficial_group_message_create.py index 46d688d158..eab94fbe7d 100644 --- a/tests/test_qqofficial_group_message_create.py +++ b/tests/test_qqofficial_group_message_create.py @@ -10,15 +10,20 @@ from botpy import ConnectionSession from astrbot.api.event import MessageChain -from astrbot.api.message_components import At, Plain +from astrbot.api.message_components import At, Plain, Record +from astrbot.core.agent.stop_policy import AgentOutputStopped from astrbot.core.message.message_event_result import ( MessageEventResult, ResultContentType, ) from astrbot.core.pipeline.respond.stage import RespondStage from astrbot.core.pipeline.result_decorate.stage import ResultDecorateStage +from astrbot.core.platform.astr_message_event import AstrMessageEvent from astrbot.core.platform.message_session import MessageSession from astrbot.core.platform.message_type import MessageType +from astrbot.core.platform.sources.qqofficial.qqofficial_message_event import ( + QQOfficialMessageEvent, +) from astrbot.core.platform.sources.qqofficial.qqofficial_platform_adapter import ( QQOfficialPlatformAdapter, _ensure_group_message_create_parser, @@ -66,6 +71,260 @@ def _dispatch_group_message(payload: dict) -> tuple[str, botpy.message.GroupMess return dispatched[0] +@pytest.mark.asyncio +async def test_qqofficial_stop_blocks_retry_and_markdown_fallback(): + event = object.__new__(QQOfficialMessageEvent) + event._extras = {} + event._force_stopped = False + event._result = None + + async def request_then_stop(*_args, **_kwargs): + event.set_extra("agent_stop_requested", True) + raise botpy.errors.SequenceNumberError("first request failed") + + request = AsyncMock(side_effect=request_then_stop) + event.bot = SimpleNamespace( + api=SimpleNamespace(_http=SimpleNamespace(request=request)) + ) + + with pytest.raises(AgentOutputStopped): + await event._send_with_markdown_fallback( + lambda payload: event.post_c2c_message("openid", **payload), + {"content": "late", "msg_id": "reply-id"}, + "late", + ) + + request.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_qqofficial_none_response_after_stop_is_not_swallowed(): + event = object.__new__(QQOfficialMessageEvent) + event._extras = {} + event._force_stopped = False + event._result = None + + async def return_none_then_stop(*_args, **_kwargs): + event.set_extra("agent_stop_requested", True) + return None + + request = AsyncMock(side_effect=return_none_then_stop) + event.bot = SimpleNamespace( + api=SimpleNamespace(_http=SimpleNamespace(request=request)) + ) + + with pytest.raises(AgentOutputStopped): + await event.post_c2c_message(openid="user", content="late") + request.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_qqofficial_streaming_stop_is_not_swallowed(): + event = object.__new__(QQOfficialMessageEvent) + event._extras = {} + event._force_stopped = False + event._result = None + event._has_send_oper = False + event.platform_meta = SimpleNamespace(name="qqofficial-test") + event.message_obj = SimpleNamespace(raw_message=object()) + event.send_buffer = None + event._post_send = AsyncMock(side_effect=AgentOutputStopped) + + async def source(): + yield MessageChain().message("late") + + with pytest.raises(AgentOutputStopped): + await event.send_streaming(source()) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("final_succeeds", [True, False]) +async def test_qqofficial_c2c_streaming_requires_terminal_delivery( + monkeypatch, + final_succeeds, +): + class FakeC2CMessage: + pass + + async def noop_base_streaming(self, generator, use_fallback=False): + return None + + monkeypatch.setattr(botpy.message, "C2CMessage", FakeC2CMessage) + monkeypatch.setattr(AstrMessageEvent, "send_streaming", noop_base_streaming) + + event = object.__new__(QQOfficialMessageEvent) + event._extras = {} + event._force_stopped = False + event._result = None + event.message_obj = SimpleNamespace(raw_message=FakeC2CMessage()) + event.send_buffer = None + states = [] + payloads = [] + + async def post_send(stream=None): + states.append(stream["state"]) + payloads.append(event.send_buffer.get_plain_text()) + event.send_buffer = None + if stream["state"] == 10 and not final_succeeds: + return None + return SimpleNamespace(id="stream-id") + + event._post_send = AsyncMock(side_effect=post_send) + + async def source(): + yield MessageChain().message("partial") + + if final_succeeds: + await event.send_streaming(source()) + else: + with pytest.raises(RuntimeError, match="streaming message delivery failed"): + await event.send_streaming(source()) + + assert states == [1, 10] + assert payloads == ["partial", "\n"] + + +@pytest.mark.asyncio +async def test_qqofficial_c2c_empty_stream_is_delivery_failure(monkeypatch): + class FakeC2CMessage: + pass + + async def noop_base_streaming(self, generator, use_fallback=False): + return None + + monkeypatch.setattr(botpy.message, "C2CMessage", FakeC2CMessage) + monkeypatch.setattr(AstrMessageEvent, "send_streaming", noop_base_streaming) + + event = object.__new__(QQOfficialMessageEvent) + event._extras = {} + event._force_stopped = False + event._result = None + event.message_obj = SimpleNamespace(raw_message=FakeC2CMessage()) + event.send_buffer = None + + async def empty_source(): + if False: + yield MessageChain().message("never") + + with pytest.raises(RuntimeError, match="produced no platform delivery"): + await event.send_streaming(empty_source()) + + assert event.send_buffer is None + + +@pytest.mark.asyncio +async def test_qqofficial_ordinary_none_response_is_delivery_failure(monkeypatch): + class FakeGroupMessage: + group_openid = "group" + + event = object.__new__(QQOfficialMessageEvent) + event._extras = {} + event._force_stopped = False + event._result = None + event._temporary_local_files = [] + event.send_buffer = None + event.message_obj = SimpleNamespace( + raw_message=FakeGroupMessage(), + message_id="message", + ) + post_group_message = AsyncMock(return_value=None) + event.bot = SimpleNamespace( + api=SimpleNamespace(post_group_message=post_group_message) + ) + + async def parsed(_message): + return "answer", None, None, None, None, None, None + + base_send = AsyncMock() + monkeypatch.setattr(botpy.message, "GroupMessage", FakeGroupMessage) + monkeypatch.setattr( + QQOfficialMessageEvent, + "_parse_to_qqofficial", + staticmethod(parsed), + ) + monkeypatch.setattr(AstrMessageEvent, "send", base_send) + + with pytest.raises(RuntimeError, match="delivery returned no response"): + await event.send(MessageChain().message("answer")) + + post_group_message.assert_awaited_once() + base_send.assert_not_awaited() + assert event.send_buffer is None + + +@pytest.mark.asyncio +async def test_qqofficial_split_delivery_stops_at_first_none_response(): + event = object.__new__(QQOfficialMessageEvent) + event._extras = {} + event._force_stopped = False + event._result = None + first = MessageChain().message("first") + second = MessageChain().message("second") + event.send_buffer = MessageChain().message("combined") + event._split_message_chain_by_media = lambda _message: [first, second] + event._post_send_one = AsyncMock(side_effect=[None, SimpleNamespace(id="second")]) + + with pytest.raises(RuntimeError, match="delivery returned no response"): + await event._post_send() + + event._post_send_one.assert_awaited_once_with(first, None) + assert event.send_buffer is None + + +@pytest.mark.asyncio +async def test_qqofficial_empty_buffer_is_delivery_failure_and_is_cleared(): + event = object.__new__(QQOfficialMessageEvent) + event._extras = {} + event._force_stopped = False + event._result = None + event.send_buffer = MessageChain() + + with pytest.raises(RuntimeError, match="buffer is empty"): + await event._post_send() + + assert event.send_buffer is None + + +@pytest.mark.asyncio +async def test_qqofficial_record_is_tracked_before_stop_unwind( + monkeypatch, + tmp_path, +): + class FakeC2CMessage: + pass + + converted_path = tmp_path / "converted.silk" + converted_path.write_bytes(b"audio") + event = object.__new__(QQOfficialMessageEvent) + event._extras = {} + event._force_stopped = False + event._result = None + event._temporary_local_files = [] + event.send_buffer = MessageChain([Record(file="source.wav")]) + event.message_obj = SimpleNamespace( + raw_message=FakeC2CMessage(), + message_id="message", + ) + + async def parsed(_message): + event.set_extra("agent_stop_requested", True) + return "", None, None, str(converted_path), None, None, None + + monkeypatch.setattr(botpy.message, "C2CMessage", FakeC2CMessage) + monkeypatch.setattr( + QQOfficialMessageEvent, + "_parse_to_qqofficial", + staticmethod(parsed), + ) + + with pytest.raises(AgentOutputStopped): + await event._post_send_one(event.send_buffer) + + assert event._temporary_local_files == [str(converted_path)] + event.cleanup_temporary_local_files() + assert not converted_path.exists() + + @pytest.mark.asyncio async def test_group_message_create_parser_is_registered_and_dispatches_group_message(): QQOfficialPlatformAdapter( diff --git a/tests/test_telegram_adapter.py b/tests/test_telegram_adapter.py index 948b84f5ba..a801c48e5d 100644 --- a/tests/test_telegram_adapter.py +++ b/tests/test_telegram_adapter.py @@ -1,11 +1,14 @@ import asyncio import importlib import sys +from types import SimpleNamespace from unittest.mock import AsyncMock, MagicMock, patch import pytest import astrbot.api.message_components as Comp +from astrbot.api.event import MessageChain +from astrbot.core.agent.stop_policy import AgentOutputStopped from astrbot.core.platform.register import unregister_platform_adapters_by_module from tests.fixtures.helpers import ( NoopAwaitable, @@ -202,6 +205,104 @@ async def test_telegram_final_segment_splits_long_markdown_messages(): assert second_call["parse_mode"] == "MarkdownV2" +@pytest.mark.asyncio +async def test_telegram_draft_stop_blocks_final_write(): + TelegramPlatformEvent = _load_telegram_platform_event() + event = object.__new__(TelegramPlatformEvent) + event._extras = {} + event._force_stopped = False + event._result = None + draft_clear_started = asyncio.Event() + release_draft_clear = asyncio.Event() + final_send = AsyncMock() + + async def draft_send(_chat_id, _draft_id, text, *_args, **_kwargs): + if text == "\u23f3": + draft_clear_started.set() + await release_draft_clear.wait() + + event._send_message_draft = draft_send + event._send_final_segment = final_send + + async def source(): + yield MessageChain().message("secret") + yield MessageChain(chain=[], type="break") + + task = asyncio.create_task( + event._send_streaming_draft( + "1", + None, + {"chat_id": "1"}, + source(), + ) + ) + await asyncio.wait_for(draft_clear_started.wait(), timeout=1) + event.set_extra("agent_stop_requested", True) + release_draft_clear.set() + + with pytest.raises(AgentOutputStopped): + await task + final_send.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_telegram_final_edit_failure_is_reported_as_delivery_failure(): + TelegramPlatformEvent = _load_telegram_platform_event() + event = object.__new__(TelegramPlatformEvent) + event._extras = {} + event._force_stopped = False + event._result = None + event._ensure_typing = AsyncMock() + event.client = SimpleNamespace( + send_message=AsyncMock(return_value=SimpleNamespace(message_id=1)), + edit_message_text=AsyncMock(side_effect=RuntimeError("edit failed")), + ) + + async def source(): + yield MessageChain().message("first") + yield MessageChain().message(" second") + + with pytest.raises(RuntimeError, match="edit failed"): + await event._send_streaming_edit( + "group", + None, + {"chat_id": "group"}, + source(), + ) + + event.client.send_message.assert_awaited_once() + assert event.client.edit_message_text.await_count == 2 + + +@pytest.mark.parametrize("stop_at", ["chat_action", "first_chunk"]) +@pytest.mark.asyncio +async def test_telegram_ordinary_send_stops_before_new_writes(stop_at): + TelegramPlatformEvent = _load_telegram_platform_event() + client = MagicMock() + event = TelegramPlatformEvent("msg", MagicMock(), MagicMock(), "session", client) + + async def request_stop(**_kwargs): + event.set_extra("agent_stop_requested", True) + + client.send_chat_action = AsyncMock( + side_effect=request_stop if stop_at == "chat_action" else None + ) + client.send_message = AsyncMock( + side_effect=request_stop if stop_at == "first_chunk" else None + ) + delta = "A" * (TelegramPlatformEvent.MAX_MESSAGE_LENGTH + 32) + + with pytest.raises(AgentOutputStopped): + await event.send_with_client( + client, + MessageChain([Comp.Plain(delta)]), + "123456", + event, + ) + + assert client.send_message.await_count == (0 if stop_at == "chat_action" else 1) + + @pytest.mark.asyncio async def test_telegram_final_segment_splits_long_plaintext_when_markdown_fails(): TelegramPlatformEvent = _load_telegram_platform_event() diff --git a/tests/test_tool_loop_agent_runner.py b/tests/test_tool_loop_agent_runner.py index 48a736e152..720db84363 100644 --- a/tests/test_tool_loop_agent_runner.py +++ b/tests/test_tool_loop_agent_runner.py @@ -30,6 +30,8 @@ TokenUsage, ) from astrbot.core.provider.provider import Provider, TTSProvider +from astrbot.core.provider.sources import request_retry +from astrbot.core.provider.sources.request_retry import retry_provider_request class MockProvider(Provider): @@ -264,28 +266,8 @@ async def text_chat(self, **kwargs) -> LLMResponse: ) -class MockStopAwareStreamProvider(MockProvider): - async def text_chat_stream(self, **kwargs): - abort_signal = kwargs.get("abort_signal") - if abort_signal is None: - return - await abort_signal.wait() - yield LLMResponse( - role="assistant", - completion_text="late chunk text", - reasoning_content="late chunk reasoning", - is_chunk=True, - ) - yield LLMResponse( - role="assistant", - completion_text="late final text", - is_chunk=False, - usage=TokenUsage(input_other=10, output=5), - ) - - class MockDelayedTTSProvider(TTSProvider): - def __init__(self, audio_path: Path): + def __init__(self, audio_path: Path | None): super().__init__({"type": "test_tts", "id": "test_tts"}, {}) self.audio_path = audio_path self.started = asyncio.Event() @@ -296,7 +278,7 @@ async def get_audio(self, text: str) -> str: self.call_count += 1 self.started.set() await self.release.wait() - return str(self.audio_path) + return str(self.audio_path) if self.audio_path else "" def meta(self) -> ProviderMeta: return ProviderMeta( @@ -435,50 +417,23 @@ async def on_agent_done(self, run_context, llm_response): self.agent_done_response = llm_response -class MockLateStopOnDoneHooks(MockHooks): - def __init__(self, event): - super().__init__() - self.event = event - - async def on_agent_done(self, run_context, llm_response): - await super().on_agent_done(run_context, llm_response) - self.event.set_extra("agent_stop_requested", True) - - -class HookPreStopRunner(ToolLoopAgentRunner): - def __init__(self, event): - super().__init__() - self.event = event - - def _is_stop_requested(self) -> bool: - if self.final_llm_resp is not None: - self.event.set_extra("agent_stop_requested", True) - return super()._is_stop_requested() - - class LateStopAfterYieldRunner: def __init__( self, event, - response_type: str, text: str, - streaming: bool, stop_after_yield: bool = True, ): self.run_context = ContextWrapper(context=MockAgentContext(event)) - self.response_type = response_type self.text = text - self.streaming = streaming + self.streaming = True self.stop_after_yield = stop_after_yield - self.stats = SimpleNamespace(to_dict=lambda: {}) self._done = False self._aborted = False - self.stop_requested = False - self.discard_late_result_called = False async def step(self): yield SimpleNamespace( - type=self.response_type, + type="streaming_delta", data={"chain": MessageChain().message(self.text)}, ) if self.stop_after_yield: @@ -489,13 +444,12 @@ def done(self): return self._done def request_stop(self): - self.stop_requested = True + pass def was_aborted(self): return self._aborted def discard_late_aborted_result(self): - self.discard_late_result_called = True self._aborted = True event = self.run_context.context.event event.set_extra("agent_user_aborted", True) @@ -505,11 +459,13 @@ def discard_late_aborted_result(self): class MockEvent: def __init__(self, umo: str, sender_id: str): self.unified_msg_origin = umo + self.plugins_name = None self._sender_id = sender_id self._extras = {} self.result = None self.trace = SimpleNamespace(record=lambda *args, **kwargs: None) self._stopped = False + self._temporary_local_files = [] def get_sender_id(self): return self._sender_id @@ -525,6 +481,9 @@ def get_extra(self, key=None, default=None): def set_result(self, result): self.result = result + def get_result(self): + return self.result + def clear_result(self): self.result = None @@ -534,8 +493,16 @@ def is_stopped(self): def stop(self): self._stopped = True - def track_temporary_local_file(self, _path): - return None + def track_temporary_local_file(self, path): + if path not in self._temporary_local_files: + self._temporary_local_files.append(path) + + def cleanup_temporary_local_files(self): + paths = list(self._temporary_local_files) + self._temporary_local_files.clear() + for path in paths: + if os.path.exists(path): + os.remove(path) def get_platform_name(self): return "test" @@ -626,35 +593,6 @@ def runner(): return ToolLoopAgentRunner() -def test_record_llm_usage_keeps_usage_reference_to_prevent_id_reuse(runner): - usage = TokenUsage(input_other=10, output=5) - runner.req = SimpleNamespace(conversation=SimpleNamespace(token_usage=0)) - runner.stats = SimpleNamespace(token_usage=TokenUsage()) - runner._recorded_usages = [] - - runner._record_llm_usage(LLMResponse(role="assistant", usage=usage)) - runner._record_llm_usage(LLMResponse(role="assistant", usage=usage)) - - assert runner.stats.token_usage.total == usage.total - assert runner.req.conversation.token_usage == usage.total - assert runner._recorded_usages == [usage] - - -def test_is_message_from_llm_response_rejects_extra_content_parts(runner): - llm_resp = LLMResponse(role="assistant", completion_text="late text") - message = Message( - role="assistant", - content=[ - TextPart(text="late text"), - ImageURLPart( - image_url=ImageURLPart.ImageURL(url="https://example.com/a.png") - ), - ], - ) - - assert runner._is_message_from_llm_response(message, llm_resp) is False - - def _make_large_tool_result_text() -> str: return "x" * 100000 @@ -1342,13 +1280,6 @@ async def test_stop_signal_returns_aborted_and_discards_delayed_response( assert final_resp is not None assert final_resp.role == "assistant" assert final_resp.completion_text == "" - assert final_resp.result_chain is None - assert final_resp.reasoning_content is None - assert final_resp.reasoning_signature is None - assert final_resp.tools_call_name == [] - assert final_resp.tools_call_args == [] - assert final_resp.tools_call_ids == [] - assert final_resp.tools_call_extra_content == {} assert ( not runner.run_context.messages or runner.run_context.messages[-1].role != "assistant" @@ -1404,193 +1335,159 @@ async def test_aborted_final_response_sanitizes_all_model_output_fields( assert "late_tool" not in serialized_messages +@pytest.mark.parametrize("stop_during", ["begin_hook", "context_process"]) @pytest.mark.asyncio -async def test_runner_step_treats_agent_user_aborted_as_stop( - runner, provider_request, mock_tool_executor, mock_hooks +async def test_stop_around_context_processing_blocks_provider_request( + stop_during, + provider_request, + mock_tool_executor, ): - event = MockEvent("test:FriendMessage:runner_user_aborted", "u1") - event.set_extra("agent_user_aborted", True) - provider = MockLeakyAbortProvider() + event = MockEvent(f"test:FriendMessage:context_{stop_during}", "u1") + provider = MockProvider() + provider.should_call_tools = False + + class Hooks(MockHooks): + async def on_agent_begin(self, run_context): + await super().on_agent_begin(run_context) + if stop_during == "begin_hook": + event.set_extra("agent_stop_requested", True) + runner = ToolLoopAgentRunner() await runner.reset( provider=provider, request=provider_request, run_context=ContextWrapper(context=MockAgentContext(event)), tool_executor=mock_tool_executor, - agent_hooks=mock_hooks, - streaming=False, + agent_hooks=Hooks(), ) + async def process(messages, **_kwargs): + event.set_extra("agent_stop_requested", True) + return messages + + runner.request_context_manager.process = AsyncMock(side_effect=process) responses = [response async for response in runner.step()] assert [response.type for response in responses] == ["aborted"] assert provider.call_count == 0 - assert runner.was_aborted() is True - final_resp = runner.get_final_llm_resp() - assert final_resp is not None - assert final_resp.completion_text == "" + if stop_during == "begin_hook": + runner.request_context_manager.process.assert_not_awaited() + else: + runner.request_context_manager.process.assert_awaited_once() @pytest.mark.asyncio -async def test_run_agent_discards_buffered_llm_result_after_abort( - runner, provider_request, mock_tool_executor, mock_hooks +async def test_stop_interrupts_pending_context_processing( + provider_request, + mock_tool_executor, + mock_hooks, ): - from astrbot.core.astr_agent_run_util import run_agent - - provider = MockDelayedTextProvider() - event = MockEvent("test:FriendMessage:buffer_abort", "u1") - + event = MockEvent("test:FriendMessage:context_pending_stop", "u1") + provider = MockProvider() + provider.should_call_tools = False + runner = ToolLoopAgentRunner() await runner.reset( provider=provider, request=provider_request, run_context=ContextWrapper(context=MockAgentContext(event)), tool_executor=mock_tool_executor, agent_hooks=mock_hooks, - streaming=False, - ) - - agent_task = asyncio.create_task( - _collect_async_iter( - run_agent( - runner, - buffer_intermediate_messages=True, - ) - ) - ) - await asyncio.wait_for(provider.started.wait(), timeout=5) - event.set_extra("agent_stop_requested", True) - provider.release.set() - - chains = await asyncio.wait_for(agent_task, timeout=5) - - assert chains == [] - assert runner.was_aborted() is True - assert event.result is None - - -@pytest.mark.asyncio -async def test_run_agent_discards_buffered_result_when_stop_arrives_before_flush(): - from astrbot.core.astr_agent_run_util import run_agent - - event = MockEvent("test:FriendMessage:late_buffer_abort", "u1") - runner = LateStopAfterYieldRunner( - event, - response_type="llm_result", - text="buffered late text", - streaming=False, - ) - - chains = await _collect_async_iter( - run_agent( - cast(Any, runner), - buffer_intermediate_messages=True, - ) ) + started = asyncio.Event() + cancelled = asyncio.Event() - assert chains == [] - assert event.result is None - assert runner.discard_late_result_called is True - assert runner.was_aborted() is True - - -@pytest.mark.asyncio -async def test_run_agent_treats_agent_user_aborted_as_stop(): - from astrbot.core.astr_agent_run_util import run_agent + async def process(_messages, **_kwargs): + started.set() + try: + await asyncio.Event().wait() + finally: + cancelled.set() - event = MockEvent("test:FriendMessage:user_aborted_abort", "u1") - event.set_extra("agent_user_aborted", True) - runner = LateStopAfterYieldRunner( - event, - response_type="llm_result", - text="aborted residual text", - streaming=False, - ) + runner.request_context_manager.process = process + task = asyncio.create_task(_collect_async_iter(runner.step())) + await asyncio.wait_for(started.wait(), timeout=1) + event.set_extra("agent_stop_requested", True) - chains = await _collect_async_iter(run_agent(cast(Any, runner))) + responses = await asyncio.wait_for(task, timeout=1) - assert chains == [] - assert event.result is None + assert [response.type for response in responses] == ["aborted"] + assert provider.call_count == 0 + assert cancelled.is_set() is True @pytest.mark.asyncio -async def test_stop_before_agent_done_hook_skips_llm_response_event( +@pytest.mark.parametrize("stop_first", [False, True]) +async def test_await_or_abort_propagates_outer_cancellation( + stop_first, + runner, provider_request, mock_tool_executor, + mock_hooks, ): - from astrbot.core import astr_agent_hooks - from astrbot.core.astr_agent_hooks import MainAgentHooks - from astrbot.core.star.star_handler import EventType - - calls = [] - - async def fake_call_event_hook(event, hook_type, *args, **kwargs): - calls.append((hook_type, args)) - return False + event = MockEvent("test:FriendMessage:outer_cancel", "u1") + await runner.reset( + provider=MockProvider(), + request=provider_request, + run_context=ContextWrapper(context=MockAgentContext(event)), + tool_executor=mock_tool_executor, + agent_hooks=mock_hooks, + ) + started = asyncio.Event() + cancelled = asyncio.Event() - event = MockEvent("test:FriendMessage:hook_race_abort", "u1") - provider = MockProvider() - provider.should_call_tools = False - runner = HookPreStopRunner(event) - monkeypatch_context = pytest.MonkeyPatch() - monkeypatch_context.setattr( - astr_agent_hooks, - "call_event_hook", - fake_call_event_hook, - ) - - try: - await runner.reset( - provider=provider, - request=provider_request, - run_context=ContextWrapper(context=MockAgentContext(event)), - tool_executor=mock_tool_executor, - agent_hooks=MainAgentHooks(), - streaming=False, - ) + async def operation(): + started.set() + try: + await asyncio.Event().wait() + finally: + cancelled.set() - responses = [response async for response in runner.step()] - finally: - monkeypatch_context.undo() + task = asyncio.create_task(runner._await_or_abort(operation())) + await asyncio.wait_for(started.wait(), timeout=1) + if stop_first: + runner.request_stop() + task.cancel() - assert [response.type for response in responses] == ["aborted"] - called_types = [hook_type for hook_type, _ in calls] - assert EventType.OnLLMResponseEvent not in called_types - assert EventType.OnAgentDoneEvent in called_types - assert event.get_extra("_llm_reasoning_content") is None - final_resp = runner.get_final_llm_resp() - assert final_resp is not None - assert final_resp.completion_text == "" - assert "这是我的最终回答" not in repr(runner.run_context.messages) + with pytest.raises(asyncio.CancelledError): + await task + assert cancelled.is_set() is True @pytest.mark.asyncio -async def test_stop_after_agent_done_hook_discards_final_assistant_message( - provider_request, - mock_tool_executor, +async def test_run_agent_discards_buffered_llm_result_after_abort( + runner, provider_request, mock_tool_executor, mock_hooks ): - event = MockEvent("test:FriendMessage:post_hook_abort", "u1") - provider = MockProvider() - provider.should_call_tools = False - hooks = MockLateStopOnDoneHooks(event) - runner = ToolLoopAgentRunner() + from astrbot.core.astr_agent_run_util import run_agent + + provider = MockDelayedTextProvider() + event = MockEvent("test:FriendMessage:buffer_abort", "u1") await runner.reset( provider=provider, request=provider_request, run_context=ContextWrapper(context=MockAgentContext(event)), tool_executor=mock_tool_executor, - agent_hooks=hooks, + agent_hooks=mock_hooks, streaming=False, ) - responses = [response async for response in runner.step()] + agent_task = asyncio.create_task( + _collect_async_iter( + run_agent( + runner, + buffer_intermediate_messages=True, + ) + ) + ) + await asyncio.wait_for(provider.started.wait(), timeout=5) + event.set_extra("agent_stop_requested", True) + provider.release.set() - assert [response.type for response in responses] == ["aborted"] + chains = await asyncio.wait_for(agent_task, timeout=5) + + assert chains == [] assert runner.was_aborted() is True - final_resp = runner.get_final_llm_resp() - assert final_resp is not None - assert final_resp.completion_text == "" - assert "这是我的最终回答" not in repr(runner.run_context.messages) + assert event.result is None @pytest.mark.asyncio @@ -1600,9 +1497,7 @@ async def test_live_agent_feeder_discards_residual_buffer_after_stop(): event = MockEvent("test:FriendMessage:live_late_buffer_abort", "u1") runner = LateStopAfterYieldRunner( event, - response_type="streaming_delta", text="partial without punctuation", - streaming=True, ) text_queue = asyncio.Queue() @@ -1620,21 +1515,24 @@ async def test_live_agent_feeder_discards_residual_buffer_after_stop(): assert text_queue.empty() +@pytest.mark.parametrize("native_stream", [False, True]) @pytest.mark.asyncio -async def test_live_agent_discards_queued_tts_audio_after_stop(tmp_path): +async def test_live_agent_discards_queued_tts_audio_after_stop(tmp_path, native_stream): from astrbot.core.astr_agent_run_util import run_live_agent + class Provider(MockDelayedTTSProvider): + def support_stream(self) -> bool: + return native_stream + audio_path = tmp_path / "late.wav" audio_path.write_bytes(b"late-audio") event = MockEvent("test:FriendMessage:live_tts_queued_abort", "u1") runner = LateStopAfterYieldRunner( event, - response_type="streaming_delta", text="complete sentence!", - streaming=True, stop_after_yield=False, ) - tts_provider = MockDelayedTTSProvider(audio_path) + tts_provider = Provider(audio_path) live_task = asyncio.create_task( _collect_async_iter(run_live_agent(cast(Any, runner), tts_provider)) @@ -1646,74 +1544,218 @@ async def test_live_agent_discards_queued_tts_audio_after_stop(tmp_path): chains = await asyncio.wait_for(live_task, timeout=5) assert chains == [] + assert event._temporary_local_files == [str(audio_path)] + event.cleanup_temporary_local_files() + assert not audio_path.exists() @pytest.mark.asyncio -async def test_stop_requested_before_stream_chunk_discards_chunk( - runner, provider_request, mock_tool_executor, mock_hooks +async def test_cancelled_genie_stream_removes_late_executor_file( + tmp_path, + monkeypatch, ): - provider = MockStopAwareStreamProvider() + """A cancelled native Genie worker must clean a file created later.""" + import threading + + from astrbot.core.astr_agent_run_util import _safe_tts_stream_wrapper + from astrbot.core.provider.sources import genie_tts + + started = threading.Event() + release = threading.Event() + finished = threading.Event() + + def generate(*, save_path, **kwargs): + started.set() + release.wait(timeout=1) + Path(save_path).write_bytes(b"late audio") + finished.set() + + monkeypatch.setattr(genie_tts, "genie", SimpleNamespace(tts=generate)) + monkeypatch.setattr(genie_tts, "get_astrbot_temp_path", lambda: str(tmp_path)) + provider = object.__new__(genie_tts.GenieTTSProvider) + provider.character_name = "test" + event = MockEvent("test:FriendMessage:genie_cancel_cleanup", "u1") + text_queue = asyncio.Queue() + audio_queue = asyncio.Queue() + await text_queue.put("sentence") - await runner.reset( - provider=provider, - request=provider_request, - run_context=ContextWrapper(context=None), - tool_executor=mock_tool_executor, - agent_hooks=mock_hooks, - streaming=True, + task = asyncio.create_task( + _safe_tts_stream_wrapper(provider, text_queue, audio_queue, event) ) + assert await asyncio.to_thread(started.wait, 1) + event.set_extra("agent_stop_requested", True) + task.cancel() + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(task, timeout=0.2) + release.set() + assert await asyncio.to_thread(finished.wait, 1) + + assert len(event._temporary_local_files) == 1 + generated_path = Path(event._temporary_local_files[0]) + for _ in range(100): + if not generated_path.exists(): + break + await asyncio.sleep(0.01) + assert not generated_path.exists() - runner.request_stop() - responses = [] - async for response in runner.step(): - responses.append(response) - assert [response.type for response in responses] == ["aborted"] - assert runner.was_aborted() is True - final_resp = runner.get_final_llm_resp() - assert final_resp is not None - assert final_resp.completion_text == "" - assert final_resp.reasoning_content is None +@pytest.mark.asyncio +async def test_cancelled_genie_get_audio_removes_late_executor_file( + tmp_path, + monkeypatch, +): + import threading + + from astrbot.core.provider.sources import genie_tts + + started = threading.Event() + release = threading.Event() + finished = threading.Event() + generated_paths = [] + + def generate(*, save_path, **kwargs): + generated_paths.append(Path(save_path)) + started.set() + release.wait(timeout=1) + Path(save_path).write_bytes(b"late audio") + finished.set() + + monkeypatch.setattr(genie_tts, "genie", SimpleNamespace(tts=generate)) + monkeypatch.setattr(genie_tts, "get_astrbot_temp_path", lambda: str(tmp_path)) + provider = object.__new__(genie_tts.GenieTTSProvider) + provider.character_name = "test" + + task = asyncio.create_task(provider.get_audio("sentence")) + assert await asyncio.to_thread(started.wait, 1) + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + release.set() + assert await asyncio.to_thread(finished.wait, 1) + + for _ in range(100): + if generated_paths and not generated_paths[0].exists(): + break + await asyncio.sleep(0.01) + assert generated_paths + assert not generated_paths[0].exists() @pytest.mark.asyncio -async def test_main_agent_hooks_skip_llm_response_hook_for_aborted_event(monkeypatch): - from astrbot.core import astr_agent_hooks - from astrbot.core.astr_agent_hooks import MainAgentHooks - from astrbot.core.star.star_handler import EventType +async def test_native_live_tts_does_not_start_queued_text_after_stop(tmp_path): + from astrbot.core.astr_agent_run_util import run_live_agent - calls = [] + class NativeTTS(MockDelayedTTSProvider): + def __init__(self): + super().__init__(tmp_path / "unused.wav") + self.texts = [] - async def fake_call_event_hook(event, hook_type, *args, **kwargs): - calls.append((hook_type, args)) - return False + def support_stream(self) -> bool: + return True - monkeypatch.setattr(astr_agent_hooks, "call_event_hook", fake_call_event_hook) + async def get_audio_stream(self, text_queue, audio_queue) -> None: + while (text := await text_queue.get()) is not None: + self.texts.append(text) + self.started.set() + await self.release.wait() + await audio_queue.put((text, b"audio")) - event = MockEvent("test:FriendMessage:hook_abort", "u1") - event.set_extra("agent_user_aborted", True) - response = LLMResponse( - role="assistant", - completion_text="", - reasoning_content="should not be stored", + event = MockEvent("test:FriendMessage:native_tts_stop", "u1") + runner = LateStopAfterYieldRunner( + event, + text="first complete sentence! second complete sentence!", + stop_after_yield=False, + ) + provider = NativeTTS() + task = asyncio.create_task( + _collect_async_iter(run_live_agent(cast(Any, runner), provider)) ) - hooks = MainAgentHooks() + await asyncio.wait_for(provider.started.wait(), timeout=1) + event.set_extra("agent_stop_requested", True) + provider.release.set() - await hooks.on_agent_done(ContextWrapper(context=MockAgentContext(event)), response) + assert await asyncio.wait_for(task, timeout=1) == [] + assert provider.texts == ["first complete sentence!"] - called_types = [hook_type for hook_type, _ in calls] - assert EventType.OnLLMResponseEvent not in called_types - assert EventType.OnAgentDoneEvent in called_types - assert event.get_extra("_llm_reasoning_content") is None - agent_done_call = next( - args for hook_type, args in calls if hook_type == EventType.OnAgentDoneEvent + +@pytest.mark.parametrize("returns_audio_path", [True, False]) +@pytest.mark.asyncio +async def test_result_decorate_stops_after_first_tts_result( + tmp_path, monkeypatch, returns_audio_path +): + """Normal result decoration must register a late TTS file before aborting.""" + from astrbot.core.message.components import Plain + from astrbot.core.message.message_event_result import ( + MessageEventResult, + ResultContentType, + ) + from astrbot.core.pipeline.result_decorate import stage as decorate_stage + + audio_path = tmp_path / "decorated-stop.wav" + audio_path.write_bytes(b"audio") + provider = MockDelayedTTSProvider(audio_path if returns_audio_path else None) + event = MockEvent("test:FriendMessage:decorate_tts_stop", "u1") + event.set_result( + MessageEventResult( + chain=[Plain("sentence"), Plain("second sentence")], + result_content_type=ResultContentType.LLM_RESULT, + ) + ) + stage = decorate_stage.ResultDecorateStage() + stage.content_safe_check_reply = False + stage.reply_prefix = "" + stage.enable_segmented_reply = False + stage.show_reasoning = False + stage.tts_trigger_probability = 1 + stage.ctx = SimpleNamespace( + plugin_manager=SimpleNamespace( + context=SimpleNamespace(get_using_tts_provider=lambda _umo: provider) + ), + astrbot_config={ + "provider_tts_settings": {"enable": True}, + }, + ) + monkeypatch.setattr( + decorate_stage.star_handlers_registry, + "get_handlers_by_event_type", + lambda *_args, **_kwargs: [], + ) + monkeypatch.setattr( + decorate_stage.SessionServiceManager, + "should_process_tts_request", + AsyncMock(return_value=True), + ) + + task = asyncio.create_task(_collect_async_iter(stage.process(event))) + await asyncio.wait_for(provider.started.wait(), timeout=1) + event.set_extra("agent_stop_requested", True) + provider.release.set() + await asyncio.wait_for(task, timeout=1) + + assert provider.call_count == 1 + assert event._temporary_local_files == ( + [str(audio_path)] if returns_audio_path else [] ) - assert agent_done_call[1].completion_text == "" - assert agent_done_call[1].reasoning_content is None + assert event.get_result() is None + event.cleanup_temporary_local_files() + if returns_audio_path: + assert not audio_path.exists() +@pytest.mark.parametrize( + ("stop_key", "expect_llm_hook"), + [ + (None, True), + ("agent_user_aborted", False), + ("agent_stop_requested", False), + ], +) @pytest.mark.asyncio -async def test_main_agent_hooks_skip_llm_response_hook_for_pending_stop(monkeypatch): +async def test_main_agent_hooks_respect_stop_state( + monkeypatch, + stop_key, + expect_llm_hook, +): from astrbot.core import astr_agent_hooks from astrbot.core.astr_agent_hooks import MainAgentHooks from astrbot.core.star.star_handler import EventType @@ -1724,68 +1766,48 @@ async def fake_call_event_hook(event, hook_type, *args, **kwargs): calls.append((hook_type, args)) return False + async def fake_call_agent_done_hook(event, run_context, llm_response): + calls.append((EventType.OnAgentDoneEvent, (run_context, llm_response))) + monkeypatch.setattr(astr_agent_hooks, "call_event_hook", fake_call_event_hook) + monkeypatch.setattr( + astr_agent_hooks, + "call_agent_done_hook", + fake_call_agent_done_hook, + ) - event = MockEvent("test:FriendMessage:hook_pending_stop", "u1") - event.set_extra("agent_stop_requested", True) + event = MockEvent("test:FriendMessage:hook_state", "u1") + if stop_key: + event.set_extra(stop_key, True) response = LLMResponse( role="assistant", - completion_text="late hook text", - reasoning_content="late hook reasoning", + completion_text="response text", + reasoning_content="response reasoning", + ) + await MainAgentHooks().on_agent_done( + ContextWrapper(context=MockAgentContext(event)), response ) - hooks = MainAgentHooks() - - await hooks.on_agent_done(ContextWrapper(context=MockAgentContext(event)), response) called_types = [hook_type for hook_type, _ in calls] - assert EventType.OnLLMResponseEvent not in called_types + assert (EventType.OnLLMResponseEvent in called_types) is expect_llm_hook assert EventType.OnAgentDoneEvent in called_types - assert event.get_extra("_llm_reasoning_content") is None - agent_done_call = next( - args for hook_type, args in calls if hook_type == EventType.OnAgentDoneEvent + assert event.get_extra("_llm_reasoning_content") == ( + "response reasoning" if expect_llm_hook else None ) - assert agent_done_call[1].completion_text == "" -def test_llm_response_hook_propagation_stops_on_pending_agent_stop(): +@pytest.mark.parametrize( + "hook_type_name", + ["OnLLMResponseEvent", "OnUsingLLMToolEvent", "OnLLMToolRespondEvent"], +) +def test_agent_output_hooks_stop_propagating_after_soft_stop(hook_type_name): from astrbot.core.pipeline.context_utils import _should_stop_hook_propagation from astrbot.core.star.star_handler import EventType event = MockEvent("test:FriendMessage:hook_propagation_stop", "u1") event.set_extra("agent_stop_requested", True) - assert _should_stop_hook_propagation(event, EventType.OnLLMResponseEvent) is True - assert _should_stop_hook_propagation(event, EventType.OnAgentDoneEvent) is False - - -@pytest.mark.asyncio -async def test_main_agent_hooks_keep_normal_llm_response_hook(monkeypatch): - from astrbot.core import astr_agent_hooks - from astrbot.core.astr_agent_hooks import MainAgentHooks - from astrbot.core.star.star_handler import EventType - - calls = [] - - async def fake_call_event_hook(event, hook_type, *args, **kwargs): - calls.append((hook_type, args)) - return False - - monkeypatch.setattr(astr_agent_hooks, "call_event_hook", fake_call_event_hook) - - event = MockEvent("test:FriendMessage:hook_normal", "u1") - response = LLMResponse( - role="assistant", - completion_text="normal text", - reasoning_content="normal reasoning", - ) - hooks = MainAgentHooks() - - await hooks.on_agent_done(ContextWrapper(context=MockAgentContext(event)), response) - - called_types = [hook_type for hook_type, _ in calls] - assert EventType.OnLLMResponseEvent in called_types - assert EventType.OnAgentDoneEvent in called_types - assert event.get_extra("_llm_reasoning_content") == "normal reasoning" + assert _should_stop_hook_propagation(event, EventType[hook_type_name]) is True @pytest.mark.asyncio @@ -1959,8 +1981,6 @@ async def test_follow_up_ticket_not_consumed_when_no_next_tool_call( @pytest.mark.asyncio async def test_skills_like_requery_passes_extra_user_content_parts(): """skills-like 模式 re-query 时应传递 extra_user_content_parts（如 image_caption）""" - from astrbot.core.agent.message import TextPart - captured_kwargs = {} class SkillsLikeProvider(MockProvider): @@ -2038,8 +2058,12 @@ async def text_chat(self, **kwargs) -> LLMResponse: @pytest.mark.asyncio -async def test_skills_like_requery_stop_discards_late_assistant_text(): - class SkillsLikeLateStopProvider(MockProvider): +async def test_skills_like_requery_stop_skips_repair_request( + provider_request, + mock_tool_executor, + mock_hooks, +): + class SkillsLikeRepairStopProvider(MockProvider): def __init__(self, event): super().__init__() self.event = event @@ -2053,66 +2077,43 @@ async def text_chat(self, **kwargs) -> LLMResponse: tools_call_name=["test_tool"], tools_call_args=[{"query": "test"}], tools_call_ids=["call_1"], - usage=TokenUsage(input_other=10, output=5), ) self.event.set_extra("agent_stop_requested", True) - return LLMResponse( - role="assistant", - completion_text="skills_like late text", - usage=TokenUsage(input_other=10, output=5), - ) + return LLMResponse(role="assistant", completion_text="") event = MockEvent(umo="test_umo", sender_id="test_sender") - provider = SkillsLikeLateStopProvider(event) - tool = FunctionTool( - name="test_tool", - description="测试", - parameters={"type": "object", "properties": {"query": {"type": "string"}}}, - handler=AsyncMock(), - ) - req = ProviderRequest( - prompt="调用工具", - func_tool=ToolSet(tools=[tool]), - contexts=[], - ) + provider = SkillsLikeRepairStopProvider(event) runner = ToolLoopAgentRunner() - hooks = MockHooks() await runner.reset( provider=provider, - request=req, + request=provider_request, run_context=ContextWrapper(context=MockAgentContext(event)), - tool_executor=cast(Any, MockToolExecutor()), - agent_hooks=hooks, + tool_executor=mock_tool_executor, + agent_hooks=mock_hooks, tool_schema_mode="skills_like", ) responses = [response async for response in runner.step()] assert [response.type for response in responses] == ["llm_result", "aborted"] - visible_text = "".join( - response.data["chain"].get_plain_text() - for response in responses - if response.type != "aborted" - ) - assert "skills_like late text" not in visible_text - assert runner.was_aborted() is True - final_resp = runner.get_final_llm_resp() - assert final_resp is not None - assert final_resp.completion_text == "" - assert hooks.agent_done_response is final_resp - assert "skills_like late text" not in repr(runner.run_context.messages) + assert provider.call_count == 2 @pytest.mark.asyncio -async def test_skills_like_requery_stop_skips_repair_request(): - class SkillsLikeRepairStopProvider(MockProvider): - def __init__(self, event): +async def test_stop_interrupts_pending_skills_like_requery( + provider_request, + mock_tool_executor, + mock_hooks, +): + class PendingRequeryProvider(MockProvider): + def __init__(self): super().__init__() - self.event = event + self.requery_started = asyncio.Event() + self.requery_cancelled = asyncio.Event() - async def text_chat(self, **kwargs) -> LLMResponse: + async def text_chat(self, **_kwargs) -> LLMResponse: self.call_count += 1 if self.call_count == 1: return LLMResponse( @@ -2123,41 +2124,39 @@ async def text_chat(self, **kwargs) -> LLMResponse: tools_call_ids=["call_1"], ) - self.event.set_extra("agent_stop_requested", True) - return LLMResponse(role="assistant", completion_text="") + self.requery_started.set() + try: + await asyncio.Event().wait() + finally: + self.requery_cancelled.set() - event = MockEvent(umo="test_umo", sender_id="test_sender") - provider = SkillsLikeRepairStopProvider(event) - tool = FunctionTool( - name="test_tool", - description="测试", - parameters={"type": "object", "properties": {"query": {"type": "string"}}}, - handler=AsyncMock(), - ) - req = ProviderRequest( - prompt="调用工具", - func_tool=ToolSet(tools=[tool]), - contexts=[], - ) + event = MockEvent("test:FriendMessage:skills_requery_pending_stop", "u1") + provider = PendingRequeryProvider() runner = ToolLoopAgentRunner() - await runner.reset( provider=provider, - request=req, + request=provider_request, run_context=ContextWrapper(context=MockAgentContext(event)), - tool_executor=cast(Any, MockToolExecutor()), - agent_hooks=MockHooks(), + tool_executor=mock_tool_executor, + agent_hooks=mock_hooks, tool_schema_mode="skills_like", ) - responses = [response async for response in runner.step()] + task = asyncio.create_task(_collect_async_iter(runner.step())) + await asyncio.wait_for(provider.requery_started.wait(), timeout=1) + event.set_extra("agent_stop_requested", True) + responses = await asyncio.wait_for(task, timeout=1) assert [response.type for response in responses] == ["llm_result", "aborted"] assert provider.call_count == 2 + assert provider.requery_cancelled.is_set() is True @pytest.mark.asyncio -async def test_skills_like_requery_fallback_checks_stop_before_yielding_text(): +async def test_skills_like_requery_fallback_checks_stop_before_yielding_text( + provider_request, + mock_tool_executor, +): class StopAfterFallbackDoneHooks(MockHooks): async def on_agent_done(self, run_context, llm_response): await super().on_agent_done(run_context, llm_response) @@ -2185,25 +2184,14 @@ async def text_chat(self, **kwargs) -> LLMResponse: event = MockEvent(umo="test_umo", sender_id="test_sender") provider = SkillsLikeFallbackProvider() - tool = FunctionTool( - name="test_tool", - description="测试", - parameters={"type": "object", "properties": {"query": {"type": "string"}}}, - handler=AsyncMock(), - ) - req = ProviderRequest( - prompt="调用工具", - func_tool=ToolSet(tools=[tool]), - contexts=[], - ) runner = ToolLoopAgentRunner() hooks = StopAfterFallbackDoneHooks() await runner.reset( provider=provider, - request=req, + request=provider_request, run_context=ContextWrapper(context=MockAgentContext(event)), - tool_executor=cast(Any, MockToolExecutor()), + tool_executor=mock_tool_executor, agent_hooks=hooks, tool_schema_mode="skills_like", ) @@ -2532,121 +2520,557 @@ async def test_follow_up_merged_into_tool_result_before_stop( @pytest.mark.asyncio -async def test_follow_up_rejected_and_runner_stops_without_execution( - runner, mock_provider, provider_request, mock_tool_executor, mock_hooks +async def test_empty_output_stop_prevents_retry( + provider_request, + mock_tool_executor, + mock_hooks, ): - """Test that when stop is requested before execution, follow-ups are rejected and runner stops gracefully.""" + """A stop raised with the first empty output must suppress later attempts.""" - mock_event = MockEvent("test:FriendMessage:stop_before_execution", "u1") - run_context = ContextWrapper(context=MockAgentContext(mock_event)) + class StopAfterEmptyProvider(MockProvider): + def __init__(self, event): + super().__init__() + self.event = event + async def text_chat(self, **kwargs) -> LLMResponse: + self.call_count += 1 + self.event.set_extra("agent_stop_requested", True) + raise EmptyModelOutputError("empty") + + event = MockEvent("test:FriendMessage:empty_retry_stop", "u1") + provider = StopAfterEmptyProvider(event) + runner = ToolLoopAgentRunner() await runner.reset( - provider=mock_provider, + provider=provider, request=provider_request, - run_context=run_context, + run_context=ContextWrapper(context=MockAgentContext(event)), tool_executor=mock_tool_executor, agent_hooks=mock_hooks, streaming=False, ) - # Request stop before any execution (simulates /stop command received at start) - runner.request_stop() - assert runner._is_stop_requested() is True + responses = [response async for response in runner.step()] - # Try to add follow-up after stop (should be rejected) - ticket_after = runner.follow_up(message_text="follow-up after stop") - assert ticket_after is None, "Post-stop follow-up should be rejected" + assert provider.call_count == 1 + assert [response.type for response in responses] == ["aborted"] - # Verify queue is empty - assert len(runner._pending_follow_ups) == 0 - # Run the agent step - should stop immediately without executing tools - async for response in runner.step(): - # Should yield an aborted response - if response.type == "aborted": - break +@pytest.mark.asyncio +async def test_streaming_output_is_not_retried_after_empty_output_error( + provider_request, + mock_tool_executor, + mock_hooks, +): + """A broken stream must not replay already published chunks.""" - # Verify runner stopped gracefully - assert runner.done() - assert runner.was_aborted() + class StreamThenEmptyProvider(MockProvider): + async def text_chat_stream(self, **kwargs): + self.call_count += 1 + chunk = LLMResponse(role="assistant", completion_text="partial") + chunk.is_chunk = True + yield chunk + raise EmptyModelOutputError("empty after stream started") + + provider = StreamThenEmptyProvider() + event = MockEvent("test:FriendMessage:stream_empty_no_retry", "u1") + runner = ToolLoopAgentRunner() + await runner.reset( + provider=provider, + request=provider_request, + run_context=ContextWrapper(context=MockAgentContext(event)), + tool_executor=mock_tool_executor, + agent_hooks=mock_hooks, + streaming=True, + ) + + responses = [response async for response in runner.step()] - # No tool execution should have occurred - assert provider_request.tool_calls_result is None + assert provider.call_count == 1 + assert [response.type for response in responses] == ["streaming_delta", "err"] @pytest.mark.asyncio -async def test_follow_up_after_stop_not_merged_into_tool_result( - runner, mock_provider, provider_request, mock_tool_executor, mock_hooks +async def test_empty_output_backoff_is_interrupted_by_event_stop( + provider_request, + mock_tool_executor, + mock_hooks, ): - """Regression test for issue #6626: verify post-stop follow-ups are not injected into tool results. + """A soft stop during retry backoff must wake without another request.""" - This test simulates the race condition where: - 1. Runner is active and executing tools - 2. A follow-up is queued (should be included in tool result) - 3. Stop is requested - 4. Another follow-up is attempted (should be rejected) - 5. Tool execution completes and merges follow-ups into result + class BackoffProvider(MockProvider): + def __init__(self): + super().__init__() + self.first_attempt_finished = asyncio.Event() - The key assertion is that only pre-stop follow-ups are merged into the tool result. - """ + async def text_chat(self, **kwargs) -> LLMResponse: + self.call_count += 1 + self.first_attempt_finished.set() + raise EmptyModelOutputError("empty") - mock_event = MockEvent("test:FriendMessage:regression_6626", "u1") - run_context = ContextWrapper(context=MockAgentContext(mock_event)) + event = MockEvent("test:FriendMessage:backoff_stop", "u1") + provider = BackoffProvider() + runner = ToolLoopAgentRunner() + await runner.reset( + provider=provider, + request=provider_request, + run_context=ContextWrapper(context=MockAgentContext(event)), + tool_executor=mock_tool_executor, + agent_hooks=mock_hooks, + streaming=False, + ) + + task = asyncio.create_task(_collect_async_iter(runner.step())) + await asyncio.wait_for(provider.first_attempt_finished.wait(), timeout=1) + started = asyncio.get_running_loop().time() + event.set_extra("agent_stop_requested", True) + responses = await asyncio.wait_for(task, timeout=1) + elapsed = asyncio.get_running_loop().time() - started + + assert provider.call_count == 1 + assert elapsed < 0.5 + assert [response.type for response in responses] == ["aborted"] + + +@pytest.mark.parametrize("streaming", [False, True]) +@pytest.mark.asyncio +async def test_stop_interrupts_provider_internal_retry_backoff( + monkeypatch, + streaming, + provider_request, + mock_tool_executor, + mock_hooks, +): + monkeypatch.setattr(request_retry, "REQUEST_RETRY_WAIT_MIN_S", 10) + monkeypatch.setattr(request_retry, "REQUEST_RETRY_WAIT_MAX_S", 10) + class InternalRetryProvider(MockProvider): + def __init__(self): + super().__init__() + self.first_attempt_finished = asyncio.Event() + + async def _request(self) -> LLMResponse: + self.call_count += 1 + if self.call_count == 1: + self.first_attempt_finished.set() + raise OSError("retryable") + return LLMResponse(role="assistant", completion_text="late") + + async def text_chat(self, **_kwargs) -> LLMResponse: + return await retry_provider_request( + "runner-stop-test", + self._request, + max_attempts=2, + ) + + async def text_chat_stream(self, **_kwargs): + yield await retry_provider_request( + "runner-stop-test", + self._request, + max_attempts=2, + ) + + event = MockEvent("test:FriendMessage:provider_internal_retry_stop", "u1") + provider = InternalRetryProvider() + runner = ToolLoopAgentRunner() await runner.reset( - provider=mock_provider, + provider=provider, request=provider_request, - run_context=run_context, + run_context=ContextWrapper(context=MockAgentContext(event)), tool_executor=mock_tool_executor, agent_hooks=mock_hooks, + streaming=streaming, + ) + + task = asyncio.create_task(_collect_async_iter(runner.step())) + await asyncio.wait_for(provider.first_attempt_finished.wait(), timeout=1) + started = asyncio.get_running_loop().time() + event.set_extra("agent_stop_requested", True) + responses = await asyncio.wait_for(task, timeout=1) + elapsed = asyncio.get_running_loop().time() - started + + assert provider.call_count == 1 + assert elapsed < 0.5 + assert [response.type for response in responses] == ["aborted"] + + +@pytest.mark.parametrize("tool_schema_mode", ["skills_like", "full"]) +@pytest.mark.asyncio +async def test_stop_after_tool_selection_yield_blocks_requery_and_executor( + tool_schema_mode, + provider_request, +): + """Resuming a tool-selection yield after stop must not start new work.""" + + class CountingExecutor: + def __init__(self): + self.calls = 0 + + def execute(self, tool, run_context, **tool_args): + self.calls += 1 + raise AssertionError("executor must not be constructed after stop") + + event = MockEvent(f"test:FriendMessage:{tool_schema_mode}_yield_stop", "u1") + provider = MockProvider() + hooks = MockHooks() + executor = CountingExecutor() + runner = ToolLoopAgentRunner() + await runner.reset( + provider=provider, + request=provider_request, + run_context=ContextWrapper(context=MockAgentContext(event)), + tool_executor=cast(Any, executor), + agent_hooks=hooks, streaming=False, + tool_schema_mode=tool_schema_mode, ) - # Add a follow-up before stop (should be included in tool result) - ticket_before = runner.follow_up(message_text="valid before stop") - assert ticket_before is not None - assert ticket_before in runner._pending_follow_ups + step = runner.step() + first = await anext(step) + event.set_extra("agent_stop_requested", True) + remaining = [response async for response in step] - # Request stop (simulates /stop command during active execution) - runner.request_stop() - assert runner._is_stop_requested() is True + assert first.type == "llm_result" + assert [response.type for response in remaining] == ["aborted"] + assert provider.call_count == 1 + assert hooks.tool_start_called is False + assert executor.calls == 0 - # Try to add follow-up after stop (should be rejected) - ticket_after = runner.follow_up(message_text="invalid after stop") - assert ticket_after is None, "Post-stop follow-up should be rejected" - # Verify queue only contains pre-stop follow-up - assert len(runner._pending_follow_ups) == 1 - assert runner._pending_follow_ups[0].text == "valid before stop" +@pytest.mark.parametrize("delivered", [False, True]) +@pytest.mark.asyncio +async def test_final_response_history_follows_delivery_boundary( + delivered, + provider_request, + mock_tool_executor, + mock_hooks, +): + """Only a response committed at delivery survives a late stop.""" + from astrbot.core.agent.stop_policy import AGENT_OUTPUT_DELIVERY_CONFIRMED_KEY + from astrbot.core.pipeline.process_stage.method.agent_sub_stages.internal import ( + InternalAgentSubStage, + ) - # Run the agent step - this will execute tool and merge follow-ups into result - async for response in runner.step(): - # The runner should execute tools and then stop - pass + event = MockEvent(f"test:FriendMessage:delivery_{delivered}", "u1") + provider = MockProvider() + provider.should_call_tools = False + runner = ToolLoopAgentRunner() + await runner.reset( + provider=provider, + request=provider_request, + run_context=ContextWrapper(context=MockAgentContext(event)), + tool_executor=mock_tool_executor, + agent_hooks=mock_hooks, + streaming=False, + ) - # Verify tool result was created with follow-up merged - # Note: When stop is requested, the tool may or may not execute depending on timing. - # The key assertion is that IF tool_calls_result exists, it only contains pre-stop follow-ups. - if provider_request.tool_calls_result is not None: - assert isinstance(provider_request.tool_calls_result, list) - assert provider_request.tool_calls_result - tool_result = str( - provider_request.tool_calls_result[0].tool_calls_result[0].content - ) + step = runner.step() + first = await anext(step) + assert first.type == "llm_result" + if delivered: + event.set_extra(AGENT_OUTPUT_DELIVERY_CONFIRMED_KEY, True) + else: + assert not InternalAgentSubStage._confirm_agent_output_delivery(event, runner) + event.set_extra("agent_stop_requested", True) + assert not InternalAgentSubStage._confirm_agent_output_delivery(event, runner) + remaining = [response async for response in step] + + assert [response.type for response in remaining] == ( + [] if delivered else ["aborted"] + ) + assert runner.was_aborted() is not delivered + assert ("这是我的最终回答" in repr(runner.run_context.messages)) is delivered + + +@pytest.mark.parametrize("stop_kind", ["mid_dispatch", "soft", "hard"]) +@pytest.mark.asyncio +async def test_agent_done_handlers_receive_isolated_responses_after_stop( + monkeypatch, stop_kind +): + """All done handlers run, with fresh safe responses after any stop.""" + from astrbot.core.pipeline import context_utils + + event = MockEvent(f"test:FriendMessage:done_{stop_kind}_stop", "u1") + if stop_kind == "soft": + event.set_extra("agent_stop_requested", True) + elif stop_kind == "hard": + event.stop() + + first_started = asyncio.Event() + release_first = asyncio.Event() + seen = [] + + async def first_handler(event, run_context, response): + seen.append(("first", response)) + if stop_kind == "mid_dispatch": + response.id = "poisoned-id" + response.usage.output = 998 + first_started.set() + await release_first.wait() + + async def second_handler(event, run_context, response): + seen.append(("second", response)) + response.completion_text = "plugin poison" + response.usage.output = 999 + event.set_extra("agent_stop_requested", False) + event.set_extra("agent_user_aborted", False) + event._stopped = False + + async def third_handler(event, run_context, response): + seen.append(("third", response)) - # Should contain the pre-stop follow-up - assert "valid before stop" in tool_result + handlers = [ + SimpleNamespace( + handler=handler, + handler_module_path=f"stop_test_plugin_{index}", + handler_name=f"handler_{index}", + ) + for index, handler in enumerate( + [first_handler, second_handler, third_handler], start=1 + ) + ] + monkeypatch.setattr( + context_utils.star_handlers_registry, + "get_handlers_by_event_type", + lambda *args, **kwargs: handlers, + ) + monkeypatch.setattr( + context_utils, + "star_map", + { + handler.handler_module_path: SimpleNamespace( + name=handler.handler_module_path + ) + for handler in handlers + }, + ) + response = LLMResponse( + role="assistant", + completion_text="raw secret", + reasoning_content="raw reasoning", + id="response-id", + usage=TokenUsage(input_other=1, output=2), + ) - # Should NOT contain the post-stop follow-up - assert "invalid after stop" not in tool_result - assert "after stop" not in tool_result or "after stop" in "valid before stop" + task = asyncio.create_task( + context_utils.call_agent_done_hook(event, object(), response) + ) + if stop_kind == "mid_dispatch": + await asyncio.wait_for(first_started.wait(), timeout=1) + event.set_extra("agent_stop_requested", True) + release_first.set() + await asyncio.wait_for(task, timeout=1) - # Ticket should be marked as consumed (merged into tool result) - assert ticket_before.consumed is True + assert [name for name, _ in seen] == ["first", "second", "third"] + if stop_kind == "mid_dispatch": + assert seen[0][1] is response + assert seen[0][1].completion_text == "raw secret" else: - # If tool execution was aborted by stop, the ticket should still be resolved - # but not consumed (since there was no tool call to merge into) - assert ticket_before.resolved.is_set() + assert seen[0][1] is not response + assert seen[0][1].completion_text == "" + assert seen[0][1] is not seen[1][1] + assert seen[1][1].completion_text == "plugin poison" + assert seen[2][1].completion_text == "" + assert seen[1][1] is not seen[2][1] + assert seen[2][1].reasoning_content is None + assert seen[2][1].id == "response-id" + assert seen[2][1].usage.output == 2 + assert response.usage.output == (998 if stop_kind == "mid_dispatch" else 2) + assert event.get_extra("agent_stop_requested") is True + + +@pytest.mark.asyncio +async def test_agent_done_handler_cancellation_is_not_swallowed(monkeypatch): + """Task cancellation must escape lifecycle hook isolation.""" + from astrbot.core.pipeline import context_utils + + started = asyncio.Event() + + async def blocking_handler(event, run_context, response): + started.set() + await asyncio.Event().wait() + + handler = SimpleNamespace( + handler=blocking_handler, + handler_module_path="cancel_test_plugin", + handler_name="blocking_handler", + ) + monkeypatch.setattr( + context_utils.star_handlers_registry, + "get_handlers_by_event_type", + lambda *args, **kwargs: [handler], + ) + monkeypatch.setattr( + context_utils, + "star_map", + {handler.handler_module_path: SimpleNamespace(name="cancel_test_plugin")}, + ) + event = MockEvent("test:FriendMessage:done_cancel", "u1") + task = asyncio.create_task( + context_utils.call_agent_done_hook( + event, + object(), + LLMResponse(role="assistant", completion_text="late"), + ) + ) + await asyncio.wait_for(started.wait(), timeout=1) + + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + + +@pytest.mark.asyncio +async def test_agent_output_hook_cancellation_is_not_swallowed(monkeypatch): + from astrbot.core.pipeline import context_utils + from astrbot.core.star.star_handler import EventType + + started = asyncio.Event() + + async def blocking_handler(event): + started.set() + await asyncio.Event().wait() + + handler = SimpleNamespace( + handler=blocking_handler, + handler_module_path="cancel_output_plugin", + handler_name="blocking_handler", + ) + monkeypatch.setattr( + context_utils.star_handlers_registry, + "get_handlers_by_event_type", + lambda *args, **kwargs: [handler], + ) + monkeypatch.setattr( + context_utils, + "star_map", + {handler.handler_module_path: SimpleNamespace(name="cancel_output_plugin")}, + ) + event = MockEvent("test:FriendMessage:output_cancel", "u1") + task = asyncio.create_task( + context_utils.call_event_hook(event, EventType.OnLLMResponseEvent) + ) + await asyncio.wait_for(started.wait(), timeout=1) + + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + + +@pytest.mark.asyncio +async def test_result_decorate_hook_cancellation_is_not_swallowed(monkeypatch): + from astrbot.core.message.components import Plain + from astrbot.core.message.message_event_result import MessageEventResult + from astrbot.core.pipeline.result_decorate import stage as decorate_stage + + started = asyncio.Event() + + async def blocking_handler(event): + started.set() + await asyncio.Event().wait() + + handler = SimpleNamespace( + handler=blocking_handler, + handler_module_path="cancel_decorate_plugin", + handler_name="blocking_handler", + ) + monkeypatch.setattr( + decorate_stage.star_handlers_registry, + "get_handlers_by_event_type", + lambda *args, **kwargs: [handler], + ) + monkeypatch.setattr( + decorate_stage, + "star_map", + {handler.handler_module_path: SimpleNamespace(name="cancel_decorate_plugin")}, + ) + event = MockEvent("test:FriendMessage:decorate_cancel", "u1") + event.set_result(MessageEventResult(chain=[Plain("answer")])) + stage = decorate_stage.ResultDecorateStage() + stage.content_safe_check_reply = False + + task = asyncio.create_task(_collect_async_iter(stage.process(event))) + await asyncio.wait_for(started.wait(), timeout=1) + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + + +@pytest.mark.asyncio +async def test_snapshot_rollback_removes_mutated_hook_history_from_next_turn(): + """Rollback and real history serialization must exclude every final secret.""" + from astrbot.core.pipeline.process_stage.method.agent_sub_stages.internal import ( + InternalAgentSubStage, + ) + + class MutatingStopHooks(MockHooks): + async def on_agent_done(self, run_context, llm_response): + llm_response.completion_text = "mutated completion secret" + llm_response.reasoning_content = "mutated reasoning secret" + llm_response.result_chain = MessageChain().message("mutated chain secret") + llm_response.raw_completion = {"secret": "mutated raw secret"} + run_context.messages.reverse() + run_context.messages.append( + Message(role="assistant", content="hook appended secret") + ) + event.set_extra("agent_stop_requested", True) + + event = MockEvent("test:FriendMessage:history_snapshot_stop", "u1") + conversation = SimpleNamespace(cid="conversation-id", token_usage=0) + request = ProviderRequest( + prompt="original question", + contexts=[], + conversation=cast(Any, conversation), + ) + provider = MockProvider() + provider.should_call_tools = False + runner = ToolLoopAgentRunner() + await runner.reset( + provider=provider, + request=request, + run_context=ContextWrapper(context=MockAgentContext(event)), + tool_executor=MockToolExecutor(), + agent_hooks=MutatingStopHooks(), + streaming=False, + ) + + responses = [response async for response in runner.step()] + + assert [response.type for response in responses] == ["aborted"] + serialized_messages = repr(runner.run_context.messages) + for secret in ( + "这是我的最终回答", + "mutated completion secret", + "mutated reasoning secret", + "mutated chain secret", + "mutated raw secret", + "hook appended secret", + ): + assert secret not in serialized_messages + + update_conversation = AsyncMock() + stage = InternalAgentSubStage() + stage.conv_manager = SimpleNamespace(update_conversation=update_conversation) + await stage._save_to_history( + event, + request, + runner.get_final_llm_resp(), + runner.run_context.messages, + runner.stats, + user_aborted=True, + ) + saved_history = update_conversation.await_args.kwargs["history"] + assert "secret" not in repr(saved_history) + + next_runner = ToolLoopAgentRunner() + await next_runner.reset( + provider=MockProvider(), + request=ProviderRequest(prompt="next question", contexts=saved_history), + run_context=ContextWrapper(context=MockAgentContext(event)), + tool_executor=MockToolExecutor(), + agent_hooks=MockHooks(), + streaming=False, + ) + assert "secret" not in repr(next_runner.run_context.messages) if __name__ == "__main__": diff --git a/tests/unit/test_aiocqhttp_reply.py b/tests/unit/test_aiocqhttp_reply.py index ac32bee6d3..cffde39d63 100644 --- a/tests/unit/test_aiocqhttp_reply.py +++ b/tests/unit/test_aiocqhttp_reply.py @@ -10,11 +10,14 @@ 协议端（napcat/Lagrange）查找消息时失败。 """ +import asyncio +from types import SimpleNamespace from unittest.mock import AsyncMock import pytest import astrbot.core.message.components as Comp +from astrbot.core.agent.stop_policy import AgentOutputStopped from astrbot.core.message.message_event_result import MessageChain from astrbot.core.pipeline.respond.stage import ( RespondStage, # noqa: F401 — 预加载避免循环导入 @@ -70,6 +73,49 @@ def test_reply_to_dict_outputs_only_id(): # ============================================================ +@pytest.mark.asyncio +async def test_onebot_stop_after_conversion_blocks_dispatch(monkeypatch): + extras = {} + converted = asyncio.Event() + release_conversion = asyncio.Event() + + async def convert(_chain): + converted.set() + await release_conversion.wait() + return [{"type": "text", "data": {"text": "late"}}] + + dispatch = AsyncMock() + monkeypatch.setattr( + AiocqhttpMessageEvent, + "_parse_onebot_json", + staticmethod(convert), + ) + monkeypatch.setattr( + AiocqhttpMessageEvent, + "_dispatch_send", + staticmethod(dispatch), + ) + stop_event = SimpleNamespace( + is_stopped=lambda: False, + get_extra=lambda key, default=None: extras.get(key, default), + ) + task = asyncio.create_task( + AiocqhttpMessageEvent.send_message( + object(), + MessageChain([Comp.Plain("late")]), + session_id="1", + stop_event=stop_event, + ) + ) + await asyncio.wait_for(converted.wait(), timeout=1) + extras["agent_stop_requested"] = True + release_conversion.set() + + with pytest.raises(AgentOutputStopped): + await task + dispatch.assert_not_awaited() + + @pytest.mark.asyncio async def test_parse_onebot_json_reply_produces_extra_fields(): """_parse_onebot_json 处理 Reply 时会输出多余字段。 diff --git a/tests/unit/test_astr_message_event.py b/tests/unit/test_astr_message_event.py index 89087d1cab..a97ff69a7c 100644 --- a/tests/unit/test_astr_message_event.py +++ b/tests/unit/test_astr_message_event.py @@ -1,10 +1,12 @@ """Tests for AstrMessageEvent class.""" +import asyncio import re from unittest.mock import AsyncMock, patch import pytest +from astrbot.core.agent.stop_policy import AgentOutputStopped from astrbot.core.message.components import ( At, AtAll, @@ -522,6 +524,39 @@ async def test_process_buffer_no_match(self, astr_message_event): assert result == "No newlines here" + @pytest.mark.asyncio + async def test_process_buffer_stops_during_segment_delay( + self, + astr_message_event, + monkeypatch, + ): + """A stop during fallback rate limiting must suppress later segments.""" + sleep_started = asyncio.Event() + release_sleep = asyncio.Event() + + async def blocked_sleep(_delay): + sleep_started.set() + await release_sleep.wait() + + monkeypatch.setattr( + "astrbot.core.platform.astr_message_event.asyncio.sleep", + blocked_sleep, + ) + pattern = re.compile(r".*\n") + with patch.object( + astr_message_event, "send", new_callable=AsyncMock + ) as mock_send: + task = asyncio.create_task( + astr_message_event.process_buffer("one\ntwo\nthree", pattern) + ) + await asyncio.wait_for(sleep_started.wait(), timeout=1) + astr_message_event.set_extra("agent_stop_requested", True) + release_sleep.set() + with pytest.raises(AgentOutputStopped): + await asyncio.wait_for(task, timeout=1) + + assert mock_send.await_count == 1 + class TestResultHelpers: """Tests for result helper methods.""" diff --git a/tests/unit/test_message_tools.py b/tests/unit/test_message_tools.py index 5853f2851b..024c1553d2 100644 --- a/tests/unit/test_message_tools.py +++ b/tests/unit/test_message_tools.py @@ -1,11 +1,18 @@ """Tests for send_message_to_user session handling.""" +import asyncio from types import SimpleNamespace from unittest.mock import AsyncMock import pytest -from astrbot.core.message.message_event_result import MessageEventResult +from astrbot.core.agent.stop_policy import AGENT_OUTPUT_DELIVERY_CONFIRMED_KEY +from astrbot.core.message.components import Plain +from astrbot.core.message.message_event_result import ( + MessageChain, + MessageEventResult, + ResultContentType, +) from astrbot.core.pipeline.respond.stage import RespondStage from astrbot.core.tools.message_tools import SendMessageToUserTool @@ -51,6 +58,8 @@ def __init__(self, result_text: str, sent_plain_texts: list[str]) -> None: self._result = MessageEventResult().message(result_text) self.send = AsyncMock() self.plugins_name = [] + self.platform_meta = SimpleNamespace(name="test") + self._has_send_oper = False def get_result(self): """Return the current message result.""" @@ -188,20 +197,164 @@ async def test_respond_stage_skips_same_text_after_send_message_to_user(): event.send.assert_not_awaited() +@pytest.mark.parametrize("send_fails", [False, True]) @pytest.mark.asyncio -async def test_respond_stage_sends_different_text_after_send_message_to_user(): - """RespondStage still sends a distinct completion after the tool call.""" +async def test_respond_stage_commits_only_successful_send(send_fails): + """Commit immediately after success, but never after a failed write.""" stage = _make_respond_stage() event = _DummyRespondEvent( result_text="I have sent the message with the tool.", sent_plain_texts=["duplicate reply"], ) + async def send_then_stop(_chain): + event.set_extra("agent_stop_requested", True) + + event.send.side_effect = ( + RuntimeError("network failed before write") if send_fails else send_then_stop + ) + result = await stage.process(event) assert result is None event.send.assert_awaited_once() assert event.get_result() is None + assert bool(event.get_extra(AGENT_OUTPUT_DELIVERY_CONFIRMED_KEY, False)) is ( + not send_fails + ) + + +@pytest.mark.asyncio +async def test_respond_stage_stops_during_segment_delay(monkeypatch): + """A stop during segmented delay must suppress every pending component.""" + stage = _make_respond_stage() + stage.enable_seg = True + stage.only_llm_result = False + stage._calc_comp_interval = AsyncMock(return_value=1) + event = _DummyRespondEvent("unused", []) + event._result.chain = [Plain("first"), Plain("second")] + sleep_started = asyncio.Event() + release_sleep = asyncio.Event() + + async def blocked_sleep(_delay): + sleep_started.set() + await release_sleep.wait() + + monkeypatch.setattr( + "astrbot.core.pipeline.respond.stage.asyncio.sleep", blocked_sleep + ) + task = asyncio.create_task(stage.process(event)) + await asyncio.wait_for(sleep_started.wait(), timeout=1) + event.set_extra("agent_stop_requested", True) + release_sleep.set() + await asyncio.wait_for(task, timeout=1) + + event.send.assert_not_awaited() + assert event.get_result() is None + assert not event.get_extra(AGENT_OUTPUT_DELIVERY_CONFIRMED_KEY, False) + + +@pytest.mark.parametrize("exit_mode", ["stop", "adapter_return"]) +@pytest.mark.asyncio +async def test_respond_stage_closes_undelivered_stream(exit_mode): + """Stopped and prematurely abandoned streams must close without a flush.""" + stage = _make_respond_stage() + event = _DummyRespondEvent("unused", []) + source_waiting = asyncio.Event() + release_source = asyncio.Event() + source_closed = asyncio.Event() + flushed = [] + + async def source_stream(): + try: + yield MessageChain().message("buffered secret") + source_waiting.set() + await release_source.wait() + finally: + source_closed.set() + + async def buffering_adapter(generator, _use_fallback): + buffered = [] + async for chain in generator: + buffered.append(chain.get_plain_text()) + if exit_mode == "adapter_return": + return + flushed.extend(buffered) + + event.send_streaming = buffering_adapter + event._result = ( + MessageEventResult() + .set_result_content_type(ResultContentType.STREAMING_RESULT) + .set_async_stream(source_stream()) + ) + task = asyncio.create_task(stage.process(event)) + if exit_mode == "stop": + await asyncio.wait_for(source_waiting.wait(), timeout=1) + event.set_extra("agent_stop_requested", True) + release_source.set() + await asyncio.wait_for(task, timeout=1) + + assert flushed == [] + assert source_closed.is_set() + assert event.get_result() is None + assert not event.get_extra(AGENT_OUTPUT_DELIVERY_CONFIRMED_KEY, False) + + +@pytest.mark.parametrize("empty_stream_kind", ["none", "break", "whitespace"]) +@pytest.mark.asyncio +async def test_respond_stage_does_not_commit_empty_stream(empty_stream_kind): + """An exhausted stream without payload must not confirm delivery.""" + stage = _make_respond_stage() + event = _DummyRespondEvent("unused", []) + + async def source_stream(): + if empty_stream_kind == "break": + yield MessageChain(chain=[], type="break") + elif empty_stream_kind == "whitespace": + yield MessageChain([Plain(" ")]) + + async def buffering_adapter(generator, _use_fallback): + async for _ in generator: + pass + + event.send_streaming = buffering_adapter + event._result = ( + MessageEventResult() + .set_result_content_type(ResultContentType.STREAMING_RESULT) + .set_async_stream(source_stream()) + ) + + await stage.process(event) + + assert event.get_result() is None + assert not event.get_extra(AGENT_OUTPUT_DELIVERY_CONFIRMED_KEY, False) + + +@pytest.mark.asyncio +async def test_respond_stage_clears_stream_after_adapter_failure(): + """A failed final flush must not leave an exhausted streaming result.""" + stage = _make_respond_stage() + event = _DummyRespondEvent("unused", []) + + async def source_stream(): + yield MessageChain([Plain("undelivered")]) + + async def failing_adapter(generator, _use_fallback): + async for _ in generator: + pass + raise RuntimeError("final flush failed") + + event.send_streaming = failing_adapter + event._result = ( + MessageEventResult() + .set_result_content_type(ResultContentType.STREAMING_RESULT) + .set_async_stream(source_stream()) + ) + + await stage.process(event) + + assert event.get_result() is None + assert not event.get_extra(AGENT_OUTPUT_DELIVERY_CONFIRMED_KEY, False) @pytest.mark.asyncio