Skip to content

fix: prevent the WebChat token ring from over-reporting context usage - #9255

Merged
Soulter merged 1 commit into
AstrBotDevs:masterfrom
Fronut:fix/9248-token-context-usage
Jul 13, 2026
Merged

Soulter merged 1 commit into
AstrBotDevs:masterfrom
Fronut:fix/9248-token-context-usage

Conversation

@Fronut

@Fronut Fronut commented Jul 13, 2026 •

Copy link
Copy Markdown
Contributor

Fixes #9248。

WebChat 输入框旁的 Token 用量圆环此前直接使用 AgentStats.token_usage 计算上下文占用率。然而,token_usage 会累计一次 Agent 运行中所有 LLM 请求的输入和输出 Token。在包含多轮工具调用的场景下,同一段历史上下文会被重复计入,导致圆环显示 200%、1000% 等明显高于实际上下文占用率的数值。

本次修改将“累计 Token 用量”和“当前上下文占用”拆分为两个独立指标:

  • token_usage 继续保留累计语义,供统计、计费和 Agent 运行信息展示使用;
  • 新增 current_context_tokens,记录最近一次完整 LLM 请求的输入 Token,即 input_other + input_cached;
  • Token 圆环改用 current_context_tokens 计算占用率,不再把多轮请求的累计用量误认为当前上下文大小。

输出 Token 没有计入 current_context_tokens,因为圆环需要反映的是最近一次请求发送给模型的输入上下文,而不是这次请求产生的输出或整个 Agent 运行的累计消耗。

Modifications / 改动点

  • 修改 astrbot/core/agent/response.py:

    • 在 AgentStats 中新增 current_context_tokens 字段;
    • 在 AgentStats.to_dict() 中序列化该字段,使其能够随 agent_stats 事件发送并保存到 WebChat 历史消息中。
  • 修改 astrbot/core/agent/runners/tool_loop_agent_runner.py:

    • 保留 token_usage 原有的累计逻辑,不影响 Provider 统计和计费用量;
    • 每次收到完整 LLM 响应后,将 current_context_tokens 更新为该次响应的输入 Token;
    • 如果最新完整响应没有提供 usage,则将 current_context_tokens 重置为 0,避免继续显示上一轮请求的陈旧上下文数据;
    • 保持原有 conversation.token_usage 更新逻辑不变。
  • 修改 dashboard/src/components/chat/Chat.vue:

    • Token 圆环优先使用 agentStats.current_context_tokens;
    • 对于修改前保存的历史消息,如果不存在该字段,则继续使用原有 token_usage 计算逻辑,保证旧数据兼容;
    • 当最新上下文用量未知或为 0 时,不显示圆环,避免展示不可靠的数据。
  • 修改 tests/test_tool_loop_agent_runner.py:

    • 增加具有不同 Token 用量的两轮工具调用回归测试;
    • 验证累计 token_usage 仍等于各轮用量之和;
    • 验证 current_context_tokens 只等于最后一次请求的输入 Token,并排除输出 Token;
    • 同时覆盖流式和非流式执行;
    • 覆盖最新响应缺少 usage 时清除陈旧上下文占用值的情况;
    • 验证 current_context_tokens 能正确序列化。
  • This is NOT a breaking change. / 这不是一个破坏性变更。

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

验证步骤

  1. 运行后端格式和静态检查:
uv run ruff format --check astrbot/core/agent/response.py astrbot/core/agent/runners/tool_loop_agent_runner.py tests/test_tool_loop_agent_runner.py
uv run ruff check astrbot/core/agent/response.py astrbot/core/agent/runners/tool_loop_agent_runner.py tests/test_tool_loop_agent_runner.py

结果:检查通过。

  1. 运行 Tool Loop Agent 和 Provider 统计相关测试:
uv run pytest tests/test_tool_loop_agent_runner.py tests/unit/test_provider_stats.py -q

结果:

...................................                                      [100%]
35 passed, 1 warning in 2.98s

其中 warning 为项目已有的 Python audioop 弃用提示,与本次修改无关。

  1. 运行前端 TypeScript 类型检查:
cd dashboard
pnpm typecheck

结果:vue-tsc --noEmit 执行通过。

  1. 额外运行了更大范围的 Agent 相关后端测试:
136 passed, 2 failed

两个失败均来自 Windows 下视频附件测试对 file:///path/... 路径分隔符的既有预期差异:测试期望 /path/to/video.mp4,实际在 Windows 上被规范化为 \path\to\video.mp4。失败用例与本次 Token 统计修改无关。

回归测试覆盖的核心场景

模拟一次包含两次 LLM 请求的工具调用流程:

请求 1:input_other=100, input_cached=10, output=1
请求 2:input_other=200, input_cached=20, output=2

修复后的预期和测试结果:

累计 token_usage:
input_other=300, input_cached=30, output=3

当前上下文 current_context_tokens:
200 + 20 = 220

这证明累计统计仍被完整保留,而 Token 圆环可以使用最后一次请求的 220 个输入 Token,不再错误使用累计值 333。


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.txt 和 pyproject.toml 文件相应位置。

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

Summary by Sourcery

Separate WebChat context-window occupancy from cumulative token usage and ensure the token ring reflects only the latest request’s input tokens.

New Features:

  • Expose current_context_tokens in AgentStats to represent input tokens of the most recent LLM request.
  • Update the WebChat token ring to use current_context_tokens for context occupancy while falling back to legacy token_usage for older messages.

Bug Fixes:

  • Prevent the WebChat token ring from over-reporting context usage by no longer using cumulative token_usage as the current context size.
  • Clear current_context_tokens when the latest LLM response omits usage data to avoid displaying stale context occupancy.

Enhancements:

  • Keep cumulative token_usage accounting unchanged for billing and statistics while adding a separate metric for display purposes.
  • Extend tool-loop agent runner stats serialization so current_context_tokens is available to the frontend.

Tests:

  • Add provider and agent runner regression tests covering multi-step tool loops, cumulative usage, latest-context tracking, missing usage handling, and serialization of current_context_tokens.

@dosubot dosubot Bot added size:S This PR changes 10-29 lines, ignoring generated files. area:core The bug / feature is about astrbot's core, backend area:webui The bug / feature is about webui(dashboard) of astrbot. labels Jul 13, 2026

@sourcery-ai sourcery-ai Bot left a comment

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.

Hey - I've found 1 issue, and left some high level feedback:

  • In Chat.vue, the computed latestTokenUsageTotal now primarily reflects latest context tokens rather than cumulative usage; consider renaming this computed property and its consumers to better match the new semantics and avoid confusion for future maintainers.
  • Using 0 to represent "unknown" current_context_tokens conflates a valid zero-length context with missing data; consider using null/None at the serialization boundary and checking explicitly for that to distinguish these cases more clearly.
Prompt for AI Agents
Please address the comments from this code review:

## Overall Comments
- In `Chat.vue`, the computed `latestTokenUsageTotal` now primarily reflects latest context tokens rather than cumulative usage; consider renaming this computed property and its consumers to better match the new semantics and avoid confusion for future maintainers.
- Using `0` to represent "unknown" `current_context_tokens` conflates a valid zero-length context with missing data; consider using `null`/`None` at the serialization boundary and checking explicitly for that to distinguish these cases more clearly.

## Individual Comments

