research-workbench/backend/agui_adapter.py

"""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,
            )
        ]