Skip to content

fix: dispatch response hooks for tool-call steps #10772

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
pei711 wants to merge 4 commits into
base: master
Choose a base branch
from
Open

fix: dispatch response hooks for tool-call steps #10772

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
d6ba95c
fix: dispatch response hooks for tool-call steps
pei711 Oct 26, 2026
a5e139e
fix(agent): preserve response hook ordering
pei711 Oct 26, 2026
df64734
fix: buffer hooked tool-call response streams
pei711 Oct 26, 2026
72e4f13
fix(agent): preserve streaming responses around hooks
pei711 Oct 26, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
11 changes: 11 additions & 0 deletions astrbot/core/agent/hooks.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,18 @@


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,
run_context: ContextWrapper[TContext],
llm_response: LLMResponse,
) -> None: ...
async def on_tool_start(
self,
run_context: ContextWrapper[TContext],
Expand Down
61 changes: 57 additions & 4 deletions astrbot/core/agent/runners/tool_loop_agent_runner.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -838,33 +838,48 @@ 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)
)
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 llm_response.reasoning_content:
yield AgentResponse(
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)

@sourcery-ai sourcery-ai Bot Oct 26, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🟡 Medium · Long streams retain every chunk

When a hooked streaming response contains a very large number of chunks, with an external response handler enabled, the runner retains a separate AgentResponse for every streamed chunk until the provider finishes. Very long outputs therefore accumulate unbounded per-chunk objects in memory, increasing memory pressure or exhausting the agent process.

Bound or compact buffered partial output while preserving the ability to flush it on failure or stop.

Prompt for AI agents 
In `astrbot/core/agent/runners/tool_loop_agent_runner.py` at line 860:

**Long streams retain every chunk**

When a hooked streaming response contains a very large number of chunks, with an external response handler enabled, the runner retains a separate `AgentResponse` for every streamed chunk until the provider finishes. Very long outputs therefore accumulate unbounded per-chunk objects in memory, increasing memory pressure or exhausting the agent process.

Bound or compact buffered partial output while preserving the ability to flush it on failure or stop.

else:
yield response
if llm_response.result_chain:
yield AgentResponse(
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:
yield AgentResponse(
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
Expand Down Expand Up @@ -894,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()
Expand All @@ -923,6 +944,38 @@ 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)
Comment thread
sourcery-ai[bot] marked this conversation as resolved.
Comment thread
sourcery-ai[bot] marked this conversation as resolved.

@sourcery-ai sourcery-ai Bot Oct 26, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🟠 High · Canceled requests still execute tools

When the runner task is canceled while an OnLLMResponseEvent handler is running, call_event_hook catches BaseException, including asyncio.CancelledError, from a plugin handler and suppresses it. The runner then continues past the response hook into tool execution, so a canceled request can still perform tool side effects.

Re-raise asyncio.CancelledError from call_event_hook so task cancellation propagates before tool execution.

Prompt for AI agents 
In `astrbot/core/agent/runners/tool_loop_agent_runner.py` at line 949:

**Canceled requests still execute tools**

When the runner task is canceled while an `OnLLMResponseEvent` handler is running, `call_event_hook` catches `BaseException`, including `asyncio.CancelledError`, from a plugin handler and suppresses it. The runner then continues past the response hook into tool execution, so a canceled request can still perform tool side effects.

Re-raise `asyncio.CancelledError` from `call_event_hook` so task cancellation propagates before tool execution.

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.
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:
Expand Down
23 changes: 22 additions & 1 deletion astrbot/core/astr_agent_hooks.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,10 +7,24 @@
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 import star_map
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
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]
) -> None:
Expand Down Expand Up @@ -40,6 +54,13 @@ async def on_agent_done(self, run_context, llm_response) -> None:
llm_response,
)

async def on_agent_response(self, run_context, llm_response) -> None:
await call_event_hook(
run_context.context.event,
EventType.OnLLMResponseEvent,
llm_response,
)

async def on_tool_start(
self,
run_context: ContextWrapper[AstrAgentContext],
Expand Down
Loading