Skip to content

## fix(agent): call on_agent_done hook before storing assistant message in history#6721

Open
Luna-channel wants to merge 1 commit intoAstrBotDevs:masterfrom
Luna-channel:fix/agent-hook-before-message-storage
Open

## fix(agent): call on_agent_done hook before storing assistant message in history#6721
Luna-channel wants to merge 1 commit intoAstrBotDevs:masterfrom
Luna-channel:fix/agent-hook-before-message-storage

Conversation

@Luna-channel
Copy link
Contributor

@Luna-channel Luna-channel commented Mar 20, 2026

Motivation / 动机

ToolLoopAgentRunner 中,assistant 消息在 on_agent_done 钩子(触发 OnLLMResponseEvent)被调用之前就已存入 run_context.messages。这意味着插件在 on_llm_response 处理器中对 completion_text 所做的任何修改,都不会反映到已存储的对话历史中。

具体影响:通过 on_llm_response 从 LLM 输出中剥离元数据标签(如 [Favour: ...][Spend: ...][Profile: ...])的插件,虽然能成功清理当前响应的文本,但未清理的原始文本已经存入了对话历史。在后续的 LLM 调用中,模型会在上下文中看到残留的元数据标签,导致重复出账、重复记忆更新等问题。很难想象这个问题之前没有人发现。我已经分析过了这个改动可能带来的后果,目前来说我认为并不会对现有的插件调用产生影响,风险极低。

Modifications / 改动点

File: astrbot/core/agent/runners/tool_loop_agent_runner.py

Reordered the on_agent_done hook call to execute before run_context.messages.append(...) in two locations:

  1. Main completion flow (~line 471): When LLM responds without tool calls
  2. _finalize_aborted_step (~line 976): When agent execution is aborted

Before (bug):

messages.append(...)   ← stores uncleaned text
on_agent_done(...)     ← plugins clean text (too late)

After (fix):

on_agent_done(...)     ← plugins clean text first
messages.append(...)   ← stores cleaned text

No new dependencies. No breaking changes. The hook contract is unchanged — on_agent_done still receives the same run_context and llm_resp objects.

  • This is NOT a breaking change. / 这不是一个破坏性变更。

Screenshots or Test Results / 运行截图或测试结果

Before fix (from real production logs):

# LLM returns response with metadata tags
completion: "...下次我一定先打报告 {memes:羞愧}\n\n[Favour: 100, Attitude: 认真讨论中, Relationship: 依然是大人最完美的作品] [Spend: 4.8, Reason: 给大人买盒蓝莓润眼片] [Profile: 备注=正在进行主动回复框架白名单测试] [ProfileDelete: 5]"

# Hooks fire and clean the text — but messages already stored the uncleaned version
# Next LLM call context contains 48 messages with stale [Favour:], [Spend:], [Profile:] tags
[BefCompact] RunCtx.messages -> [48] system,user,assistant,...,user

# LLM sees old tags and generates duplicate operations
[Refund: 0.5, Reason: ...] [Favour: 100, ...] [Profile: 备注=被大人抓到偷偷花钱反省中] [ProfileDelete: 5]

After fix: The on_agent_done hook fires first, plugins clean completion_text, and the cleaned text is stored in run_context.messages. Subsequent LLM calls see clean history without metadata tags.

Syntax validation: python -c "import ast; ast.parse(...)" → OK
Lint: ruff check → All checks passed


Checklist / 检查清单

  • 😊 If there are new features added in the PR, I have discussed it with the authors through issues/emails, etc.
    / 如果 PR 中有新加入的功能,已经通过 Issue / 邮件等方式和作者讨论过。

  • 👀 My changes have been well-tested, and "Verification Steps" and "Screenshots" have been provided above.
    / 我的更改经过了良好的测试,并已在上方提供了"验证步骤"和"运行截图"

  • 🤓 I have ensured that no new dependencies are introduced, OR if new dependencies are introduced, they have been added to the appropriate locations in requirements.txt and pyproject.toml.
    / 我确保没有引入新依赖库,或者引入了新依赖库的同时将其添加到 requirements.txtpyproject.toml 文件相应位置。

  • 😮 My changes do not introduce malicious code.
    / 我的更改没有引入恶意代码。

