furyhawk/deer-flow
1
0
Fork 0
mirror of https://github.com/bytedance/deer-flow.git synced 2026-06-10 17:35:57 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat(agent): add ToolOutputBudgetMiddleware for oversized tool output protection (#3303)

Browse Source 
* 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>
This commit is contained in:
Xinmin Zeng
2026-05-29 22:59:26 +08:00
committed by GitHub
parent e683ed6a76
commit ca487578a4
7 changed files with 1481 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backend/packages/harness/deerflow/agents/middlewares/tool_error_handling_middleware.py
+3 -1
View File
@@ -77,9 +77,11 @@ def _build_runtime_middlewares(
"""Build shared base middlewares for agent execution."""
from deerflow.agents.middlewares.llm_error_handling_middleware import LLMErrorHandlingMiddleware
from deerflow.agents.middlewares.thread_data_middleware import ThreadDataMiddleware
from deerflow.agents.middlewares.tool_output_budget_middleware import ToolOutputBudgetMiddleware
from deerflow.sandbox.middleware import SandboxMiddleware
middlewares: list[AgentMiddleware] = [
ToolOutputBudgetMiddleware.from_app_config(app_config),
ThreadDataMiddleware(lazy_init=lazy_init),
SandboxMiddleware(lazy_init=lazy_init),
]
@@ -87,7 +89,7 @@ def _build_runtime_middlewares(
if include_uploads:
from deerflow.agents.middlewares.uploads_middleware import UploadsMiddleware
middlewares.insert(1, UploadsMiddleware())
middlewares.insert(2, UploadsMiddleware())
if include_dangling_tool_call_patch:
from deerflow.agents.middlewares.dangling_tool_call_middleware import DanglingToolCallMiddleware
backend/packages/harness/deerflow/agents/middlewares/tool_output_budget_middleware.py
+489
View File
@@ -0,0 +1,489 @@
"""Middleware that enforces a per-result budget on tool outputs.
Oversized tool results are persisted to disk and replaced with a compact
preview containing a file reference. When disk persistence is
unavailable the middleware falls back to head+tail truncation so the
model context is never blown by a single large tool return.
"""
from __future__ import annotations
import asyncio
import logging
import os
import uuid
from collections.abc import Awaitable, Callable
from dataclasses import replace as dc_replace
from typing import Any, override
from langchain.agents import AgentState
from langchain.agents.middleware import AgentMiddleware
from langchain.agents.middleware.types import ModelCallResult, ModelRequest, ModelResponse
from langchain_core.messages import ToolMessage
from langgraph.prebuilt.tool_node import ToolCallRequest
from langgraph.types import Command
from deerflow.config.tool_output_config import ToolOutputConfig
logger = logging.getLogger(__name__)
def _default_config() -> ToolOutputConfig:
return ToolOutputConfig()
# ---------------------------------------------------------------------------
# Text helpers
# ---------------------------------------------------------------------------
def _message_text(content: Any) -> str | None:
"""Extract a plain-text representation from a ToolMessage content field.
Returns ``None`` for non-string / multimodal content so the caller
can skip budget enforcement (images, structured blocks, etc.).
"""
if isinstance(content, str):
return content
if content is None:
return None
if isinstance(content, list):
pieces: list[str] = []
for part in content:
if isinstance(part, str):
pieces.append(part)
elif isinstance(part, dict) and isinstance(part.get("text"), str):
pieces.append(part["text"])
else:
return None
return "\n".join(pieces) if pieces else None
return None
def _snap_to_line_boundary(text: str, pos: int) -> int:
"""Return *pos* or the nearest preceding newline+1, whichever is closer.
Used so that previews and truncations end on a complete line when
possible. If no newline exists in the second half of ``text[:pos]``
the original *pos* is returned unchanged.
"""
if pos <= 0 or pos >= len(text):
return pos
half = pos // 2
nl = text.rfind("\n", half, pos)
if nl >= 0:
return nl + 1
return pos
# ---------------------------------------------------------------------------
# Disk persistence
# ---------------------------------------------------------------------------
_EXT_MAP: dict[str, str] = {
"bash": "log",
"bash_tool": "log",
"web_fetch": "log",
}
def _sanitize_tool_name(name: str) -> str:
"""Strip path separators and traversal components from a tool name."""
base = os.path.basename(name)
safe = base.replace("..", "").replace("/", "_").replace("\\", "_")
return safe or "unknown"
def _externalize(
content: str,
*,
tool_name: str,
tool_call_id: str,
outputs_path: str,
storage_subdir: str,
) -> str | None:
"""Write *content* to disk and return the virtual path, or ``None`` on failure."""
if os.path.isabs(storage_subdir) or ".." in storage_subdir:
return None
storage_dir = os.path.join(outputs_path, storage_subdir)
try:
os.makedirs(storage_dir, exist_ok=True)
except OSError:
return None
safe_name = _sanitize_tool_name(tool_name)
ext = _EXT_MAP.get(tool_name, "txt")
short_id = uuid.uuid4().hex[:12]
filename = f"{safe_name}-{short_id}.{ext}"
filepath = os.path.join(storage_dir, filename)
if not os.path.abspath(filepath).startswith(os.path.abspath(storage_dir)):
return None
try:
with open(filepath, "w", encoding="utf-8") as f:
f.write(content)
except OSError:
return None
virtual_base = "/mnt/user-data/outputs"
return f"{virtual_base}/{storage_subdir}/{filename}"
# ---------------------------------------------------------------------------
# Preview / fallback builders
# ---------------------------------------------------------------------------
def _build_preview(
content: str,
*,
tool_name: str,
virtual_path: str,
head_chars: int,
tail_chars: int,
) -> str:
"""Build a preview with a file reference for externalized output."""
total = len(content)
head_end = _snap_to_line_boundary(content, min(head_chars, total))
tail_start = max(head_end, total - tail_chars)
tail_start_snapped = _snap_to_line_boundary(content, tail_start)
if tail_start_snapped > head_end:
tail_start = tail_start_snapped
head = content[:head_end]
tail = content[tail_start:] if tail_start < total else ""
omitted = total - len(head) - len(tail)
ref = f"\n\n[Full {tool_name} output saved to {virtual_path} ({total} chars, ~{total // 4} tokens). Use read_file with start_line and end_line to access specific sections. {omitted} chars omitted from this preview.]\n\n"
parts = [head, ref]
if tail:
parts.append(tail)
return "".join(parts)
def _build_fallback(
content: str,
*,
tool_name: str,
max_chars: int,
head_chars: int,
tail_chars: int,
) -> str:
"""Build a head+tail truncation when disk persistence is unavailable.
The returned string is guaranteed to be no longer than *max_chars*.
"""
total = len(content)
if max_chars <= 0 or total <= max_chars:
return content
marker_template = "\n\n[... {n} chars omitted from {tn} output. Persistent storage unavailable. Consider narrowing the query or using more specific parameters.]\n\n"
marker_overhead = len(marker_template.format(n=total, tn=tool_name))
if marker_overhead >= max_chars:
return content[:max_chars]
budget = max_chars - marker_overhead
effective_head = min(head_chars, budget)
effective_tail = min(tail_chars, max(0, budget - effective_head))
head_end = _snap_to_line_boundary(content, min(effective_head, total))
tail_start = max(head_end, total - effective_tail)
tail_start_snapped = _snap_to_line_boundary(content, tail_start)
if tail_start_snapped > head_end:
tail_start = tail_start_snapped
head = content[:head_end]
tail = content[tail_start:] if tail_start < total else ""
omitted = total - len(head) - len(tail)
marker = marker_template.format(n=omitted, tn=tool_name)
parts = [head, marker]
if tail:
parts.append(tail)
return "".join(parts)
# ---------------------------------------------------------------------------
# Core budget logic
# ---------------------------------------------------------------------------
def _resolve_outputs_path(request: ToolCallRequest) -> str | None:
"""Best-effort extraction of the thread outputs path."""
runtime = getattr(request, "runtime", None)
if runtime is None:
return None
state = getattr(runtime, "state", None)
if state is None:
return None
thread_data = state.get("thread_data")
if not isinstance(thread_data, dict):
return None
outputs_path = thread_data.get("outputs_path")
return outputs_path if isinstance(outputs_path, str) else None
def _budget_content(
content: str,
*,
tool_name: str,
tool_call_id: str,
outputs_path: str | None,
config: ToolOutputConfig,
) -> str | None:
"""Apply budget to *content*. Returns ``None`` if no change needed."""
threshold = config.tool_overrides.get(tool_name, config.externalize_min_chars)
if threshold <= 0 and config.fallback_max_chars <= 0:
return None
if len(content) <= threshold and len(content) <= config.fallback_max_chars:
return None
if threshold > 0 and len(content) > threshold and outputs_path:
virtual_path = _externalize(
content,
tool_name=tool_name,
tool_call_id=tool_call_id,
outputs_path=outputs_path,
storage_subdir=config.storage_subdir,
)
if virtual_path is not None:
logger.info(
"Externalized %s output (%d chars) to %s",
tool_name,
len(content),
virtual_path,
)
return _build_preview(
content,
tool_name=tool_name,
virtual_path=virtual_path,
head_chars=config.preview_head_chars,
tail_chars=config.preview_tail_chars,
)
if config.fallback_max_chars > 0 and len(content) > config.fallback_max_chars:
logger.warning(
"Fallback-truncating %s output: %d chars → %d max",
tool_name,
len(content),
config.fallback_max_chars,
)
return _build_fallback(
content,
tool_name=tool_name,
max_chars=config.fallback_max_chars,
head_chars=config.fallback_head_chars,
tail_chars=config.fallback_tail_chars,
)
return None
# ---------------------------------------------------------------------------
# Result patchers
# ---------------------------------------------------------------------------
def _patch_tool_message(msg: ToolMessage, config: ToolOutputConfig, outputs_path: str | None) -> ToolMessage:
"""Apply budget to a single ToolMessage. Returns the original if unchanged."""
tool_name = msg.name or "unknown"
if tool_name in config.exempt_tools:
return msg
text = _message_text(msg.content)
if text is None:
return msg
replacement = _budget_content(
text,
tool_name=tool_name,
tool_call_id=msg.tool_call_id or "",
outputs_path=outputs_path,
config=config,
)
if replacement is None:
return msg
update: dict[str, Any] = {"content": replacement}
if getattr(msg, "response_metadata", None):
update["response_metadata"] = dict(msg.response_metadata)
if getattr(msg, "additional_kwargs", None):
update["additional_kwargs"] = dict(msg.additional_kwargs)
return msg.model_copy(update=update)
def _effective_trigger(tool_name: str, config: ToolOutputConfig) -> int:
"""Smallest content length that could trigger budgeting for *tool_name*.
Mirrors the trigger conditions in :func:`_budget_content` (per-tool
externalize threshold OR global fallback), so the pre-scan never produces
a false negative. Returns ``-1`` when nothing could ever trigger.
"""
candidates: list[int] = []
externalize = config.tool_overrides.get(tool_name, config.externalize_min_chars)
if externalize > 0:
candidates.append(externalize)
if config.fallback_max_chars > 0:
candidates.append(config.fallback_max_chars)
return min(candidates) if candidates else -1
def _tool_message_over_budget(msg: ToolMessage, config: ToolOutputConfig) -> bool:
"""Cheap, per-tool-aware check: is this ToolMessage non-exempt and over its trigger?"""
if (msg.name or "") in config.exempt_tools:
return False
trigger = _effective_trigger(msg.name or "", config)
if trigger < 0:
return False
text = _message_text(msg.content)
return text is not None and len(text) > trigger
def _needs_budget(result: ToolMessage | Command, config: ToolOutputConfig) -> bool:
"""Fast check whether *result* could need budgeting (avoids thread offload for small outputs)."""
if isinstance(result, ToolMessage):
return _tool_message_over_budget(result, config)
update = getattr(result, "update", None)
if isinstance(update, dict):
for msg in update.get("messages", []):
if isinstance(msg, ToolMessage) and _tool_message_over_budget(msg, config):
return True
return False
def _patch_result(result: ToolMessage | Command, config: ToolOutputConfig, outputs_path: str | None) -> ToolMessage | Command:
"""Apply budget to a tool call result (ToolMessage or Command)."""
if isinstance(result, ToolMessage):
return _patch_tool_message(result, config, outputs_path)
update = getattr(result, "update", None)
if not isinstance(update, dict):
return result
messages = update.get("messages")
if not isinstance(messages, list):
return result
new_messages: list[Any] = []
changed = False
for msg in messages:
if isinstance(msg, ToolMessage):
patched = _patch_tool_message(msg, config, outputs_path)
if patched is not msg:
changed = True
new_messages.append(patched)
else:
new_messages.append(msg)
if not changed:
return result
return dc_replace(result, update={**update, "messages": new_messages})
def _patch_model_messages(messages: list[Any], config: ToolOutputConfig) -> list[Any] | None:
"""Apply budget to historical ToolMessages in a model request. Returns ``None`` if unchanged.
A cheap pre-scan bails out before allocating a new list when no historical
ToolMessage exceeds the budget — the common case once every result has
already been budgeted at tool-call time, so a long history is not rebuilt
on every model call.
"""
if not any(isinstance(msg, ToolMessage) and _tool_message_over_budget(msg, config) for msg in messages):
return None
updated: list[Any] = []
changed = False
for msg in messages:
if isinstance(msg, ToolMessage):
patched = _patch_tool_message(msg, config, outputs_path=None)
if patched is not msg:
changed = True
updated.append(patched)
else:
updated.append(msg)
return updated if changed else None
# ---------------------------------------------------------------------------
# Middleware class
# ---------------------------------------------------------------------------
class ToolOutputBudgetMiddleware(AgentMiddleware[AgentState]):
"""Enforce per-result budget on tool outputs via externalization or truncation."""
def __init__(self, config: ToolOutputConfig | None = None) -> None:
super().__init__()
self._config = config if config is not None else _default_config()
@classmethod
def from_app_config(cls, app_config: Any) -> ToolOutputBudgetMiddleware:
tool_output = getattr(app_config, "tool_output", None)
if isinstance(tool_output, ToolOutputConfig):
return cls(config=tool_output)
return cls()
# -- tool call hooks ---------------------------------------------------
@override
def wrap_tool_call(
self,
request: ToolCallRequest,
handler: Callable[[ToolCallRequest], ToolMessage | Command],
) -> ToolMessage | Command:
result = handler(request)
if not self._config.enabled:
return result
if not _needs_budget(result, self._config):
return result
outputs_path = _resolve_outputs_path(request)
return _patch_result(result, self._config, outputs_path)
@override
async def awrap_tool_call(
self,
request: ToolCallRequest,
handler: Callable[[ToolCallRequest], Awaitable[ToolMessage | Command]],
) -> ToolMessage | Command:
result = await handler(request)
if not self._config.enabled:
return result
if not _needs_budget(result, self._config):
return result
outputs_path = _resolve_outputs_path(request)
return await asyncio.to_thread(_patch_result, result, self._config, outputs_path)
# -- model call hooks (historical message truncation) ------------------
@override
def wrap_model_call(
self,
request: ModelRequest,
handler: Callable[[ModelRequest], ModelResponse],
) -> ModelCallResult:
if self._config.enabled:
messages = getattr(request, "messages", None)
if isinstance(messages, list):
patched = _patch_model_messages(messages, self._config)
if patched is not None:
request = request.override(messages=patched)
return handler(request)
@override
async def awrap_model_call(
self,
request: ModelRequest,
handler: Callable[[ModelRequest], Awaitable[ModelResponse]],
) -> ModelCallResult:
if self._config.enabled:
messages = getattr(request, "messages", None)
if isinstance(messages, list):
patched = _patch_model_messages(messages, self._config)
if patched is not None:
request = request.override(messages=patched)
return await handler(request)
backend/packages/harness/deerflow/config/app_config.py
+2
View File
@@ -30,6 +30,7 @@ from deerflow.config.summarization_config import SummarizationConfig, load_summa
from deerflow.config.title_config import TitleConfig, load_title_config_from_dict
from deerflow.config.token_usage_config import TokenUsageConfig
from deerflow.config.tool_config import ToolConfig, ToolGroupConfig
from deerflow.config.tool_output_config import ToolOutputConfig
from deerflow.config.tool_search_config import ToolSearchConfig, load_tool_search_config_from_dict
load_dotenv()
@@ -93,6 +94,7 @@ class AppConfig(BaseModel):
skills: SkillsConfig = Field(default_factory=SkillsConfig, description="Skills configuration")
skill_evolution: SkillEvolutionConfig = Field(default_factory=SkillEvolutionConfig, description="Agent-managed skill evolution configuration")
extensions: ExtensionsConfig = Field(default_factory=ExtensionsConfig, description="Extensions configuration (MCP servers and skills state)")
tool_output: ToolOutputConfig = Field(default_factory=ToolOutputConfig, description="Tool output budget protection configuration")
tool_search: ToolSearchConfig = Field(default_factory=ToolSearchConfig, description="Tool search / deferred loading configuration")
title: TitleConfig = Field(default_factory=TitleConfig, description="Automatic title generation configuration")
summarization: SummarizationConfig = Field(default_factory=SummarizationConfig, description="Conversation summarization configuration")
backend/packages/harness/deerflow/config/tool_output_config.py
+62
View File
@@ -0,0 +1,62 @@
"""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.",
)