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_acceptance" validation_result: dict[str, Any] | None = None class StubAgentService(AgentService): def __init__(self, *, fail: bool = False) -> None: super().__init__() self.fail = fail self.calls: list[dict[str, Any]] = [] async def process_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}", ) 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}", ) class DirectModeOnlyAgentService(StubAgentService): async def submit_direct(self, message: str, **kwargs: Any) -> StubRunResult: # type: ignore[override] raise RuntimeError("submit_direct should not be used when service is not running") 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", "prompt_locale": "zh-Hant", "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, "gateway_user_id": None, "title": None, "execution_context": None, "prompt_locale": "zh-Hant", "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_acceptance" assert message["evidence_status"] == "recorded" assert message["validation_result"] is None assert "validation_status" not in message 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_message_uses_direct_processing_when_loop_is_not_running() -> None: service = DirectModeOnlyAgentService() 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() assert service.calls == [ { "message": "hello", "session_id": "web:alpha", "source": "websocket", "user_id": None, "gateway_user_id": None, "title": None, "execution_context": None, "prompt_locale": None, "model": None, "provider_name": None, "embedding_model": None, "max_tool_iterations": None, } ] assert message["type"] == "message" assert message["content"] == "echo:hello" def test_rest_chat_uses_direct_processing_when_loop_is_not_running() -> None: service = DirectModeOnlyAgentService() app = create_app(service=service, manage_service_lifecycle=False) with TestClient(app) as client: response = client.post( "/api/chat", json={"session_id": "web:alpha", "message": "hello", "prompt_locale": "en"}, ) assert response.status_code == 200 assert service.calls == [ { "message": "hello", "session_id": "web:alpha", "source": "web", "user_id": None, "gateway_user_id": None, "title": None, "execution_context": None, "prompt_locale": "en", "model": None, "provider_name": None, "embedding_model": None, "temperature": None, "max_tokens": None, "max_tool_iterations": None, "fallback_target": None, "auxiliary_target": None, "embedding_target": None, } ] assert response.json()["output_text"] == "echo:hello" def test_rest_chat_uses_authenticated_user_for_gateway_identity() -> None: service = DirectModeOnlyAgentService() app = create_app(service=service, manage_service_lifecycle=False) app.state.auth_tokens["token-1"] = "tom" with TestClient(app) as client: response = client.post( "/api/chat", headers={"Authorization": "Bearer token-1"}, json={"session_id": "web:alpha", "message": "hello", "user_id": "other"}, ) assert response.status_code == 200 assert service.calls == [ { "message": "hello", "session_id": "web:alpha", "source": "web", "user_id": "other", "gateway_user_id": "tom", "title": None, "execution_context": None, "prompt_locale": None, "model": None, "provider_name": None, "embedding_model": None, "temperature": None, "max_tokens": None, "max_tool_iterations": None, "fallback_target": None, "auxiliary_target": None, "embedding_target": None, } ] def test_websocket_uses_authenticated_user_for_gateway_identity() -> None: service = StubAgentService() app = create_app(service=service, manage_service_lifecycle=False) app.state.auth_tokens["token-1"] = "tom" with TestClient(app) as client: with client.websocket_connect("/ws/web:alpha?token=token-1") as websocket: websocket.send_json({"type": "message", "content": "hello", "user_id": "other"}) assert websocket.receive_json() == {"type": "status", "status": "thinking"} websocket.receive_json() websocket.receive_json() assert service.calls == [ { "message": "hello", "session_id": "web:alpha", "source": "websocket", "user_id": "other", "gateway_user_id": "tom", "title": None, "execution_context": None, "prompt_locale": None, "model": None, "provider_name": None, "embedding_model": None, "max_tool_iterations": None, } ] 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"}