copilotkit-intelligence-runtime 0.1.0__py3-none-any.whl

copilotkit_runtime/finalizer.py

@@ -0,0 +1,75 @@
+"""Track unfinished AG-UI streams without retaining message or tool content."""
+
+import json
+
+from .models import Json
+
+
+class EventFinalizer:
+    """Append TS-compatible closers only when the source omitted a terminal event."""
+
+    def __init__(self) -> None:
+        self.messages: dict[str, None] = {}
+        self.tools: dict[str, tuple[bool, bool]] = {}
+        self.terminal = False
+
+    def observe(self, event: Json) -> None:
+        """Retain only open identifiers and terminal status from an emitted event."""
+        kind = event.get("type")
+        message = event.get("messageId")
+        tool = event.get("toolCallId")
+        if kind in ("RUN_FINISHED", "RUN_ERROR"):
+            self.terminal = True
+        elif kind == "TEXT_MESSAGE_START" and isinstance(message, str):
+            self.messages[message] = None
+        elif kind == "TEXT_MESSAGE_END" and isinstance(message, str):
+            self.messages.pop(message, None)
+        elif kind == "TOOL_CALL_START" and isinstance(tool, str):
+            self.tools[tool] = (False, False)
+        elif isinstance(tool, str) and tool in self.tools:
+            ended, result = self.tools[tool]
+            self.tools[tool] = (
+                ended or kind == "TOOL_CALL_END",
+                result or kind == "TOOL_CALL_RESULT",
+            )
+
+    def finish(self, *, stop_requested: bool = False) -> list[Json]:
+        """Close missing streams and emit a clean stop or an incomplete-stream error."""
+        if self.terminal:
+            return []
+        message = (
+            "Run stopped by user"
+            if stop_requested
+            else "Run ended without emitting a terminal event"
+        )
+        result: list[Json] = [
+            {"type": "TEXT_MESSAGE_END", "messageId": key} for key in self.messages
+        ]
+        for key, (ended, has_result) in self.tools.items():
+            if not ended:
+                result.append({"type": "TOOL_CALL_END", "toolCallId": key})
+            if not has_result:
+                result.append(
+                    {
+                        "type": "TOOL_CALL_RESULT",
+                        "toolCallId": key,
+                        "messageId": f"{key}-result",
+                        "role": "tool",
+                        "content": json.dumps(
+                            {
+                                "status": "stopped" if stop_requested else "error",
+                                "reason": "stop_requested"
+                                if stop_requested
+                                else "missing_terminal_event",
+                                "message": message,
+                            },
+                            separators=(",", ":"),
+                        ),
+                    }
+                )
+        result.append(
+            {"type": "RUN_FINISHED"}
+            if stop_requested
+            else {"type": "RUN_ERROR", "message": message, "code": "INCOMPLETE_STREAM"}
+        )
+        return result

copilotkit_runtime/gateway.py

