feat: AG-UI adapter translating amplifierd SSE events to AG-UI protocol (with MCP App metadata passthrough)

backend/agui_adapter.py

@@ -0,0 +1,179 @@
+"""AG-UI adapter: translates amplifierd SSE events to AG-UI protocol events.
+
+Handles content_block:start/delta/end, tool:result, and orchestrator:complete
+events from amplifierd and converts them into the AG-UI event types that
+CopilotKit expects. Passes through MCP App metadata (_meta and structuredContent)
+from tool results for iframe rendering.
+"""
+
+import uuid
+from typing import Any
+
+from ag_ui.core import (
+    RunFinishedEvent,
+    TextMessageContentEvent,
+    TextMessageEndEvent,
+    TextMessageStartEvent,
+    ToolCallArgsEvent,
+    ToolCallEndEvent,
+    ToolCallResultEvent,
+    ToolCallStartEvent,
+)
+
+
+class AmplifierdEventTranslator:
+    """Translates amplifierd SSE events to AG-UI protocol events.
+
+    Args:
+        run_id: The AG-UI run ID for this conversation run.
+        thread_id: The AG-UI thread ID for the current thread.
+    """
+
+    def __init__(self, run_id: str, thread_id: str) -> None:
+        self._run_id = run_id
+        self._thread_id = thread_id
+        # Tracks active blocks by index: {index: {type, id?, name?, message_id?}}
+        self._blocks: dict[int, dict[str, Any]] = {}
+        self._current_message_id: str | None = None
+
+    def translate(self, event: dict[str, Any]) -> list[Any]:
+        """Translate a single amplifierd SSE event into zero or more AG-UI events.
+
+        Args:
+            event: Dict with 'event_type' and 'data' keys.
+
+        Returns:
+            List of ag_ui.core event objects (may be empty for unknown events).
+        """
+        event_type = event.get("event_type", "")
+        data = event.get("data", {})
+
+        if event_type == "content_block:start":
+            return self._handle_block_start(data)
+        elif event_type == "content_block:delta":
+            return self._handle_block_delta(data)
+        elif event_type == "content_block:end":
+            return self._handle_block_end(data)
+        elif event_type == "tool:result":
+            return self._handle_tool_result(data)
+        elif event_type == "orchestrator:complete":
+            return self._handle_complete()
+        else:
+            return []
+
+    def _handle_block_start(self, data: dict[str, Any]) -> list[Any]:
+        """Handle content_block:start — emit TEXT_MESSAGE_START or TOOL_CALL_START."""
+        index = data.get("index", 0)
+        block_type = data.get("type", "text")
+
+        if block_type == "tool_use":
+            tool_id = data.get("id", "")
+            tool_name = data.get("name", "")
+            self._blocks[index] = {
+                "type": "tool_use",
+                "tool_id": tool_id,
+                "tool_name": tool_name,
+            }
+            return [
+                ToolCallStartEvent(
+                    tool_call_id=tool_id,
+                    tool_call_name=tool_name,
+                    parent_message_id=self._current_message_id,
+                )
+            ]
+        else:
+            # Text block
+            message_id = str(uuid.uuid4())
+            self._blocks[index] = {"type": "text", "message_id": message_id}
+            self._current_message_id = message_id
+            return [
+                TextMessageStartEvent(
+                    message_id=message_id,
+                    role="assistant",
+                )
+            ]
+
+    def _handle_block_delta(self, data: dict[str, Any]) -> list[Any]:
+        """Handle content_block:delta — emit TEXT_MESSAGE_CONTENT or TOOL_CALL_ARGS."""
+        index = data.get("index", 0)
+        delta = data.get("delta", {})
+        delta_type = delta.get("type", "")
+
+        block_info = self._blocks.get(index)
+
+        if block_info is None:
+            # Graceful handling: no matching start — treat as tool args if input_json_delta
+            if delta_type == "input_json_delta":
+                return [
+                    ToolCallArgsEvent(
+                        tool_call_id="",
+                        delta=delta.get("partial_json", ""),
+                    )
+                ]
+            return []
+
+        if block_info["type"] == "tool_use":
+            return [
+                ToolCallArgsEvent(
+                    tool_call_id=block_info["tool_id"],
+                    delta=delta.get("partial_json", ""),
+                )
+            ]
+        else:
+            # Text block
+            token = delta.get("text", "")
+            if not token:
+                return []
+            return [
+                TextMessageContentEvent(
+                    message_id=block_info["message_id"],
+                    delta=token,
+                )
+            ]
+
+    def _handle_block_end(self, data: dict[str, Any]) -> list[Any]:
+        """Handle content_block:end — emit TEXT_MESSAGE_END or TOOL_CALL_END."""
+        index = data.get("index", 0)
+        block_info = self._blocks.pop(index, None)
+
+        if block_info is None:
+            return []
+
+        if block_info["type"] == "tool_use":
+            return [ToolCallEndEvent(tool_call_id=block_info["tool_id"])]
+        else:
+            return [TextMessageEndEvent(message_id=block_info["message_id"])]
+
+    def _handle_tool_result(self, data: dict[str, Any]) -> list[Any]:
+        """Handle tool:result — emit TOOL_CALL_RESULT with optional MCP App metadata."""
+        tool_call_id = data.get("tool_use_id", "")
+        content = data.get("content", "")
+        message_id = str(uuid.uuid4())
+
+        event = ToolCallResultEvent(
+            message_id=message_id,
+            tool_call_id=tool_call_id,
+            content=content if isinstance(content, str) else str(content),
+            role="tool",
+        )
+
+        # Conditionally pass through MCP App metadata for iframe rendering
+        extra: dict[str, Any] = {}
+        if "_meta" in data:
+            extra["_meta"] = data["_meta"]
+        if "structuredContent" in data:
+            extra["structuredContent"] = data["structuredContent"]
+
+        if extra:
+            event.__pydantic_extra__ = extra
+
+        return [event]
+
+    def _handle_complete(self) -> list[Any]:
+        """Handle orchestrator:complete — emit RUN_FINISHED."""
+        return [
+            RunFinishedEvent(
+                thread_id=self._thread_id,
+                run_id=self._run_id,
+            )
+        ]

