fix: stabilize tool injection for LLM prefix cache hits
Two changes to make the tool schema sent to the LLM deterministic: 1. ToolSet.normalize() — sort tools by name before serialization. Called at the end of build_main_agent() after all injection passes. Eliminates ordering drift from plugin load order, MCP reconnection, and persona tool list differences. 2. Always inject full sandbox tool set — ComputerToolProvider now returns get_default_sandbox_tools() unconditionally, regardless of sandbox boot state. Browser tools are always in the schema even if the sandbox profile lacks browser capability. The executor rejects calls to unavailable browser tools with a descriptive error instead of silently omitting them from the schema. This eliminates the pre-boot/post-boot tool set jump that caused prefix cache misses on the second request of a conversation.
This commit is contained in:
@@ -106,6 +106,15 @@ class ToolSet:
|
||||
"""Remove a tool by its name."""
|
||||
self.tools = [tool for tool in self.tools if tool.name != name]
|
||||
|
||||
def normalize(self) -> None:
|
||||
"""Sort tools by name for deterministic serialization.
|
||||
|
||||
This ensures the serialized tool schema sent to the LLM is
|
||||
identical across requests regardless of registration/injection
|
||||
order, enabling LLM provider prefix cache hits.
|
||||
"""
|
||||
self.tools.sort(key=lambda t: t.name)
|
||||
|
||||
def get_tool(self, name: str) -> FunctionTool | None:
|
||||
"""Get a tool by its name."""
|
||||
for tool in self.tools:
|
||||
|
||||
@@ -169,10 +169,61 @@ class FunctionToolExecutor(BaseFunctionToolExecutor[AstrAgentContext]):
|
||||
|
||||
return
|
||||
else:
|
||||
# Guard: reject sandbox tools whose capability is unavailable.
|
||||
# Tools are always injected (for schema stability / prefix caching),
|
||||
# but execution is blocked when the sandbox lacks the capability.
|
||||
rejection = cls._check_sandbox_capability(tool, run_context)
|
||||
if rejection is not None:
|
||||
yield rejection
|
||||
return
|
||||
|
||||
async for r in cls._execute_local(tool, run_context, **tool_args):
|
||||
yield r
|
||||
return
|
||||
|
||||
# Browser tool names that require the "browser" sandbox capability.
|
||||
_BROWSER_TOOL_NAMES: frozenset[str] = frozenset({
|
||||
"astrbot_execute_browser",
|
||||
"astrbot_execute_browser_batch",
|
||||
"astrbot_run_browser_skill",
|
||||
})
|
||||
|
||||
@classmethod
|
||||
def _check_sandbox_capability(
|
||||
cls,
|
||||
tool: FunctionTool,
|
||||
run_context: ContextWrapper[AstrAgentContext],
|
||||
) -> mcp.types.CallToolResult | None:
|
||||
"""Return a rejection result if the tool requires a sandbox capability
|
||||
that is not available, or None if the tool may proceed."""
|
||||
if tool.name not in cls._BROWSER_TOOL_NAMES:
|
||||
return None
|
||||
|
||||
from astrbot.core.computer.computer_client import get_sandbox_capabilities
|
||||
|
||||
session_id = run_context.context.event.unified_msg_origin
|
||||
caps = get_sandbox_capabilities(session_id)
|
||||
|
||||
# Sandbox not yet booted — allow through (boot will happen on first
|
||||
# shell/python call; browser tools will fail naturally if truly unavailable).
|
||||
if caps is None:
|
||||
return None
|
||||
|
||||
if "browser" not in caps:
|
||||
msg = (
|
||||
f"Tool '{tool.name}' requires browser capability, but the current "
|
||||
f"sandbox profile does not include it (capabilities: {list(caps)}). "
|
||||
"Please ask the administrator to switch to a sandbox profile with "
|
||||
"browser support, or use shell/python tools instead."
|
||||
)
|
||||
logger.warning("[ToolExec] capability_rejected tool=%s caps=%s", tool.name, list(caps))
|
||||
return mcp.types.CallToolResult(
|
||||
content=[mcp.types.TextContent(type="text", text=msg)],
|
||||
isError=True,
|
||||
)
|
||||
|
||||
return None
|
||||
|
||||
@classmethod
|
||||
def _get_runtime_computer_tools(
|
||||
cls,
|
||||
|
||||
@@ -1074,6 +1074,10 @@ async def build_main_agent(
|
||||
asyncio.create_task(_handle_webchat(event, req, provider))
|
||||
|
||||
if req.func_tool and req.func_tool.tools:
|
||||
# Sort tools by name for deterministic serialization so that
|
||||
# LLM provider prefix caching can match across requests.
|
||||
req.func_tool.normalize()
|
||||
|
||||
tool_prompt = (
|
||||
TOOL_CALL_PROMPT
|
||||
if config.tool_schema_mode == "full"
|
||||
|
||||
@@ -181,15 +181,17 @@ class ComputerToolProvider:
|
||||
def _sandbox_tools(self, ctx: ToolProviderContext) -> list[FunctionTool]:
|
||||
"""Collect tools for sandbox mode.
|
||||
|
||||
Prefers the precise post-boot tool list from the running session's
|
||||
booter instance (``get_sandbox_tools``). When the sandbox has not
|
||||
yet been booted, falls back to the conservative pre-boot default
|
||||
declared by the booter class (``get_default_sandbox_tools``).
|
||||
Always returns the full (pre-boot default) tool set declared by the
|
||||
booter class, regardless of whether the sandbox is already booted.
|
||||
|
||||
This ensures the tool schema sent to the LLM is stable across the
|
||||
entire conversation lifecycle (pre-boot and post-boot produce the
|
||||
same set), enabling LLM prefix cache hits. Tools whose underlying
|
||||
capability is unavailable at runtime are rejected by the executor
|
||||
with a descriptive error message instead of being omitted from the
|
||||
schema.
|
||||
"""
|
||||
from astrbot.core.computer.computer_client import (
|
||||
get_default_sandbox_tools,
|
||||
get_sandbox_tools,
|
||||
)
|
||||
from astrbot.core.computer.computer_client import get_default_sandbox_tools
|
||||
|
||||
booter_type = ctx.sandbox_cfg.get("booter", "shipyard_neo")
|
||||
|
||||
@@ -203,12 +205,7 @@ class ComputerToolProvider:
|
||||
os.environ["SHIPYARD_ENDPOINT"] = ep
|
||||
os.environ["SHIPYARD_ACCESS_TOKEN"] = at
|
||||
|
||||
# Prefer precise post-boot tools from the running session
|
||||
booted_tools = get_sandbox_tools(ctx.session_id)
|
||||
if booted_tools:
|
||||
return booted_tools
|
||||
|
||||
# Pre-boot: conservative default from booter class
|
||||
# Always return the full tool set for schema stability
|
||||
return get_default_sandbox_tools(ctx.sandbox_cfg)
|
||||
|
||||
def _sandbox_prompt_addon(self, ctx: ToolProviderContext) -> str:
|
||||
|
||||
Reference in New Issue
Block a user