aicc-py/tests/test_integration.py

"""Integration test: end-to-end AICC session via in-process transport."""

from __future__ import annotations

import pytest

from aicc import AICCClient, Bridge, ToolClass
from aicc.protocol import ModelClass, TickMode, WorldKind
from aicc.transport.in_process import InProcessTransport


@pytest.mark.asyncio
async def test_handshake_and_call_tool():
    bridge = Bridge(
        name="test-room",
        kind=WorldKind.TEXT,
        tick_rate_hz=10.0,
        tick_mode=TickMode.FIXED,
        agent_model=ModelClass.EDGE_MEDIUM,
    )

    @bridge.tool(description="Get the agent's current position and rotation.")
    async def proprioception() -> dict:
        return {
            "position": {"x": 1.0, "y": 0.5, "z": 2.0},
            "rotation": {"yaw": 0.0, "pitch": 0.0},
            "velocity": {"x": 0.0, "y": 0.0, "z": 0.0},
            "health": 100,
        }

    @bridge.tool(description="Move the agent forward by the given distance.")
    async def move(forward: float = 0.0) -> dict:
        return {"moved": forward}

    @bridge.tool(cls=ToolClass.GENERATOR, description="Place an object in the world.")
    async def place_object(prefab: str) -> dict:
        return {"spawned": prefab}

    t = InProcessTransport.start(bridge)
    async with t:
        client = AICCClient(t)
        async with client:
            manifest = await client.handshake()
            assert manifest.world.name == "test-room"
            assert manifest.tick_rate_hz == 10.0
            tool_ids = {t.id for t in manifest.tools}
            assert {"proprioception", "move", "place_object"} <= tool_ids

            res = await client.call_tool("proprioception", {})
            assert res.ok
            assert res.output["position"] == {"x": 1.0, "y": 0.5, "z": 2.0}

            res = await client.call_tool("move", {"forward": 2.5})
            assert res.ok
            assert res.output == {"moved": 2.5}

            res = await client.call_tool("place_object", {"prefab": "crate_01"})
            assert res.ok
            assert res.output == {"spawned": "crate_01"}


@pytest.mark.asyncio
async def test_unknown_tool_returns_tool_error():
    bridge = Bridge(name="x")

    @bridge.tool(description="noop")
    async def ping() -> dict:
        return {"pong": True}

    t = InProcessTransport.start(bridge)
    async with t:
        client = AICCClient(t)
        async with client:
            await client.handshake()
            from aicc.errors import ToolError
            from aicc.protocol import ErrorCode

            with pytest.raises(ToolError) as ei:
                await client.call_tool("does_not_exist", {})
            assert ei.value.code == ErrorCode.TOOL_UNKNOWN
            assert not ei.value.retryable


@pytest.mark.asyncio
async def test_protocol_mismatch_raises():
    from aicc.errors import ProtocolError

    bridge = Bridge(name="x")

    @bridge.tool(description="noop")
    async def ping() -> dict:
        return {"pong": True}

    t = InProcessTransport.start(bridge)
    async with t:
        client = AICCClient(t, protocol="aicc/0.2")
        async with client:
            with pytest.raises(ProtocolError):
                await client.handshake()


@pytest.mark.asyncio
async def test_capability_required_tool_is_unavailable():
    from aicc.errors import ToolError
    from aicc.protocol import ErrorCode

    bridge = Bridge(name="x")

    @bridge.tool(
        description="place",
        cls=ToolClass.GENERATOR,
        requires_capability="can_modify_world",
    )
    async def place_object() -> dict:
        return {"spawned": True}

    t = InProcessTransport.start(bridge)
    async with t:
        client = AICCClient(t)
        async with client:
            await client.handshake()
            with pytest.raises(ToolError) as ei:
                await client.call_tool("place_object", {})
            assert ei.value.code == ErrorCode.TOOL_UNAVAILABLE


@pytest.mark.asyncio
async def test_execution_failure_is_surfaced():
    from aicc.errors import ToolError
    from aicc.protocol import ErrorCode

    bridge = Bridge(name="x")

    @bridge.tool(description="always fails")
    async def boom() -> dict:
        raise RuntimeError("kaboom")

    t = InProcessTransport.start(bridge)
    async with t:
        client = AICCClient(t)
        async with client:
            await client.handshake()
            with pytest.raises(ToolError) as ei:
                await client.call_tool("boom", {})
            assert ei.value.code == ErrorCode.EXECUTION_FAILED


@pytest.mark.asyncio
async def test_schema_generation_basic_types():
    bridge = Bridge(name="x")

    @bridge.tool(description="basic types")
    async def fn(a: int, b: str, c: float = 1.0) -> dict:
        return {"a": a, "b": b, "c": c}

    schema = bridge.tools()[0].input_schema
    assert schema["properties"]["a"]["type"] == "integer"
    assert schema["properties"]["b"]["type"] == "string"
    assert schema["properties"]["c"]["type"] == "number"
    assert "a" in schema["required"]
    assert "b" in schema["required"]
    assert "c" not in schema["required"]