From ceb99bd4dd8a884176649878931f8161ea358f43 Mon Sep 17 00:00:00 2001 From: w33d Date: Sun, 12 Jul 2026 17:05:48 +0800 Subject: [PATCH 1/2] fix(agent): preserve tool metadata after skills-like requery --- .../agent/runners/tool_loop_agent_runner.py | 3 +++ tests/test_tool_loop_agent_runner.py | 22 +++++++++++++++++-- 2 files changed, 23 insertions(+), 2 deletions(-) diff --git a/astrbot/core/agent/runners/tool_loop_agent_runner.py b/astrbot/core/agent/runners/tool_loop_agent_runner.py index 88038473af..4f48cd76a0 100644 --- a/astrbot/core/agent/runners/tool_loop_agent_runner.py +++ b/astrbot/core/agent/runners/tool_loop_agent_runner.py @@ -871,6 +871,9 @@ async def step(self): llm_resp.tools_call_name = requery_resp.tools_call_name llm_resp.tools_call_args = requery_resp.tools_call_args llm_resp.tools_call_ids = requery_resp.tools_call_ids + llm_resp.tools_call_extra_content = ( + requery_resp.tools_call_extra_content + ) tool_call_result_blocks = [] cached_images = [] # Collect cached images for LLM visibility diff --git a/tests/test_tool_loop_agent_runner.py b/tests/test_tool_loop_agent_runner.py index b4464680fb..7bd231db13 100644 --- a/tests/test_tool_loop_agent_runner.py +++ b/tests/test_tool_loop_agent_runner.py @@ -1284,8 +1284,8 @@ async def test_follow_up_ticket_not_consumed_when_no_next_tool_call( @pytest.mark.asyncio -async def test_skills_like_requery_passes_extra_user_content_parts(): - """skills-like 模式 re-query 时应传递 extra_user_content_parts（如 image_caption）""" +async def test_skills_like_requery_preserves_request_and_tool_metadata(): + """skills-like re-query preserves request parts and provider tool metadata.""" from astrbot.core.agent.message import TextPart captured_kwargs = {} @@ -1301,6 +1301,11 @@ async def text_chat(self, **kwargs) -> LLMResponse: tools_call_name=["test_tool"], tools_call_args=[{"query": "test"}], tools_call_ids=["call_1"], + tools_call_extra_content={ + "call_1": { + "google": {"thought_signature": "selection-signature"} + } + }, usage=TokenUsage(input_other=10, output=5), ) if self.call_count == 2: @@ -1312,6 +1317,9 @@ async def text_chat(self, **kwargs) -> LLMResponse: tools_call_name=["test_tool"], tools_call_args=[{"query": "actual"}], tools_call_ids=["call_2"], + tools_call_extra_content={ + "call_2": {"google": {"thought_signature": "requery-signature"}} + }, usage=TokenUsage(input_other=10, output=5), ) # 后续调用：正常回复 @@ -1363,6 +1371,16 @@ async def text_chat(self, **kwargs) -> LLMResponse: assert len(parts) == 1 assert parts[0].text == "一张猫的照片" + assistant_tool_message = next( + message + for message in run_context.messages + if message.role == "assistant" and message.tool_calls + ) + assert assistant_tool_message.tool_calls[0].id == "call_2" + assert assistant_tool_message.tool_calls[0].extra_content == { + "google": {"thought_signature": "requery-signature"} + } + @pytest.mark.asyncio async def test_follow_up_accepted_when_active_and_not_stopping( From 1125c4229b32176f41f997dfb8d19dce6abcee56 Mon Sep 17 00:00:00 2001 From: w33d Date: Sun, 12 Jul 2026 18:42:58 +0800 Subject: [PATCH 2/2] fix(agent): count skills-like requery usage --- .../agent/runners/tool_loop_agent_runner.py | 8 + tests/test_tool_loop_agent_runner.py | 139 +++++++++++++++++- 2 files changed, 143 insertions(+), 4 deletions(-) diff --git a/astrbot/core/agent/runners/tool_loop_agent_runner.py b/astrbot/core/agent/runners/tool_loop_agent_runner.py index 4f48cd76a0..387c79c898 100644 --- a/astrbot/core/agent/runners/tool_loop_agent_runner.py +++ b/astrbot/core/agent/runners/tool_loop_agent_runner.py @@ -1331,6 +1331,10 @@ async def _resolve_tool_exec( abort_signal=self._abort_signal, request_max_retries=self.request_max_retries, ) + if requery_resp and requery_resp.usage: + self.stats.token_usage += requery_resp.usage + if self.req.conversation: + self.req.conversation.token_usage += requery_resp.usage.total if requery_resp: llm_resp = requery_resp self._sanitize_malformed_tool_calls(llm_resp) @@ -1359,6 +1363,10 @@ async def _resolve_tool_exec( abort_signal=self._abort_signal, request_max_retries=self.request_max_retries, ) + if repair_resp and repair_resp.usage: + self.stats.token_usage += repair_resp.usage + if self.req.conversation: + self.req.conversation.token_usage += repair_resp.usage.total if repair_resp: llm_resp = repair_resp self._sanitize_malformed_tool_calls(llm_resp) diff --git a/tests/test_tool_loop_agent_runner.py b/tests/test_tool_loop_agent_runner.py index 7bd231db13..c3d91aaac4 100644 --- a/tests/test_tool_loop_agent_runner.py +++ b/tests/test_tool_loop_agent_runner.py @@ -1284,8 +1284,8 @@ async def test_follow_up_ticket_not_consumed_when_no_next_tool_call( @pytest.mark.asyncio -async def test_skills_like_requery_preserves_request_and_tool_metadata(): - """skills-like re-query preserves request parts and provider tool metadata.""" +async def test_skills_like_requery_preserves_request_tool_metadata_and_usage(): + """skills-like re-query preserves metadata and counts its token usage.""" from astrbot.core.agent.message import TextPart captured_kwargs = {} @@ -1306,7 +1306,7 @@ async def text_chat(self, **kwargs) -> LLMResponse: "google": {"thought_signature": "selection-signature"} } }, - usage=TokenUsage(input_other=10, output=5), + usage=TokenUsage(input_other=10, input_cached=2, output=5), ) if self.call_count == 2: # 第二次调用：re-query with param schema @@ -1320,7 +1320,7 @@ async def text_chat(self, **kwargs) -> LLMResponse: tools_call_extra_content={ "call_2": {"google": {"thought_signature": "requery-signature"}} }, - usage=TokenUsage(input_other=10, output=5), + usage=TokenUsage(input_other=20, input_cached=3, output=7), ) # 后续调用：正常回复 return LLMResponse( @@ -1339,11 +1339,13 @@ async def text_chat(self, **kwargs) -> LLMResponse: tool_set = ToolSet(tools=[tool]) caption_part = TextPart(text="一张猫的照片") + conversation = SimpleNamespace(cid="test-conversation", token_usage=0) req = ProviderRequest( prompt="看看这张图", func_tool=tool_set, contexts=[], extra_user_content_parts=[caption_part], + conversation=cast(Any, conversation), ) event = MockEvent(umo="test_umo", sender_id="test_sender") @@ -1380,6 +1382,135 @@ async def text_chat(self, **kwargs) -> LLMResponse: assert assistant_tool_message.tool_calls[0].extra_content == { "google": {"thought_signature": "requery-signature"} } + assert runner.stats.token_usage == TokenUsage( + input_other=30, + input_cached=5, + output=12, + ) + assert conversation.token_usage == 47 + + +@pytest.mark.asyncio +async def test_skills_like_requery_repair_counts_both_extra_requests(): + """skills-like repair counts both provider requests after tool selection.""" + + class RepairProvider(MockProvider): + async def text_chat(self, **kwargs) -> LLMResponse: + self.call_count += 1 + if self.call_count == 1: + return LLMResponse( + role="assistant", + tools_call_name=["test_tool"], + tools_call_args=[{"query": "test"}], + tools_call_ids=["call_1"], + usage=TokenUsage(input_other=2, input_cached=3, output=5), + ) + if self.call_count == 2: + return LLMResponse( + role="assistant", + completion_text="", + usage=TokenUsage(input_other=7, input_cached=11, output=13), + ) + if self.call_count == 3: + return LLMResponse( + role="assistant", + tools_call_name=["test_tool"], + tools_call_args=[{"query": "repaired"}], + tools_call_ids=["call_3"], + usage=TokenUsage(input_other=17, input_cached=19, output=23), + ) + raise AssertionError("Unexpected provider request") + + provider = RepairProvider() + tool = FunctionTool( + name="test_tool", + description="test", + parameters={"type": "object", "properties": {"query": {"type": "string"}}}, + handler=AsyncMock(), + ) + conversation = SimpleNamespace(cid="test-conversation", token_usage=0) + request = ProviderRequest( + prompt="run tool", + func_tool=ToolSet(tools=[tool]), + contexts=[], + conversation=cast(Any, conversation), + ) + run_context = ContextWrapper( + context=MockAgentContext(MockEvent("test_umo", "test_sender")) + ) + runner = ToolLoopAgentRunner() + + await runner.reset( + provider=provider, + request=request, + run_context=run_context, + tool_executor=cast(Any, MockToolExecutor()), + agent_hooks=MockHooks(), + tool_schema_mode="skills_like", + ) + + async for _ in runner.step(): + pass + + assert provider.call_count == 3 + assert runner.stats.token_usage == TokenUsage( + input_other=26, + input_cached=33, + output=41, + ) + assert conversation.token_usage == 100 + + +@pytest.mark.asyncio +async def test_skills_like_without_requery_does_not_double_count_usage(): + """A tool selection without a matching schema is counted only once.""" + + class UnknownToolProvider(MockProvider): + async def text_chat(self, **kwargs) -> LLMResponse: + self.call_count += 1 + return LLMResponse( + role="assistant", + tools_call_name=["missing_tool"], + tools_call_args=[{}], + tools_call_ids=["call_missing"], + usage=TokenUsage(input_other=29, input_cached=31, output=37), + ) + + provider = UnknownToolProvider() + tool = FunctionTool( + name="test_tool", + description="test", + parameters={"type": "object", "properties": {}}, + handler=AsyncMock(), + ) + conversation = SimpleNamespace(cid="test-conversation", token_usage=0) + request = ProviderRequest( + prompt="run tool", + func_tool=ToolSet(tools=[tool]), + contexts=[], + conversation=cast(Any, conversation), + ) + runner = ToolLoopAgentRunner() + + await runner.reset( + provider=provider, + request=request, + run_context=ContextWrapper(context=None), + tool_executor=cast(Any, MockToolExecutor()), + agent_hooks=MockHooks(), + tool_schema_mode="skills_like", + ) + + async for _ in runner.step(): + pass + + assert provider.call_count == 1 + assert runner.stats.token_usage == TokenUsage( + input_other=29, + input_cached=31, + output=37, + ) + assert conversation.token_usage == 97 @pytest.mark.asyncio