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

2026-05-26 18:29:26 +00:00
parent 8e065266ea
commit 026825ce02
2 changed files with 399 additions and 0 deletions
@@ -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 == []