nanobot/tests/agent/test_turn_delivery.py

from pathlib import Path

import pytest

from nanobot.agent.turn_delivery import TurnDeliveryFactory
from nanobot.bus.events import InboundMessage
from nanobot.bus.queue import MessageBus
from nanobot.bus.runtime_events import RuntimeEventBus
from nanobot.session.manager import SessionManager
from nanobot.session.webui_turns import WebuiTurnRoutePolicy
from nanobot.webui.metadata import (
    WEBSOCKET_TURN_OWNER_METADATA_KEY,
    WEBUI_TURN_METADATA_KEY,
)


def test_websocket_lifecycles_get_distinct_internal_owners(tmp_path: Path) -> None:
    factory = TurnDeliveryFactory(
        MessageBus(),
        RuntimeEventBus(),
        route_policy=WebuiTurnRoutePolicy(SessionManager(tmp_path / "sessions")),
    )
    first_msg = InboundMessage(
        channel="websocket",
        sender_id="user",
        chat_id="chat-a",
        content="first",
        metadata={WEBSOCKET_TURN_OWNER_METADATA_KEY: "attacker-reused-owner"},
    )
    second_msg = InboundMessage(
        channel="websocket",
        sender_id="user",
        chat_id="chat-a",
        content="second",
        metadata={WEBSOCKET_TURN_OWNER_METADATA_KEY: "attacker-reused-owner"},
    )

    first = factory.create(first_msg, first_msg.session_key)
    second = factory.create(second_msg, second_msg.session_key)
    first_owner = first.lifecycle_message.metadata[WEBSOCKET_TURN_OWNER_METADATA_KEY]
    second_owner = second.lifecycle_message.metadata[WEBSOCKET_TURN_OWNER_METADATA_KEY]

    assert first_owner == first.delivery_message.metadata[WEBSOCKET_TURN_OWNER_METADATA_KEY]
    assert first_owner != second_owner
    assert first_owner != "attacker-reused-owner"
    assert second_owner != "attacker-reused-owner"
    assert WEBUI_TURN_METADATA_KEY not in first.lifecycle_message.metadata
    assert first_msg.metadata[WEBSOCKET_TURN_OWNER_METADATA_KEY] == first_owner
    assert second_msg.metadata[WEBSOCKET_TURN_OWNER_METADATA_KEY] == second_owner


def test_websocket_lifecycle_reuses_registered_ingress_owner(tmp_path: Path) -> None:
    from nanobot.session import webui_turns as wth

    owner = wth.register_queued_websocket_turn_if_idle("chat-queued", "turn-queued")
    assert owner is not None
    msg = InboundMessage(
        channel="websocket",
        sender_id="user",
        chat_id="chat-queued",
        content="queued",
        metadata={
            WEBSOCKET_TURN_OWNER_METADATA_KEY: owner,
            WEBUI_TURN_METADATA_KEY: "turn-queued",
        },
    )
    factory = TurnDeliveryFactory(
        MessageBus(),
        RuntimeEventBus(),
        route_policy=WebuiTurnRoutePolicy(SessionManager(tmp_path / "sessions")),
    )

    try:
        delivery = factory.create(msg, msg.session_key)

        assert delivery.lifecycle_message.metadata[WEBSOCKET_TURN_OWNER_METADATA_KEY] == owner
        assert msg.metadata[WEBSOCKET_TURN_OWNER_METADATA_KEY] == owner
    finally:
        wth.clear_websocket_turn_if_current("chat-queued", owner)