Summary by Sourcery

Bug Fixes:

  • Call the on_agent_done hook before recording final assistant messages in normal and aborted agent flows to prevent uncleaned LLM outputs from being stored in run_context.messages.

@dosubot dosubot bot added the size:S This PR changes 10-29 lines, ignoring generated files. label Mar 20, 2026
@gemini-code-assist
Copy link
Contributor

Summary of Changes

Hello, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request addresses a critical bug in the ToolLoopAgentRunner where assistant messages were stored in the conversation history before the on_agent_done hook was invoked. This timing issue prevented plugins from effectively cleaning or modifying the LLM's completion_text before it became a permanent part of the run_context.messages. The fix reorders these operations, ensuring that plugins can process and clean the response, particularly by stripping metadata tags, before the message is recorded, thereby preventing the persistence of uncleaned text and avoiding downstream issues like redundant operations in subsequent LLM interactions.

Highlights

  • Hook Execution Order: The on_agent_done hook is now called before the assistant message is appended to the run_context.messages. This ensures that any modifications made by plugins to the completion_text (e.g., stripping metadata tags) are reflected in the stored conversation history.
  • Impact on Message History: Previously, uncleaned LLM responses, including metadata tags, were stored in the conversation history, leading to issues like duplicate operations in subsequent LLM calls. This fix prevents such uncleaned text from being persisted.
  • Affected Scenarios: The reordering has been applied to both the main LLM completion flow (when no tool calls are made) and the _finalize_aborted_step method (when agent execution is aborted), covering all relevant paths where assistant messages are stored.
Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature Command Description
Code Review /gemini review Performs a code review for the current pull request in its current state.
Pull Request Summary /gemini summary Provides a summary of the current pull request in its current state.
Comment @gemini-code-assist Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help /gemini help Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for GitHub and other Google products, sign up here.

Footnotes

  1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution.

Copy link
Contributor

@sourcery-ai sourcery-ai bot left a comment

Choose a reason for hiding this comment

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

Hey - I've left some high level feedback:

  • Since on_agent_done is now invoked before run_context.messages.append(...), consider auditing any existing hooks that may implicitly rely on the final assistant message already being present in run_context.messages and, if needed, add a defensive check or explicit comment near the hook call to clarify this ordering contract.
  • The try/except wrapper around agent_hooks.on_agent_done is now duplicated in multiple places; consider extracting this into a small helper (e.g., _safe_on_agent_done(llm_resp)) to keep the control flow consistent and easier to maintain.
  • When logging hook errors (Error in on_agent_done hook), consider including additional context such as the current run ID or agent name to make it easier to trace failures in multi-run or multi-agent environments.
Prompt for AI Agents
Please address the comments from this code review:

## Overall Comments
- Since `on_agent_done` is now invoked before `run_context.messages.append(...)`, consider auditing any existing hooks that may implicitly rely on the final assistant message already being present in `run_context.messages` and, if needed, add a defensive check or explicit comment near the hook call to clarify this ordering contract.
- The try/except wrapper around `agent_hooks.on_agent_done` is now duplicated in multiple places; consider extracting this into a small helper (e.g., `_safe_on_agent_done(llm_resp)`) to keep the control flow consistent and easier to maintain.
- When logging hook errors (`Error in on_agent_done hook`), consider including additional context such as the current run ID or agent name to make it easier to trace failures in multi-run or multi-agent environments.

Sourcery is free for open source - if you like our reviews please consider sharing them ✨
Help me be more useful! Please click 👍 or 👎 on each comment and I'll use the feedback to improve your reviews.

@dosubot dosubot bot added the area:provider The bug / feature is about AI Provider, Models, LLM Agent, LLM Agent Runner. label Mar 20, 2026
Copy link
Contributor

@gemini-code-assist gemini-code-assist bot left a comment

Choose a reason for hiding this comment

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

Code Review

