aicc-capsule/testbed/bridge.py

"""The testbed bridge: connects the AICC protocol to the room world.

Registers the world tools (proprioception, vision, depth, hear, world_query,
move, turn, look_at, interact) plus the conformance tools the scenario suite
expects (echo, boom, bump). Subclasses ``aicc.Bridge`` only to learn which
session is handling the current tool call, so tools can push async events
(collision, audio) via ``emit_event``.
"""

from __future__ import annotations

import asyncio
import base64
import copy
import io
import math
from typing import Any

from aicc.bridge import Bridge
from aicc.protocol import TickMode, ToolClass
from pydantic import BaseModel

from testbed.room.render import Raycaster
from testbed.room.world import Room

VISION_WIDTH = 160
VISION_HEIGHT = 120
DEPTH_WIDTH = VISION_WIDTH // 4
DEPTH_HEIGHT = VISION_HEIGHT // 4


# ---------------------------------------------------------------------------
# Typed tool outputs (advertised as output_schema in the manifest)
# ---------------------------------------------------------------------------


class Vec3(BaseModel):
    x: float
    y: float
    z: float


class Rotation(BaseModel):
    yaw_deg: float
    pitch_deg: float


class ProprioceptionOutput(BaseModel):
    position: Vec3
    rotation: Rotation
    velocity: Vec3
    health: float
    tick: int


class VisionOutput(BaseModel):
    png_b64: str
    width: int
    height: int
    tick: int


class DepthOutput(BaseModel):
    width: int
    height: int
    max_depth: float
    depth: list[list[float]]
    tick: int


class Sound(BaseModel):
    kind: str
    direction_deg: float
    intensity: float


class HearOutput(BaseModel):
    sounds: list[Sound]
    tick: int


class MoveOutput(BaseModel):
    moved: float
    collision: bool
    collision_normal: Vec3 | None
    position: Vec3
    tick: int


class TurnOutput(BaseModel):
    rotation: Rotation
    tick: int


class LookAtOutput(BaseModel):
    target: str
    rotation: Rotation
    tick: int


class InteractOutput(BaseModel):
    target: str
    success: bool
    message: str
    tick: int


class WorldQueryOutput(BaseModel):
    room: dict[str, Any]
    obstacles: list[dict[str, Any]]
    beacon: dict[str, Any]
    tick: int


class EchoOutput(BaseModel):
    value: str


class BumpOutput(BaseModel):
    bumped: bool


# ---------------------------------------------------------------------------
# Bridge
# ---------------------------------------------------------------------------