@@ -0,0 +1,287 @@
+"""Phoenix ingestion with independent control reception and acknowledged delivery."""
+
+import asyncio
+import base64
+import json
+from urllib.parse import parse_qsl, urlencode, urlsplit, urlunsplit
+from uuid import uuid4
+
+from websockets.asyncio.client import ClientConnection, connect
+from websockets.exceptions import WebSocketException
+from websockets.typing import Subprotocol
+
+from .models import Json, RuntimeConfig
+from .telemetry import Telemetry
+
+
+class DeliveryRejected(ConnectionError):
+    """A permanent gateway rejection cannot be repaired through transport retries."""
+
+
+class Gateway:
+    """One independent receiver and one ordered, bounded delivery lane per run."""
+
+    def __init__(
+        self, config: RuntimeConfig, thread_id: str, run_id: str, telemetry: Telemetry
+    ) -> None:
+        self.config = config
+        self.thread_id = thread_id
+        self.run_id = run_id
+        self.telemetry = telemetry
+        self.topic = f"ingestion:{run_id}"
+        self.socket: ClientConnection | None = None
+        self.ref = 0
+        self.join_ref = ""
+        self.sequence = 0
+        self.failed = False
+        self.supports_batch = False
+        self.stop_requested = asyncio.Event()
+        self._disconnected = asyncio.Event()
+        self._reader: asyncio.Task[None] | None = None
+        self._pending: dict[str, tuple[str, asyncio.Future[Json]]] = {}
+        self.lock = asyncio.Lock()
+        self._closed = False
+        self._deliveries: set[asyncio.Task[None]] = set()
+
+    async def join(self) -> None:
+        """Wait for authenticated join, retrying only recoverable startup failures."""
+        for attempt in range(self.config.max_delivery_attempts):
+            try:
+                await self._join_once()
+                return
+            except DeliveryRejected:
+                await self._disconnect()
+                raise
+            except (ConnectionError, TimeoutError, OSError, WebSocketException):
+                await self._disconnect()
+                if attempt + 1 == self.config.max_delivery_attempts:
+                    raise
+                await asyncio.sleep(min(0.1 * 2**attempt, 2))
+
+    async def _join_once(self) -> None:
+        """Open a fresh socket, start its receiver, and negotiate batching capability."""
+        if self._closed:
+            raise DeliveryRejected("Gateway is closed")
+        parts = urlsplit(self.config.runner_url)
+        path = parts.path.rstrip("/")
+        if not path.endswith("/websocket"):
+            path += "/websocket"
+        query = dict(parse_qsl(parts.query))
+        query["vsn"] = "2.0.0"
+        url = urlunsplit((parts.scheme, parts.netloc, path, urlencode(query), ""))
+        token = base64.b64encode(self.config.api_key.encode()).decode().rstrip("=")
+        self.socket = await connect(
+            url,
+            subprotocols=[Subprotocol("phoenix"), Subprotocol("base64url.bearer.phx." + token)],
+            open_timeout=self.config.ack_timeout,
+            close_timeout=1,
+            max_size=self.config.max_body_bytes,
+        )
+        self._disconnected.clear()
+        self._reader = asyncio.create_task(
+            self._receive(self.socket), name="copilotkit-gateway-receiver"
+        )
+        try:
+            reply = await self._push(
+                "phx_join", {"thread_id": self.thread_id, "run_id": self.run_id}
+            )
+            self.supports_batch = "runner_event_batch_v1" in reply.get("capabilities", [])
+        except BaseException:
+            await self._disconnect()
+            raise
+
+    async def _receive(self, socket: ClientConnection) -> None:
+        """Read control frames while the agent is idle or a push waits for its ACK."""
+        try:
+            async for raw in socket:
+                frame = json.loads(raw)
+                if not isinstance(frame, list) or len(frame) != 5:
+                    raise ConnectionError("Invalid Phoenix frame")
+                _, ref, topic, event, payload = frame
+                if event == "ag-ui" and topic == self.topic and isinstance(payload, dict):
+                    if payload.get("type") == "CUSTOM" and payload.get("name") == "stop":
+                        self.stop_requested.set()
+                elif event == "phx_reply" and ref in self._pending:
+                    expected_topic, future = self._pending[ref]
+                    if topic == expected_topic and not future.done():
+                        if not isinstance(payload, dict):
+                            future.set_exception(ConnectionError("Invalid Phoenix reply"))
+                        else:
+                            future.set_result(payload)
+                elif event in ("phx_close", "phx_error") and topic == self.topic:
+                    raise ConnectionError("Phoenix channel closed")
+        except (ValueError, ConnectionError, OSError, WebSocketException):
+            pass
+        finally:
+            if self.socket is socket:
+                self._disconnected.set()
+                for _, future in self._pending.values():
+                    if not future.done():
+                        future.set_exception(ConnectionError("Gateway connection closed"))
+
+    async def _push(self, event: str, payload: Json, topic: str | None = None) -> Json:
+        """Match one reply reference without consuming another push's control messages."""
+        if self.socket is None or self._disconnected.is_set():
+            raise ConnectionError("Gateway is disconnected")
+        self.ref += 1
+        ref, target = str(self.ref), topic or self.topic
+        if event == "phx_join":
+            self.join_ref = ref
+        future: asyncio.Future[Json] = asyncio.get_running_loop().create_future()
+        self._pending[ref] = (target, future)
+        try:
+            async with asyncio.timeout(self.config.ack_timeout):
+                await self.socket.send(
+                    json.dumps([self.join_ref, ref, target, event, payload], allow_nan=False)
+                )
+                reply = await future
+            response = reply.get("response", {})
+            response = response if isinstance(response, dict) else {}
+            if reply.get("status") != "ok":
+                retryable = (
+                    response.get("retryable") is True
+                    or response.get("reason") == "gateway_draining"
+                )
+                if response.get("retryable") is False or (event == "phx_join" and not retryable):
+                    raise DeliveryRejected("Phoenix push permanently rejected")
+                raise ConnectionError("Phoenix push rejected")
+            return response
+        finally:
+            self._pending.pop(ref, None)
+            if future.done() and not future.cancelled():
+                future.exception()
+            else:
+                future.cancel()
+
+    async def send(self, event: Json) -> None:
+        """Send one immutable event, with backpressure until its durable ACK."""
+        await self.send_many([event])
+
+    async def send_many(self, events: list[Json]) -> None:
+        """Defer cancellation until the active immutable group reaches its ACK boundary."""
+        detached: list[Json] = json.loads(json.dumps(events, allow_nan=False))
+        task = asyncio.create_task(self._deliver(detached))
+        self._deliveries.add(task)
+        cancelled = False
+        try:
+            while not task.done():
+                try:
+                    await asyncio.shield(task)
+                except asyncio.CancelledError:
+                    # Each stop/lease cancellation must respect the same ACK boundary.
+                    # Explicit abort() cancels the delivery itself, ending this loop.
+                    cancelled = True
+            task.result()
+            if cancelled:
+                raise asyncio.CancelledError
+        finally:
+            self._deliveries.discard(task)
+
+    async def _deliver(self, events: list[Json]) -> None:
+        """Stamp at most 32 events atomically and replay the same payloads after disconnect."""
+        if not events or len(events) > 32:
+            raise ValueError("Gateway batches require between 1 and 32 events")
+        # send_many already detached and validated this snapshot.
+        async with self.lock:
+            if self.failed or self._closed:
+                raise DeliveryRejected("Gateway delivery is unavailable")
+            payloads = []
+            for index, event in enumerate(events):
+                if not isinstance(event, dict) or not isinstance(event.get("type"), str):
+                    raise ValueError("Invalid AG-UI event")
+                metadata = event.get("metadata") or {}
+                if not isinstance(metadata, dict):
+                    raise ValueError("Invalid event metadata")
+                payloads.append(
+                    {
+                        **event,
+                        "threadId": self.thread_id,
+                        "runId": self.run_id,
+                        "thread_id": self.thread_id,
+                        "run_id": self.run_id,
+                        "metadata": {
+                            **metadata,
+                            "cpki_event_id": str(uuid4()),
+                            "cpki_event_seq": self.sequence + index + 1,
+                        },
+                    }
+                )
+            self.sequence += len(payloads)
+            position = 0
+            attempts = 0
+            try:
+                while position < len(payloads):
+                    try:
+                        if self.socket is None or self._disconnected.is_set():
+                            await self._disconnect()
+                            await self._join_once()
+                        batch = (
+                            payloads[position:]
+                            if self.supports_batch
+                            else payloads[position : position + 1]
+                        )
+                        await self._push(
+                            "events" if self.supports_batch else "event",
+                            {"events": batch} if self.supports_batch else batch[0],
+                        )
+                        position += len(batch)
+                        attempts = 0
+                    except DeliveryRejected:
+                        raise
+                    except (ConnectionError, TimeoutError, OSError, WebSocketException):
+                        attempts += 1
+                        await self._disconnect()
+                        if attempts >= self.config.max_delivery_attempts:
+                            raise ConnectionError("Event delivery retry budget exhausted")
+                        await asyncio.sleep(min(0.1 * 2 ** (attempts - 1), 2))
+            except BaseException:
+                self.failed = True
+                raise
+
+    async def keepalive(self) -> None:
+        """Reconnect idle channels and maintain Phoenix heartbeat frames without rerunning agents."""
+        while not self._closed:
+            try:
+                await asyncio.wait_for(self._disconnected.wait(), timeout=15)
+            except TimeoutError:
+                pass
+            async with self.lock:
+                if self._closed:
+                    return
+                if self._disconnected.is_set():
+                    await self._disconnect()
+                    await self.join()
+                else:
+                    try:
+                        await self._push("heartbeat", {}, "phoenix")
+                    except (ConnectionError, TimeoutError, OSError, WebSocketException):
+                        await self._disconnect()
+                        await self.join()
+
+    async def _disconnect(self) -> None:
+        """Close one transport and settle its receiver before replacing it."""
+        socket, reader = self.socket, self._reader
+        self.socket = None
+        self._reader = None
+        if reader:
+            reader.cancel()
+            await asyncio.gather(reader, return_exceptions=True)
+        if socket:
+            await socket.close()
+
+    async def aclose(self) -> None:
+        """Close only after delivery completes, or after a terminal failure/shutdown."""
+        self._closed = True
+        for task in self._deliveries:
+            task.cancel()
+        await self._disconnect()
+
+    def abort(self) -> None:
+        """Force-close transport at the host shutdown deadline; no ACK is claimed."""
+        self._closed = self.failed = True
+        for task in self._deliveries:
+            task.cancel()
+        if self._reader:
+            self._reader.cancel()
+        if self.socket:
+            self.socket.transport.abort()

