mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-08-11 23:38:44 +00:00
e7b88a97ed
InputSanitizationMiddleware: escapes blocked XML tags in user messages (<system> -> <system>) — de-identify, don't reject (AWS PII-style). Wraps user input in plain-text boundary markers (--- BEGIN/END USER INPUT ---) per OWASP structured-prompt guidance. System-Context Confidentiality clause prevents LLM from revealing internal instructions. Boundary-marker layer hardened against delimiter injection (self-suppression + break-out). Multimodal content blocks preserved during sanitization. Addresses review feedback from @fancyboi999 and @WillemJiang: 1. Boundary-marker injection: strict startswith+endswith idempotency, neutralize user-supplied BEGIN/END tokens, forged-idempotency bypass fixed 2. Multimodal data loss: _rebuild_content preserves interleaved non-text blocks 3. CI fixes: replay_provider strips boundary markers before hashing, middleware chain assertions updated 4. Lint: ruff format applied on Linux (slice spacing + f-string collapse) Test: 71 input_sanitization + 1 replay_golden + 14 tool_error + 93 tool_output_budget = 179/179 passed
304 lines
12 KiB
Python
304 lines
12 KiB
Python
import sys
|
|
from types import ModuleType, SimpleNamespace
|
|
|
|
import pytest
|
|
from langchain_core.messages import ToolMessage
|
|
from langgraph.errors import GraphInterrupt
|
|
|
|
from deerflow.agents.middlewares.tool_error_handling_middleware import (
|
|
ToolErrorHandlingMiddleware,
|
|
build_subagent_runtime_middlewares,
|
|
)
|
|
from deerflow.agents.middlewares.view_image_middleware import ViewImageMiddleware
|
|
from deerflow.config.app_config import AppConfig, CircuitBreakerConfig
|
|
from deerflow.config.guardrails_config import GuardrailsConfig
|
|
from deerflow.config.model_config import ModelConfig
|
|
from deerflow.config.sandbox_config import SandboxConfig
|
|
|
|
|
|
def _request(name: str = "web_search", tool_call_id: str | None = "tc-1"):
|
|
tool_call = {"name": name}
|
|
if tool_call_id is not None:
|
|
tool_call["id"] = tool_call_id
|
|
return SimpleNamespace(tool_call=tool_call)
|
|
|
|
|
|
def _module(name: str, **attrs):
|
|
module = ModuleType(name)
|
|
for key, value in attrs.items():
|
|
setattr(module, key, value)
|
|
return module
|
|
|
|
|
|
def _make_app_config(*, supports_vision: bool = False) -> AppConfig:
|
|
return AppConfig(
|
|
models=[
|
|
ModelConfig(
|
|
name="test-model",
|
|
display_name="test-model",
|
|
description=None,
|
|
use="langchain_openai:ChatOpenAI",
|
|
model="test-model",
|
|
supports_vision=supports_vision,
|
|
)
|
|
],
|
|
sandbox=SandboxConfig(use="test"),
|
|
guardrails=GuardrailsConfig(enabled=False),
|
|
circuit_breaker=CircuitBreakerConfig(failure_threshold=7, recovery_timeout_sec=11),
|
|
)
|
|
|
|
|
|
def _stub_runtime_middleware_imports(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
class FakeMiddleware:
|
|
def __init__(self, *args, **kwargs):
|
|
self.args = args
|
|
self.kwargs = kwargs
|
|
|
|
class FakeLLMErrorHandlingMiddleware:
|
|
def __init__(self, *, app_config):
|
|
self.app_config = app_config
|
|
|
|
monkeypatch.setitem(
|
|
sys.modules,
|
|
"deerflow.agents.middlewares.llm_error_handling_middleware",
|
|
_module(
|
|
"deerflow.agents.middlewares.llm_error_handling_middleware",
|
|
LLMErrorHandlingMiddleware=FakeLLMErrorHandlingMiddleware,
|
|
),
|
|
)
|
|
monkeypatch.setitem(
|
|
sys.modules,
|
|
"deerflow.agents.middlewares.thread_data_middleware",
|
|
_module("deerflow.agents.middlewares.thread_data_middleware", ThreadDataMiddleware=FakeMiddleware),
|
|
)
|
|
monkeypatch.setitem(
|
|
sys.modules,
|
|
"deerflow.sandbox.middleware",
|
|
_module("deerflow.sandbox.middleware", SandboxMiddleware=FakeMiddleware),
|
|
)
|
|
monkeypatch.setitem(
|
|
sys.modules,
|
|
"deerflow.agents.middlewares.dangling_tool_call_middleware",
|
|
_module("deerflow.agents.middlewares.dangling_tool_call_middleware", DanglingToolCallMiddleware=FakeMiddleware),
|
|
)
|
|
monkeypatch.setitem(
|
|
sys.modules,
|
|
"deerflow.agents.middlewares.sandbox_audit_middleware",
|
|
_module("deerflow.agents.middlewares.sandbox_audit_middleware", SandboxAuditMiddleware=FakeMiddleware),
|
|
)
|
|
|
|
|
|
def test_build_subagent_runtime_middlewares_threads_app_config_to_llm_middleware(monkeypatch: pytest.MonkeyPatch):
|
|
captured: dict[str, object] = {}
|
|
|
|
class FakeMiddleware:
|
|
def __init__(self, *args, **kwargs):
|
|
self.args = args
|
|
self.kwargs = kwargs
|
|
|
|
class FakeLLMErrorHandlingMiddleware:
|
|
def __init__(self, *, app_config):
|
|
captured["app_config"] = app_config
|
|
|
|
app_config = _make_app_config()
|
|
|
|
monkeypatch.setitem(
|
|
sys.modules,
|
|
"deerflow.agents.middlewares.llm_error_handling_middleware",
|
|
_module(
|
|
"deerflow.agents.middlewares.llm_error_handling_middleware",
|
|
LLMErrorHandlingMiddleware=FakeLLMErrorHandlingMiddleware,
|
|
),
|
|
)
|
|
monkeypatch.setitem(
|
|
sys.modules,
|
|
"deerflow.agents.middlewares.thread_data_middleware",
|
|
_module("deerflow.agents.middlewares.thread_data_middleware", ThreadDataMiddleware=FakeMiddleware),
|
|
)
|
|
monkeypatch.setitem(
|
|
sys.modules,
|
|
"deerflow.sandbox.middleware",
|
|
_module("deerflow.sandbox.middleware", SandboxMiddleware=FakeMiddleware),
|
|
)
|
|
monkeypatch.setitem(
|
|
sys.modules,
|
|
"deerflow.agents.middlewares.dangling_tool_call_middleware",
|
|
_module("deerflow.agents.middlewares.dangling_tool_call_middleware", DanglingToolCallMiddleware=FakeMiddleware),
|
|
)
|
|
monkeypatch.setitem(
|
|
sys.modules,
|
|
"deerflow.agents.middlewares.sandbox_audit_middleware",
|
|
_module("deerflow.agents.middlewares.sandbox_audit_middleware", SandboxAuditMiddleware=FakeMiddleware),
|
|
)
|
|
monkeypatch.setitem(
|
|
sys.modules,
|
|
"deerflow.agents.middlewares.input_sanitization_middleware",
|
|
_module("deerflow.agents.middlewares.input_sanitization_middleware", InputSanitizationMiddleware=FakeMiddleware),
|
|
)
|
|
|
|
middlewares = build_subagent_runtime_middlewares(app_config=app_config, lazy_init=False)
|
|
|
|
assert captured["app_config"] is app_config
|
|
# 8 baseline (InputSanitization, ToolOutputBudget, ThreadData, Sandbox,
|
|
# DanglingToolCall, LLMErrorHandling, SandboxAudit, ToolErrorHandling)
|
|
# + 1 SafetyFinishReasonMiddleware (enabled by default).
|
|
from deerflow.agents.middlewares.safety_finish_reason_middleware import SafetyFinishReasonMiddleware
|
|
from deerflow.agents.middlewares.tool_output_budget_middleware import ToolOutputBudgetMiddleware
|
|
|
|
assert len(middlewares) == 9
|
|
assert isinstance(middlewares[0], FakeMiddleware) # InputSanitizationMiddleware stub
|
|
assert isinstance(middlewares[1], ToolOutputBudgetMiddleware)
|
|
assert any(isinstance(m, ToolErrorHandlingMiddleware) for m in middlewares)
|
|
assert isinstance(middlewares[-1], SafetyFinishReasonMiddleware)
|
|
|
|
|
|
def test_wrap_tool_call_passthrough_on_success():
|
|
middleware = ToolErrorHandlingMiddleware()
|
|
req = _request()
|
|
expected = ToolMessage(content="ok", tool_call_id="tc-1", name="web_search")
|
|
|
|
result = middleware.wrap_tool_call(req, lambda _req: expected)
|
|
|
|
assert result is expected
|
|
|
|
|
|
def test_wrap_tool_call_returns_error_tool_message_on_exception():
|
|
middleware = ToolErrorHandlingMiddleware()
|
|
req = _request(name="web_search", tool_call_id="tc-42")
|
|
|
|
def _boom(_req):
|
|
raise RuntimeError("network down")
|
|
|
|
result = middleware.wrap_tool_call(req, _boom)
|
|
|
|
assert isinstance(result, ToolMessage)
|
|
assert result.tool_call_id == "tc-42"
|
|
assert result.name == "web_search"
|
|
assert result.status == "error"
|
|
assert "Tool 'web_search' failed" in result.text
|
|
assert "network down" in result.text
|
|
|
|
|
|
def test_wrap_tool_call_uses_fallback_tool_call_id_when_missing():
|
|
middleware = ToolErrorHandlingMiddleware()
|
|
req = _request(name="mcp_tool", tool_call_id=None)
|
|
|
|
def _boom(_req):
|
|
raise ValueError("bad request")
|
|
|
|
result = middleware.wrap_tool_call(req, _boom)
|
|
|
|
assert isinstance(result, ToolMessage)
|
|
assert result.tool_call_id == "missing_tool_call_id"
|
|
assert result.name == "mcp_tool"
|
|
assert result.status == "error"
|
|
|
|
|
|
def test_wrap_tool_call_reraises_graph_interrupt():
|
|
middleware = ToolErrorHandlingMiddleware()
|
|
req = _request(name="ask_clarification", tool_call_id="tc-int")
|
|
|
|
def _interrupt(_req):
|
|
raise GraphInterrupt(())
|
|
|
|
with pytest.raises(GraphInterrupt):
|
|
middleware.wrap_tool_call(req, _interrupt)
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_awrap_tool_call_returns_error_tool_message_on_exception():
|
|
middleware = ToolErrorHandlingMiddleware()
|
|
req = _request(name="mcp_tool", tool_call_id="tc-async")
|
|
|
|
async def _boom(_req):
|
|
raise TimeoutError("request timed out")
|
|
|
|
result = await middleware.awrap_tool_call(req, _boom)
|
|
|
|
assert isinstance(result, ToolMessage)
|
|
assert result.tool_call_id == "tc-async"
|
|
assert result.name == "mcp_tool"
|
|
assert result.status == "error"
|
|
assert "request timed out" in result.text
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_awrap_tool_call_reraises_graph_interrupt():
|
|
middleware = ToolErrorHandlingMiddleware()
|
|
req = _request(name="ask_clarification", tool_call_id="tc-int-async")
|
|
|
|
async def _interrupt(_req):
|
|
raise GraphInterrupt(())
|
|
|
|
with pytest.raises(GraphInterrupt):
|
|
await middleware.awrap_tool_call(req, _interrupt)
|
|
|
|
|
|
def test_subagent_runtime_middlewares_include_view_image_for_vision_model(monkeypatch):
|
|
app_config = _make_app_config(supports_vision=True)
|
|
_stub_runtime_middleware_imports(monkeypatch)
|
|
|
|
middlewares = build_subagent_runtime_middlewares(app_config=app_config, model_name="test-model")
|
|
|
|
assert any(isinstance(middleware, ViewImageMiddleware) for middleware in middlewares)
|
|
|
|
|
|
def test_subagent_runtime_middlewares_include_view_image_for_default_vision_model(monkeypatch):
|
|
app_config = _make_app_config(supports_vision=True)
|
|
_stub_runtime_middleware_imports(monkeypatch)
|
|
|
|
middlewares = build_subagent_runtime_middlewares(app_config=app_config, model_name=None)
|
|
|
|
assert any(isinstance(middleware, ViewImageMiddleware) for middleware in middlewares)
|
|
|
|
|
|
def test_subagent_runtime_middlewares_skip_view_image_for_text_model(monkeypatch):
|
|
app_config = _make_app_config(supports_vision=False)
|
|
_stub_runtime_middleware_imports(monkeypatch)
|
|
|
|
middlewares = build_subagent_runtime_middlewares(app_config=app_config, model_name="test-model")
|
|
|
|
assert not any(isinstance(middleware, ViewImageMiddleware) for middleware in middlewares)
|
|
|
|
|
|
def test_subagent_runtime_middlewares_attach_deferred_filter_when_setup_has_names(monkeypatch):
|
|
"""A subagent built with deferred MCP tools gets DeferredToolFilterMiddleware, positioned before SafetyFinishReasonMiddleware (mirrors the lead ordering)."""
|
|
from langchain_core.tools import tool as as_tool
|
|
|
|
from deerflow.agents.middlewares.deferred_tool_filter_middleware import DeferredToolFilterMiddleware
|
|
from deerflow.agents.middlewares.safety_finish_reason_middleware import SafetyFinishReasonMiddleware
|
|
from deerflow.tools.builtins.tool_search import build_deferred_tool_setup
|
|
from deerflow.tools.mcp_metadata import tag_mcp_tool
|
|
|
|
app_config = _make_app_config()
|
|
_stub_runtime_middleware_imports(monkeypatch)
|
|
|
|
@as_tool
|
|
def mcp_thing(x: str) -> str:
|
|
"deferred mcp tool"
|
|
return x
|
|
|
|
setup = build_deferred_tool_setup([tag_mcp_tool(mcp_thing)], enabled=True)
|
|
assert setup.deferred_names # sanity: populated setup
|
|
|
|
middlewares = build_subagent_runtime_middlewares(app_config=app_config, deferred_setup=setup)
|
|
|
|
filters = [m for m in middlewares if isinstance(m, DeferredToolFilterMiddleware)]
|
|
assert len(filters) == 1
|
|
filter_idx = next(i for i, m in enumerate(middlewares) if isinstance(m, DeferredToolFilterMiddleware))
|
|
safety_idx = next(i for i, m in enumerate(middlewares) if isinstance(m, SafetyFinishReasonMiddleware))
|
|
assert filter_idx < safety_idx
|
|
|
|
|
|
def test_subagent_runtime_middlewares_skip_deferred_filter_without_names(monkeypatch):
|
|
"""No deferred setup (disabled / no MCP tool) -> no DeferredToolFilterMiddleware."""
|
|
from deerflow.agents.middlewares.deferred_tool_filter_middleware import DeferredToolFilterMiddleware
|
|
from deerflow.tools.builtins.tool_search import DeferredToolSetup
|
|
|
|
app_config = _make_app_config()
|
|
_stub_runtime_middleware_imports(monkeypatch)
|
|
|
|
for setup in (None, DeferredToolSetup(None, frozenset(), None)):
|
|
middlewares = build_subagent_runtime_middlewares(app_config=app_config, deferred_setup=setup)
|
|
assert not any(isinstance(m, DeferredToolFilterMiddleware) for m in middlewares)
|