### Comment 1
<location path="tests/test_tool_loop_agent_runner.py" line_range="654-678" />
<code_context>
+    assert runner.stats.to_dict()["current_context_tokens"] == 220
+
+
+@pytest.mark.asyncio
+async def test_stats_clear_current_context_when_latest_usage_is_missing(
+    runner, provider_request, mock_tool_executor, mock_hooks
+):
+    """Do not expose stale context occupancy when the latest usage is unknown."""
+    provider = MissingFinalUsageProvider()
+    await runner.reset(
+        provider=provider,
+        request=provider_request,
+        run_context=ContextWrapper(context=None),
+        tool_executor=mock_tool_executor,
+        agent_hooks=mock_hooks,
+        streaming=False,
+    )
+
+    async for _ in runner.step_until_done(3):
+        pass
+
+    assert runner.stats.token_usage == TokenUsage(
+        input_other=100,
+        input_cached=10,
+        output=1,
+    )
+    assert runner.stats.current_context_tokens == 0
+    assert runner.stats.to_dict()["current_context_tokens"] == 0
+
+
</code_context>
<issue_to_address>
**suggestion (testing):** Extend the missing-usage test to also cover streaming mode.

This regression test currently only exercises `streaming=False`, even though the underlying stats logic is shared between streaming and non-streaming flows. To ensure we also catch any discrepancies in how final responses are emitted in streaming runs, please parameterize `test_stats_clear_current_context_when_latest_usage_is_missing` on the `streaming` flag (similar to the previous test) so both modes are covered and we verify stale context is never exposed during streaming.

```suggestion
@pytest.mark.parametrize("streaming", [False, True])
@pytest.mark.asyncio
async def test_stats_clear_current_context_when_latest_usage_is_missing(
    runner, provider_request, mock_tool_executor, mock_hooks, streaming
):
    """Do not expose stale context occupancy when the latest usage is unknown."""
    provider = MissingFinalUsageProvider()
    await runner.reset(
        provider=provider,
        request=provider_request,
        run_context=ContextWrapper(context=None),
        tool_executor=mock_tool_executor,
        agent_hooks=mock_hooks,
        streaming=streaming,
    )

    async for _ in runner.step_until_done(3):
        pass

    assert runner.stats.token_usage == TokenUsage(
        input_other=100,
        input_cached=10,
        output=1,
    )
    assert runner.stats.current_context_tokens == 0
    assert runner.stats.to_dict()["current_context_tokens"] == 0
```
</issue_to_address>

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.

Comment thread tests/test_tool_loop_agent_runner.py

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

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.

Code Review

This pull request introduces tracking for the most recent LLM request's input tokens (current_context_tokens) separately from cumulative token usage, allowing the UI to display the current context-window occupancy. The changes span the agent stats model, the tool loop runner, the chat dashboard component, and corresponding unit tests. Feedback on the changes suggests removing a redundant is_chunk check in the runner since chunks are already filtered earlier, and simplifying the null/undefined check in the Vue component using a loose inequality check (!= null).

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

Comment thread astrbot/core/agent/runners/tool_loop_agent_runner.py Outdated
Comment thread dashboard/src/components/chat/Chat.vue Outdated
@Fronut

Fronut commented Jul 13, 2026

Copy link
Copy Markdown
Contributor Author

/gemini review

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

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.

Code Review

This pull request introduces a current_context_tokens field to track the input tokens sent in the most recent LLM request separately from cumulative token usage, allowing the dashboard to display the latest context-window occupancy. The frontend has been updated to use this new field, and corresponding unit tests have been added to verify the behavior. No review comments were provided, so there is no additional feedback.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

@dosubot dosubot Bot added the lgtm This PR has been approved by a maintainer label Jul 13, 2026
@Soulter Soulter changed the title fix: 将“累计 Token 用量”和“当前上下文占用”拆分为两个独立指标 fix: prevent the WebChat token ring from over-reporting context usage Jul 13, 2026
@Soulter
Soulter merged commit cc0ec5e into AstrBotDevs:master Jul 13, 2026
21 checks passed
@Fronut
Fronut deleted the fix/9248-token-context-usage branch July 13, 2026 13:23
BegoniaHe pushed a commit to Xero-Team/AstrBot that referenced this pull request Jul 14, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

area:core The bug / feature is about astrbot's core, backend area:webui The bug / feature is about webui(dashboard) of astrbot. lgtm This PR has been approved by a maintainer size:S This PR changes 10-29 lines, ignoring generated files.

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[Bug] token用量圆环累加导致百分比虚高

2 participants