From d6ba95c9a8bd9dc2684908e7e7c25afa02283d6e Mon Sep 17 00:00:00 2001 From: pei711 <199601925+pei711@users.noreply.github.com> Date: Mon, 26 Oct 2026 10:53:27 +0800 Subject: [PATCH 1/4] fix: dispatch response hooks for tool-call steps --- astrbot/core/agent/hooks.py | 5 +++ .../agent/runners/tool_loop_agent_runner.py | 5 +++ astrbot/core/astr_agent_hooks.py | 12 +++++++ tests/test_tool_loop_agent_runner.py | 33 +++++++++++++++++++ 4 files changed, 55 insertions(+) diff --git a/astrbot/core/agent/hooks.py b/astrbot/core/agent/hooks.py index 74ca6335b31..baf00f843ae 100644 --- a/astrbot/core/agent/hooks.py +++ b/astrbot/core/agent/hooks.py @@ -10,6 +10,11 @@ class BaseAgentRunHooks(Generic[TContext]): async def on_agent_begin(self, run_context: ContextWrapper[TContext]) -> None: ... + async def on_agent_response( + self, + run_context: ContextWrapper[TContext], + llm_response: LLMResponse, + ) -> None: ... async def on_tool_start( self, run_context: ContextWrapper[TContext], diff --git a/astrbot/core/agent/runners/tool_loop_agent_runner.py b/astrbot/core/agent/runners/tool_loop_agent_runner.py index f8418fe4205..fc14105ea24 100644 --- a/astrbot/core/agent/runners/tool_loop_agent_runner.py +++ b/astrbot/core/agent/runners/tool_loop_agent_runner.py @@ -923,6 +923,11 @@ async def step(self): if not llm_resp.tools_call_name: await self._complete_with_assistant_response(llm_resp) + else: + try: + await self.agent_hooks.on_agent_response(self.run_context, llm_resp) + except Exception as e: + logger.error(f"Error in on_agent_response hook: {e}", exc_info=True) # 返回 LLM 结果 if llm_resp.reasoning_content: diff --git a/astrbot/core/astr_agent_hooks.py b/astrbot/core/astr_agent_hooks.py index 3155257d712..35edeb074a3 100644 --- a/astrbot/core/astr_agent_hooks.py +++ b/astrbot/core/astr_agent_hooks.py @@ -40,6 +40,18 @@ async def on_agent_done(self, run_context, llm_response) -> None: llm_response, ) + async def on_agent_response(self, run_context, llm_response) -> None: + if llm_response and llm_response.reasoning_content: + run_context.context.event.set_extra( + "_llm_reasoning_content", llm_response.reasoning_content + ) + + await call_event_hook( + run_context.context.event, + EventType.OnLLMResponseEvent, + llm_response, + ) + async def on_tool_start( self, run_context: ContextWrapper[AstrAgentContext], diff --git a/tests/test_tool_loop_agent_runner.py b/tests/test_tool_loop_agent_runner.py index ef8e818bb5c..51b36e3ceec 100644 --- a/tests/test_tool_loop_agent_runner.py +++ b/tests/test_tool_loop_agent_runner.py @@ -398,6 +398,7 @@ class MockHooks(BaseAgentRunHooks): def __init__(self): self.agent_begin_called = False self.agent_done_called = False + self.agent_responses = [] self.tool_start_called = False self.tool_end_called = False @@ -413,6 +414,9 @@ 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 + async def on_agent_response(self, run_context, llm_response): + self.agent_responses.append(llm_response) + class ClearingAgentBeginHooks(MockHooks): async def on_agent_begin(self, run_context): @@ -912,6 +916,35 @@ async def test_hooks_called_with_max_step( assert mock_hooks.tool_end_called, "on_tool_end应该被调用" +@pytest.mark.asyncio +async def test_tool_call_response_invokes_response_hook( + runner, mock_provider, provider_request, mock_tool_executor, mock_hooks +): + mock_provider.should_call_tools = True + mock_provider.max_calls_before_normal_response = 100 + + await runner.reset( + provider=mock_provider, + request=provider_request, + run_context=ContextWrapper(context=None), + tool_executor=mock_tool_executor, + agent_hooks=mock_hooks, + streaming=False, + ) + + responses = [response async for response in runner.step()] + + assert len(mock_hooks.agent_responses) == 1 + assert mock_hooks.agent_responses[0].completion_text == "我需要使用工具来帮助您" + assert mock_hooks.agent_responses[0].tools_call_name == ["test_tool"] + assert not mock_hooks.agent_done_called + assert any( + response.type == "llm_result" + and response.data["chain"].get_plain_text() == "我需要使用工具来帮助您" + for response in responses + ) + + @pytest.mark.asyncio async def test_tool_result_includes_all_calltoolresult_content( runner, mock_provider, provider_request, mock_hooks, monkeypatch From a5e139e1f85b69cdf373cdb6dcecc925c7e26cf9 Mon Sep 17 00:00:00 2001 From: pei711 <199601925+pei711@users.noreply.github.com> Date: Mon, 26 Oct 2026 11:07:19 +0800 Subject: [PATCH 2/4] fix(agent): preserve response hook ordering --- astrbot/core/astr_agent_hooks.py | 5 ---- tests/test_tool_loop_agent_runner.py | 38 +++++++++++++++++++++++++++- 2 files changed, 37 insertions(+), 6 deletions(-) diff --git a/astrbot/core/astr_agent_hooks.py b/astrbot/core/astr_agent_hooks.py index 35edeb074a3..28d060f40e7 100644 --- a/astrbot/core/astr_agent_hooks.py +++ b/astrbot/core/astr_agent_hooks.py @@ -41,11 +41,6 @@ async def on_agent_done(self, run_context, llm_response) -> None: ) async def on_agent_response(self, run_context, llm_response) -> None: - if llm_response and llm_response.reasoning_content: - run_context.context.event.set_extra( - "_llm_reasoning_content", llm_response.reasoning_content - ) - await call_event_hook( run_context.context.event, EventType.OnLLMResponseEvent, diff --git a/tests/test_tool_loop_agent_runner.py b/tests/test_tool_loop_agent_runner.py index 51b36e3ceec..7843276ba94 100644 --- a/tests/test_tool_loop_agent_runner.py +++ b/tests/test_tool_loop_agent_runner.py @@ -12,6 +12,7 @@ sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))) from astrbot.core.agent.agent import Agent +from astrbot.core.astr_agent_hooks import MainAgentHooks from astrbot.core.agent.handoff import HandoffTool from astrbot.core.agent.hooks import BaseAgentRunHooks from astrbot.core.agent.message import ImageURLPart, Message, TextPart @@ -24,6 +25,7 @@ from astrbot.core.message.message_event_result import MessageChain from astrbot.core.provider.entities import LLMResponse, ProviderRequest, TokenUsage from astrbot.core.provider.provider import Provider +from astrbot.core.star.star_handler import EventType class MockProvider(Provider): @@ -399,6 +401,7 @@ def __init__(self): self.agent_begin_called = False self.agent_done_called = False self.agent_responses = [] + self.event_order = [] self.tool_start_called = False self.tool_end_called = False @@ -416,6 +419,7 @@ async def on_agent_done(self, run_context, llm_response): async def on_agent_response(self, run_context, llm_response): self.agent_responses.append(llm_response) + self.event_order.append("hook") class ClearingAgentBeginHooks(MockHooks): @@ -932,12 +936,20 @@ async def test_tool_call_response_invokes_response_hook( streaming=False, ) - responses = [response async for response in runner.step()] + responses = [] + async for response in runner.step(): + responses.append(response) + if ( + response.type == "llm_result" + and response.data["chain"].get_plain_text() == "我需要使用工具来帮助您" + ): + mock_hooks.event_order.append("text") assert len(mock_hooks.agent_responses) == 1 assert mock_hooks.agent_responses[0].completion_text == "我需要使用工具来帮助您" assert mock_hooks.agent_responses[0].tools_call_name == ["test_tool"] assert not mock_hooks.agent_done_called + assert mock_hooks.event_order == ["hook", "text"] assert any( response.type == "llm_result" and response.data["chain"].get_plain_text() == "我需要使用工具来帮助您" @@ -945,6 +957,30 @@ async def test_tool_call_response_invokes_response_hook( ) +@pytest.mark.asyncio +async def test_main_agent_response_hook_dispatches_llm_response_event(monkeypatch): + event = MagicMock() + llm_response = LLMResponse( + role="assistant", + completion_text="intermediate response", + tools_call_name=["test_tool"], + tools_call_args=[{"query": "test"}], + tools_call_ids=["call_123"], + ) + call_event_hook = AsyncMock() + monkeypatch.setattr( + "astrbot.core.astr_agent_hooks.call_event_hook", call_event_hook + ) + + await MainAgentHooks().on_agent_response( + ContextWrapper(context=MockAgentContext(event)), llm_response + ) + + call_event_hook.assert_awaited_once_with( + event, EventType.OnLLMResponseEvent, llm_response + ) + + @pytest.mark.asyncio async def test_tool_result_includes_all_calltoolresult_content( runner, mock_provider, provider_request, mock_hooks, monkeypatch From df647341b95c2b349c505306568fc139f83eb6ec Mon Sep 17 00:00:00 2001 From: pei711 <199601925+pei711@users.noreply.github.com> Date: Mon, 26 Oct 2026 11:08:32 +0800 Subject: [PATCH 3/4] fix: buffer hooked tool-call response streams --- astrbot/core/agent/hooks.py | 6 ++ .../agent/runners/tool_loop_agent_runner.py | 72 +++++++++++++------ astrbot/core/astr_agent_hooks.py | 18 ++++- tests/test_tool_loop_agent_runner.py | 35 +++++++-- 4 files changed, 103 insertions(+), 28 deletions(-) diff --git a/astrbot/core/agent/hooks.py b/astrbot/core/agent/hooks.py index baf00f843ae..76c67df4bfe 100644 --- a/astrbot/core/agent/hooks.py +++ b/astrbot/core/agent/hooks.py @@ -9,6 +9,12 @@ class BaseAgentRunHooks(Generic[TContext]): + def should_buffer_streaming_response( + self, run_context: ContextWrapper[TContext] + ) -> bool: + """Whether response hooks may modify streamed response content.""" + return False + async def on_agent_begin(self, run_context: ContextWrapper[TContext]) -> None: ... async def on_agent_response( self, diff --git a/astrbot/core/agent/runners/tool_loop_agent_runner.py b/astrbot/core/agent/runners/tool_loop_agent_runner.py index fc14105ea24..2cc89f06653 100644 --- a/astrbot/core/agent/runners/tool_loop_agent_runner.py +++ b/astrbot/core/agent/runners/tool_loop_agent_runner.py @@ -838,33 +838,38 @@ async def step(self): return self.run_context.messages = processed_messages self._simple_print_message_role("[AftCompact]", self.run_context.messages) + buffer_streaming_response = self.streaming and bool( + self.agent_hooks.should_buffer_streaming_response(self.run_context) + ) async for llm_response in self._iter_llm_responses_with_fallback(): 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 - - if llm_response.reasoning_content: - yield AgentResponse( - type="streaming_delta", - data=AgentResponseData( - chain=MessageChain(type="reasoning").message( - llm_response.reasoning_content, + if not buffer_streaming_response: + if llm_response.reasoning_content: + yield AgentResponse( + type="streaming_delta", + data=AgentResponseData( + chain=MessageChain(type="reasoning").message( + llm_response.reasoning_content, + ), ), - ), - ) - if llm_response.result_chain: - yield AgentResponse( - type="streaming_delta", - data=AgentResponseData(chain=llm_response.result_chain), - ) - elif llm_response.completion_text: - yield AgentResponse( - type="streaming_delta", - data=AgentResponseData( - chain=MessageChain().message(llm_response.completion_text), - ), - ) + ) + if llm_response.result_chain: + yield AgentResponse( + type="streaming_delta", + data=AgentResponseData(chain=llm_response.result_chain), + ) + elif llm_response.completion_text: + yield AgentResponse( + type="streaming_delta", + data=AgentResponseData( + chain=MessageChain().message( + llm_response.completion_text + ), + ), + ) if self._is_stop_requested(): break continue @@ -929,6 +934,31 @@ async def step(self): except Exception as e: logger.error(f"Error in on_agent_response hook: {e}", exc_info=True) + # Response hooks can modify the complete response. Emit streaming deltas + # only after they run so clients never receive content the hook changed. + if buffer_streaming_response: + if llm_resp.reasoning_content: + yield AgentResponse( + type="streaming_delta", + data=AgentResponseData( + chain=MessageChain(type="reasoning").message( + llm_resp.reasoning_content, + ), + ), + ) + if llm_resp.result_chain: + yield AgentResponse( + type="streaming_delta", + data=AgentResponseData(chain=llm_resp.result_chain), + ) + elif llm_resp.completion_text: + yield AgentResponse( + type="streaming_delta", + data=AgentResponseData( + chain=MessageChain().message(llm_resp.completion_text), + ), + ) + # 返回 LLM 结果 if llm_resp.reasoning_content: yield AgentResponse( diff --git a/astrbot/core/astr_agent_hooks.py b/astrbot/core/astr_agent_hooks.py index 28d060f40e7..148be3cb694 100644 --- a/astrbot/core/astr_agent_hooks.py +++ b/astrbot/core/astr_agent_hooks.py @@ -7,10 +7,21 @@ 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.star.star_handler import EventType +from astrbot.core.star.star_handler import EventType, star_handlers_registry class MainAgentHooks(BaseAgentRunHooks[AstrAgentContext]): + def should_buffer_streaming_response( + self, run_context: ContextWrapper[AstrAgentContext] + ) -> bool: + event = run_context.context.event + return bool( + star_handlers_registry.get_handlers_by_event_type( + EventType.OnLLMResponseEvent, + plugins_name=event.plugins_name, + ) + ) + async def on_agent_begin( self, run_context: ContextWrapper[AstrAgentContext] ) -> None: @@ -78,7 +89,10 @@ async def on_tool_end( class EmptyAgentHooks(BaseAgentRunHooks[AstrAgentContext]): - pass + def should_buffer_streaming_response( + self, run_context: ContextWrapper[AstrAgentContext] + ) -> bool: + return False MAIN_AGENT_HOOKS = MainAgentHooks() diff --git a/tests/test_tool_loop_agent_runner.py b/tests/test_tool_loop_agent_runner.py index 7843276ba94..9c0572d9a56 100644 --- a/tests/test_tool_loop_agent_runner.py +++ b/tests/test_tool_loop_agent_runner.py @@ -12,13 +12,13 @@ sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))) from astrbot.core.agent.agent import Agent -from astrbot.core.astr_agent_hooks import MainAgentHooks from astrbot.core.agent.handoff import HandoffTool 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.tool_loop_agent_runner import ToolLoopAgentRunner from astrbot.core.agent.tool import FunctionTool, ToolSet +from astrbot.core.astr_agent_hooks import MainAgentHooks from astrbot.core.astr_agent_run_util import run_agent from astrbot.core.astr_agent_tool_exec import FunctionToolExecutor from astrbot.core.exceptions import EmptyModelOutputError @@ -402,9 +402,15 @@ def __init__(self): self.agent_done_called = False self.agent_responses = [] self.event_order = [] + self.events = [] + self.buffer_streaming_response = False + self.rewrite_response = False self.tool_start_called = False self.tool_end_called = False + def should_buffer_streaming_response(self, run_context): + return self.buffer_streaming_response + async def on_agent_begin(self, run_context): self.agent_begin_called = True @@ -418,6 +424,9 @@ async def on_agent_done(self, run_context, llm_response): self.agent_done_called = True async def on_agent_response(self, run_context, llm_response): + self.events.append("response_hook") + if self.rewrite_response: + llm_response.completion_text = f"hooked: {llm_response.completion_text}" self.agent_responses.append(llm_response) self.event_order.append("hook") @@ -926,6 +935,8 @@ async def test_tool_call_response_invokes_response_hook( ): mock_provider.should_call_tools = True mock_provider.max_calls_before_normal_response = 100 + mock_hooks.buffer_streaming_response = True + mock_hooks.rewrite_response = True await runner.reset( provider=mock_provider, @@ -933,7 +944,7 @@ async def test_tool_call_response_invokes_response_hook( run_context=ContextWrapper(context=None), tool_executor=mock_tool_executor, agent_hooks=mock_hooks, - streaming=False, + streaming=True, ) responses = [] @@ -941,20 +952,34 @@ async def test_tool_call_response_invokes_response_hook( responses.append(response) if ( response.type == "llm_result" - and response.data["chain"].get_plain_text() == "我需要使用工具来帮助您" + and response.data["chain"].get_plain_text() + == "hooked: 我需要使用工具来帮助您" ): mock_hooks.event_order.append("text") + if response.type == "streaming_delta": + mock_hooks.events.append("streaming_delta") assert len(mock_hooks.agent_responses) == 1 - assert mock_hooks.agent_responses[0].completion_text == "我需要使用工具来帮助您" + assert ( + mock_hooks.agent_responses[0].completion_text + == "hooked: 我需要使用工具来帮助您" + ) assert mock_hooks.agent_responses[0].tools_call_name == ["test_tool"] assert not mock_hooks.agent_done_called assert mock_hooks.event_order == ["hook", "text"] assert any( response.type == "llm_result" - and response.data["chain"].get_plain_text() == "我需要使用工具来帮助您" + and response.data["chain"].get_plain_text() == "hooked: 我需要使用工具来帮助您" for response in responses ) + assert any( + response.type == "streaming_delta" + and response.data["chain"].get_plain_text() == "hooked: 我需要使用工具来帮助您" + for response in responses + ) + assert mock_hooks.events.index("response_hook") < mock_hooks.events.index( + "streaming_delta" + ) @pytest.mark.asyncio From 72e4f1399466564ed3b9e9f0169381471f9936f8 Mon Sep 17 00:00:00 2001 From: pei711 <199601925+pei711@users.noreply.github.com> Date: Tue, 27 Oct 2026 00:33:48 +0800 Subject: [PATCH 4/4] fix(agent): preserve streaming responses around hooks --- .../agent/runners/tool_loop_agent_runner.py | 64 ++++++--- astrbot/core/astr_agent_hooks.py | 20 +-- tests/test_tool_loop_agent_runner.py | 131 ++++++++++++++++++ 3 files changed, 182 insertions(+), 33 deletions(-) diff --git a/astrbot/core/agent/runners/tool_loop_agent_runner.py b/astrbot/core/agent/runners/tool_loop_agent_runner.py index 2cc89f06653..bd8f993bd39 100644 --- a/astrbot/core/agent/runners/tool_loop_agent_runner.py +++ b/astrbot/core/agent/runners/tool_loop_agent_runner.py @@ -841,35 +841,45 @@ async def step(self): buffer_streaming_response = self.streaming and bool( self.agent_hooks.should_buffer_streaming_response(self.run_context) ) + buffered_streaming_responses: list[AgentResponse] = [] async for llm_response in self._iter_llm_responses_with_fallback(): 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 - if not buffer_streaming_response: - if llm_response.reasoning_content: - yield AgentResponse( - type="streaming_delta", - data=AgentResponseData( - chain=MessageChain(type="reasoning").message( - llm_response.reasoning_content, - ), - ), - ) - if llm_response.result_chain: - yield AgentResponse( - type="streaming_delta", - data=AgentResponseData(chain=llm_response.result_chain), - ) - elif llm_response.completion_text: - yield AgentResponse( - type="streaming_delta", - data=AgentResponseData( - chain=MessageChain().message( - llm_response.completion_text - ), + if llm_response.reasoning_content: + response = AgentResponse( + type="streaming_delta", + data=AgentResponseData( + chain=MessageChain(type="reasoning").message( + llm_response.reasoning_content, ), - ) + ), + ) + if buffer_streaming_response: + buffered_streaming_responses.append(response) + else: + yield response + if llm_response.result_chain: + response = AgentResponse( + type="streaming_delta", + data=AgentResponseData(chain=llm_response.result_chain), + ) + if buffer_streaming_response: + buffered_streaming_responses.append(response) + else: + yield response + elif llm_response.completion_text: + response = AgentResponse( + type="streaming_delta", + data=AgentResponseData( + chain=MessageChain().message(llm_response.completion_text), + ), + ) + if buffer_streaming_response: + buffered_streaming_responses.append(response) + else: + yield response if self._is_stop_requested(): break continue @@ -899,16 +909,22 @@ async def step(self): break # got final response if self._is_stop_requested(): + for response in buffered_streaming_responses: + yield response yield await self._finalize_aborted_step() return if not llm_resp_result: + for response in buffered_streaming_responses: + yield response return # 处理 LLM 响应 llm_resp = llm_resp_result if llm_resp.role == "err": + for response in buffered_streaming_responses: + yield response # 如果 LLM 响应错误，转换到错误状态 self.final_llm_resp = llm_resp self.stats.end_time = time.time() @@ -933,6 +949,8 @@ async def step(self): await self.agent_hooks.on_agent_response(self.run_context, llm_resp) except Exception as e: logger.error(f"Error in on_agent_response hook: {e}", exc_info=True) + if not llm_resp.tools_call_name: + await self._complete_with_assistant_response(llm_resp) # Response hooks can modify the complete response. Emit streaming deltas # only after they run so clients never receive content the hook changed. diff --git a/astrbot/core/astr_agent_hooks.py b/astrbot/core/astr_agent_hooks.py index 148be3cb694..d53dcef869e 100644 --- a/astrbot/core/astr_agent_hooks.py +++ b/astrbot/core/astr_agent_hooks.py @@ -7,6 +7,7 @@ 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.star.star import star_map from astrbot.core.star.star_handler import EventType, star_handlers_registry @@ -15,12 +16,14 @@ def should_buffer_streaming_response( self, run_context: ContextWrapper[AstrAgentContext] ) -> bool: event = run_context.context.event - return bool( - star_handlers_registry.get_handlers_by_event_type( - EventType.OnLLMResponseEvent, - plugins_name=event.plugins_name, - ) - ) + for handler in star_handlers_registry.get_handlers_by_event_type( + EventType.OnLLMResponseEvent, + plugins_name=event.plugins_name, + ): + plugin = star_map.get(handler.handler_module_path) + if plugin and not plugin.reserved: + return True + return False async def on_agent_begin( self, run_context: ContextWrapper[AstrAgentContext] @@ -89,10 +92,7 @@ async def on_tool_end( class EmptyAgentHooks(BaseAgentRunHooks[AstrAgentContext]): - def should_buffer_streaming_response( - self, run_context: ContextWrapper[AstrAgentContext] - ) -> bool: - return False + pass MAIN_AGENT_HOOKS = MainAgentHooks() diff --git a/tests/test_tool_loop_agent_runner.py b/tests/test_tool_loop_agent_runner.py index 9c0572d9a56..f4adb441f35 100644 --- a/tests/test_tool_loop_agent_runner.py +++ b/tests/test_tool_loop_agent_runner.py @@ -405,6 +405,7 @@ def __init__(self): self.events = [] self.buffer_streaming_response = False self.rewrite_response = False + self.clear_tool_calls = False self.tool_start_called = False self.tool_end_called = False @@ -425,6 +426,8 @@ async def on_agent_done(self, run_context, llm_response): async def on_agent_response(self, run_context, llm_response): self.events.append("response_hook") + if self.clear_tool_calls: + llm_response.tools_call_name = [] if self.rewrite_response: llm_response.completion_text = f"hooked: {llm_response.completion_text}" self.agent_responses.append(llm_response) @@ -982,6 +985,112 @@ async def test_tool_call_response_invokes_response_hook( ) +@pytest.mark.asyncio +async def test_cleared_tool_calls_finalize_as_assistant_response( + runner, mock_provider, provider_request, mock_tool_executor, mock_hooks +): + mock_provider.should_call_tools = True + mock_provider.max_calls_before_normal_response = 100 + mock_hooks.clear_tool_calls = True + + await runner.reset( + provider=mock_provider, + request=provider_request, + run_context=ContextWrapper(context=None), + tool_executor=mock_tool_executor, + agent_hooks=mock_hooks, + streaming=False, + ) + + responses = [response async for response in runner.step()] + + assert runner.done() + assert runner.get_final_llm_resp().tools_call_name == [] + assert mock_hooks.agent_done_called + assert not mock_hooks.tool_start_called + assert any(response.type == "llm_result" for response in responses) + + +@pytest.mark.asyncio +async def test_buffered_stream_chunks_are_flushed_when_provider_fails( + runner, provider_request, mock_tool_executor, mock_hooks +): + class FailingStreamingProvider(MockProvider): + async def text_chat_stream(self, **kwargs): + yield LLMResponse( + role="assistant", + completion_text="partial response", + is_chunk=True, + ) + raise RuntimeError("stream failed") + + mock_hooks.buffer_streaming_response = True + + await runner.reset( + provider=FailingStreamingProvider(), + request=provider_request, + run_context=ContextWrapper(context=None), + tool_executor=mock_tool_executor, + agent_hooks=mock_hooks, + streaming=True, + ) + + responses = [response async for response in runner.step()] + + assert any( + response.type == "streaming_delta" + and response.data["chain"].get_plain_text() == "partial response" + for response in responses + ) + assert any(response.type == "err" for response in responses) + + +@pytest.mark.asyncio +async def test_buffered_partial_stream_is_flushed_when_stopped( + runner, provider_request, mock_tool_executor, mock_hooks +): + class PausedStreamingProvider(MockProvider): + def __init__(self): + super().__init__() + self.chunk_yielded = asyncio.Event() + + async def text_chat_stream(self, **kwargs): + yield LLMResponse( + role="assistant", + completion_text="partial response", + is_chunk=True, + ) + self.chunk_yielded.set() + await asyncio.Future() + + provider = PausedStreamingProvider() + mock_hooks.buffer_streaming_response = True + + await runner.reset( + provider=provider, + request=provider_request, + run_context=ContextWrapper(context=None), + tool_executor=mock_tool_executor, + agent_hooks=mock_hooks, + streaming=True, + ) + + async def collect_responses(): + return [response async for response in runner.step()] + + step_task = asyncio.create_task(collect_responses()) + await asyncio.wait_for(provider.chunk_yielded.wait(), timeout=5) + runner.request_stop() + responses = await asyncio.wait_for(step_task, timeout=5) + + assert any( + response.type == "streaming_delta" + and response.data["chain"].get_plain_text() == "partial response" + for response in responses + ) + assert any(response.type == "aborted" for response in responses) + + @pytest.mark.asyncio async def test_main_agent_response_hook_dispatches_llm_response_event(monkeypatch): event = MagicMock() @@ -1006,6 +1115,28 @@ async def test_main_agent_response_hook_dispatches_llm_response_event(monkeypatc ) +@pytest.mark.parametrize( + ("reserved", "expected"), + [(True, False), (False, True)], +) +def test_main_agent_buffers_only_for_external_response_handlers( + monkeypatch, reserved, expected +): + from astrbot.core.star.star import star_map + + module_path = "tests.response_handler" + handler = SimpleNamespace(handler_module_path=module_path) + event = SimpleNamespace(plugins_name=None) + run_context = ContextWrapper(context=SimpleNamespace(event=event)) + monkeypatch.setattr( + "astrbot.core.astr_agent_hooks.star_handlers_registry.get_handlers_by_event_type", + lambda *args, **kwargs: [handler], + ) + monkeypatch.setitem(star_map, module_path, SimpleNamespace(reserved=reserved)) + + assert MainAgentHooks().should_buffer_streaming_response(run_context) is expected + + @pytest.mark.asyncio async def test_tool_result_includes_all_calltoolresult_content( runner, mock_provider, provider_request, mock_hooks, monkeypatch