129 lines
4.9 KiB
Python
129 lines
4.9 KiB
Python
"""Pausing audit writes must preserve tool results and runtime error handling."""
|
|
|
|
from functools import partial
|
|
from types import SimpleNamespace
|
|
from unittest.mock import Mock
|
|
|
|
import pytest
|
|
from langchain_core.messages import ToolMessage
|
|
from langgraph.errors import GraphBubbleUp
|
|
|
|
from deerflow.agents.middlewares import tool_metrics_middleware as metrics
|
|
from deerflow.agents.middlewares.llm_metrics_middleware import LlmMetricsMiddleware
|
|
from deerflow.agents.middlewares.tool_error_handling_middleware import (
|
|
ToolErrorHandlingMiddleware,
|
|
build_lead_runtime_middlewares,
|
|
build_subagent_runtime_middlewares,
|
|
)
|
|
from deerflow.config.app_config import AppConfig
|
|
|
|
BUILDERS = [build_lead_runtime_middlewares, build_subagent_runtime_middlewares]
|
|
|
|
|
|
def _config(enabled=None):
|
|
options = {} if enabled is None else {"tool_metrics_enabled": enabled}
|
|
return AppConfig(
|
|
sandbox={"use": "deerflow.sandbox.local:LocalSandboxProvider"},
|
|
runtime_performance=options,
|
|
)
|
|
|
|
|
|
@pytest.mark.parametrize("builder", BUILDERS)
|
|
@pytest.mark.parametrize("enabled", [None, False])
|
|
def test_paused_runtime_never_constructs_a_recorder(builder, enabled, monkeypatch):
|
|
constructor = Mock(side_effect=AssertionError("paused runtime constructed a DB recorder"))
|
|
monkeypatch.setattr(metrics, "ToolMetricsMiddleware", constructor)
|
|
|
|
middlewares = builder(app_config=_config(enabled))
|
|
|
|
constructor.assert_not_called()
|
|
assert any(isinstance(m, ToolErrorHandlingMiddleware) for m in middlewares)
|
|
assert any(isinstance(m, LlmMetricsMiddleware) for m in middlewares)
|
|
|
|
|
|
def _tool_wrappers(builder, enabled):
|
|
# Exercise the actual error/audit wrappers selected by both runtime builders.
|
|
return [m for m in builder(app_config=_config(enabled)) if isinstance(m, (ToolErrorHandlingMiddleware, metrics.ToolMetricsMiddleware))]
|
|
|
|
|
|
@pytest.mark.parametrize("builder", BUILDERS)
|
|
@pytest.mark.parametrize("enabled", [False, True])
|
|
@pytest.mark.parametrize("fails", [False, True])
|
|
def test_sync_tool_execution_and_error_conversion_survive_pause(builder, enabled, fails, monkeypatch):
|
|
flushed = Mock()
|
|
monkeypatch.setattr(metrics.ToolMetricsMiddleware, "_flush", lambda _, record: flushed(record))
|
|
request = SimpleNamespace(tool_call={"id": "call-1", "name": "bash", "args": {}})
|
|
expected = ToolMessage(content="actual tool output", tool_call_id="call-1", name="bash")
|
|
|
|
def tool(req):
|
|
assert req is request
|
|
if fails:
|
|
raise ValueError("tool failed")
|
|
return expected
|
|
|
|
handler = tool
|
|
wrappers = _tool_wrappers(builder, enabled)
|
|
if enabled:
|
|
assert isinstance(wrappers[-1], metrics.ToolMetricsMiddleware)
|
|
for middleware in reversed(wrappers):
|
|
handler = partial(middleware.wrap_tool_call, handler=handler)
|
|
result = handler(request)
|
|
|
|
assert result.status == ("error" if fails else "success")
|
|
if not fails:
|
|
assert result is expected
|
|
assert flushed.call_count == int(enabled)
|
|
if enabled:
|
|
assert flushed.call_args.args[0].status == result.status
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize("builder", BUILDERS)
|
|
@pytest.mark.parametrize("enabled", [False, True])
|
|
@pytest.mark.parametrize("fails", [False, True])
|
|
async def test_async_tool_execution_and_error_conversion_survive_pause(builder, enabled, fails, monkeypatch):
|
|
flushed = Mock()
|
|
monkeypatch.setattr(metrics.ToolMetricsMiddleware, "_flush", lambda _, record: flushed(record))
|
|
request = SimpleNamespace(tool_call={"id": "call-1", "name": "bash", "args": {}})
|
|
expected = ToolMessage(content="actual tool output", tool_call_id="call-1", name="bash")
|
|
|
|
async def tool(req):
|
|
assert req is request
|
|
if fails:
|
|
raise ValueError("tool failed")
|
|
return expected
|
|
|
|
handler = tool
|
|
for middleware in reversed(_tool_wrappers(builder, enabled)):
|
|
handler = partial(middleware.awrap_tool_call, handler=handler)
|
|
result = await handler(request)
|
|
|
|
assert result.status == ("error" if fails else "success")
|
|
if not fails:
|
|
assert result is expected
|
|
assert flushed.call_count == int(enabled)
|
|
if enabled:
|
|
assert flushed.call_args.args[0].status == result.status
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize("builder", BUILDERS)
|
|
@pytest.mark.parametrize("enabled", [False, True])
|
|
async def test_pause_does_not_intercept_graph_control_flow(builder, enabled, monkeypatch):
|
|
flushed = Mock()
|
|
monkeypatch.setattr(metrics.ToolMetricsMiddleware, "_flush", lambda _, record: flushed(record))
|
|
request = SimpleNamespace(tool_call={"id": "call-1", "name": "bash", "args": {}})
|
|
signal = GraphBubbleUp("pause")
|
|
|
|
async def tool(_):
|
|
raise signal
|
|
|
|
handler = tool
|
|
for middleware in reversed(_tool_wrappers(builder, enabled)):
|
|
handler = partial(middleware.awrap_tool_call, handler=handler)
|
|
|
|
with pytest.raises(GraphBubbleUp) as exc:
|
|
await handler(request)
|
|
assert exc.value is signal
|
|
flushed.assert_not_called()
|