from __future__ import annotations import asyncio from collections.abc import Callable from pathlib import Path from unittest.mock import AsyncMock import pytest from nanobot.agent.tools.session_messages import SendSessionMessageTool from nanobot.bus.events import InboundMessage from nanobot.bus.outbound_events import SessionMessageInputEvent from nanobot.bus.queue import MessageBus from nanobot.security.workspace_access import WORKSPACE_SCOPE_METADATA_KEY from nanobot.session.manager import SessionManager from nanobot.session.session_handles import SessionHandle, SessionHandleDirectory from nanobot.session.session_messages import ( SESSION_MESSAGE_METADATA_KEY, SESSION_REPLY_TIMEOUT_METADATA_KEY, SessionMessageError, session_message_envelope, session_message_inbound, session_reply_timeout_envelope, session_reply_timeout_inbound, ) from nanobot.session.webui_turns import project_session_message_input from nanobot.webui.transcript import read_transcript_lines class _FakeTimer: def __init__(self) -> None: self.cancelled = False def cancel(self) -> None: self.cancelled = True class _FakeScheduler: def __init__(self) -> None: self.calls: list[tuple[float, Callable[[], None], _FakeTimer]] = [] def __call__(self, delay: float, callback: Callable[[], None]) -> _FakeTimer: timer = _FakeTimer() self.calls.append((delay, callback, timer)) return timer class _FakeClock: def __init__(self) -> None: self.now = 0.0 def __call__(self) -> float: return self.now def advance(self, seconds: float) -> None: self.now += seconds class FakeSessionHandleDirectory: def __init__(self, identities: list[SessionHandle]) -> None: self._by_key = {identity.session_key: identity for identity in identities} self._by_name = {identity.name.casefold(): identity for identity in identities} def ensure(self, session_key: str) -> SessionHandle: identity = self.handle_for_session(session_key) if identity is None: raise ValueError(f"unknown session: {session_key}") return identity def resolve(self, name: str) -> SessionHandle | None: return self._by_name.get(name.casefold()) def handle_for_session(self, key: str) -> SessionHandle | None: return self._by_key.get(key) def _identity(name: str, session_key: str, workspace: Path) -> SessionHandle: color_slot = 1 if name == "lead" else 2 return SessionHandle( id=f"handle_{color_slot:032x}", name=name, color_slot=color_slot, session_key=session_key, workspace=workspace, ) def _persist(sessions: SessionManager, key: str) -> None: session = sessions.get_or_create(key) session.metadata["webui"] = True sessions.save(session) def _service( tmp_path: Path, *, max_messages_per_minute: int = 6, schedule_later: Callable[[float, Callable[[], None]], _FakeTimer] | None = None, clock: Callable[[], float] | None = None, ) -> tuple[SendSessionMessageTool, MessageBus, SessionManager]: workspace = tmp_path / "project" workspace.mkdir(exist_ok=True) sessions = SessionManager(tmp_path / "state") source = _identity("lead", "websocket:lead", workspace) target = _identity("reviewer", "websocket:reviewer", workspace) for identity in (source, target): _persist(sessions, identity.session_key) sessions.invalidate(source.session_key) sessions.invalidate(target.session_key) bus = MessageBus() return ( SendSessionMessageTool( sessions=sessions, bus=bus, directory=FakeSessionHandleDirectory([source, target]), max_messages_per_minute=max_messages_per_minute, schedule_later=schedule_later, clock=clock, ), bus, sessions, ) @pytest.mark.asyncio async def test_webui_input_is_persisted_and_projected( tmp_path: Path, monkeypatch: pytest.MonkeyPatch, ) -> None: monkeypatch.setattr("nanobot.webui.transcript.get_webui_dir", lambda: tmp_path / "webui") service, bus, _sessions = _service(tmp_path) target_handle = await service.enqueue( source_session_key="websocket:lead", target_handle="@reviewer", content="Review the implementation.", expect_reply=True, reply_timeout_seconds=60, ) assert target_handle == "@reviewer" inbound = bus.inbound.get_nowait() assert inbound.channel == "websocket" assert inbound.sender_id == "session" assert inbound.chat_id == "reviewer" assert inbound.session_key_override == "websocket:reviewer" assert inbound.require_existing_session is True assert inbound.content == "Review the implementation." assert set(inbound.metadata) == {SESSION_MESSAGE_METADATA_KEY} envelope = session_message_envelope(inbound.metadata) assert envelope is not None assert session_message_inbound(inbound) == envelope message_id = envelope["message_id"] assert envelope["expect_reply"] is True assert set(envelope) == { "message_id", "created_at_ms", "expect_reply", "source", "target", } assert envelope["source"] == { "name": "lead", "session_key": "websocket:lead", "handle_id": "handle_00000000000000000000000000000001", "color_slot": 1, } assert envelope["target"] == { "name": "reviewer", "session_key": "websocket:reviewer", } assert bus.outbound.empty() await project_session_message_input(bus, inbound, "websocket:reviewer") live = bus.outbound.get_nowait() assert (live.channel, live.chat_id) == ("websocket", "reviewer") assert isinstance(live.event, SessionMessageInputEvent) assert live.event.content == "Review the implementation." assert live.event.session_message == { "direction": "incoming", "message_id": message_id, "session": { "id": "handle_00000000000000000000000000000001", "name": "lead", "color_slot": 1, }, } assert "websocket:" not in str(live.event.session_message) assert bus.outbound.empty() transcript = read_transcript_lines("websocket:reviewer") assert len(transcript) == 1 assert transcript[0]["text"] == "Review the implementation." assert transcript[0]["session_message"] == live.event.session_message @pytest.mark.asyncio async def test_projection_publishes_when_transcript_persistence_fails( tmp_path: Path, monkeypatch: pytest.MonkeyPatch, ) -> None: monkeypatch.setattr("nanobot.webui.transcript.get_webui_dir", lambda: tmp_path / "webui") service, bus, _sessions = _service(tmp_path) await service.enqueue( source_session_key="websocket:lead", target_handle="@reviewer", content="Review the implementation.", expect_reply=False, ) inbound = bus.inbound.get_nowait() attempts = 0 def fail_append(*args: object, **kwargs: object) -> None: nonlocal attempts attempts += 1 raise OSError("write failed") monkeypatch.setattr( "nanobot.session.webui_turns.append_session_message_input", fail_append, ) await project_session_message_input(bus, inbound, "websocket:reviewer") assert attempts == 1 assert bus.outbound.qsize() == 1 assert read_transcript_lines("websocket:reviewer") == [] @pytest.mark.asyncio async def test_enqueue_accepts_persisted_target_that_is_not_cached(tmp_path: Path) -> None: service, bus, sessions = _service(tmp_path) assert sessions.get_cached("websocket:reviewer") is None await service.enqueue( source_session_key="websocket:lead", target_handle="reviewer", content="Ping", expect_reply=False, ) assert bus.inbound_size == 1 assert sessions.get_cached("websocket:reviewer") is None @pytest.mark.asyncio async def test_enqueue_supports_non_webui_sessions( tmp_path: Path, ) -> None: workspace = tmp_path / "project" workspace.mkdir() sessions = SessionManager(tmp_path / "state") source_key = "telegram:source" source = sessions.get_or_create(source_key) source.metadata[WORKSPACE_SCOPE_METADATA_KEY] = { "project_path": str(workspace), "access_mode": "restricted", } sessions.save(source) target_key = "telegram:target" target_session = sessions.get_or_create(target_key) target_session.metadata[WORKSPACE_SCOPE_METADATA_KEY] = { "project_path": str(workspace), "access_mode": "restricted", } sessions.save(target_session) directory = SessionHandleDirectory(sessions) target = directory.ensure_many([target_key])[target_key] bus = MessageBus() service = SendSessionMessageTool( sessions=sessions, bus=bus, directory=directory, ) await service.enqueue( source_session_key=source_key, target_handle=target.name, content="Ping", expect_reply=False, ) inbound = bus.inbound.get_nowait() envelope = session_message_inbound(inbound) assert envelope is not None assert envelope["source"]["session_key"] == source_key assert envelope["target"]["session_key"] == target_key assert directory.handle_for_session(source_key) is not None @pytest.mark.asyncio async def test_enqueue_uses_persistent_session_handles(tmp_path: Path) -> None: workspace = tmp_path / "project" workspace.mkdir() sessions = SessionManager(tmp_path / "state") for key, title in ( ("websocket:lead", "Lead"), ("websocket:reviewer", "Reviewer"), ): session = sessions.get_or_create(key) session.metadata.update({ "title": title, "webui": True, WORKSPACE_SCOPE_METADATA_KEY: { "project_path": str(workspace), "access_mode": "restricted", }, }) sessions.save(session) directory = SessionHandleDirectory(sessions) handles = directory.ensure_many(["websocket:lead", "websocket:reviewer"]) source = handles["websocket:lead"] target = handles["websocket:reviewer"] bus = MessageBus() service = SendSessionMessageTool(sessions=sessions, bus=bus, directory=directory) target_handle = await service.enqueue( source_session_key=source.session_key, target_handle=f"@{target.name}", content="Please review this.", expect_reply=True, reply_timeout_seconds=60, ) assert target_handle == f"@{target.name}" inbound = bus.inbound.get_nowait() envelope = session_message_envelope(inbound.metadata) assert envelope is not None assert envelope["source"]["handle_id"] == source.id assert envelope["target"]["session_key"] == target.session_key @pytest.mark.asyncio async def test_enqueue_rejects_stale_target_before_bus_mutation(tmp_path: Path) -> None: service, bus, sessions = _service(tmp_path) sessions.delete_session("websocket:reviewer") with pytest.raises(SessionMessageError, match="not persisted") as exc_info: await service.enqueue( source_session_key="websocket:lead", target_handle="reviewer", content="Ping", expect_reply=False, ) assert exc_info.value.code == "target_not_found" assert bus.inbound_size == 0 @pytest.mark.asyncio async def test_enqueue_allows_self_send(tmp_path: Path) -> None: workspace = tmp_path / "project" workspace.mkdir() sessions = SessionManager(tmp_path / "state") source = _identity("lead", "websocket:lead", workspace) _persist(sessions, source.session_key) bus = MessageBus() service = SendSessionMessageTool( sessions=sessions, bus=bus, directory=FakeSessionHandleDirectory([source]), ) await service.enqueue( source_session_key=source.session_key, target_handle="@lead", content="Loop", expect_reply=False, ) assert bus.inbound_size == 1 @pytest.mark.asyncio async def test_enqueue_accepts_cross_workspace_target(tmp_path: Path) -> None: workspace = tmp_path / "project" other = tmp_path / "other" workspace.mkdir() other.mkdir() sessions = SessionManager(tmp_path / "state") source = _identity("lead", "websocket:lead", workspace) target = _identity("reviewer", "websocket:reviewer", other) for identity in (source, target): _persist(sessions, identity.session_key) bus = MessageBus() scheduler = _FakeScheduler() service = SendSessionMessageTool( sessions=sessions, bus=bus, directory=FakeSessionHandleDirectory([source, target]), schedule_later=scheduler, ) await service.enqueue( source_session_key=source.session_key, target_handle="reviewer", content="Ping", expect_reply=True, reply_timeout_seconds=60, ) inbound = bus.inbound.get_nowait() envelope = session_message_inbound(inbound) assert envelope is not None assert envelope["source"]["session_key"] == source.session_key assert envelope["target"]["session_key"] == target.session_key scheduler.calls[0][1]() await asyncio.sleep(0) await asyncio.sleep(0) timeout = bus.inbound.get_nowait() assert session_reply_timeout_inbound(timeout) is not None @pytest.mark.asyncio async def test_enqueue_enforces_per_session_minute_limit(tmp_path: Path) -> None: clock = _FakeClock() service, bus, _sessions = _service( tmp_path, max_messages_per_minute=2, clock=clock, ) for index in range(2): await service.enqueue( source_session_key="websocket:lead", target_handle="reviewer", content=f"Ping {index}", expect_reply=False, ) with pytest.raises(SessionMessageError) as exc_info: await service.enqueue( source_session_key="websocket:lead", target_handle="reviewer", content="Too many", expect_reply=False, ) assert exc_info.value.code == "rate_limited" await service.enqueue( source_session_key="websocket:reviewer", target_handle="lead", content="Independent sender", expect_reply=False, ) clock.advance(60) await service.enqueue( source_session_key="websocket:lead", target_handle="reviewer", content="New window", expect_reply=False, ) assert bus.inbound_size == 4 @pytest.mark.asyncio async def test_enqueue_publish_failure_does_not_consume_rate_limit( tmp_path: Path, ) -> None: service, bus, _sessions = _service(tmp_path, max_messages_per_minute=1) original_publish = bus.publish_inbound bus.publish_inbound = AsyncMock(side_effect=RuntimeError("bus unavailable")) with pytest.raises(RuntimeError, match="bus unavailable"): await service.enqueue( source_session_key="websocket:lead", target_handle="reviewer", content="first attempt", expect_reply=False, ) bus.publish_inbound = original_publish target_handle = await service.enqueue( source_session_key="websocket:lead", target_handle="reviewer", content="retry", expect_reply=False, ) assert target_handle == "@reviewer" assert bus.inbound_size == 1 @pytest.mark.asyncio async def test_enqueue_requires_timeout_only_for_requested_replies(tmp_path: Path) -> None: service, bus, _sessions = _service(tmp_path) with pytest.raises(SessionMessageError) as exc_info: await service.enqueue( source_session_key="websocket:lead", target_handle="reviewer", content="Please reply", expect_reply=True, ) assert exc_info.value.code == "invalid_reply_timeout" with pytest.raises(SessionMessageError) as exc_info: await service.enqueue( source_session_key="websocket:lead", target_handle="reviewer", content="Please reply", expect_reply=True, reply_timeout_seconds=61, ) assert exc_info.value.code == "invalid_reply_timeout" with pytest.raises(SessionMessageError) as exc_info: await service.enqueue( source_session_key="websocket:lead", target_handle="reviewer", content="No reply needed", expect_reply=False, reply_timeout_seconds=60, ) assert exc_info.value.code == "unexpected_reply_timeout" assert bus.inbound_size == 0 @pytest.mark.asyncio async def test_requested_reply_timeout_resumes_the_waiting_session(tmp_path: Path) -> None: scheduler = _FakeScheduler() service, bus, _sessions = _service(tmp_path, schedule_later=scheduler) await service.enqueue( source_session_key="websocket:lead", target_handle="reviewer", content="Please reply", expect_reply=True, reply_timeout_seconds=60, ) bus.inbound.get_nowait() assert len(scheduler.calls) == 1 delay, expire, timer = scheduler.calls[0] assert delay == 60 assert timer.cancelled is False expire() await asyncio.sleep(0) await asyncio.sleep(0) timeout_message = bus.inbound.get_nowait() timeout = session_reply_timeout_envelope(timeout_message.metadata) assert timeout is not None assert session_reply_timeout_inbound(timeout_message) == timeout assert timeout["timeout_seconds"] == 60 assert timeout["source"]["session_key"] == "websocket:lead" assert timeout["target"]["session_key"] == "websocket:reviewer" @pytest.mark.asyncio async def test_session_reply_cancels_its_pending_timeout(tmp_path: Path) -> None: scheduler = _FakeScheduler() service, bus, _sessions = _service(tmp_path, schedule_later=scheduler) await service.enqueue( source_session_key="websocket:lead", target_handle="reviewer", content="Please reply", expect_reply=True, reply_timeout_seconds=60, ) bus.inbound.get_nowait() timer = scheduler.calls[0][2] await service.enqueue( source_session_key="websocket:reviewer", target_handle="lead", content="Here is the answer", expect_reply=False, ) assert timer.cancelled is True reply = bus.inbound.get_nowait() assert session_message_inbound(reply) is not None @pytest.mark.asyncio async def test_new_reply_wait_replaces_the_previous_wait_for_the_same_session_pair( tmp_path: Path, ) -> None: scheduler = _FakeScheduler() service, bus, _sessions = _service(tmp_path, schedule_later=scheduler) await service.enqueue( source_session_key="websocket:reviewer", target_handle="lead", content="First question", expect_reply=True, reply_timeout_seconds=60, ) await service.enqueue( source_session_key="websocket:reviewer", target_handle="lead", content="Second question", expect_reply=True, reply_timeout_seconds=30, ) assert bus.inbound_size == 2 assert len(scheduler.calls) == 2 assert scheduler.calls[0][2].cancelled is True assert scheduler.calls[1][2].cancelled is False scheduler.calls[0][1]() await asyncio.sleep(0) await asyncio.sleep(0) assert bus.inbound_size == 2 def test_session_message_envelope_rejects_dynamic_boundary_violations() -> None: assert session_message_envelope(None) is None assert session_message_envelope({SESSION_MESSAGE_METADATA_KEY: {}}) is None metadata = _session_message_metadata() envelope = metadata[SESSION_MESSAGE_METADATA_KEY] assert isinstance(envelope, dict) envelope.pop("expect_reply") assert session_message_envelope(metadata) is None metadata = _session_message_metadata() envelope = metadata[SESSION_MESSAGE_METADATA_KEY] assert isinstance(envelope, dict) envelope["expect_reply"] = 1 assert session_message_envelope(metadata) is None def test_session_inbound_checks_sender_and_route_not_lifecycle_policy() -> None: metadata = _session_message_metadata() internal = InboundMessage( channel="system", sender_id="session", chat_id="websocket:reviewer", content="Review this", metadata=metadata, session_key_override="websocket:reviewer", ) forged = InboundMessage( channel="websocket", sender_id="user", chat_id="reviewer", content="/stop", metadata=metadata, session_key_override="websocket:reviewer", require_existing_session=True, ) wrong_target = InboundMessage( channel="system", sender_id="session", chat_id="websocket:other", content="/stop", metadata=metadata, session_key_override="websocket:other", require_existing_session=True, ) assert session_message_inbound(internal) is not None assert session_message_inbound(forged) is None assert session_message_inbound(wrong_target) is None def test_session_reply_timeout_checks_sender_and_route_not_lifecycle_policy() -> None: metadata = _session_message_metadata() request = metadata[SESSION_MESSAGE_METADATA_KEY] assert isinstance(request, dict) timeout_metadata = { SESSION_REPLY_TIMEOUT_METADATA_KEY: { **request, "timeout_seconds": 60, }, } internal = InboundMessage( channel="system", sender_id="session_timeout", chat_id="websocket:lead", content="", metadata=timeout_metadata, session_key_override="websocket:lead", ) forged = InboundMessage( channel="websocket", sender_id="user", chat_id="lead", content="", metadata=timeout_metadata, session_key_override="websocket:lead", require_existing_session=True, ) assert session_reply_timeout_envelope(timeout_metadata) is not None assert session_reply_timeout_inbound(internal) is not None assert session_reply_timeout_inbound(forged) is None request["expect_reply"] = False assert session_reply_timeout_envelope({ SESSION_REPLY_TIMEOUT_METADATA_KEY: { **request, "timeout_seconds": 60, }, }) is None def _session_message_metadata() -> dict[str, object]: return { SESSION_MESSAGE_METADATA_KEY: { "message_id": "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": "websocket:reviewer", }, } }