copilotkit_runtime/mcp_apps.py

@@ -0,0 +1,299 @@
+"""MCP Apps middleware with native authenticated Streamable HTTP sessions."""
+
+import asyncio
+import hashlib
+import json
+from collections.abc import AsyncIterator
+from dataclasses import dataclass, field
+from datetime import timedelta
+from urllib.parse import urlsplit
+from uuid import uuid4
+
+from mcp import ClientSession, types
+from mcp.client.streamable_http import streamablehttp_client
+from mcp.shared.message import SessionMessage
+
+from .models import Json
+
+
+@dataclass(frozen=True)
+class MCPServer:
+    """Server-owned endpoint/authentication. Browser input cannot alter this configuration."""
+
+    url: str
+    server_id: str | None = None
+    agent_id: str | None = None
+    headers: dict[str, str] = field(default_factory=dict, repr=False)
+
+    def __post_init__(self) -> None:
+        """Reject non-HTTP endpoints and credentials embedded in URLs."""
+        parsed = urlsplit(self.url)
+        if parsed.scheme not in ("http", "https") or not parsed.hostname or parsed.username:
+            raise ValueError("MCP server requires an HTTP URL without embedded credentials")
+
+    @property
+    def server_hash(self) -> str:
+        """Match the opaque reference used by the TypeScript MCP Apps renderer."""
+        serialized = json.dumps({"type": "http", "url": self.url}, separators=(",", ":"))
+        return hashlib.md5(serialized.encode(), usedforsecurity=False).hexdigest()
+
+
+@dataclass(frozen=True)
+class MCPAppsConfig:
+    """Bound MCP request duration, discovery pages, and buffered arguments."""
+
+    servers: tuple[MCPServer, ...] = ()
+    timeout: float = 30
+    max_tool_pages: int = 20
+    max_argument_bytes: int = 2 * 1024 * 1024
+
+
+class MCPAppsMiddleware:
+    """Expose UI-enabled tools and execute iframe requests against configured servers."""
+
+    def __init__(self, config: MCPAppsConfig) -> None:
+        self.config = config
+
+    def _servers(self, agent_id: str) -> tuple[MCPServer, ...]:
+        """Select the allowed MCP servers for one agent."""
+        return tuple(
+            server for server in self.config.servers if server.agent_id in (None, agent_id)
+        )
+
+    async def _request(self, server: MCPServer, method: str, params: Json | None = None) -> Json:
+        """Negotiate an MCP session and close its transport after one bounded operation."""
+        async with asyncio.timeout(self.config.timeout):
+            async with streamablehttp_client(
+                server.url,
+                headers=dict(server.headers),
+                timeout=self.config.timeout,
+                sse_read_timeout=self.config.timeout,
+            ) as (read, write, _):
+                async with ClientSession(
+                    read, write, read_timeout_seconds=timedelta(seconds=self.config.timeout)
+                ) as session:
+                    # Public protocol calls allow the UI extension without an SDK subclass.
+                    initialized = await session.send_request(
+                        types.ClientRequest(
+                            types.InitializeRequest(
+                                params=types.InitializeRequestParams(
+                                    protocolVersion="2025-06-18",
+                                    capabilities=types.ClientCapabilities.model_validate(
+                                        {
+                                            "extensions": {
+                                                "io.modelcontextprotocol/ui": {
+                                                    "mimeTypes": ["text/html;profile=mcp-app"]
+                                                }
+                                            }
+                                        }
+                                    ),
+                                    clientInfo=types.Implementation(
+                                        name="copilotkit-runtime-python", version="0.1.0"
+                                    ),
+                                )
+                            )
+                        ),
+                        types.InitializeResult,
+                    )
+                    if initialized.protocolVersion not in (
+                        "2024-11-05",
+                        "2025-03-26",
+                        "2025-06-18",
+                        "2025-11-25",
+                    ):
+                        raise ValueError("Unsupported MCP protocol version")
+                    await session.send_notification(
+                        types.ClientNotification(types.InitializedNotification())
+                    )
+                    if method == "tools/list":
+                        tools: list[Json] = []
+                        cursor = None
+                        cursors: set[str] = set()
+                        for _page in range(self.config.max_tool_pages):
+                            page = await session.list_tools(cursor=cursor)
+                            tools.extend(
+                                tool.model_dump(mode="json", by_alias=True, exclude_none=True)
+                                for tool in page.tools
+                            )
+                            cursor = page.nextCursor
+                            if not cursor:
+                                return {"tools": tools}
+                            if cursor in cursors:
+                                raise ValueError("MCP discovery repeated its cursor")
+                            cursors.add(cursor)
+                        raise ValueError("MCP discovery page limit exceeded")
+                    if method == "tools/call":
+                        result = await session.send_request(
+                            types.ClientRequest(
+                                types.CallToolRequest(
+                                    params=types.CallToolRequestParams.model_validate(params or {})
+                                )
+                            ),
+                            types.CallToolResult,
+                        )
+                        return result.model_dump(mode="json", by_alias=True, exclude_none=True)
+                    if method == "resources/read":
+                        resource = await session.read_resource((params or {})["uri"])
+                        return resource.model_dump(mode="json", by_alias=True, exclude_none=True)
+                    if method == "ping":
+                        pong = await session.send_ping()
+                        return pong.model_dump(mode="json", by_alias=True, exclude_none=True)
+                    if method == "notifications/message":
+                        await write.send(
+                            SessionMessage(
+                                types.JSONRPCMessage(
+                                    types.JSONRPCNotification(
+                                        jsonrpc="2.0", method=method, params=params
+                                    )
+                                )
+                            )
+                        )
+                        return {"success": True}
+                    raise ValueError("MCP method not allowed")
+
+    async def discover(
+        self, input: Json, agent_id: str
+    ) -> tuple[Json, dict[str, tuple[MCPServer, str]]]:
+        """Inject tools with UI resources and reject ambiguous server tool names."""
+        tools = list(input.get("tools", []))
+        by_name: dict[str, tuple[MCPServer, str]] = {}
+        existing = {tool.get("name") for tool in tools}
+        for server in self._servers(agent_id):
+            result = await self._request(server, "tools/list")
+            for tool in result["tools"]:
+                meta = tool.get("_meta")
+                meta = meta if isinstance(meta, dict) else {}
+                ui = meta.get("ui")
+                if isinstance(ui, dict) and "visibility" in ui:
+                    visibility = ui["visibility"]
+                    if not isinstance(visibility, list) or "model" not in visibility:
+                        continue
+                uri = ui.get("resourceUri") if isinstance(ui, dict) else None
+                if not isinstance(uri, str):
+                    uri = meta.get("ui/resourceUri")
+                if not isinstance(uri, str):
+                    continue
+                name = tool["name"]
+                if name in existing:
+                    raise ValueError("MCP tool name collision")
+                existing.add(name)
+                by_name[name] = (server, uri)
+                tools.append(
+                    {
+                        "name": name,
+                        "description": tool.get("description", "") + f"\n[UI Resource: {uri}]",
+                        "parameters": tool.get("inputSchema", {"type": "object", "properties": {}}),
+                    }
+                )
+        return {**input, "tools": tools}, by_name
+
+    async def proxy(self, request: Json, agent_id: str) -> AsyncIterator[Json]:
+        """Bypass the agent for allowlisted browser iframe requests."""
+        yield {"type": "RUN_STARTED"}
+        servers = self._servers(agent_id)
+        server_id = request.get("serverId")
+        matches = [
+            item
+            for item in servers
+            if (
+                item.server_id == server_id
+                if server_id
+                else item.server_hash == request.get("serverHash")
+            )
+        ]
+        server = matches[0] if len(matches) == 1 else None
+        if server is None:
+            result: Json = {"error": "Unknown MCP server"}
+        elif request.get("method") not in (
+            "tools/call",
+            "resources/read",
+            "notifications/message",
+            "ping",
+        ):
+            result = {"error": "MCP method not allowed for UI proxy"}
+        elif "params" in request and not isinstance(request["params"], dict):
+            result = {"error": "Invalid MCP params"}
+        else:
+            try:
+                result = await self._request(server, request["method"], request.get("params"))
+            except Exception:
+                result = {"error": "MCP request failed"}
+        yield {"type": "RUN_FINISHED", "result": result}
+
+    async def transform(
+        self, source: AsyncIterator[Json], input: Json, tools: dict[str, tuple[MCPServer, str]]
+    ) -> AsyncIterator[Json]:
+        """Execute unresolved UI calls before the final run event and emit their activities."""
+        calls: dict[str, Json] = {}
+        resolved: set[str] = set()
+        for message in input.get("messages", []):
+            if message.get("role") == "assistant":
+                for call in message.get("toolCalls", []):
+                    calls[call["id"]] = {
+                        "name": call["function"]["name"],
+                        "args": call["function"].get("arguments", ""),
+                    }
+            elif message.get("role") == "tool":
+                resolved.add(message["toolCallId"])
+        terminal = None
+        async for event in source:
+            kind, call_id = event.get("type"), event.get("toolCallId", "")
+            if kind == "TOOL_CALL_START":
+                calls[call_id] = {"name": event["toolCallName"], "args": ""}
+            elif kind == "TOOL_CALL_ARGS" and call_id in calls:
+                calls[call_id]["args"] += event.get("delta", "")
+                if len(calls[call_id]["args"].encode()) > self.config.max_argument_bytes:
+                    raise ValueError("MCP arguments exceed configured limit")
+            elif kind == "TOOL_CALL_RESULT":
+                resolved.add(call_id)
+            if kind == "RUN_FINISHED":
+                terminal = event
+            else:
+                yield event
+        if terminal:
+            for call_id, call in calls.items():
+                if call_id in resolved or call["name"] not in tools:
+                    continue
+                server, resource_uri = tools[call["name"]]
+                try:
+                    arguments = json.loads(call["args"] or "{}")
+                    if not isinstance(arguments, dict):
+                        raise ValueError("MCP tool arguments must be an object")
+                    result = await self._request(
+                        server, "tools/call", {"name": call["name"], "arguments": arguments}
+                    )
+                    content = result.get("content", [])
+                    text = "\n".join(
+                        item["text"]
+                        for item in content
+                        if item.get("type") == "text" and isinstance(item.get("text"), str)
+                    )
+                    yield {
+                        "type": "TOOL_CALL_RESULT",
+                        "toolCallId": call_id,
+                        "messageId": str(uuid4()),
+                        "content": text or json.dumps(content),
+                    }
+                    activity: Json = {
+                        "result": result,
+                        "resourceUri": resource_uri,
+                        "serverHash": server.server_hash,
+                        "toolInput": arguments,
+                    }
+                    if server.server_id:
+                        activity["serverId"] = server.server_id
+                    yield {
+                        "type": "ACTIVITY_SNAPSHOT",
+                        "messageId": str(uuid4()),
+                        "activityType": "mcp-apps",
+                        "content": activity,
+                        "replace": True,
+                    }
+                except Exception:
+                    yield {
+                        "type": "TOOL_CALL_RESULT",
+                        "toolCallId": call_id,
+                        "messageId": str(uuid4()),
+                        "content": json.dumps({"error": "MCP tool execution failed"}),
+                    }
+            yield terminal