Skip to content

fix: prevent late LLM responses after stop #9154

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
Sisyphbaous-DT-Project wants to merge 2 commits into
base: master
Choose a base branch
from
+1,186 −81
Open

fix: prevent late LLM responses after stop #9154

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
b471c1e
fix: 修复 stop 后迟到 LLM 回复漏发
Sisyphbaous-DT-Project Jul 5, 2026
53dc22e
fix: 处理 stop 修复的 review 风险点
Sisyphbaous-DT-Project Jul 6, 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
184 changes: 144 additions & 40 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 @@ -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 "
Expand Down Expand Up @@ -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:
Comment thread
Sisyphbaous-DT-Project marked this conversation as resolved.
"""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()
Expand All @@ -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(
Expand Down Expand Up @@ -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_usages: list[T.Any] = []

# These two are used for tool schema mode handling
# We now have two modes:
Expand Down Expand Up @@ -323,6 +337,16 @@ 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
Comment thread
Sisyphbaous-DT-Project marked this conversation as resolved.

def _read_tool_hint(self) -> str:
if self.read_tool is not None:
return f"`{self.read_tool.name}`"
Expand Down Expand Up @@ -488,6 +512,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", "<unknown>")
is_last_candidate = idx == total_candidates - 1
if idx > 0:
Expand Down Expand Up @@ -526,6 +553,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.",
Expand Down Expand Up @@ -734,6 +763,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
Expand All @@ -747,34 +780,36 @@ 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",
data=AgentResponseData(
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:
Expand Down Expand Up @@ -809,7 +844,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:
Expand Down Expand Up @@ -838,11 +878,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",
Expand All @@ -864,8 +914,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
Expand Down Expand Up @@ -1332,6 +1380,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.
Expand Down Expand Up @@ -1370,46 +1421,99 @@ 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

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):
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
Comment thread
Sisyphbaous-DT-Project marked this conversation as resolved.

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)

Expand Down
Loading
Loading