class RoomBridge(Bridge):
    """AICC bridge over one Room. Event-driven ticks: the world advances
    only on tool calls (protocol tick_mode = "event")."""

    def __init__(
        self,
        world: Room | None = None,
        name: str = "testbed_room_01",
        **kwargs: Any,
    ) -> None:
        self.world = world or Room()
        self.renderer = Raycaster(VISION_WIDTH, VISION_HEIGHT)
        self._session_id: str | None = None
        kwargs.setdefault("tick_mode", TickMode.EVENT)
        kwargs.setdefault("tick_rate_hz", 10.0)
        super().__init__(name=name, kind="3d", **kwargs)
        self._register_tools()

    # -- session plumbing for event emission from tool handlers ------------

    async def _handle_tool_call(self, session_id: str, call) -> dict[str, Any]:
        self._session_id = session_id
        return await super()._handle_tool_call(session_id, call)

    def emit(
        self,
        topic: str,
        payload: dict[str, Any] | None = None,
        *,
        tick: int | None = None,
    ) -> None:
        """Push an async event to the session currently handling a tool call."""
        if self._session_id is not None:
            self.emit_event(self._session_id, topic, payload, tick=tick)

    # -- world helpers ------------------------------------------------------

    def _tick(self) -> int:
        return self.world.advance_tick()

    def _vec3(self, x: float, y: float, z: float) -> Vec3:
        return Vec3(x=round(x, 3), y=round(y, 3), z=round(z, 3))

    def _rotation(self) -> Rotation:
        c = self.world.capsule
        return Rotation(
            yaw_deg=round(c.yaw_deg % 360.0, 2), pitch_deg=round(c.pitch_deg, 2)
        )

    def _bearing_to(self, dx: float, dz: float) -> float:
        """Compass bearing of (dx, dz) relative to the capsule's heading."""
        abs_bearing = math.degrees(math.atan2(dx, dz))
        return (abs_bearing - self.world.capsule.yaw_deg) % 360.0

    # -- tool registration --------------------------------------------------

    def _register_tools(self) -> None:
        world = self.world
        renderer = self.renderer

        @self.tool(
            cls=ToolClass.SENSOR,
            description="Return the capsule's own state: position, rotation, velocity, health.",
        )
        async def proprioception() -> ProprioceptionOutput:
            tick = self._tick()
            c = world.capsule
            return ProprioceptionOutput(
                position=self._vec3(c.x, 0.0, c.z),
                rotation=self._rotation(),
                velocity=self._vec3(0.0, 0.0, c.speed),
                health=round(c.health, 1),
                tick=tick,
            ).model_dump()

        @self.tool(
            cls=ToolClass.SENSOR,
            description="Return the current first-person RGB frame as base64 PNG (160x120).",
            limits={"calls_per_minute": 60},
        )
        async def vision() -> VisionOutput:
            tick = self._tick()
            # Render off the event loop so frames never stall the bridge under
            # load; deepcopy keeps the snapshot consistent with this tool call.
            snapshot = copy.deepcopy(world)
            img, _ = await asyncio.to_thread(renderer.render, snapshot)

            def _encode() -> str:
                buf = io.BytesIO()
                img.save(buf, format="PNG")
                return base64.b64encode(buf.getvalue()).decode("ascii")

            return VisionOutput(
                png_b64=await asyncio.to_thread(_encode),
                width=VISION_WIDTH,
                height=VISION_HEIGHT,
                tick=tick,
            ).model_dump()

        @self.tool(
            cls=ToolClass.SENSOR,
            description="Return a depth map (meters, 40x30) aligned to the vision frame.",
        )
        async def depth(max_depth: float = 10.0) -> DepthOutput:
            tick = self._tick()
            snapshot = copy.deepcopy(world)
            _, grid = await asyncio.to_thread(renderer.render, snapshot, max_depth)
            return DepthOutput(
                width=DEPTH_WIDTH,
                height=DEPTH_HEIGHT,
                max_depth=round(max_depth, 2),
                depth=grid,
                tick=tick,
            ).model_dump()

        @self.tool(
            cls=ToolClass.SENSOR,
            description="Return audio events since the last call: beacon hum when near, collision thuds.",
        )
        async def hear() -> HearOutput:
            tick = self._tick()
            sounds = world.hear_now()
            return HearOutput(
                sounds=[Sound(**s) for s in sounds],
                tick=tick,
            ).model_dump()

        @self.tool(
            cls=ToolClass.SENSOR,
            description="Return static room layout: bounds, obstacle positions, beacon location.",
        )
        async def world_query() -> WorldQueryOutput:
            tick = self._tick()
            d = world.describe()
            return WorldQueryOutput(**d, tick=tick).model_dump()

        @self.tool(
            cls=ToolClass.ACTUATOR,
            description="Move the capsule forward along its heading by the given distance in meters.",
        )
        async def move(forward: float = 1.0) -> MoveOutput:
            tick = self._tick()
            if forward < 0.0 or forward > 5.0:
                raise ValueError(f"forward must be in [0, 5] m, got {forward}")
            moved, hit = world.move_forward(forward)
            c = world.capsule
            normal = None
            if hit is not None:
                normal = self._vec3(hit.normal_x, 0.0, hit.normal_z)
                self.emit(
                    "collision",
                    {
                        "other": hit.other,
                        "normal": normal.model_dump(),
                        "impulse": round(hit.impulse, 3),
                    },
                    tick=tick,
                )
                world.queue_audio(
                    "thud",
                    self._bearing_to(-hit.normal_x, -hit.normal_z),
                    min(1.0, 0.4 + hit.impulse * 0.15),
                )
            return MoveOutput(
                moved=round(moved, 3),
                collision=hit is not None,
                collision_normal=normal,
                position=self._vec3(c.x, 0.0, c.z),
                tick=tick,
            ).model_dump()

        @self.tool(
            cls=ToolClass.ACTUATOR,
            description="Rotate the camera: yaw_deg turns left/right, pitch_deg tilts up/down.",
        )
        async def turn(yaw_deg: float = 0.0, pitch_deg: float = 0.0) -> TurnOutput:
            tick = self._tick()
            world.turn(yaw_deg, pitch_deg)
            return TurnOutput(rotation=self._rotation(), tick=tick).model_dump()

        @self.tool(
            cls=ToolClass.ACTUATOR,
            description="Orient the camera toward a named target (e.g. 'beacon').",
        )
        async def look_at(target: str) -> LookAtOutput:
            tick = self._tick()
            if target != "beacon":
                raise ValueError(
                    f"unknown target {target!r}; available targets: beacon"
                )
            yaw, pitch = world.look_at_beacon()
            return LookAtOutput(
                target=target,
                rotation=Rotation(
                    yaw_deg=round(yaw % 360.0, 2), pitch_deg=round(pitch, 2)
                ),
                tick=tick,
            ).model_dump()

        @self.tool(
            cls=ToolClass.ACTUATOR,
            description="Use the named interactable object (e.g. 'beacon'). Must be within reach.",
        )
        async def interact(target: str = "beacon") -> InteractOutput:
            tick = self._tick()
            if target != "beacon":
                raise ValueError(
                    f"unknown target {target!r}; available targets: beacon"
                )
            success, message = world.interact_beacon()
            if success and world.beacon.active:
                world.queue_audio(
                    "beacon_chime",
                    self._bearing_to(
                        world.beacon.x - world.capsule.x,
                        world.beacon.z - world.capsule.z,
                    ),
                    1.0,
                )
            return InteractOutput(
                target=target, success=success, message=message, tick=tick
            ).model_dump()

        # -- conformance tools (design doc: register alongside world tools) --

        @self.tool(
            cls=ToolClass.SENSOR,
            description="Echo a value back to the caller.",
        )
        async def echo(value: str = "") -> EchoOutput:
            return EchoOutput(value=value).model_dump()

        @self.tool(
            cls=ToolClass.ACTUATOR,
            description="Always raises an exception.",
        )
        async def boom() -> dict[str, Any]:
            raise RuntimeError("kaboom")

        @self.tool(
            cls=ToolClass.ACTUATOR,
            description="Returns ok and emits a collision event.",
        )
        async def bump() -> BumpOutput:
            tick = self._tick()
            self.emit(
                "collision",
                {
                    "other": "wall",
                    "normal": {"x": 0.0, "y": 0.0, "z": 1.0},
                    "impulse": 1.0,
                },
                tick=tick,
            )
            return BumpOutput(bumped=True).model_dump()


def build_bridge(world: Room | None = None, **kwargs: Any) -> RoomBridge:
    """Construct a fully-wired testbed bridge."""
    return RoomBridge(world=world, **kwargs)


async def serve(host: str = "127.0.0.1", port: int = 8765) -> None:
    """Run the bridge over WebSocket until interrupted."""
    from aicc.transport.websocket import WebSocketServer

    bridge = build_bridge()
    async with WebSocketServer(bridge, host=host, port=port):
        print(f"AICC capsule testbed listening on ws://{host}:{port}")
        await bridge.serve_forever()


if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser(
        description="Run the AICC capsule testbed bridge over WebSocket."
    )
    parser.add_argument("--host", default="127.0.0.1")
    parser.add_argument("--port", type=int, default=8765)
    args = parser.parse_args()
    try:
        asyncio.run(serve(args.host, args.port))
    except KeyboardInterrupt:
        pass