mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-06-10 01:15:58 +00:00
Xinmin Zeng
ca487578a4
* feat(agent): add ToolOutputBudgetMiddleware for oversized tool output protection Closes #3289. Adds a unified middleware that enforces per-result budgets on ALL tool outputs (MCP, sandbox, community, custom), preventing oversized external tool results from blowing the model context window. Design informed by claude-code (persistToolResult), hermes-agent (tool_result_storage), and pi (OutputAccumulator) — the three most mature implementations in production coding-agent frameworks. Key features: - Disk externalization: oversized outputs written to thread-local .tool-results/ directory, replaced with compact preview + file reference. Model can read full output via read_file with offset/limit. - Fallback truncation: head+tail truncation when disk is unavailable (no thread_data, write failure), ensuring the context is always protected. - read_file exemption: prevents persist-read-persist infinite loops (independently discovered by claude-code, hermes-agent, and pi). - Per-tool threshold overrides via config. - Line-boundary-aware truncation (no partial lines in previews). - Multimodal content passthrough (images/structured blocks skip budget). - Historical ToolMessage patching in wrap_model_call for checkpoint recovery scenarios. Related: #3222 (design RFC), #1844 (comprehensive context management), #3137 (write_file args compaction), #1677 (sandbox tool truncation). * test: add MCP content_and_artifact format coverage Add 5 tests for MCP tool output format (list of content blocks): - text content blocks are extracted and budgeted - multiple text blocks are joined and budgeted - image content blocks are skipped (multimodal passthrough) - mixed text+image blocks are skipped - small text blocks pass through unchanged Total test count: 59 (was 54). * fix(agent): address Codex review findings for ToolOutputBudgetMiddleware Three issues identified by Codex code review, all fixed: 1. `enabled` config field was unused — middleware now checks `config.enabled` and skips all processing when disabled. 2. `_build_fallback` could exceed `fallback_max_chars` — the marker text itself (~139 chars) was not deducted from the budget. Now pre-computes marker overhead and falls back to hard slice when max_chars is smaller than the marker. 3. Sync file I/O in async path — `awrap_tool_call` now delegates `_patch_result` to `asyncio.to_thread` to avoid blocking the event loop during disk writes. Tests updated to use realistic fallback_max_chars values (500+) that can accommodate the marker overhead, plus two new tests: - `test_result_never_exceeds_max_chars` (parametric across sizes) - `test_very_small_max_chars_does_not_crash` * fix(agent): address Copilot review — path traversal, async perf, shared config 1. Path traversal defense: sanitize tool_name via _sanitize_tool_name() (strips separators, .., absolute paths), validate storage_subdir is relative, and verify resolved filepath stays inside storage_dir. 2. Async hot-path optimization: add _needs_budget() cheap check before asyncio.to_thread offload — small outputs (99% of calls) skip the thread overhead entirely. 3. Replace shared module-level _DEFAULT_CONFIG with _default_config() factory to prevent cross-instance mutation of mutable fields. 12 new tests: TestSanitizeToolName (5), TestExternalizePathTraversal (3), TestNeedsBudget (4). * fix(agent): correct preview hint to match read_file actual API read_file uses start_line/end_line (1-indexed line numbers), not offset/limit. The previous wording was copied from hermes-agent which has a different read_file interface. * perf(agent): hoist hot-path imports, add model-call pre-scan (review #3303) Address maintainer review feedback: 1. Hoist inline imports to module level — `import asyncio` (was in awrap_tool_call hot path) and `from dataclasses import replace` (was in _patch_result) now live at module top. 2. Add a cheap pre-scan to _patch_model_messages so the historical message list is not rebuilt on every model call when nothing is oversized (the common case once results are budgeted at tool-call time). Also adds the same _needs_budget gate to the sync wrap_tool_call for symmetry with awrap_tool_call. The pre-scan is refactored into per-tool-aware helpers (_effective_trigger / _tool_message_over_budget) that mirror the exact trigger conditions in _budget_content — including tool_overrides — so the fast-path can never produce a false negative (silently skipping budgeting for a tool with a low per-tool threshold). 7 new regression tests lock the per-tool-override-through-pre-scan path and the model-call early return. --------- Co-authored-by: Willem Jiang <willem.jiang@gmail.com>
63 lines
2.3 KiB
Python
63 lines
2.3 KiB
Python
"""Configuration for tool output budget protection."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from pydantic import BaseModel, Field
|
|
|
|
|
|
class ToolOutputConfig(BaseModel):
|
|
"""Config section for tool-result output budget enforcement.
|
|
|
|
When a tool returns more than ``externalize_min_chars`` characters,
|
|
the full output is persisted to disk and replaced with a compact
|
|
preview + file reference. If disk persistence is unavailable the
|
|
output falls back to head+tail truncation.
|
|
"""
|
|
|
|
enabled: bool = Field(
|
|
default=True,
|
|
description="Enable the tool output budget middleware.",
|
|
)
|
|
externalize_min_chars: int = Field(
|
|
default=12_000,
|
|
ge=0,
|
|
description="Character threshold to trigger disk externalization. Outputs below this pass through unchanged. Set to 0 to disable externalization (fallback truncation still applies when output exceeds fallback_max_chars).",
|
|
)
|
|
preview_head_chars: int = Field(
|
|
default=2_000,
|
|
ge=0,
|
|
description="Characters to keep from the head of the output in the preview.",
|
|
)
|
|
preview_tail_chars: int = Field(
|
|
default=1_000,
|
|
ge=0,
|
|
description="Characters to keep from the tail of the output in the preview.",
|
|
)
|
|
fallback_max_chars: int = Field(
|
|
default=30_000,
|
|
ge=0,
|
|
description="Maximum characters when disk persistence is unavailable. 0 disables fallback truncation.",
|
|
)
|
|
fallback_head_chars: int = Field(
|
|
default=8_000,
|
|
ge=0,
|
|
description="Head characters for fallback truncation.",
|
|
)
|
|
fallback_tail_chars: int = Field(
|
|
default=3_000,
|
|
ge=0,
|
|
description="Tail characters for fallback truncation.",
|
|
)
|
|
storage_subdir: str = Field(
|
|
default=".tool-results",
|
|
description="Subdirectory under the thread outputs path for persisted tool results.",
|
|
)
|
|
exempt_tools: list[str] = Field(
|
|
default_factory=lambda: ["read_file", "read_file_tool"],
|
|
description="Tool names exempt from budget enforcement (prevents persist→read→persist loops).",
|
|
)
|
|
tool_overrides: dict[str, int] = Field(
|
|
default_factory=dict,
|
|
description="Per-tool externalize_min_chars overrides. Keys are tool names, values are char thresholds. Use 0 to disable externalization for a specific tool.",
|
|
)
|