wehub-resource-sync
109 lines
3.4 KiB
Python
109 lines
3.4 KiB
Python
"""Tests for task exception metadata produced by ToolErrorHandlingMiddleware."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
|
|
from langchain_core.messages import ToolMessage
|
|
from langgraph.types import Command
|
|
|
|
from deerflow.agents.middlewares.tool_error_handling_middleware import (
|
|
ToolErrorHandlingMiddleware,
|
|
)
|
|
from deerflow.subagents.status_contract import (
|
|
SUBAGENT_ERROR_KEY,
|
|
SUBAGENT_STATUS_KEY,
|
|
)
|
|
|
|
|
|
class _FakeRequest:
|
|
"""Stand-in for ``ToolCallRequest`` used by the middleware."""
|
|
|
|
def __init__(self, tool_name: str, tool_call_id: str = "call-1") -> None:
|
|
self.tool_call = {"name": tool_name, "id": tool_call_id}
|
|
|
|
|
|
def test_task_tool_exception_returns_failed_metadata():
|
|
middleware = ToolErrorHandlingMiddleware()
|
|
request = _FakeRequest("task")
|
|
|
|
def handler(_req):
|
|
raise RuntimeError("blew up during execution")
|
|
|
|
result = middleware.wrap_tool_call(request, handler)
|
|
|
|
assert isinstance(result, ToolMessage)
|
|
assert result.additional_kwargs.get(SUBAGENT_STATUS_KEY) == "failed"
|
|
assert "RuntimeError" in result.additional_kwargs.get(SUBAGENT_ERROR_KEY, "")
|
|
|
|
|
|
def test_async_task_tool_exception_returns_failed_metadata():
|
|
middleware = ToolErrorHandlingMiddleware()
|
|
request = _FakeRequest("task")
|
|
|
|
async def handler(_req):
|
|
raise RuntimeError("async boom")
|
|
|
|
result = asyncio.run(middleware.awrap_tool_call(request, handler))
|
|
|
|
assert isinstance(result, ToolMessage)
|
|
assert result.additional_kwargs.get(SUBAGENT_STATUS_KEY) == "failed"
|
|
assert "RuntimeError" in result.additional_kwargs.get(SUBAGENT_ERROR_KEY, "")
|
|
|
|
|
|
def test_successful_plain_task_tool_message_is_not_stamped_from_content():
|
|
middleware = ToolErrorHandlingMiddleware()
|
|
request = _FakeRequest("task")
|
|
|
|
def handler(_req):
|
|
return ToolMessage(content="Task Succeeded. Result: ok", tool_call_id="call-1", name="task")
|
|
|
|
result = middleware.wrap_tool_call(request, handler)
|
|
|
|
assert isinstance(result, ToolMessage)
|
|
assert SUBAGENT_STATUS_KEY not in (result.additional_kwargs or {})
|
|
|
|
|
|
def test_does_not_stamp_non_task_tool_exception():
|
|
middleware = ToolErrorHandlingMiddleware()
|
|
request = _FakeRequest("bash")
|
|
|
|
def handler(_req):
|
|
raise RuntimeError("command failed")
|
|
|
|
result = middleware.wrap_tool_call(request, handler)
|
|
|
|
assert isinstance(result, ToolMessage)
|
|
assert SUBAGENT_STATUS_KEY not in (result.additional_kwargs or {})
|
|
|
|
|
|
def test_task_command_with_metadata_bypasses_middleware_stamp():
|
|
middleware = ToolErrorHandlingMiddleware()
|
|
request = _FakeRequest("task")
|
|
command = Command(
|
|
update={
|
|
"messages": [
|
|
ToolMessage(
|
|
"Task Succeeded. Result: ok",
|
|
tool_call_id="call-1",
|
|
name="task",
|
|
additional_kwargs={"subagent_status": "completed", "subagent_result_brief": "ok"},
|
|
)
|
|
]
|
|
}
|
|
)
|
|
|
|
assert middleware.wrap_tool_call(request, lambda _req: command) is command
|
|
|
|
|
|
def test_additional_kwargs_round_trip_via_json():
|
|
msg = ToolMessage(
|
|
content="Task Succeeded. Result: ok",
|
|
tool_call_id="call-1",
|
|
name="task",
|
|
additional_kwargs={SUBAGENT_STATUS_KEY: "completed", SUBAGENT_ERROR_KEY: ""},
|
|
)
|
|
serialised = msg.model_dump_json()
|
|
restored = ToolMessage.model_validate_json(serialised)
|
|
assert restored.additional_kwargs.get(SUBAGENT_STATUS_KEY) == "completed"
|