tests/test_agui_adapter.py

@@ -0,0 +1,220 @@
+"""Tests for agui_adapter module — AmplifierdEventTranslator."""
+
+import sys
+import os
+
+# Add backend to sys.path so we can import agui_adapter
+sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "backend"))
+
+from agui_adapter import AmplifierdEventTranslator
+
+
+class TestTextBlocks:
+    """Tests for text content_block events → AG-UI text message events."""
+
+    def setup_method(self):
+        self.translator = AmplifierdEventTranslator(
+            run_id="run_1", thread_id="thread_1"
+        )
+
+    def test_text_block_start(self):
+        """text_block_start → TEXT_MESSAGE_START with role='assistant' and messageId."""
+        event = {
+            "event_type": "content_block:start",
+            "data": {"index": 0, "type": "text"},
+        }
+        result = self.translator.translate(event)
+        assert len(result) == 1
+        ev = result[0]
+        assert ev.type.value == "TEXT_MESSAGE_START"
+        assert ev.role == "assistant"
+        assert ev.message_id  # must have a non-empty message_id
+
+    def test_text_delta(self):
+        """text_delta → TEXT_MESSAGE_CONTENT with delta='Hello '."""
+        # Start the text block first
+        start_event = {
+            "event_type": "content_block:start",
+            "data": {"index": 0, "type": "text"},
+        }
+        self.translator.translate(start_event)
+
+        delta_event = {
+            "event_type": "content_block:delta",
+            "data": {
+                "index": 0,
+                "delta": {"type": "text_delta", "text": "Hello "},
+            },
+        }
+        result = self.translator.translate(delta_event)
+        assert len(result) == 1
+        ev = result[0]
+        assert ev.type.value == "TEXT_MESSAGE_CONTENT"
+        assert ev.delta == "Hello "
+
+    def test_text_block_end(self):
+        """text_block_end → TEXT_MESSAGE_END."""
+        # Start the block first
+        start_event = {
+            "event_type": "content_block:start",
+            "data": {"index": 0, "type": "text"},
+        }
+        self.translator.translate(start_event)
+
+        end_event = {
+            "event_type": "content_block:end",
+            "data": {"index": 0},
+        }
+        result = self.translator.translate(end_event)
+        assert len(result) == 1
+        ev = result[0]
+        assert ev.type.value == "TEXT_MESSAGE_END"
+
+
+class TestToolBlocks:
+    """Tests for tool_use content_block events and tool results → AG-UI tool call events."""
+
+    def setup_method(self):
+        self.translator = AmplifierdEventTranslator(
+            run_id="run_1", thread_id="thread_1"
+        )
+
+    def test_tool_use_start(self):
+        """tool_use_start → TOOL_CALL_START with toolCallName='bash', toolCallId='tool_123'."""
+        event = {
+            "event_type": "content_block:start",
+            "data": {
+                "index": 0,
+                "type": "tool_use",
+                "id": "tool_123",
+                "name": "bash",
+            },
+        }
+        result = self.translator.translate(event)
+        assert len(result) == 1
+        ev = result[0]
+        assert ev.type.value == "TOOL_CALL_START"
+        assert ev.tool_call_name == "bash"
+        assert ev.tool_call_id == "tool_123"
+
+    def test_tool_use_delta(self):
+        """tool_use_delta → TOOL_CALL_ARGS with delta='{"command": "ls"}'."""
+        # Start the tool block first
+        start_event = {
+            "event_type": "content_block:start",
+            "data": {
+                "index": 0,
+                "type": "tool_use",
+                "id": "tool_123",
+                "name": "bash",
+            },
+        }
+        self.translator.translate(start_event)
+
+        delta_event = {
+            "event_type": "content_block:delta",
+            "data": {
+                "index": 0,
+                "delta": {
+                    "type": "input_json_delta",
+                    "partial_json": '{"command": "ls"}',
+                },
+            },
+        }
+        result = self.translator.translate(delta_event)
+        assert len(result) == 1
+        ev = result[0]
+        assert ev.type.value == "TOOL_CALL_ARGS"
+        assert ev.delta == '{"command": "ls"}'
+
+    def test_tool_use_end(self):
+        """tool_use_end → TOOL_CALL_END."""
+        # Start first
+        start_event = {
+            "event_type": "content_block:start",
+            "data": {
+                "index": 0,
+                "type": "tool_use",
+                "id": "tool_123",
+                "name": "bash",
+            },
+        }
+        self.translator.translate(start_event)
+
+        end_event = {
+            "event_type": "content_block:end",
+            "data": {"index": 0},
+        }
+        result = self.translator.translate(end_event)
+        assert len(result) == 1
+        ev = result[0]
+        assert ev.type.value == "TOOL_CALL_END"
+
+    def test_tool_result(self):
+        """tool_result → TOOL_CALL_RESULT with toolCallId."""
+        event = {
+            "event_type": "tool:result",
+            "data": {"tool_use_id": "tool_123", "content": "output"},
+        }
+        result = self.translator.translate(event)
+        assert len(result) == 1
+        ev = result[0]
+        assert ev.type.value == "TOOL_CALL_RESULT"
+        assert ev.tool_call_id == "tool_123"
+
+    def test_tool_result_with_mcp_app_metadata(self):
+        """tool_result_with_mcp_app_metadata → TOOL_CALL_RESULT preserves _meta.ui.resourceUri='app://chart' and structuredContent.chartType='bar'."""
+        event = {
+            "event_type": "tool:result",
+            "data": {
+                "tool_use_id": "tool_123",
+                "content": "output",
+                "_meta": {"ui": {"resourceUri": "app://chart"}},
+                "structuredContent": {"chartType": "bar"},
+            },
+        }
+        result = self.translator.translate(event)
+        assert len(result) == 1
+        ev = result[0]
+        assert ev.type.value == "TOOL_CALL_RESULT"
+        d = ev.model_dump()
+        assert d.get("_meta", {}).get("ui", {}).get("resourceUri") == "app://chart"
+        assert d.get("structuredContent", {}).get("chartType") == "bar"
+
+    def test_tool_result_without_mcp_app_metadata(self):
+        """tool_result_without_mcp_app_metadata → no _meta or structuredContent keys."""
+        event = {
+            "event_type": "tool:result",
+            "data": {"tool_use_id": "tool_123", "content": "output"},
+        }
+        result = self.translator.translate(event)
+        assert len(result) == 1
+        ev = result[0]
+        d = ev.model_dump()
+        assert "_meta" not in d
+        assert "structuredContent" not in d
+
+
+class TestLifecycle:
+    """Tests for lifecycle events → AG-UI run lifecycle events."""
+
+    def setup_method(self):
+        self.translator = AmplifierdEventTranslator(
+            run_id="run_1", thread_id="thread_1"
+        )
+
+    def test_orchestrator_complete(self):
+        """orchestrator_complete → RUN_FINISHED with threadId and runId."""
+        event = {"event_type": "orchestrator:complete", "data": {}}
+        result = self.translator.translate(event)
+        assert len(result) == 1
+        ev = result[0]
+        assert ev.type.value == "RUN_FINISHED"
+        assert ev.thread_id == "thread_1"
+        assert ev.run_id == "run_1"
+
+    def test_unknown_event(self):
+        """unknown event → empty list."""
+        event = {"event_type": "some:unknown_event", "data": {}}
+        result = self.translator.translate(event)
+        assert result == []