diff --git a/src/agent_engine/engine/langgraph/tools/tool_result.py b/src/agent_engine/engine/langgraph/tools/tool_result.py new file mode 100644 index 00000000..1f31933a --- /dev/null +++ b/src/agent_engine/engine/langgraph/tools/tool_result.py @@ -0,0 +1,103 @@ +"""Provider-agnostic normalized representation for tool execution outcomes.""" + +from __future__ import annotations + +import json +import logging +from dataclasses import asdict, dataclass, is_dataclass +from typing import Any + +from langchain_core.messages import ToolMessage + +logger = logging.getLogger(__name__) + + +def _as_text(content: Any) -> str: + if isinstance(content, str): + return content + if isinstance(content, list): + parts: list[str] = [] + for item in content: + if isinstance(item, str): + parts.append(item) + elif isinstance(item, dict) and "text" in item: + parts.append(str(item["text"])) + return "".join(parts) + return str(content) if content is not None else "" + + +@dataclass(frozen=True) +class NormalizedToolResult: + """Provider-agnostic representation of a tool execution outcome. + + Separates model-facing text content from machine-readable structured output + and optional artifact metadata, preserving MCP structuredContent without + collapsing everything to a plain string. + """ + + text: str + structured: Any | None = None + artifact: Any | None = None + + +def normalize_tool_result(raw_result: Any) -> NormalizedToolResult: + """Normalize a raw tool execution return value into a NormalizedToolResult. + + Supports: + - String output (local tools) + - LangChain ToolMessage (carrying MCP content blocks and artifact.structuredContent) + - Dicts with content / structuredContent + - Local Python objects / dicts / Pydantic models + """ + if isinstance(raw_result, NormalizedToolResult): + return raw_result + + if isinstance(raw_result, str): + return NormalizedToolResult(text=raw_result) + + text = "" + structured: Any | None = None + artifact: Any | None = None + + try: + if isinstance(raw_result, ToolMessage): + text = _as_text(raw_result.content) + artifact = getattr(raw_result, "artifact", None) + if artifact is not None: + if isinstance(artifact, dict) and "structuredContent" in artifact: + structured = artifact["structuredContent"] + elif hasattr(artifact, "structuredContent"): + structured = artifact.structuredContent + elif isinstance(raw_result, dict): + artifact = raw_result.get("artifact") + if "content" in raw_result: + text = _as_text(raw_result["content"]) + if "structuredContent" in raw_result: + structured = raw_result["structuredContent"] + elif isinstance(artifact, dict) and "structuredContent" in artifact: + structured = artifact["structuredContent"] + elif "content" not in raw_result and "artifact" not in raw_result: + structured = raw_result + elif is_dataclass(raw_result) and not isinstance(raw_result, type): + structured = asdict(raw_result) + elif hasattr(raw_result, "model_dump") and callable(raw_result.model_dump): + structured = raw_result.model_dump() + elif hasattr(raw_result, "dict") and callable(raw_result.dict): + structured = raw_result.dict() + elif isinstance(raw_result, (bytes, bytearray)): + text = raw_result.decode("utf-8", errors="replace") + else: + text = "" + except Exception as exc: + logger.warning("Failed to extract structured content from tool result: %s", exc) + text = "" + structured = None + artifact = None + + if not text.strip() and structured is not None: + try: + text = json.dumps(structured, default=str) + except Exception: + text = "" + + return NormalizedToolResult(text=text, structured=structured, artifact=artifact) diff --git a/src/agent_engine/runtime/hooks/models.py b/src/agent_engine/runtime/hooks/models.py index 50cd5e30..014bb0d1 100644 --- a/src/agent_engine/runtime/hooks/models.py +++ b/src/agent_engine/runtime/hooks/models.py @@ -230,6 +230,9 @@ class ToolCallContext: metadata: dict[str, object] = field(default_factory=dict) +_UNSET = object() + + @dataclass(frozen=True) class ToolResultContext: """A successful tool call's result, passed to ``transform_tool_result`` hooks @@ -250,15 +253,26 @@ class ToolResultContext: tool_name: str provider: ToolProvider result: str + structured_result: JsonValue | None = None + artifact: JsonValue | None = None server_id: str | None = None latency_ms: int | None = None metadata: dict[str, object] = field(default_factory=dict) - structured_result: JsonValue | None = None - artifact: JsonValue | None = None - def with_result(self, result: str) -> ToolResultContext: - """Replace text while preserving structured output and artifacts.""" - return dataclasses.replace(self, result=result) + @property + def structured(self) -> JsonValue | None: + return self.structured_result + + def with_result( + self, + result: str, + structured: Any | None = _UNSET, + artifact: Any | None = _UNSET, + ) -> ToolResultContext: + """Return a copy with ``result`` replaced (immutable update).""" + st = self.structured_result if structured is _UNSET else structured + art = self.artifact if artifact is _UNSET else artifact + return dataclasses.replace(self, result=result, structured_result=st, artifact=art) def with_structured_result(self, structured_result: JsonValue | None) -> ToolResultContext: """Return a copy with the machine-readable result replaced.""" diff --git a/tests/engine/test_tool_result_normalization.py b/tests/engine/test_tool_result_normalization.py new file mode 100644 index 00000000..3f6027b9 --- /dev/null +++ b/tests/engine/test_tool_result_normalization.py @@ -0,0 +1,171 @@ +"""Tests for tool result normalization and MCP structuredContent support.""" + +from __future__ import annotations + +from dataclasses import dataclass + +import pytest +from langchain_core.messages import ToolMessage + +from agent_engine.approvals.in_memory_tool_execution_repository import ( + InMemoryToolExecutionRepository, +) +from agent_engine.approvals.models import ToolExecutionStatus +from agent_engine.approvals.tool_execution_manager import ToolExecutionManager +from agent_engine.engine.langgraph.tools.tool_result import ( + normalize_tool_result, +) +from agent_engine.runtime.hooks.models import ToolResultContext +from agent_engine.runtime.tool_results import NormalizedToolResult + + +def test_normalize_plain_string_tool() -> None: + res = normalize_tool_result("hello world") + assert res.text == "hello world" + assert res.structured is None + assert res.artifact is None + + +def test_normalize_mcp_text_and_structured_content() -> None: + msg = ToolMessage( + content=[{"type": "text", "text": "Found 2 invoices"}], + tool_call_id="call_123", + artifact={ + "structuredContent": { + "invoices": [ + {"id": "INV-123", "amount": 500}, + {"id": "INV-456", "amount": 800}, + ] + } + }, + ) + res = normalize_tool_result(msg) + assert res.text == "Found 2 invoices" + assert res.structured == { + "invoices": [ + {"id": "INV-123", "amount": 500}, + {"id": "INV-456", "amount": 800}, + ] + } + assert res.artifact == { + "structuredContent": { + "invoices": [ + {"id": "INV-123", "amount": 500}, + {"id": "INV-456", "amount": 800}, + ] + } + } + + +def test_normalize_multiple_text_blocks_and_structured_content() -> None: + msg = ToolMessage( + content=[ + {"type": "text", "text": "Header line.\n"}, + {"type": "text", "text": "Detail line."}, + ], + tool_call_id="call_456", + artifact={"structuredContent": {"count": 42}}, + ) + res = normalize_tool_result(msg) + assert res.text == "Header line.\nDetail line." + assert res.structured == {"count": 42} + + +def test_normalize_structured_only_result() -> None: + data = {"structuredContent": {"balance": 1250}} + res = normalize_tool_result(data) + assert res.structured == {"balance": 1250} + assert "1250" in res.text + + +def test_normalize_local_structured_dict() -> None: + local_data = {"status": "ok", "items": [1, 2, 3]} + res = normalize_tool_result(local_data) + assert res.structured == {"status": "ok", "items": [1, 2, 3]} + assert '{"status": "ok", "items": [1, 2, 3]}' in res.text + + +@dataclass +class CustomToolOutput: + status: str + count: int + + +def test_normalize_local_dataclass() -> None: + output = CustomToolOutput(status="success", count=5) + res = normalize_tool_result(output) + assert res.structured == {"status": "success", "count": 5} + assert "success" in res.text + + +@pytest.mark.asyncio +async def test_execution_manager_persists_and_restores_structured_and_artifact() -> None: + repo = InMemoryToolExecutionRepository() + manager = ToolExecutionManager(execution_repository=repo) + + exec_id = "exec_test_123" + await manager.begin_execution( + exec_id, tool_call_id="tc_1", run_id="run_1", tool_name="test_tool" + ) + result = NormalizedToolResult( + text="Found items", + structured={"items": ["a", "b"]}, + artifact={"meta": "v1"}, + ) + await manager.finish_execution( + exec_id, + status=ToolExecutionStatus.SUCCEEDED, + result=result, + ) + + restored = await manager.restored_result(exec_id) + assert restored is not None + assert restored.text == "Found items" + assert restored.structured == {"items": ["a", "b"]} + assert restored.artifact == {"meta": "v1"} + + +def test_normalize_structured_only_tool_message() -> None: + msg = ToolMessage( + content=[], + tool_call_id="call_structured_only", + artifact={"structuredContent": {"balance": 1250}}, + ) + res = normalize_tool_result(msg) + assert res.structured == {"balance": 1250} + assert res.text == '{"balance": 1250}' + + +def test_normalize_auxiliary_artifact_not_conflated_with_structured() -> None: + msg = ToolMessage( + content="Generated report PDF", + tool_call_id="call_aux_artifact", + artifact={"file_id": "abc_123", "mime_type": "application/pdf"}, + ) + res = normalize_tool_result(msg) + assert res.text == "Generated report PDF" + assert res.structured is None + assert res.artifact == {"file_id": "abc_123", "mime_type": "application/pdf"} + + +def test_tool_result_context_with_result_preserves_and_clears_structured() -> None: + ctx = ToolResultContext( + agent_id="a1", + tool_name="t1", + provider="mcp", + result="Long raw output text", + structured_result={"data": [1, 2, 3]}, + artifact={"meta": "v1"}, + ) + + # Truncate text while preserving structured and artifact + updated = ctx.with_result("Truncated text") + assert updated.result == "Truncated text" + assert updated.structured == {"data": [1, 2, 3]} + assert updated.artifact == {"meta": "v1"} + + # Explicitly clear structured and artifact + cleared = ctx.with_result("Cleared text", structured=None, artifact=None) + assert cleared.result == "Cleared text" + assert cleared.structured is None + assert cleared.artifact is None