livekit_agents/test_beaver_terminal_client.py

import asyncio
import json
import sys
from pathlib import Path

import aiohttp
import pytest
from aiohttp import web

if __name__ == "__main__":
    sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
    raise SystemExit(pytest.main([__file__]))

try:
    from custom.beaver_terminal_client import (
        BeaverTerminalClient,
        BeaverTerminalError,
        MessageIdGenerator,
        build_connect_frame,
        build_message_frame,
    )
except ModuleNotFoundError:
    from beaver_terminal_client import (
        BeaverTerminalClient,
        BeaverTerminalError,
        MessageIdGenerator,
        build_connect_frame,
        build_message_frame,
    )


def test_build_connect_frame_uses_stable_peer_id() -> None:
    frame = build_connect_frame(peer_id="device-001", device_name="desk-terminal")

    assert frame == {
        "type": "connect",
        "peer_id": "device-001",
        "device_name": "desk-terminal",
        "capabilities": ["text"],
    }


def test_build_message_frame_uses_message_id_and_text() -> None:
    frame = build_message_frame(message_id="device-001-000001", text="hello")

    assert frame == {
        "type": "message",
        "message_id": "device-001-000001",
        "text": "hello",
    }


def test_message_id_generator_uses_monotonic_peer_counter() -> None:
    generator = MessageIdGenerator(peer_id="device-001", initial_counter=7)

    assert generator.next_id() == "device-001-000008"
    assert generator.next_id() == "device-001-000009"
    assert generator.counter == 9


async def test_client_connects_sends_text_and_returns_assistant_reply(
    unused_tcp_port: int,
) -> None:
    received: list[dict[str, object]] = []

    async def websocket_handler(request: web.Request) -> web.WebSocketResponse:
        ws = web.WebSocketResponse()
        await ws.prepare(request)

        async for message in ws:
            assert message.type == aiohttp.WSMsgType.TEXT
            frame = json.loads(message.data)
            received.append(frame)

            if frame["type"] == "connect":
                await ws.send_json(
                    {
                        "type": "connected",
                        "channel_id": "terminal-dev",
                        "session_id": "terminal-dev:local:device-001",
                    }
                )
            elif frame["type"] == "message":
                await ws.send_json(
                    {
                        "type": "ack",
                        "message_id": frame["message_id"],
                        "session_id": "terminal-dev:local:device-001",
                        "accepted": True,
                    }
                )
                await ws.send_json(
                    {
                        "type": "message",
                        "role": "assistant",
                        "message_id": frame["message_id"],
                        "run_id": "run-1",
                        "text": "assistant reply",
                        "finish_reason": "stop",
                    }
                )

        return ws

    app = web.Application()
    app.router.add_get("/api/channels/terminal-dev/ws", websocket_handler)
    runner = web.AppRunner(app)
    await runner.setup()
    site = web.TCPSite(runner, "127.0.0.1", unused_tcp_port)
    await site.start()

    client = BeaverTerminalClient(
        url=f"http://127.0.0.1:{unused_tcp_port}/api/channels/terminal-dev/ws",
        peer_id="device-001",
        device_name="desk-terminal",
    )

    try:
        await client.connect()
        reply = await client.send_text("hello")
    finally:
        await client.close()
        await runner.cleanup()

    assert client.session_id == "terminal-dev:local:device-001"
    assert reply == "assistant reply"
    assert received == [
        {
            "type": "connect",
            "peer_id": "device-001",
            "device_name": "desk-terminal",
            "capabilities": ["text"],
        },
        {
            "type": "message",
            "message_id": "device-001-000001",
            "text": "hello",
        },
    ]


async def test_client_returns_cached_duplicate_reply(unused_tcp_port: int) -> None:
    async def websocket_handler(request: web.Request) -> web.WebSocketResponse:
        ws = web.WebSocketResponse()
        await ws.prepare(request)

        async for message in ws:
            frame = json.loads(message.data)
            if frame["type"] == "connect":
                await ws.send_json(
                    {
                        "type": "connected",
                        "channel_id": "terminal-dev",
                        "session_id": "terminal-dev:local:device-001",
                    }
                )
            elif frame["type"] == "message":
                await ws.send_json(
                    {
                        "type": "ack",
                        "message_id": frame["message_id"],
                        "session_id": "terminal-dev:local:device-001",
                        "accepted": False,
                        "duplicate": True,
                        "pending": False,
                        "reply": "cached assistant reply",
                    }
                )

        return ws

    app = web.Application()
    app.router.add_get("/api/channels/terminal-dev/ws", websocket_handler)
    runner = web.AppRunner(app)
    await runner.setup()
    site = web.TCPSite(runner, "127.0.0.1", unused_tcp_port)
    await site.start()

    client = BeaverTerminalClient(
        url=f"http://127.0.0.1:{unused_tcp_port}/api/channels/terminal-dev/ws",
        peer_id="device-001",
        device_name="desk-terminal",
    )

    try:
        await client.connect()
        reply = await client.send_text("hello")
    finally:
        await client.close()
        await runner.cleanup()

    assert reply == "cached assistant reply"


