Skip to content

fix(agent): preserve tool metadata after skills-like requery #9220

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
Last-emo-boy wants to merge 2 commits into
base: master
Choose a base branch
from
+164 −4
Open

fix(agent): preserve tool metadata after skills-like requery #9220

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
ceb99bd
fix(agent): preserve tool metadata after skills-like requery
Last-emo-boy Jul 12, 2026
1125c42
fix(agent): count skills-like requery usage
Last-emo-boy Jul 12, 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
11 changes: 11 additions & 0 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 @@ -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
)
Comment on lines +874 to +876

@gemini-code-assist gemini-code-assist Bot Jul 12, 2026

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.

medium

The token usage of the skills-like re-query response (requery_resp) is currently not tracked or accumulated into the agent's total token usage statistics (self.stats.token_usage) or the conversation's token usage (self.req.conversation.token_usage). This leads to inaccurate token usage reporting when the skills-like mode is active.

We should accumulate the token usage from requery_resp when it is available.

                    llm_resp.tools_call_extra_content = (
                        requery_resp.tools_call_extra_content
                    )
                    if requery_resp.usage:
                        self.stats.token_usage += requery_resp.usage
                        if self.req.conversation:
                            self.req.conversation.token_usage += requery_resp.usage.total


tool_call_result_blocks = []
cached_images = [] # Collect cached images for LLM visibility
Expand Down Expand Up @@ -1328,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)
Expand Down Expand Up @@ -1356,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)
Expand Down
157 changes: 153 additions & 4 deletions tests/test_tool_loop_agent_runner.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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_tool_metadata_and_usage():
"""skills-like re-query preserves metadata and counts its token usage."""
from astrbot.core.agent.message import TextPart

captured_kwargs = {}
Expand All @@ -1301,7 +1301,12 @@ async def text_chat(self, **kwargs) -> LLMResponse:
tools_call_name=["test_tool"],
tools_call_args=[{"query": "test"}],
tools_call_ids=["call_1"],
usage=TokenUsage(input_other=10, output=5),
tools_call_extra_content={
"call_1": {
"google": {"thought_signature": "selection-signature"}
}
},
usage=TokenUsage(input_other=10, input_cached=2, output=5),
)
if self.call_count == 2:
# 第二次调用：re-query with param schema
Expand All @@ -1312,7 +1317,10 @@ async def text_chat(self, **kwargs) -> LLMResponse:
tools_call_name=["test_tool"],
tools_call_args=[{"query": "actual"}],
tools_call_ids=["call_2"],
usage=TokenUsage(input_other=10, output=5),
tools_call_extra_content={
"call_2": {"google": {"thought_signature": "requery-signature"}}
},
usage=TokenUsage(input_other=20, input_cached=3, output=7),
)
# 后续调用：正常回复
return LLMResponse(
Expand All @@ -1331,11 +1339,13 @@ async def text_chat(self, **kwargs) -> LLMResponse:
tool_set = ToolSet(tools=[tool])

caption_part = TextPart(text="<image_caption>一张猫的照片</image_caption>")
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")
Expand Down Expand Up @@ -1363,6 +1373,145 @@ async def text_chat(self, **kwargs) -> LLMResponse:
assert len(parts) == 1
assert parts[0].text == "<image_caption>一张猫的照片</image_caption>"

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"}
}
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
async def test_follow_up_accepted_when_active_and_not_stopping(
Expand Down
Loading