本次改动修正了 on_agent_done 钩子调用时机的问题,确保在将助手消息存入历史记录之前执行,从而允许插件正确清理和修改 LLM 的响应。这个修正是合理的,并且解决了 PR 描述中提到的数据一致性问题。

此外,我发现这个改动在两个地方引入了重复的 try...except 代码块。我提供了一个代码审查评论,建议通过提取辅助方法来重构这部分代码,以提高代码的可维护性。

Comment on lines +471 to +476
# call the on_agent_done hook BEFORE recording the message,
# so that plugins can clean metadata tags from completion_text first.
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)
Copy link
Contributor

Choose a reason for hiding this comment

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

medium

为了提高代码的可维护性并避免重复,建议将这个 try...except 块提取到一个私有的辅助方法中。此逻辑在 _finalize_aborted_step 方法(第 976-979 行)中也被重复使用。

您可以考虑在 ToolLoopAgentRunner 类中添加一个类似这样的方法:

    async def _safe_run_on_agent_done(self, llm_resp: LLMResponse):
        """Safely runs the on_agent_done hook and logs any errors."""
        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)

然后,您就可以用一行 await self._safe_run_on_agent_done(llm_resp) 来替换当前的 try...except 块,并在 _finalize_aborted_step 中也进行同样的替换。

@dosubot
Copy link

dosubot bot commented Mar 20, 2026

Related Documentation

1 document(s) may need updating based on files changed in this PR:

AstrBotTeam's Space

pr4697的改动
View Suggested Changes
@@ -740,7 +740,8 @@
 - **[PR #5850](https://github.com/AstrBotDevs/AstrBot/pull/5850) 增强**:停止信号现在会传递到 LLM 调用和工具执行器,通过 `abort_signal` 参数使停止机制在整个执行堆栈中一致传播
 - **工具执行中断**:长时间运行的工具等待(包括子代理 handoff)现在可以被停止请求中断,`_iter_tool_executor_results()` 方法监控 `abort_signal`,使用 `asyncio.wait()` 与 `FIRST_COMPLETED` 在工具结果和中止信号之间竞争
 - **部分输出保留**:当工具执行期间发生中断时,`_finalize_aborted_step()` 方法会保留已生成的部分 LLM 响应
-- Agent 转换为 DONE 状态,并触发 `on_agent_done` 钩子
+- **钩子执行顺序(PR #6721)**:在 Agent 完成或中止时,`on_agent_done` 钩子会在 assistant 消息存入对话历史之前触发。这确保插件可以在 `on_llm_response` 处理器中修改 `completion_text`(例如剥离元数据标签),修改后的内容会被保存到历史记录中,而非原始的 LLM 输出
+- Agent 转换为 DONE 状态
 - 对话历史和会话状态得以保留,响应类型标记为 "aborted"
 - `was_aborted()` 方法返回 `true`,表明任务被用户主动中止
 
@@ -769,6 +770,7 @@
 - 历史保存条件从 `not event.is_stopped()` 改为 `not event.is_stopped() or agent_runner.was_aborted()`
 - 确保即使会话被中止,已生成的部分内容也能正确保存到历史记录中
 - 空响应在用户中止时会被妥善处理,不会导致历史记录丢失
+- **钩子与历史记录顺序(PR #6721)**:`on_agent_done` 钩子现在在 assistant 消息存入历史记录之前触发。这确保插件对 `completion_text` 的修改(例如通过 `on_llm_response` 处理器剥离元数据标签)能够反映到存储的对话历史中,防止未清理的 LLM 输出泄漏到后续上下文中
 
 ##### 前端实现
 - 状态管理(`useMessages.ts`):新增 `currentRequestController`、`currentReader`、`currentRunningSessionId` 和 `userStopRequested` 状态

[Accept] [Decline]

Note: You must be authenticated to accept/decline updates.

How did I do? Any feedback?  Join Discord

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

area:provider The bug / feature is about AI Provider, Models, LLM Agent, LLM Agent Runner. size:S This PR changes 10-29 lines, ignoring generated files.

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant