diff --git a/astrbot/core/agent/response.py b/astrbot/core/agent/response.py index 9e61fa8c7..531a08f6d 100644 --- a/astrbot/core/agent/response.py +++ b/astrbot/core/agent/response.py @@ -18,6 +18,8 @@ class AgentResponse: @dataclass class AgentStats: token_usage: TokenUsage = field(default_factory=TokenUsage) + current_context_tokens: int = 0 + """Input tokens sent in the most recent LLM request.""" start_time: float = 0.0 end_time: float = 0.0 time_to_first_token: float = 0.0 @@ -29,6 +31,7 @@ class AgentStats: def to_dict(self) -> dict: return { "token_usage": self.token_usage.__dict__, + "current_context_tokens": self.current_context_tokens, "start_time": self.start_time, "end_time": self.end_time, "time_to_first_token": self.time_to_first_token, diff --git a/astrbot/core/agent/runners/tool_loop_agent_runner.py b/astrbot/core/agent/runners/tool_loop_agent_runner.py index 88038473a..43f6beb99 100644 --- a/astrbot/core/agent/runners/tool_loop_agent_runner.py +++ b/astrbot/core/agent/runners/tool_loop_agent_runner.py @@ -770,9 +770,14 @@ class ToolLoopAgentRunner(BaseAgentRunner[TContext]): 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 + # Chunk responses have already continued above. A missing usage report + # means the latest context occupancy is unknown. + self.stats.current_context_tokens = 0 + if llm_response.usage: + # Keep cumulative usage for billing and expose the latest request + # input separately for context-window occupancy displays. self.stats.token_usage += llm_response.usage + self.stats.current_context_tokens = llm_response.usage.input if self.req.conversation: self.req.conversation.token_usage = llm_response.usage.total break # got final response diff --git a/dashboard/src/components/chat/Chat.vue b/dashboard/src/components/chat/Chat.vue index 38e0e7c49..560fdbffd 100644 --- a/dashboard/src/components/chat/Chat.vue +++ b/dashboard/src/components/chat/Chat.vue @@ -867,11 +867,16 @@ const currentTokenMetadata = computed(() => { const model = currentTokenProvider.value?.model; return model ? tokenModelMetadata.value[model] || null : null; }); -const latestTokenUsageTotal = computed(() => { +const latestContextTokens = computed(() => { for (let index = activeMessages.value.length - 1; index >= 0; index -= 1) { const message = activeMessages.value[index]; if (isUserMessage(message)) continue; - const usage = message.content?.agentStats?.token_usage; + const stats = message.content?.agentStats; + if (!stats) continue; + if (stats.current_context_tokens != null) { + return readTokenCount(stats.current_context_tokens); + } + const usage = stats.token_usage; if (!usage) continue; return ( readTokenCount(usage.input_other) + @@ -882,7 +887,7 @@ const latestTokenUsageTotal = computed(() => { return 0; }); const tokenUsageIndicator = computed(() => { - const used = latestTokenUsageTotal.value; + const used = latestContextTokens.value; const limit = contextLimit(currentTokenProvider.value, currentTokenMetadata.value); if (used <= 0 || limit <= 0) return null; diff --git a/tests/test_tool_loop_agent_runner.py b/tests/test_tool_loop_agent_runner.py index b4464680f..1fe2d819f 100644 --- a/tests/test_tool_loop_agent_runner.py +++ b/tests/test_tool_loop_agent_runner.py @@ -142,6 +142,41 @@ class MockMixedContentToolExecutor: return generator() +class VaryingUsageProvider(MockProvider): + """Return distinct token usage values for each tool-loop request.""" + + async def text_chat(self, **kwargs) -> LLMResponse: + self.call_count += 1 + usage = TokenUsage( + input_other=self.call_count * 100, + input_cached=self.call_count * 10, + output=self.call_count, + ) + if self.call_count == 1: + return LLMResponse( + role="assistant", + tools_call_name=["test_tool"], + tools_call_args=[{"query": "test"}], + tools_call_ids=["call_varying_usage"], + usage=usage, + ) + return LLMResponse( + role="assistant", + completion_text="final", + usage=usage, + ) + + +class MissingFinalUsageProvider(VaryingUsageProvider): + """Omit usage from the final response after reporting an earlier request.""" + + async def text_chat(self, **kwargs) -> LLMResponse: + response = await super().text_chat(**kwargs) + if self.call_count == 2: + response.usage = None + return response + + class MockFailingProvider(MockProvider): async def text_chat(self, **kwargs) -> LLMResponse: self.call_count += 1 @@ -587,6 +622,63 @@ async def test_normal_completion_without_max_step( assert runner.req.func_tool is not None, "正常完成时工具不应该被禁用" +@pytest.mark.asyncio +@pytest.mark.parametrize("streaming", [False, True]) +async def test_stats_separate_latest_context_from_cumulative_usage( + runner, provider_request, mock_tool_executor, mock_hooks, streaming +): + """Context occupancy uses the latest input while usage remains cumulative.""" + provider = VaryingUsageProvider() + 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 provider.call_count == 2 + assert runner.stats.token_usage == TokenUsage( + input_other=300, + input_cached=30, + output=3, + ) + assert runner.stats.current_context_tokens == 220 + assert runner.stats.to_dict()["current_context_tokens"] == 220 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("streaming", [False, True]) +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 + + @pytest.mark.asyncio async def test_max_step_with_streaming( runner, mock_provider, provider_request, mock_tool_executor, mock_hooks