@pytest.mark.asyncio
async def test_same_chat_different_sessions_restore_previous_active_projection(
    tmp_path: Path,
) -> None:
    from unittest.mock import AsyncMock, MagicMock

    from nanobot.session import webui_turns as wth

    factory = TurnDeliveryFactory(
        MessageBus(),
        RuntimeEventBus(),
        route_policy=WebuiTurnRoutePolicy(SessionManager(tmp_path / "sessions")),
    )
    first_msg = InboundMessage(
        channel="websocket",
        sender_id="user",
        chat_id="shared-chat",
        content="first",
        metadata={WEBUI_TURN_METADATA_KEY: "turn-first"},
        session_key_override="websocket:session-first",
    )
    second_msg = InboundMessage(
        channel="websocket",
        sender_id="user",
        chat_id="shared-chat",
        content="second",
        metadata={WEBUI_TURN_METADATA_KEY: "turn-second"},
        session_key_override="websocket:session-second",
    )
    first = factory.create(first_msg, first_msg.session_key)
    second = factory.create(second_msg, second_msg.session_key)
    first_owner = first.lifecycle_message.metadata[WEBSOCKET_TURN_OWNER_METADATA_KEY]
    second_owner = second.lifecycle_message.metadata[WEBSOCKET_TURN_OWNER_METADATA_KEY]
    bus = MagicMock()
    bus.publish_outbound = AsyncMock()

    try:
        await wth.publish_turn_run_status(
            bus,
            first.lifecycle_message,
            "running",
            started_at=100.0,
        )
        await wth.publish_turn_run_status(
            bus,
            second.lifecycle_message,
            "running",
            started_at=200.0,
        )

        assert wth.websocket_turn_wall_started_at("shared-chat") == 200.0
        assert wth.websocket_turn_id("shared-chat") == "turn-second"
        assert wth.clear_websocket_turn_if_current("shared-chat", second_owner) is True
        assert wth.websocket_turn_wall_started_at("shared-chat") == 100.0
        assert wth.websocket_turn_id("shared-chat") == "turn-first"
        assert wth._WEBSOCKET_TURN_OWNERS["shared-chat"] == first_owner
        assert wth.clear_websocket_turn_if_current("shared-chat", first_owner) is True
        assert wth.websocket_turn_wall_started_at("shared-chat") is None
    finally:
        wth._WEBSOCKET_ACTIVE_TURNS.pop("shared-chat", None)
        wth._WEBSOCKET_TURN_WALL_STARTED_AT.pop("shared-chat", None)
        wth._WEBSOCKET_TURN_IDS.pop("shared-chat", None)
        wth._WEBSOCKET_TURN_OWNERS.pop("shared-chat", None)


def test_late_subagent_route_requires_webui_owned_session(tmp_path: Path) -> None:
    sessions = SessionManager(tmp_path)
    factory = TurnDeliveryFactory(
        MessageBus(),
        RuntimeEventBus(),
        route_policy=WebuiTurnRoutePolicy(sessions),
    )
    session_key = "websocket:chat-a"
    msg = InboundMessage(
        channel="system",
        sender_id="subagent",
        chat_id=session_key,
        content="Background research completed",
        session_key_override=session_key,
        metadata={
            "injected_event": "subagent_result",
            "subagent_task_id": "sub-1",
        },
    )

    hidden_route = factory.create(msg, session_key).route

    assert hidden_route.channel == "websocket"
    assert hidden_route.chat_id == "chat-a"
    assert hidden_route.metadata == {}
    assert hidden_route.publish_lifecycle is False

    session = sessions.get_or_create(session_key)
    session.metadata["webui"] = True
    first_visible_route = factory.create(msg, session_key).route
    second_visible_route = factory.create(msg, session_key).route

    assert first_visible_route.publish_lifecycle is True
    assert set(first_visible_route.metadata) == {
        "webui",
        "_wants_stream",
        WEBSOCKET_TURN_OWNER_METADATA_KEY,
        WEBUI_TURN_METADATA_KEY,
    }
    assert first_visible_route.metadata["webui"] is True
    assert first_visible_route.metadata["_wants_stream"] is True
    first_turn_id = first_visible_route.metadata[WEBUI_TURN_METADATA_KEY]
    second_turn_id = second_visible_route.metadata[WEBUI_TURN_METADATA_KEY]
    assert first_turn_id.startswith("subagent:")
    assert second_turn_id.startswith("subagent:")
    assert first_turn_id != second_turn_id
    assert (
        first_visible_route.metadata[WEBSOCKET_TURN_OWNER_METADATA_KEY]
        != second_visible_route.metadata[WEBSOCKET_TURN_OWNER_METADATA_KEY]
    )
    assert msg.metadata == {
        "injected_event": "subagent_result",
        "subagent_task_id": "sub-1",
    }