"""Black-box smoke test for the real gateway WebUI transport.""" from __future__ import annotations import asyncio import json import socket import subprocess import sys import time from pathlib import Path from urllib.parse import quote import httpx import pytest import websockets from nanobot.session.manager import SessionManager from nanobot.session.recovery import PENDING_USER_TURN_KEY, RUNTIME_CHECKPOINT_KEY _BOOTSTRAP_SECRET = "smoke-secret" def _free_port() -> int: with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: sock.bind(("127.0.0.1", 0)) return int(sock.getsockname()[1]) def _write_smoke_config(path: Path, *, workspace: Path, ws_port: int, gateway_port: int) -> None: config = { "agents": { "defaults": { "workspace": str(workspace), "provider": "custom", "model": "custom/smoke-model", "maxToolIterations": 1, "dream": {"enabled": False}, } }, "providers": { "custom": { "apiKey": "smoke-no-external-call", "apiBase": "http://127.0.0.1:9/v1", } }, "channels": { "websocket": { "enabled": True, "host": "127.0.0.1", "port": ws_port, "allowFrom": ["*"], "tokenIssueSecret": _BOOTSTRAP_SECRET, } }, "gateway": { "host": "127.0.0.1", "port": gateway_port, "heartbeat": {"enabled": False}, }, } path.write_text(json.dumps(config), encoding="utf-8") def _start_gateway(config_path: Path, log_path: Path) -> subprocess.Popen[bytes]: log_file = log_path.open("wb") try: process = subprocess.Popen( [ sys.executable, "-m", "nanobot", "gateway", "--config", str(config_path), ], cwd=Path(__file__).resolve().parents[2], stdout=log_file, stderr=subprocess.STDOUT, ) finally: log_file.close() return process def _stop_gateway(process: subprocess.Popen[bytes]) -> None: if process.poll() is not None: return process.terminate() try: process.wait(timeout=10) except subprocess.TimeoutExpired: process.kill() process.wait(timeout=10) def _get_json(url: str, *, token: str | None = None) -> dict: headers = {"Authorization": f"Bearer {token}"} if token else {} response = httpx.get(url, headers=headers, timeout=5.0, trust_env=False) response.raise_for_status() return response.json() def _get_bootstrap(url: str) -> dict: response = httpx.get( url, headers={"X-Nanobot-Auth": _BOOTSTRAP_SECRET}, timeout=5.0, trust_env=False, ) response.raise_for_status() return response.json() def _wait_for_bootstrap(base_url: str, process: subprocess.Popen[bytes], log_path: Path) -> dict: deadline = time.monotonic() + 20 last_error: Exception | None = None while time.monotonic() < deadline: if process.poll() is not None: break try: return _get_bootstrap(f"{base_url}/webui/bootstrap") except (httpx.HTTPError, OSError) as exc: last_error = exc time.sleep(0.2) logs = log_path.read_text(encoding="utf-8", errors="replace") raise AssertionError(f"gateway did not start; last_error={last_error!r}\n{logs}") async def _recv_until(ws: websockets.WebSocketClientProtocol, event: str) -> dict: deadline = time.monotonic() + 20 while time.monotonic() < deadline: raw = await asyncio.wait_for(ws.recv(), timeout=5) payload = json.loads(raw) if payload.get("event") == event: return payload raise AssertionError(f"websocket event {event!r} was not received") @pytest.mark.asyncio async def test_gateway_webui_bootstrap_message_and_thread_hydration(tmp_path: Path) -> None: ws_port = _free_port() gateway_port = _free_port() workspace = tmp_path / "workspace" workspace.mkdir() config_path = tmp_path / "config.json" log_path = tmp_path / "gateway.log" _write_smoke_config( config_path, workspace=workspace, ws_port=ws_port, gateway_port=gateway_port, ) process = _start_gateway(config_path, log_path) base_url = f"http://127.0.0.1:{ws_port}" try: bootstrap = _wait_for_bootstrap(base_url, process, log_path) assert bootstrap["model_name"] == "custom/smoke-model" ws_url = f'{bootstrap["ws_url"]}?token={bootstrap["token"]}&client_id=smoke' async with websockets.connect(ws_url) as ws: ready = await _recv_until(ws, "ready") assert ready["client_id"] == "smoke" await ws.send(json.dumps({"type": "new_chat"})) attached = await _recv_until(ws, "attached") chat_id = attached["chat_id"] await _recv_until(ws, "session_updated") await ws.send(json.dumps({ "type": "message", "chat_id": chat_id, "content": "/model", "webui": True, "turn_id": "smoke-turn", })) answer = await _recv_until(ws, "message") assert "Current model: `custom/smoke-model`" in answer["text"] await _recv_until(ws, "turn_end") await ws.send(json.dumps({ "type": "message", "chat_id": chat_id, "content": "!printf shell-ok", "webui": True, "user_shell": True, "turn_id": "shell-turn", })) shell = await _recv_until(ws, "message") assert "shell-ok" in shell["text"] assert shell["turn_id"] == "shell-turn" await _recv_until(ws, "turn_end") api_token = _wait_for_bootstrap(base_url, process, log_path)["api_token"] sessions = _get_json(f"{base_url}/api/sessions", token=api_token) key = f"websocket:{chat_id}" assert key in {row["key"] for row in sessions["sessions"]} encoded_key = quote(key, safe="") thread = _get_json( f"{base_url}/api/sessions/{encoded_key}/webui-thread", token=api_token, ) contents = [str(message.get("content") or "") for message in thread["messages"]] assert "/model" in contents assert any("Current model: `custom/smoke-model`" in text for text in contents) assert "!printf shell-ok" in contents assert any("shell-ok" in text for text in contents) finally: _stop_gateway(process) def test_gateway_restart_restores_a_completed_answer_without_replaying_model( tmp_path: Path, ) -> None: """Exercise recovery through two real gateway processes and durable files.""" ws_port = _free_port() gateway_port = _free_port() workspace = tmp_path / "workspace" workspace.mkdir() config_path = tmp_path / "config.json" first_log = tmp_path / "gateway-first.log" second_log = tmp_path / "gateway-second.log" _write_smoke_config( config_path, workspace=workspace, ws_port=ws_port, gateway_port=gateway_port, ) base_url = f"http://127.0.0.1:{ws_port}" first = _start_gateway(config_path, first_log) try: _wait_for_bootstrap(base_url, first, first_log) finally: _stop_gateway(first) sessions_root = tmp_path / "sessions" sessions = SessionManager(workspace, sessions_root=sessions_root) session = sessions.get_or_create("websocket:recovery-smoke") session.messages.append({"role": "user", "content": "recover this answer"}) session.metadata["webui"] = True session.metadata[PENDING_USER_TURN_KEY] = True session.metadata[RUNTIME_CHECKPOINT_KEY] = { "phase": "final_response", "assistant_message": { "role": "assistant", "content": "restored without another model request", }, "completed_tool_results": [], "pending_tool_calls": [], } sessions.save(session, fsync=True) second = _start_gateway(config_path, second_log) try: bootstrap = _wait_for_bootstrap(base_url, second, second_log) deadline = time.monotonic() + 20 restored = None while time.monotonic() < deadline: restored = SessionManager( workspace, sessions_root=sessions_root, ).get_or_create("websocket:recovery-smoke") if any( message.get("content") == "restored without another model request" for message in restored.messages ): break time.sleep(0.1) else: logs = second_log.read_text(encoding="utf-8", errors="replace") raise AssertionError(f"answer was not recovered after restart\n{logs}") assert restored is not None assert PENDING_USER_TURN_KEY not in restored.metadata assert RUNTIME_CHECKPOINT_KEY not in restored.metadata assert restored.metadata["webui_recovery"]["reason"] == "answer_restored" async def assert_attach_state() -> None: ws_url = f'{bootstrap["ws_url"]}?token={bootstrap["token"]}&client_id=recovery-smoke' async with websockets.connect(ws_url) as ws: await _recv_until(ws, "ready") await ws.send(json.dumps({"type": "attach", "chat_id": "recovery-smoke"})) attached = await _recv_until(ws, "attached") assert attached["recovery_state"]["status"] == "recovered" assert attached["recovery_state"]["reason"] == "answer_restored" asyncio.run(assert_attach_state()) finally: _stop_gateway(second)