from __future__ import annotations from dataclasses import dataclass, field from typing import Any from fastapi.testclient import TestClient from beaver.interfaces.web.app import create_app from beaver.services.agent_service import AgentService @dataclass(slots=True) class StubRunResult: session_id: str run_id: str = "run-1" output_text: str = "ok" finish_reason: str = "stop" tool_iterations: int = 0 provider_name: str | None = "stub" model: str | None = "stub-model" usage: dict[str, Any] = field(default_factory=lambda: {"total_tokens": 3}) task_id: str | None = "task-1" task_status: str | None = "awaiting_feedback" validation_result: dict[str, Any] | None = field(default_factory=lambda: {"accepted": True}) class StubAgentService(AgentService): def __init__(self, *, fail: bool = False) -> None: super().__init__() self.fail = fail self.calls: list[dict[str, Any]] = [] async def submit_direct(self, message: str, **kwargs: Any) -> StubRunResult: # type: ignore[override] self.calls.append({"message": message, **kwargs}) if self.fail: raise RuntimeError("boom") return StubRunResult( session_id=kwargs.get("session_id") or "web:default", output_text=f"echo:{message}", ) def test_websocket_ping_pong() -> None: app = create_app(service=StubAgentService(), manage_service_lifecycle=False) with TestClient(app) as client: with client.websocket_connect("/ws/web:alpha") as websocket: websocket.send_json({"type": "ping"}) assert websocket.receive_json() == {"type": "pong"} def test_websocket_message_returns_chat_metadata_and_session_updated() -> None: service = StubAgentService() app = create_app(service=service, manage_service_lifecycle=False) with TestClient(app) as client: with client.websocket_connect("/ws/web:alpha") as websocket: websocket.send_json( { "type": "message", "content": "hello", "metadata": {"source": "test"}, "attachments": [{"file_id": "file-1", "name": "a.txt"}], } ) assert websocket.receive_json() == {"type": "status", "status": "thinking"} message = websocket.receive_json() session_updated = websocket.receive_json() assert service.calls == [ { "message": "hello", "session_id": "web:alpha", "source": "websocket", "user_id": None, "title": None, "execution_context": None, "model": None, "provider_name": None, "embedding_model": None, "max_tool_iterations": None, } ] assert message["type"] == "message" assert message["role"] == "assistant" assert message["content"] == "echo:hello" assert message["session_id"] == "web:alpha" assert message["run_id"] == "run-1" assert message["task_id"] == "task-1" assert message["task_status"] == "awaiting_feedback" assert message["validation_result"] == {"accepted": True} assert message["validation_status"] == "passed" assert message["metadata"]["input_metadata"] == { "source": "test", "attachments": [{"file_id": "file-1", "name": "a.txt"}], } assert session_updated == { "type": "session_updated", "session_id": "web:alpha", "source": "websocket", } def test_websocket_empty_content_returns_error_without_runtime_call() -> None: service = StubAgentService() app = create_app(service=service, manage_service_lifecycle=False) with TestClient(app) as client: with client.websocket_connect("/ws/web:alpha") as websocket: websocket.send_json({"type": "message", "content": " "}) assert websocket.receive_json() == {"type": "error", "error": "'content' is required"} assert service.calls == [] def test_websocket_runtime_error_returns_assistant_error_message() -> None: service = StubAgentService(fail=True) app = create_app(service=service, manage_service_lifecycle=False) with TestClient(app) as client: with client.websocket_connect("/ws/web:alpha") as websocket: websocket.send_json({"type": "message", "content": "hello"}) assert websocket.receive_json() == {"type": "status", "status": "thinking"} message = websocket.receive_json() websocket.send_json({"type": "ping"}) pong = websocket.receive_json() assert message["type"] == "message" assert message["role"] == "assistant" assert message["session_id"] == "web:alpha" assert message["finish_reason"] == "error" assert message["tool_iterations"] == 0 assert "boom" in message["content"] assert pong == {"type": "pong"}