"""Bridge: environment side of the AICC protocol.""" from __future__ import annotations import asyncio import uuid from collections.abc import Callable from dataclasses import dataclass, field from typing import Any from aicc.protocol import ( PROTOCOL_VERSION, Capabilities, ErrorCode, ErrorMessage, ErrorPayload, Heartbeat, MessageType, ModelClass, ModelDecl, SessionClose, SessionInit, TickMode, ToolCall, ToolResult, ToolSpec, WorldKind, WorldSpec, ) from aicc.tool import ToolDefinition, call_tool_impl @dataclass class Session: """A live bridge session with one connected agent.""" session_id: str = field(default_factory=lambda: str(uuid.uuid4())) manifest: SessionInit | None = None send_queue: asyncio.Queue[dict[str, Any]] = field(default_factory=asyncio.Queue) closed: bool = False class Bridge: """Environment-side bridge. Holds the world state, registered tools, and per-session state. Bridges run inside transports (in-process, websocket, etc.) and do not perform I/O directly. Example: bridge = Bridge(name="capsule", kind=WorldKind.THREE_D) @bridge.tool(description="Get current position.") async def proprioception() -> dict: return {"position": {"x": 0, "y": 0, "z": 0}, ...} bridge.register_tool(...) """ def __init__( self, name: str, kind: WorldKind = WorldKind.THREE_D, tick_rate_hz: float = 10.0, tick_mode: TickMode = TickMode.FIXED, agent_model: ModelClass = ModelClass.EDGE_MEDIUM, expected_first_token_ms: int | None = None, expected_full_response_ms: int | None = None, ): self.name = name self.kind = kind self.tick_rate_hz = tick_rate_hz self.tick_mode = tick_mode self.agent_model = agent_model self.expected_first_token_ms = expected_first_token_ms self.expected_full_response_ms = expected_full_response_ms self._tools: dict[str, ToolDefinition] = {} self._sessions: dict[str, Session] = {} # ---------- Tool registration ---------- def tool( self, _fn: Callable[..., Any] | None = None, *, id: str | None = None, cls=None, description: str | None = None, requires_capability: str | None = None, limits: dict[str, Any] | None = None, strict_input: bool = False, ): """Register a tool. Use as @bridge.tool or @bridge.tool(cls=..., ...).""" from aicc.tool import tool as _tool if _fn is not None and callable(_fn): defn = _tool(_fn) self.register_tool(defn) return _fn def wrap(fn: Callable[..., Any]) -> Callable[..., Any]: defn = _tool( fn, id=id, cls=cls, description=description, requires_capability=requires_capability, limits=limits, strict_input=strict_input, ) self.register_tool(defn) return fn return wrap def register_tool(self, definition: ToolDefinition) -> None: if definition.id in self._tools: raise ValueError(f"Tool id '{definition.id}' is already registered.") self._tools[definition.id] = definition def tools(self) -> list[ToolDefinition]: return list(self._tools.values()) # ---------- Capability computation ---------- def capabilities(self) -> Capabilities: return Capabilities( sensors=[t.id for t in self._tools.values() if t.cls.value == "sensor"], actuators=[t.id for t in self._tools.values() if t.cls.value == "actuator"], generators=[t.id for t in self._tools.values() if t.cls.value == "generator"], ) # ---------- Session handling ---------- def open_session(self) -> Session: session = Session() manifest = self._build_manifest(session.session_id) session.manifest = manifest self._sessions[session.session_id] = session return session def close_session(self, session_id: str) -> None: sess = self._sessions.pop(session_id, None) if sess is not None: sess.closed = True def emit_event( self, session_id: str, topic: str, payload: dict[str, Any] | None = None, *, tick: int | None = None, ) -> None: """Queue an async event for delivery to a connected agent. The transport is responsible for draining `session.send_queue` and forwarding events to the client. Safe to call from tool handlers. """ sess = self._sessions.get(session_id) if sess is None or sess.closed: return from aicc.protocol import EventMessage event = EventMessage( session_id=session_id, topic=topic, payload=payload or {}, meta={"tick": tick} if tick is not None else None, ) sess.send_queue.put_nowait(event.model_dump(by_alias=True)) def session(self, session_id: str) -> Session: sess = self._sessions.get(session_id) if sess is None or sess.closed: raise KeyError(session_id) return sess async def serve_forever(self) -> None: """Run the bridge event loop until cancelled. Useful with a server transport (e.g. WebSocketServer) already running: async with WebSocketServer(bridge, port=8765): await bridge.serve_forever() """ while True: await asyncio.sleep(3600) def _build_manifest(self, session_id: str) -> SessionInit: return SessionInit( session_id=session_id, tick_rate_hz=self.tick_rate_hz, tick_mode=self.tick_mode, world=WorldSpec(name=self.name, kind=self.kind), agent_model=ModelDecl( class_=self.agent_model, expected_first_token_ms=self.expected_first_token_ms, expected_full_response_ms=self.expected_full_response_ms, ), capabilities=self.capabilities(), tools=[ ToolSpec( id=t.id, class_=t.cls, description=t.description, input_schema=t.input_schema, output_schema=t.output_schema, requires_capability=t.requires_capability, limits=t.limits, strict_input=t.strict_input, ) for t in self._tools.values() ], ) # ---------- Message handling ---------- async def handle_message(self, session_id: str, raw: dict[str, Any]) -> dict[str, Any] | None: """Handle a single incoming message from an agent. Returns the response message (or None for fire-and-forget events/heartbeats). May also push async events to the session queue. """ msg_type = raw.get("type") # Protocol version check BEFORE any model parsing: a wrong version # must be answered with protocol_mismatch, not a parse error. proto = raw.get("protocol") if proto is not None and proto != PROTOCOL_VERSION: return ErrorMessage( session_id=session_id, error=ErrorPayload( code=ErrorCode.PROTOCOL_MISMATCH, message=( f"Protocol mismatch: bridge='{PROTOCOL_VERSION}', " f"message='{proto}'" ), retryable=False, ), ).model_dump(by_alias=True) if msg_type == MessageType.MANIFEST_REQUEST.value: try: sess = self.session(session_id) except KeyError: return self._session_expired(session_id) return sess.manifest.model_dump(by_alias=True) if sess.manifest else None if msg_type == MessageType.TOOL_CALL.value: try: call = ToolCall.model_validate(raw) except Exception: # noqa: BLE001 - malformed tool_call return ToolResult( session_id=session_id, call_id=raw.get("call_id") or "", ok=False, error=ErrorPayload( code=ErrorCode.INVALID_INPUT, message="Malformed tool_call message.", retryable=False, ), ).model_dump(by_alias=True) try: self.session(session_id) except KeyError: return self._session_expired(session_id) return await self._handle_tool_call(session_id, call) if msg_type == MessageType.SESSION_CLOSE.value: self.close_session(session_id) return SessionClose( session_id=session_id, reason=(raw.get("reason") or "client_close"), ).model_dump(by_alias=True) if msg_type == MessageType.HEARTBEAT.value: return Heartbeat(session_id=session_id).model_dump(by_alias=True) if msg_type == MessageType.SESSION_RESUME.value: try: sess = self.session(raw.get("session_id") or session_id) except KeyError: return self._session_expired(session_id) return sess.manifest.model_dump(by_alias=True) if sess.manifest else None return None def _session_expired(self, session_id: str) -> dict[str, Any]: return ErrorMessage( session_id=session_id, error=ErrorPayload( code=ErrorCode.SESSION_EXPIRED, message=f"Session '{session_id}' is unknown or closed.", retryable=False, ), ).model_dump(by_alias=True) async def _handle_tool_call(self, session_id: str, call: ToolCall) -> dict[str, Any]: defn = self._tools.get(call.tool) if defn is None: return ToolResult( session_id=session_id, call_id=call.call_id, ok=False, error=ErrorPayload( code=ErrorCode.TOOL_UNKNOWN, message=f"Tool '{call.tool}' is not registered.", retryable=False, ), ).model_dump(by_alias=True) if defn.requires_capability and defn.requires_capability not in self._sessions[session_id].manifest.capabilities.sensors + self._sessions[session_id].manifest.capabilities.actuators + self._sessions[session_id].manifest.capabilities.generators: # type: ignore[union-attr] return ToolResult( session_id=session_id, call_id=call.call_id, ok=False, error=ErrorPayload( code=ErrorCode.TOOL_UNAVAILABLE, message=f"Tool '{call.tool}' requires capability '{defn.requires_capability}'.", retryable=False, ), ).model_dump(by_alias=True) try: output = await call_tool_impl(defn, call.input) except asyncio.CancelledError: raise except Exception as exc: # noqa: BLE001 return ToolResult( session_id=session_id, call_id=call.call_id, ok=False, error=ErrorPayload( code=ErrorCode.EXECUTION_FAILED, message=f"{type(exc).__name__}: {exc}", retryable=False, ), ).model_dump(by_alias=True) if not isinstance(output, dict): output = {"value": output} return ToolResult( session_id=session_id, call_id=call.call_id, ok=True, output=output, ).model_dump(by_alias=True)