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.session_messages import ( SESSION_MESSAGE_METADATA_KEY, SESSION_REPLY_TIMEOUT_METADATA_KEY, ) from nanobot.session.webui_turns import WebuiTurnRoutePolicy from nanobot.webui.metadata import ( WEBSOCKET_TURN_OWNER_METADATA_KEY, WEBUI_MESSAGE_SOURCE_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", } def test_session_route_targets_its_webui_session_with_source_provenance( tmp_path: Path, monkeypatch: pytest.MonkeyPatch, ) -> None: monkeypatch.setattr("nanobot.config.paths.get_data_dir", lambda: tmp_path) sessions = SessionManager(tmp_path) session_key = "websocket:reviewer" session = sessions.get_or_create(session_key) session.metadata["webui"] = True sessions.save(session) envelope = { "message_id": "handle-message-1", "created_at_ms": 1, "expect_reply": True, "source": { "name": "lead", "session_key": "websocket:lead", "handle_id": "handle_00000000000000000000000000000001", "color_slot": 1, }, "target": { "name": "reviewer", "session_key": session_key, }, } msg = InboundMessage( channel="system", sender_id="session", chat_id=session_key, content="Review this", session_key_override=session_key, require_existing_session=True, metadata={SESSION_MESSAGE_METADATA_KEY: envelope}, ) factory = TurnDeliveryFactory( MessageBus(), RuntimeEventBus(), route_policy=WebuiTurnRoutePolicy(sessions), ) route = factory.create(msg, session_key).route assert route.channel == "websocket" assert route.chat_id == "reviewer" assert route.publish_lifecycle is True assert route.metadata[SESSION_MESSAGE_METADATA_KEY] == envelope assert route.metadata[WEBUI_MESSAGE_SOURCE_METADATA_KEY] == { "kind": "session", "label": "@lead", } assert route.metadata[WEBUI_TURN_METADATA_KEY].startswith("session-message:") assert msg.metadata == {SESSION_MESSAGE_METADATA_KEY: envelope} def test_session_timeout_route_resumes_its_webui_session(tmp_path: Path) -> None: sessions = SessionManager(tmp_path) session_key = "websocket:lead" session = sessions.get_or_create(session_key) session.metadata["webui"] = True sessions.save(session) envelope = { "message_id": "handle-message-1", "created_at_ms": 1, "expect_reply": True, "timeout_seconds": 60, "source": { "name": "lead", "session_key": session_key, "handle_id": "handle_00000000000000000000000000000001", "color_slot": 1, }, "target": { "name": "reviewer", "session_key": "websocket:reviewer", }, } msg = InboundMessage( channel="system", sender_id="session_timeout", chat_id=session_key, content="", session_key_override=session_key, require_existing_session=True, metadata={SESSION_REPLY_TIMEOUT_METADATA_KEY: envelope}, ) factory = TurnDeliveryFactory( MessageBus(), RuntimeEventBus(), route_policy=WebuiTurnRoutePolicy(sessions), ) route = factory.create(msg, session_key).route assert route.channel == "websocket" assert route.chat_id == "lead" assert route.publish_lifecycle is True assert route.metadata[WEBUI_TURN_METADATA_KEY].startswith("session-reply-timeout:") assert WEBUI_MESSAGE_SOURCE_METADATA_KEY not in route.metadata def test_session_route_does_not_create_a_missing_target_session( tmp_path: Path, monkeypatch: pytest.MonkeyPatch, ) -> None: monkeypatch.setattr("nanobot.config.paths.get_data_dir", lambda: tmp_path) sessions = SessionManager(tmp_path) session_key = "websocket:deleted" envelope = { "message_id": "handle-message-deleted", "created_at_ms": 1, "expect_reply": True, "source": { "name": "lead", "session_key": "websocket:lead", "handle_id": "handle_00000000000000000000000000000001", "color_slot": 1, }, "target": { "name": "deleted", "session_key": session_key, }, } msg = InboundMessage( channel="system", sender_id="session", chat_id=session_key, content="Review this", session_key_override=session_key, require_existing_session=True, metadata={SESSION_MESSAGE_METADATA_KEY: envelope}, ) factory = TurnDeliveryFactory( MessageBus(), RuntimeEventBus(), route_policy=WebuiTurnRoutePolicy(sessions), ) route = factory.create(msg, session_key).route assert route.publish_lifecycle is False assert sessions.get_cached(session_key) is None assert sessions.read_session_metadata(session_key) is None