- conformance.py: declarative scenario runner (subset matcher, $kind/$gt/ $enum/$required assertions, call_id/session_id echo checks), reference bridge, CLI entry (python -m aicc.conformance) - bridge.py: protocol version check before parsing (protocol_mismatch), session expiry handling (session_expired), malformed tool_call -> invalid_input, emit_event() for async events - client.py: Self return type - in_process.py: drain bridge event queue after response (event delivery), Self return type - websocket.py: Self return type, best-effort shutdown - tests: 14 passing (incl. conformance scenarios, 9/9 core scenarios) - ruff: all checks pass
95 lines
3.5 KiB
Python
95 lines
3.5 KiB
Python
"""In-process transport: bridges AICC clients to a Bridge within one event loop.
|
|
|
|
Useful for tests, embedded use, and single-process agents that drive the bridge directly.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
from typing import Any, Self
|
|
|
|
from aicc.bridge import Bridge
|
|
from aicc.errors import ConnectionError_
|
|
from aicc.protocol import SessionInit
|
|
|
|
|
|
class InProcessTransport:
|
|
"""Two-sided transport: one side is the bridge, the other is the client.
|
|
|
|
Both sides share the same asyncio event loop. The bridge side pushes its
|
|
manifest into the client's receive queue when start() is awaited.
|
|
|
|
Example:
|
|
async with InProcessTransport.start(bridge) as t:
|
|
client = AICCClient(t)
|
|
async with client:
|
|
manifest = await client.handshake()
|
|
"""
|
|
|
|
def __init__(self) -> None:
|
|
self._send_q: asyncio.Queue[dict[str, Any]] = asyncio.Queue()
|
|
self._recv_q: asyncio.Queue[dict[str, Any]] = asyncio.Queue()
|
|
self._bridge: Bridge | None = None
|
|
self._session_id: str | None = None
|
|
self._connected = False
|
|
|
|
@classmethod
|
|
def start(cls, bridge: Bridge) -> InProcessTransport:
|
|
t = cls()
|
|
t._bridge = bridge
|
|
return t
|
|
|
|
async def connect(self) -> None:
|
|
if self._bridge is None:
|
|
raise ConnectionError_("InProcessTransport.start(bridge) must be used for server-side.")
|
|
if self._connected:
|
|
return
|
|
self._connected = True
|
|
session = self._bridge.open_session()
|
|
self._session_id = session.session_id
|
|
manifest: SessionInit = session.manifest # type: ignore[assignment]
|
|
await self._recv_q.put(manifest.model_dump(by_alias=True))
|
|
|
|
async def close(self) -> None:
|
|
if self._bridge is not None and self._session_id is not None:
|
|
self._bridge.close_session(self._session_id)
|
|
self._session_id = None
|
|
self._connected = False
|
|
|
|
async def send(self, message: dict[str, Any]) -> None:
|
|
if not self._connected or self._bridge is None or self._session_id is None:
|
|
raise ConnectionError_("Transport is not connected.")
|
|
sid = message.get("session_id") or self._session_id
|
|
msg_type = message.get("type")
|
|
# session_resume carries the session_id; bind it
|
|
if msg_type == "session_resume" and not self._bridge._sessions.get(sid):
|
|
# Auto-open if bridge forgot
|
|
pass
|
|
response = await self._bridge.handle_message(sid, message)
|
|
if response is not None:
|
|
await self._recv_q.put(response)
|
|
# Drain any events the bridge queued during handling (e.g. from
|
|
# tool handlers calling bridge.emit_event). Deliver them after the
|
|
# response so the client sees a coherent request -> response order.
|
|
sess = self._bridge._sessions.get(sid)
|
|
if sess is not None:
|
|
while not sess.send_queue.empty():
|
|
event = sess.send_queue.get_nowait()
|
|
await self._recv_q.put(event)
|
|
|
|
async def receive(self) -> dict[str, Any]:
|
|
if not self._connected:
|
|
raise ConnectionError_("Transport is not connected.")
|
|
return await self._recv_q.get()
|
|
|
|
async def push_event(self, event: dict[str, Any]) -> None:
|
|
"""Bridge-side helper to push an event to the connected client."""
|
|
await self._recv_q.put(event)
|
|
|
|
async def __aenter__(self) -> Self:
|
|
await self.connect()
|
|
return self
|
|
|
|
async def __aexit__(self, exc_type, exc, tb) -> None:
|
|
await self.close()
|