feat: implement handling for large tool results with overflow file writing and read tool integration

This commit is contained in:
Soulter
2026-04-10 17:26:17 +08:00
parent 3acda6f77a
commit 1745e9c4fb
8 changed files with 356 additions and 46 deletions

View File

@@ -4,9 +4,11 @@ import sys
import time
import traceback
import typing as T
import uuid
from collections.abc import AsyncIterator
from contextlib import suppress
from dataclasses import dataclass, field
from pathlib import Path
from mcp.types import (
BlobResourceContents,
@@ -25,7 +27,7 @@ from tenacity import (
from astrbot import logger
from astrbot.core.agent.message import ImageURLPart, TextPart, ThinkPart
from astrbot.core.agent.tool import ToolSet
from astrbot.core.agent.tool import FunctionTool, ToolSet
from astrbot.core.agent.tool_image_cache import tool_image_cache
from astrbot.core.exceptions import EmptyModelOutputError
from astrbot.core.message.components import Json
@@ -45,7 +47,7 @@ from astrbot.core.provider.provider import Provider
from ..context.compressor import ContextCompressor
from ..context.config import ContextConfig
from ..context.manager import ContextManager
from ..context.token_counter import TokenCounter
from ..context.token_counter import EstimateTokenCounter, TokenCounter
from ..hooks import BaseAgentRunHooks
from ..message import AssistantMessageSegment, Message, ToolCallMessageSegment
from ..response import AgentResponseData, AgentStats
@@ -97,6 +99,8 @@ ToolExecutorResultT = T.TypeVar("ToolExecutorResultT")
class ToolLoopAgentRunner(BaseAgentRunner[TContext]):
TOOL_RESULT_MAX_ESTIMATED_TOKENS = 27_500
TOOL_RESULT_PREVIEW_MAX_ESTIMATED_TOKENS = 7000
EMPTY_OUTPUT_RETRY_ATTEMPTS = 3
EMPTY_OUTPUT_RETRY_WAIT_MIN_S = 1
EMPTY_OUTPUT_RETRY_WAIT_MAX_S = 4
@@ -151,6 +155,11 @@ class ToolLoopAgentRunner(BaseAgentRunner[TContext]):
"Otherwise, change strategy, adjust arguments, or explain the limitation "
"to the user."
)
TOOL_RESULT_OVERFLOW_NOTICE_TEMPLATE = (
"Truncated tool output preview shown above. "
"The tool output was too large to include directly and was written to "
"`{overflow_path}`. Use {read_tool_hint} with a narrower window to inspect it."
)
def _get_persona_custom_error_message(self) -> str | None:
"""Read persona-level custom error message from event extras when available."""
@@ -206,6 +215,8 @@ class ToolLoopAgentRunner(BaseAgentRunner[TContext]):
custom_compressor: ContextCompressor | None = None,
tool_schema_mode: str | None = "full",
fallback_providers: list[Provider] | None = None,
tool_result_overflow_dir: str | None = None,
read_tool: FunctionTool | None = None,
**kwargs: T.Any,
) -> None:
self.req = request
@@ -217,6 +228,9 @@ class ToolLoopAgentRunner(BaseAgentRunner[TContext]):
self.truncate_turns = truncate_turns
self.custom_token_counter = custom_token_counter
self.custom_compressor = custom_compressor
self.tool_result_overflow_dir = tool_result_overflow_dir
self.read_tool = read_tool
self._tool_result_token_counter = EstimateTokenCounter()
# we will do compress when:
# 1. before requesting LLM
# TODO: 2. after LLM output a tool call
@@ -298,6 +312,103 @@ class ToolLoopAgentRunner(BaseAgentRunner[TContext]):
self.stats = AgentStats()
self.stats.start_time = time.time()
def _read_tool_hint(self) -> str:
if self.read_tool is not None:
return f"`{self.read_tool.name}`"
return "the available file-read tool"
async def _write_tool_result_overflow_file(
self,
*,
tool_call_id: str,
content: str,
) -> str:
if self.tool_result_overflow_dir is None:
raise ValueError("tool_result_overflow_dir is not configured")
overflow_dir = Path(self.tool_result_overflow_dir).resolve(strict=False)
safe_tool_call_id = (
"".join(
ch if ch.isalnum() or ch in {"-", "_", "."} else "_"
for ch in tool_call_id
).strip("._")
or "tool_call"
)
file_name = f"{safe_tool_call_id}_{uuid.uuid4().hex[:8]}.txt"
overflow_path = overflow_dir / file_name
def _run() -> str:
overflow_dir.mkdir(parents=True, exist_ok=True)
overflow_path.write_text(content, encoding="utf-8")
return str(overflow_path)
return await asyncio.to_thread(_run)
async def _materialize_large_tool_result(
self,
*,
tool_call_id: str,
content: str,
) -> str:
if self.tool_result_overflow_dir is None or self.read_tool is None:
return content
estimated_tokens = self._tool_result_token_counter.count_tokens(
[Message(role="tool", content=content, tool_call_id=tool_call_id)]
)
if estimated_tokens <= self.TOOL_RESULT_MAX_ESTIMATED_TOKENS:
return content
preview = self._truncate_tool_result_preview(content, tool_call_id=tool_call_id)
try:
overflow_path = await self._write_tool_result_overflow_file(
tool_call_id=tool_call_id,
content=content,
)
except Exception as exc:
logger.warning(
"Failed to spill oversized tool result for %s: %s",
tool_call_id,
exc,
exc_info=True,
)
error_notice = (
"Tool output exceeded the inline result limit "
f"({estimated_tokens} estimated tokens > "
f"{self.TOOL_RESULT_MAX_ESTIMATED_TOKENS}) and could not be written "
f"to `{self.tool_result_overflow_dir}`: {exc}"
)
if not preview:
return error_notice
return f"{preview}\n\n{error_notice}"
notice = self.TOOL_RESULT_OVERFLOW_NOTICE_TEMPLATE.format(
overflow_path=overflow_path,
read_tool_hint=self._read_tool_hint(),
)
if not preview:
return notice
return f"{preview}\n\n{notice}"
def _truncate_tool_result_preview(
self,
content: str,
*,
tool_call_id: str,
) -> str:
preview = content
while preview:
estimated_tokens = self._tool_result_token_counter.count_tokens(
[Message(role="tool", content=preview, tool_call_id=tool_call_id)]
)
if estimated_tokens <= self.TOOL_RESULT_PREVIEW_MAX_ESTIMATED_TOKENS:
return preview
next_len = len(preview) // 2
if next_len <= 0:
break
preview = preview[:next_len]
return preview
async def _iter_llm_responses(
self, *, include_model: bool = True
) -> T.AsyncGenerator[LLMResponse, None]:
@@ -933,9 +1044,14 @@ class ToolLoopAgentRunner(BaseAgentRunner[TContext]):
"The tool has returned a data type that is not supported."
)
if result_parts:
inline_result = "\n\n".join(result_parts)
inline_result = await self._materialize_large_tool_result(
tool_call_id=func_tool_id,
content=inline_result,
)
_append_tool_call_result(
func_tool_id,
"\n\n".join(result_parts)
inline_result
+ self._build_repeated_tool_call_guidance(
func_tool_name, tool_call_streak
),

View File

@@ -81,7 +81,10 @@ from astrbot.core.tools.web_search_tools import (
TavilyWebSearchTool,
normalize_legacy_web_search_config,
)
from astrbot.core.utils.astrbot_path import get_astrbot_workspaces_path
from astrbot.core.utils.astrbot_path import (
get_astrbot_system_tmp_path,
get_astrbot_workspaces_path,
)
from astrbot.core.utils.file_extract import extract_file_moonshotai
from astrbot.core.utils.llm_metadata import LLM_METADATAS
from astrbot.core.utils.media_utils import (
@@ -1471,6 +1474,14 @@ async def build_main_agent(
fallback_providers=_get_fallback_chat_providers(
provider, plugin_context, config.provider_settings
),
tool_result_overflow_dir=(
get_astrbot_system_tmp_path()
if req.func_tool and req.func_tool.get_tool("astrbot_file_read_tool")
else None
),
read_tool=(
req.func_tool.get_tool("astrbot_file_read_tool") if req.func_tool else None
),
)
if apply_reset:

View File

@@ -182,8 +182,22 @@ def detect_text_encoding(sample: bytes) -> str | None:
for encoding in _TEXT_ENCODINGS:
try:
decoded = sample.decode(encoding)
except UnicodeDecodeError:
continue
except UnicodeDecodeError as exc:
# Probe samples can end in the middle of a multibyte sequence.
# When the decode failure only happens at the sample tail, trim a few
# bytes and retry so UTF-8 text is not misclassified as binary.
if exc.start >= len(sample) - 4:
decoded = ""
for trim_bytes in range(1, min(4, len(sample)) + 1):
try:
decoded = sample[:-trim_bytes].decode(encoding)
break
except UnicodeDecodeError:
continue
if not decoded:
continue
else:
continue
if _looks_like_text(decoded):
return encoding

View File

@@ -36,6 +36,7 @@ from astrbot.core.star.filter.platform_adapter_type import (
PlatformAdapterType,
)
from astrbot.core.subagent_orchestrator import SubAgentOrchestrator
from astrbot.core.utils.astrbot_path import get_astrbot_system_tmp_path
from ..exceptions import ProviderNotFoundError
from .filter.command import CommandFilter
@@ -232,6 +233,13 @@ class Context:
for k, v in kwargs.items()
if k not in ["stream", "agent_hooks", "agent_context"]
}
if request.func_tool and request.func_tool.get_tool("astrbot_file_read_tool"):
other_kwargs.setdefault(
"tool_result_overflow_dir", get_astrbot_system_tmp_path()
)
other_kwargs.setdefault(
"read_tool", request.func_tool.get_tool("astrbot_file_read_tool")
)
await agent_runner.reset(
provider=prov,

View File

@@ -46,6 +46,7 @@ from astrbot.core.computer.file_read_utils import read_file_tool_result
from astrbot.core.message.components import File
from astrbot.core.utils.astrbot_path import (
get_astrbot_skills_path,
get_astrbot_system_tmp_path,
get_astrbot_temp_path,
)
@@ -71,7 +72,7 @@ def _restricted_env_path_labels(umo: str) -> list[str]:
return [
"data/skills",
f"data/workspaces/{normalized_umo}",
"/tmp/.astrbot",
get_astrbot_system_tmp_path(),
]
@@ -91,7 +92,7 @@ def _read_allowed_roots(umo: str) -> tuple[Path, ...]:
return (
Path(get_astrbot_skills_path()).resolve(strict=False),
_workspace_root(umo),
Path("/tmp/.astrbot").resolve(strict=False),
Path(get_astrbot_system_tmp_path()).resolve(strict=False),
)

View File

@@ -1,33 +1,33 @@
"""Astrbot统一路径获取
"""Centralized AstrBot path helpers.
项目路径:固定为源码所在路径
根目录路径:默认为当前工作目录,可通过环境变量 ASTRBOT_ROOT 指定
数据目录路径:固定为根目录下的 data 目录
配置文件路径:固定为数据目录下的 config 目录
插件目录路径:固定为数据目录下的 plugins 目录
插件数据目录路径:固定为数据目录下的 plugin_data 目录
T2I 模板目录路径:固定为数据目录下的 t2i_templates 目录
WebChat 数据目录路径:固定为数据目录下的 webchat 目录
临时文件目录路径:固定为数据目录下的 temp 目录
Skills 目录路径:固定为数据目录下的 skills 目录
Workspaces 目录路径:固定为数据目录下的 workspaces 目录
第三方依赖目录路径:固定为数据目录下的 site-packages 目录
Project path:
- Fixed to the source tree location.
Root path:
- Defaults to the current working directory.
- Can be overridden with the ``ASTRBOT_ROOT`` environment variable.
Data subdirectories:
- Most runtime data lives under ``<root>/data``.
- A few tool-runtime files intentionally live under the system temporary
directory as ``.astrbot``.
"""
import os
import tempfile
from astrbot.core.utils.runtime_env import is_packaged_desktop_runtime
def get_astrbot_path() -> str:
"""获取Astrbot项目路径"""
"""Return the AstrBot project source path."""
return os.path.realpath(
os.path.join(os.path.dirname(os.path.abspath(__file__)), "../../../"),
)
def get_astrbot_root() -> str:
"""获取Astrbot根目录路径"""
"""Return the AstrBot root directory."""
if path := os.environ.get("ASTRBOT_ROOT"):
return os.path.realpath(path)
if is_packaged_desktop_runtime():
@@ -36,60 +36,65 @@ def get_astrbot_root() -> str:
def get_astrbot_data_path() -> str:
"""获取Astrbot数据目录路径"""
"""Return the AstrBot data directory path."""
return os.path.realpath(os.path.join(get_astrbot_root(), "data"))
def get_astrbot_config_path() -> str:
"""获取Astrbot配置文件路径"""
"""Return the AstrBot config directory path."""
return os.path.realpath(os.path.join(get_astrbot_data_path(), "config"))
def get_astrbot_plugin_path() -> str:
"""获取Astrbot插件目录路径"""
"""Return the AstrBot plugin directory path."""
return os.path.realpath(os.path.join(get_astrbot_data_path(), "plugins"))
def get_astrbot_plugin_data_path() -> str:
"""获取Astrbot插件数据目录路径"""
"""Return the AstrBot plugin data directory path."""
return os.path.realpath(os.path.join(get_astrbot_data_path(), "plugin_data"))
def get_astrbot_t2i_templates_path() -> str:
"""获取Astrbot T2I 模板目录路径"""
"""Return the AstrBot T2I templates directory path."""
return os.path.realpath(os.path.join(get_astrbot_data_path(), "t2i_templates"))
def get_astrbot_webchat_path() -> str:
"""获取Astrbot WebChat 数据目录路径"""
"""Return the AstrBot WebChat data directory path."""
return os.path.realpath(os.path.join(get_astrbot_data_path(), "webchat"))
def get_astrbot_temp_path() -> str:
"""获取Astrbot临时文件目录路径"""
"""Return the AstrBot temporary data directory path."""
return os.path.realpath(os.path.join(get_astrbot_data_path(), "temp"))
def get_astrbot_skills_path() -> str:
"""获取Astrbot Skills 目录路径"""
"""Return the AstrBot skills directory path."""
return os.path.realpath(os.path.join(get_astrbot_data_path(), "skills"))
def get_astrbot_workspaces_path() -> str:
"""获取Astrbot Workspaces 目录路径"""
"""Return the AstrBot workspaces directory path."""
return os.path.realpath(os.path.join(get_astrbot_data_path(), "workspaces"))
def get_astrbot_system_tmp_path() -> str:
"""Return the shared system temporary directory used by local tools."""
return os.path.realpath(os.path.join(tempfile.gettempdir(), ".astrbot"))
def get_astrbot_site_packages_path() -> str:
"""获取Astrbot第三方依赖目录路径"""
"""Return the AstrBot third-party site-packages directory path."""
return os.path.realpath(os.path.join(get_astrbot_data_path(), "site-packages"))
def get_astrbot_knowledge_base_path() -> str:
"""获取Astrbot知识库根目录路径"""
"""Return the AstrBot knowledge base root path."""
return os.path.realpath(os.path.join(get_astrbot_data_path(), "knowledge_base"))
def get_astrbot_backups_path() -> str:
"""获取Astrbot备份目录路径"""
"""Return the AstrBot backups directory path."""
return os.path.realpath(os.path.join(get_astrbot_data_path(), "backups"))

View File

@@ -92,6 +92,12 @@ def _make_large_text() -> str:
return "".join(f"line-{index:05d}-{'x' * 48}\n" for index in range(6000))
def test_detect_text_encoding_allows_utf8_probe_cut_mid_character():
sample = '{"results": ["中文内容"]}'.encode()[:-1]
assert file_read_utils.detect_text_encoding(sample) in {"utf-8", "utf-8-sig"}
@pytest.mark.asyncio
async def test_file_read_tool_rejects_large_full_text_read_before_local_stream_read(
monkeypatch: pytest.MonkeyPatch,
@@ -249,7 +255,7 @@ async def test_file_read_tool_stores_long_converted_document_in_workspace(
assert len(converted_files) == 1
assert converted_files[0].read_text(encoding="utf-8") == long_text
assert str(converted_files[0]) in result
assert "narrow `offset`/`limit` window" in result
assert "Read or grep that file with a narrow window." in result
@pytest.mark.asyncio

View File

@@ -1,6 +1,7 @@
import asyncio
import os
import sys
from pathlib import Path
from types import SimpleNamespace
from typing import Any, cast
from unittest.mock import AsyncMock
@@ -97,6 +98,26 @@ class MockToolExecutor:
return generator()
class LargeTextToolExecutor:
"""模拟返回超长文本的工具执行器"""
def __init__(self, text: str):
self.text = text
@classmethod
def from_text(cls, text: str) -> "LargeTextToolExecutor":
return cls(text)
def execute(self, tool, run_context, **tool_args):
async def generator():
from mcp.types import CallToolResult, TextContent
result = CallToolResult(content=[TextContent(type="text", text=self.text)])
yield result
return generator()
class MockMixedContentToolExecutor:
"""模拟返回图片 + 文本的工具执行器"""
@@ -193,6 +214,32 @@ class MockToolCallProvider(MockProvider):
)
class SingleToolThenFinalProvider(MockProvider):
def __init__(self, tool_name: str, tool_args: dict[str, str] | None = None):
super().__init__()
self.tool_name = tool_name
self.tool_args = tool_args or {}
async def text_chat(self, **kwargs) -> LLMResponse:
self.call_count += 1
func_tool = kwargs.get("func_tool")
if func_tool is None or self.call_count > 1:
return LLMResponse(
role="assistant",
completion_text="最终回复",
usage=TokenUsage(input_other=10, output=5),
)
return LLMResponse(
role="assistant",
completion_text="",
tools_call_name=[self.tool_name],
tools_call_args=[self.tool_args],
tools_call_ids=["call_large_result"],
usage=TokenUsage(input_other=10, output=5),
)
class SequentialToolProvider(MockProvider):
def __init__(self, tool_sequence: list[str]):
super().__init__()
@@ -334,6 +381,10 @@ def runner():
return ToolLoopAgentRunner()
def _make_large_tool_result_text() -> str:
return "x" * 100000
@pytest.mark.asyncio
async def test_max_step_limit_functionality(
runner, mock_provider, provider_request, mock_tool_executor, mock_hooks
@@ -1124,18 +1175,116 @@ async def test_follow_up_accepted_when_active_and_not_stopping(
streaming=False,
)
# Runner is active (not done) and stop is not requested
assert not runner.done()
assert runner._is_stop_requested() is False
ticket = runner.follow_up(message_text="valid follow-up message")
assert ticket is not None, (
"Follow-up should be accepted when runner is active and not stopping"
@pytest.mark.asyncio
async def test_large_tool_result_is_spilled_to_file_and_replaced_with_read_notice(
tmp_path,
):
tool = FunctionTool(
name="test_tool",
description="测试工具",
parameters={"type": "object", "properties": {"query": {"type": "string"}}},
handler=AsyncMock(),
)
assert ticket.text == "valid follow-up message"
assert ticket.consumed is False
assert ticket in runner._pending_follow_ups
read_tool = FunctionTool(
name="astrbot_file_read_tool",
description="read file",
parameters={"type": "object", "properties": {"path": {"type": "string"}}},
handler=AsyncMock(),
)
tool_set = ToolSet(tools=[tool, read_tool])
provider = SingleToolThenFinalProvider(tool.name, {"query": "large"})
request = ProviderRequest(prompt="run tool", func_tool=tool_set, contexts=[])
runner = ToolLoopAgentRunner()
await runner.reset(
provider=provider,
request=request,
run_context=ContextWrapper(context=None),
tool_executor=cast(
Any,
LargeTextToolExecutor.from_text(_make_large_tool_result_text()),
),
agent_hooks=MockHooks(),
streaming=False,
tool_result_overflow_dir=str(tmp_path),
read_tool=read_tool,
)
responses = []
async for response in runner.step_until_done(3):
responses.append(response)
tool_messages = [m for m in runner.run_context.messages if m.role == "tool"]
assert len(tool_messages) == 1
tool_message_content = str(tool_messages[0].content)
assert "xxxxxxxxxx" in tool_message_content
assert "Truncated tool output preview shown above." in tool_message_content
assert "The tool output was too large to include directly" in tool_message_content
assert "`astrbot_file_read_tool`" in tool_message_content
assert "Use `astrbot_file_read_tool` to inspect it." in tool_message_content
overflow_files = list(Path(tmp_path).glob("call_large_result_*.txt"))
assert len(overflow_files) == 1
assert (
overflow_files[0].read_text(encoding="utf-8") == _make_large_tool_result_text()
)
assert str(overflow_files[0]) in tool_message_content
llm_results = [resp for resp in responses if resp.type == "llm_result"]
assert llm_results
@pytest.mark.asyncio
async def test_large_tool_result_keeps_preview_when_spill_fails(
tmp_path,
monkeypatch: pytest.MonkeyPatch,
):
tool = FunctionTool(
name="test_tool",
description="测试工具",
parameters={"type": "object", "properties": {"query": {"type": "string"}}},
handler=AsyncMock(),
)
read_tool = FunctionTool(
name="astrbot_file_read_tool",
description="read file",
parameters={"type": "object", "properties": {"path": {"type": "string"}}},
handler=AsyncMock(),
)
tool_set = ToolSet(tools=[tool, read_tool])
provider = SingleToolThenFinalProvider(tool.name, {"query": "large"})
request = ProviderRequest(prompt="run tool", func_tool=tool_set, contexts=[])
runner = ToolLoopAgentRunner()
async def _raise_spill_error(*, tool_call_id: str, content: str) -> str:
raise OSError("disk full")
monkeypatch.setattr(runner, "_write_tool_result_overflow_file", _raise_spill_error)
await runner.reset(
provider=provider,
request=request,
run_context=ContextWrapper(context=None),
tool_executor=cast(
Any,
LargeTextToolExecutor.from_text(_make_large_tool_result_text()),
),
agent_hooks=MockHooks(),
streaming=False,
tool_result_overflow_dir=str(tmp_path),
read_tool=read_tool,
)
async for _ in runner.step_until_done(3):
pass
tool_messages = [m for m in runner.run_context.messages if m.role == "tool"]
assert len(tool_messages) == 1
tool_message_content = str(tool_messages[0].content)
assert "xxxxxxxxxx" in tool_message_content
assert "Tool output exceeded the inline result limit" in tool_message_content
assert "disk full" in tool_message_content
@pytest.mark.asyncio