"""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. The 'data' value is the raw SSE payload from amplifierd, which wraps the event-specific fields under a nested ``"data"`` key (e.g. ``{"event": "...", "data": {...}, "session_id": ..., ...}``). Returns: List of ag_ui.core event objects (may be empty for unknown events). """ event_type = event.get("event_type", "") outer = event.get("data", {}) # amplifierd wraps event-specific fields in a nested "data" key. # Fall back to the outer dict for tests that pass data directly. data = outer.get("data", outer) 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 — register block; defer TEXT_MESSAGE_START to first delta. Supports both amplifierd wire format (``block_index`` / ``block_type``) and the legacy test format (``index`` / ``type``) via fallbacks. """ index = data.get("block_index", data.get("index", 0)) block_type = data.get("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, ) ] elif block_type == "text": # Lazy start: don't emit TEXT_MESSAGE_START until we know there's actual content. message_id = str(uuid.uuid4()) self._blocks[index] = {"type": "text", "message_id": message_id, "started": False} self._current_message_id = message_id return [] else: # thinking, image, or any other block type — skip entirely self._blocks[index] = {"type": "skip"} return [] def _handle_block_delta(self, data: dict[str, Any]) -> list[Any]: """Handle content_block:delta — emit TEXT_MESSAGE_CONTENT or TOOL_CALL_ARGS. For text blocks, TEXT_MESSAGE_START is emitted lazily here on the first non-empty token so that thinking/empty text blocks produce zero AG-UI events. Supports both amplifierd wire format (``block_index``) and legacy test format (``index``) via fallback. """ index = data.get("block_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"] == "skip": return [] if block_info["type"] == "tool_use": return [ ToolCallArgsEvent( tool_call_id=block_info["tool_id"], delta=delta.get("partial_json", ""), ) ] else: # Text block — lazy start: emit TEXT_MESSAGE_START on the first real token token = delta.get("text", "") if not token: return [] if not block_info.get("started"): block_info["started"] = True return [ TextMessageStartEvent( message_id=block_info["message_id"], role="assistant", ), TextMessageContentEvent( message_id=block_info["message_id"], delta=token, ), ] 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. Supports two modes: - **Streaming mode** (delta events were received): the block was already opened via ``_handle_block_delta``; just close it with TEXT_MESSAGE_END. - **Batch mode** (no delta events, full text in ``block.text``): emit TEXT_MESSAGE_START + TEXT_MESSAGE_CONTENT + TEXT_MESSAGE_END in one shot. This is the format used by amplifierd, which sends complete blocks at end time. Text blocks with no content in either mode are silently dropped. Supports both amplifierd wire format (``block_index``) and legacy test format (``index``) via fallback. """ index = data.get("block_index", data.get("index", 0)) block_info = self._blocks.pop(index, None) if block_info is None: return [] if block_info["type"] == "skip": return [] if block_info["type"] == "tool_use": return [ToolCallEndEvent(tool_call_id=block_info["tool_id"])] else: message_id = block_info["message_id"] if block_info.get("started"): # Streaming mode: deltas already opened the message — just close it. return [TextMessageEndEvent(message_id=message_id)] # Batch mode: full text is in the "block" field (amplifierd wire format). block = data.get("block", {}) text = block.get("text", "") if text: return [ TextMessageStartEvent(message_id=message_id, role="assistant"), TextMessageContentEvent(message_id=message_id, delta=text), TextMessageEndEvent(message_id=message_id), ] return [] 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, ) ]