"""Unit tests for WebSocket handler internals and reconnect logic."""
import asyncio
import json
from unittest.mock import AsyncMock, MagicMock
import pytest
from fastapi import WebSocketDisconnect
from navi.api import websocket as ws_mod
from navi.core.orchestrator import AgentSessionOrchestrator, SessionRun
@pytest.fixture(autouse=True)
def _clear_state(monkeypatch):
"""Clear global state before every WS test."""
yield
@pytest.fixture
def mock_websocket():
ws = AsyncMock()
ws.accept = AsyncMock()
ws.close = AsyncMock()
ws.send_json = AsyncMock()
return ws
@pytest.fixture
def mock_session():
session = MagicMock()
session.user_id = "test-user-id"
return session
@pytest.fixture
def mock_user():
user = MagicMock()
user.id = "test-user-id"
user.role = "admin"
return user
@pytest.fixture
def fake_orchestrator():
container = MagicMock()
container.profile_registry = None
container.tool_registry = None
container.backend_registry = None
container.cp_registry = None
container.workers = []
container.memory_store = None
container.mcp_manager = None
return AgentSessionOrchestrator(container)
# ── SessionRun buffer tests ─────────────────────────────────────────────────
@pytest.mark.anyio
async def test_event_buffer_appended_and_replayed():
"""Broadcast stores serialised events; oldest evicted when cap exceeded."""
run = SessionRun()
class FakeEvent:
def __init__(self, idx: int) -> None:
self.idx = idx
def to_wire(self) -> dict:
return {"type": "stream_delta", "delta": str(self.idx)}
for i in range(3):
await run.broadcast(("event", FakeEvent(i)))
assert len(run.events) == 3
assert run.events[0] == {"type": "stream_delta", "delta": "0"}
assert run.events[2] == {"type": "stream_delta", "delta": "2"}
# Fill buffer past limit
run.events.clear()
for i in range(500 + 5):
await run.broadcast(("event", FakeEvent(i)))
assert len(run.events) == 500
assert run.events[0] == {"type": "stream_delta", "delta": "5"}
assert run.events[-1] == {"type": "stream_delta", "delta": str(500 + 4)}
# ── Reconnect / replay tests ─────────────────────────────────────────────────
@pytest.mark.anyio
async def test_reconnect_replays_buffered_events(mock_websocket, mock_session, mock_user, monkeypatch):
"""Re-attach to active run yields replay_start, buffered events, replay_end, then session_sync."""
monkeypatch.setattr(ws_mod, "get_current_user_ws", AsyncMock(return_value=mock_user))
mock_store = MagicMock()
mock_store.get = AsyncMock(return_value=mock_session)
monkeypatch.setattr(ws_mod, "get_session_store", lambda: mock_store)
monkeypatch.setattr(ws_mod, "_stream_to_client", AsyncMock(return_value=True))
fake_container = MagicMock()
fake_container.profile_registry = None
fake_container.tool_registry = None
fake_container.backend_registry = None
fake_container.cp_registry = None
fake_container.orchestrator = AgentSessionOrchestrator(fake_container)
monkeypatch.setattr("navi.api.deps._resolve_container", lambda: fake_container)
orchestrator = fake_container.orchestrator
run = orchestrator.create_run("s1")
run.events = [
{"type": "stream_delta", "delta": "hello"},
{"type": "thinking_delta", "delta": "hmm"},
]
mock_websocket.receive_text = AsyncMock(side_effect=WebSocketDisconnect())
await ws_mod.websocket_session("s1", mock_websocket)
calls = [c.args[0] for c in mock_websocket.send_json.call_args_list]
types = [c["type"] for c in calls]
assert types == [
"stream_start",
"replay_start",
"stream_delta",
"thinking_delta",
"replay_end",
"session_sync",
]
assert calls[1]["count"] == 2
orchestrator._sessions.pop("s1", None)
@pytest.mark.anyio
async def test_session_sync_after_reconnect_when_done(mock_websocket, mock_session, mock_user, monkeypatch):
"""Reconnect when no run is active → only session_sync, no replay."""
monkeypatch.setattr(ws_mod, "get_current_user_ws", AsyncMock(return_value=mock_user))
mock_store = MagicMock()
mock_store.get = AsyncMock(return_value=mock_session)
monkeypatch.setattr(ws_mod, "get_session_store", lambda: mock_store)
monkeypatch.setattr(ws_mod, "_stream_to_client", AsyncMock(return_value=True))
fake_container = MagicMock()
fake_container.profile_registry = None
fake_container.tool_registry = None
fake_container.backend_registry = None
fake_container.cp_registry = None
fake_container.orchestrator = AgentSessionOrchestrator(fake_container)
monkeypatch.setattr("navi.api.deps._resolve_container", lambda: fake_container)
mock_websocket.receive_text = AsyncMock(side_effect=WebSocketDisconnect())
await ws_mod.websocket_session("s1", mock_websocket)
calls = [c.args[0] for c in mock_websocket.send_json.call_args_list]
assert len(calls) == 1
assert calls[0]["type"] == "session_sync"
@pytest.mark.anyio
async def test_session_sync_after_recall_run(mock_websocket, mock_session, mock_user, monkeypatch):
"""Reconnect while a headless recall run is active → session_sync (no replay, no error)."""
monkeypatch.setattr(ws_mod, "get_current_user_ws", AsyncMock(return_value=mock_user))
mock_store = MagicMock()
mock_store.get = AsyncMock(return_value=mock_session)
monkeypatch.setattr(ws_mod, "get_session_store", lambda: mock_store)
monkeypatch.setattr(ws_mod, "_stream_to_client", AsyncMock(return_value=True))
fake_container = MagicMock()
fake_container.profile_registry = None
fake_container.tool_registry = None
fake_container.backend_registry = None
fake_container.cp_registry = None
fake_container.orchestrator = AgentSessionOrchestrator(fake_container)
monkeypatch.setattr("navi.api.deps._resolve_container", lambda: fake_container)
orchestrator = fake_container.orchestrator
orchestrator.mark_busy("s1")
mock_websocket.receive_text = AsyncMock(side_effect=WebSocketDisconnect())
await ws_mod.websocket_session("s1", mock_websocket)
calls = [c.args[0] for c in mock_websocket.send_json.call_args_list]
types = [c["type"] for c in calls]
assert types == ["session_sync"]
await orchestrator.clear_busy("s1")
# ── Concurrent run guard ─────────────────────────────────────────────────────
@pytest.mark.anyio
async def test_concurrent_run_guard_rejects_second_message(mock_websocket, mock_session, mock_user, monkeypatch):
"""Sending a second message while a run is active yields a WebSocket error."""
monkeypatch.setattr(ws_mod, "get_current_user_ws", AsyncMock(return_value=mock_user))
mock_store = MagicMock()
mock_store.get = AsyncMock(return_value=mock_session)
monkeypatch.setattr(ws_mod, "get_session_store", lambda: mock_store)
monkeypatch.setattr(ws_mod, "_stream_to_client", AsyncMock(return_value=True))
fake_container = MagicMock()
fake_container.profile_registry = None
fake_container.tool_registry = None
fake_container.backend_registry = None
fake_container.cp_registry = None
fake_container.orchestrator = AgentSessionOrchestrator(fake_container)
monkeypatch.setattr("navi.api.deps._resolve_container", lambda: fake_container)
orchestrator = fake_container.orchestrator
# _run_agent sleeps so the run stays registered
async def fake_run_agent(*a, **kw):
await asyncio.sleep(3600)
monkeypatch.setattr(orchestrator, "run_agent", fake_run_agent)
message_count = 0
async def fake_receive_text():
nonlocal message_count
message_count += 1
if message_count == 1:
return json.dumps({"type": "message", "content": "first"})
if message_count == 2:
return json.dumps({"type": "message", "content": "second"})
raise WebSocketDisconnect()
mock_websocket.receive_text = fake_receive_text
await ws_mod.websocket_session("s1", mock_websocket)
calls = [c.args[0] for c in mock_websocket.send_json.call_args_list]
error_calls = [c for c in calls if c["type"] == "error"]
assert len(error_calls) == 1
assert "already running" in error_calls[0]["message"]
# Cleanup background task
state = orchestrator._sessions.get("s1")
if state and state.run and state.run.task:
state.run.task.cancel()
try:
await state.run.task
except asyncio.CancelledError:
pass
orchestrator._sessions.pop("s1", None)
# ── /compact (forced compression) ───────────────────────────────────────────
@pytest.mark.anyio
async def test_compact_message_runs_forced_compression(mock_websocket, mock_session, mock_user, monkeypatch):
"""A {"type": "compact"} message dispatches to run_compact (not run_agent),
and sends NO stream_start — the client distinguishes a forced compact from an
in-turn auto-compress by the absence of a streaming turn."""
monkeypatch.setattr(ws_mod, "get_current_user_ws", AsyncMock(return_value=mock_user))
mock_store = MagicMock()
mock_store.get = AsyncMock(return_value=mock_session)
monkeypatch.setattr(ws_mod, "get_session_store", lambda: mock_store)
monkeypatch.setattr(ws_mod, "_stream_to_client", AsyncMock(return_value=True))
fake_container = MagicMock()
fake_container.profile_registry = None
fake_container.tool_registry = None
fake_container.backend_registry = None
fake_container.cp_registry = None
fake_container.orchestrator = AgentSessionOrchestrator(fake_container)
monkeypatch.setattr("navi.api.deps._resolve_container", lambda: fake_container)
orchestrator = fake_container.orchestrator
compact_calls: list = []
async def fake_run_compact(sid, store):
compact_calls.append((sid, store))
monkeypatch.setattr(orchestrator, "run_compact", fake_run_compact)
received: list[str] = []
async def fake_receive_text():
if not received:
received.append(json.dumps({"type": "compact"}))
return received[-1]
raise WebSocketDisconnect()
mock_websocket.receive_text = fake_receive_text
await ws_mod.websocket_session("s1", mock_websocket)
# The handler schedules run_compact via create_task but the mocked
# _stream_to_client / receive_text never yield to it, so drive it here.
state = orchestrator._sessions.get("s1")
if state and state.run and state.run.task:
try:
await state.run.task
except asyncio.CancelledError:
pass
assert len(compact_calls) == 1
assert compact_calls[0][0] == "s1"
# No stream_start is sent for a compact.
calls = [c.args[0] for c in mock_websocket.send_json.call_args_list]
assert all(c.get("type") != "stream_start" for c in calls)
orchestrator._sessions.pop("s1", None)
@pytest.mark.anyio
async def test_compact_while_agent_running_is_rejected(mock_websocket, mock_session, mock_user, monkeypatch):
"""A /compact while an agent turn is active is rejected — compressing
session.context mid-turn would race the running agent."""
monkeypatch.setattr(ws_mod, "get_current_user_ws", AsyncMock(return_value=mock_user))
mock_store = MagicMock()
mock_store.get = AsyncMock(return_value=mock_session)
monkeypatch.setattr(ws_mod, "get_session_store", lambda: mock_store)
monkeypatch.setattr(ws_mod, "_stream_to_client", AsyncMock(return_value=True))
fake_container = MagicMock()
fake_container.profile_registry = None
fake_container.tool_registry = None
fake_container.backend_registry = None
fake_container.cp_registry = None
fake_container.orchestrator = AgentSessionOrchestrator(fake_container)
monkeypatch.setattr("navi.api.deps._resolve_container", lambda: fake_container)
orchestrator = fake_container.orchestrator
async def fake_run_agent(*a, **kw):
await asyncio.sleep(3600)
monkeypatch.setattr(orchestrator, "run_agent", fake_run_agent)
message_count = 0
async def fake_receive_text():
nonlocal message_count
message_count += 1
if message_count == 1:
return json.dumps({"type": "message", "content": "first"})
if message_count == 2:
return json.dumps({"type": "compact"})
raise WebSocketDisconnect()
mock_websocket.receive_text = fake_receive_text
await ws_mod.websocket_session("s1", mock_websocket)
calls = [c.args[0] for c in mock_websocket.send_json.call_args_list]
error_calls = [c for c in calls if c["type"] == "error"]
assert len(error_calls) == 1
assert "wait for it to finish" in error_calls[0]["message"]
state = orchestrator._sessions.get("s1")
if state and state.run and state.run.task:
state.run.task.cancel()
try:
await state.run.task
except asyncio.CancelledError:
pass
orchestrator._sessions.pop("s1", None)
# ── run_compact event/error broadcasting ────────────────────────────────────
@pytest.mark.anyio
async def test_run_compact_broadcasts_events_and_done(fake_orchestrator, monkeypatch):
"""run_compact forwards agent.compact_stream events to subscribers and ends
with a done marker so _stream_to_client returns."""
from navi.core.events import CompressionStarted, ContextCompressed
async def fake_compact_stream(session_id):
yield CompressionStarted(context_tokens=100, max_context_tokens=4096)
yield ContextCompressed(messages_before=20, messages_after=5, context_tokens=80, max_context_tokens=4096)
class _FakeAgent:
async def compact_stream(self, session_id):
async for ev in fake_compact_stream(session_id):
yield ev
def fake_build_agent(self, _store):
return _FakeAgent()
monkeypatch.setattr(AgentSessionOrchestrator, "_build_agent", fake_build_agent)
run = fake_orchestrator.create_run("s1")
queue = run.subscribe()
await fake_orchestrator.run_compact("s1", session_store=MagicMock())
items = []
while not queue.empty():
items.append(queue.get_nowait())
kinds = [item[0] for item in items]
assert kinds == ["event", "event", "done"]
assert isinstance(items[0][1], CompressionStarted)
assert isinstance(items[1][1], ContextCompressed)
# run_compact's finally clears state.run, and with no websockets attached
# _cleanup then removes the session entry entirely.
assert "s1" not in fake_orchestrator._sessions
@pytest.mark.anyio
async def test_run_compact_surfaces_nothing_to_compact_as_error(fake_orchestrator, monkeypatch):
"""NothingToCompactError is broadcast as an error event (user feedback),
followed by done — not raised into the WS handler."""
from navi.exceptions import NothingToCompactError
class _FakeAgent:
async def compact_stream(self, session_id):
raise NothingToCompactError("nothing")
yield # pragma: no cover - makes this an async generator
def fake_build_agent(self, _store):
return _FakeAgent()
monkeypatch.setattr(AgentSessionOrchestrator, "_build_agent", fake_build_agent)
run = fake_orchestrator.create_run("s1")
queue = run.subscribe()
await fake_orchestrator.run_compact("s1", session_store=MagicMock())
items = []
while not queue.empty():
items.append(queue.get_nowait())
kinds = [item[0] for item in items]
assert kinds == ["error", "done"]
assert "nothing" in items[0][1]
assert "s1" not in fake_orchestrator._sessions