async def test_client_raises_on_error_frames(unused_tcp_port: int) -> None:
    async def websocket_handler(request: web.Request) -> web.WebSocketResponse:
        ws = web.WebSocketResponse()
        await ws.prepare(request)

        async for message in ws:
            frame = json.loads(message.data)
            if frame["type"] == "connect":
                await ws.send_json(
                    {
                        "type": "connected",
                        "channel_id": "terminal-dev",
                        "session_id": "terminal-dev:local:device-001",
                    }
                )
            elif frame["type"] == "message":
                await ws.send_json({"type": "error", "error": "text is required"})

        return ws

    app = web.Application()
    app.router.add_get("/api/channels/terminal-dev/ws", websocket_handler)
    runner = web.AppRunner(app)
    await runner.setup()
    site = web.TCPSite(runner, "127.0.0.1", unused_tcp_port)
    await site.start()

    client = BeaverTerminalClient(
        url=f"http://127.0.0.1:{unused_tcp_port}/api/channels/terminal-dev/ws",
        peer_id="device-001",
        device_name="desk-terminal",
    )

    try:
        await client.connect()
        with pytest.raises(BeaverTerminalError, match="text is required"):
            await client.send_text("hello")
    finally:
        await client.close()
        await runner.cleanup()


async def test_client_treats_assistant_finish_reason_error_as_failed_turn(
    unused_tcp_port: int,
) -> None:
    async def websocket_handler(request: web.Request) -> web.WebSocketResponse:
        ws = web.WebSocketResponse()
        await ws.prepare(request)

        async for message in ws:
            frame = json.loads(message.data)
            if frame["type"] == "connect":
                await ws.send_json(
                    {
                        "type": "connected",
                        "channel_id": "terminal-dev",
                        "session_id": "terminal-dev:local:device-001",
                    }
                )
            elif frame["type"] == "message":
                await ws.send_json(
                    {
                        "type": "ack",
                        "message_id": frame["message_id"],
                        "session_id": "terminal-dev:local:device-001",
                        "accepted": True,
                    }
                )
                await ws.send_json(
                    {
                        "type": "message",
                        "role": "assistant",
                        "message_id": frame["message_id"],
                        "run_id": "run-1",
                        "text": "failed turn",
                        "finish_reason": "error",
                    }
                )

        return ws

    app = web.Application()
    app.router.add_get("/api/channels/terminal-dev/ws", websocket_handler)
    runner = web.AppRunner(app)
    await runner.setup()
    site = web.TCPSite(runner, "127.0.0.1", unused_tcp_port)
    await site.start()

    client = BeaverTerminalClient(
        url=f"http://127.0.0.1:{unused_tcp_port}/api/channels/terminal-dev/ws",
        peer_id="device-001",
        device_name="desk-terminal",
    )

    try:
        await client.connect()
        with pytest.raises(BeaverTerminalError, match="failed turn"):
            await client.send_text("hello")
    finally:
        await client.close()
        await runner.cleanup()


async def test_client_ping_sends_ping_and_waits_for_pong(unused_tcp_port: int) -> None:
    async def websocket_handler(request: web.Request) -> web.WebSocketResponse:
        ws = web.WebSocketResponse()
        await ws.prepare(request)

        async for message in ws:
            frame = json.loads(message.data)
            if frame["type"] == "connect":
                await ws.send_json(
                    {
                        "type": "connected",
                        "channel_id": "terminal-dev",
                        "session_id": "terminal-dev:local:device-001",
                    }
                )
            elif frame["type"] == "ping":
                await ws.send_json({"type": "pong"})

        return ws

    app = web.Application()
    app.router.add_get("/api/channels/terminal-dev/ws", websocket_handler)
    runner = web.AppRunner(app)
    await runner.setup()
    site = web.TCPSite(runner, "127.0.0.1", unused_tcp_port)
    await site.start()

    client = BeaverTerminalClient(
        url=f"http://127.0.0.1:{unused_tcp_port}/api/channels/terminal-dev/ws",
        peer_id="device-001",
        device_name="desk-terminal",
    )

    try:
        await client.connect()
        assert await client.ping()
    finally:
        await client.close()
        await runner.cleanup()


async def test_client_reconnects_with_same_peer_id_when_socket_closes_before_send(
    unused_tcp_port: int,
) -> None:
    connect_peer_ids: list[str] = []
    connection_count = 0

    async def websocket_handler(request: web.Request) -> web.WebSocketResponse:
        nonlocal connection_count
        connection_count += 1
        current_connection = connection_count
        ws = web.WebSocketResponse()
        await ws.prepare(request)

        async for message in ws:
            frame = json.loads(message.data)
            if frame["type"] == "connect":
                connect_peer_ids.append(frame["peer_id"])
                await ws.send_json(
                    {
                        "type": "connected",
                        "channel_id": "terminal-dev",
                        "session_id": "terminal-dev:local:device-001",
                    }
                )
                if current_connection == 1:
                    await ws.close()
            elif frame["type"] == "message":
                await ws.send_json(
                    {
                        "type": "ack",
                        "message_id": frame["message_id"],
                        "session_id": "terminal-dev:local:device-001",
                        "accepted": True,
                    }
                )
                await ws.send_json(
                    {
                        "type": "message",
                        "role": "assistant",
                        "message_id": frame["message_id"],
                        "run_id": "run-2",
                        "text": "reply after reconnect",
                        "finish_reason": "stop",
                    }
                )

        return ws

    app = web.Application()
    app.router.add_get("/api/channels/terminal-dev/ws", websocket_handler)
    runner = web.AppRunner(app)
    await runner.setup()
    site = web.TCPSite(runner, "127.0.0.1", unused_tcp_port)
    await site.start()

    client = BeaverTerminalClient(
        url=f"http://127.0.0.1:{unused_tcp_port}/api/channels/terminal-dev/ws",
        peer_id="device-001",
        device_name="desk-terminal",
    )

    try:
        await client.connect()
        await asyncio.sleep(0.01)
        reply = await client.send_text("hello")
    finally:
        await client.close()
        await runner.cleanup()

    assert reply == "reply after reconnect"
    assert connect_peer_ids == ["device-001", "device-001"]