feat: refactor ChatSessionManager to be application-specific and integrate httpx2 for websocket transport

tests/api/test_chat_ws.py

@@ -1,81 +1,106 @@
 """Tests for the websocket chat service using the local stub agent.
 
-Uses Starlette's synchronous ``TestClient`` websocket support (rather than
-the ``httpx.AsyncClient`` fixture used elsewhere) since it triggers the
-FastAPI lifespan on ``__enter__``/``__exit__``, wiring up
-``app.chat_agent`` / ``app.state.chat_sessions`` exactly like a real
-server run.
+Uses HTTPX2's async ASGI WebSocket transport. The app lifespan is managed
+explicitly because the transport does not trigger lifespan events itself.
 """
 
+from collections.abc import AsyncIterator
+from contextlib import asynccontextmanager
+
 import pytest
-from fastapi.testclient import TestClient
+from httpx2 import AsyncClient
+from httpx2.websockets import ASGIWebSocketTransport
 
 from app.main import app
 
 pytestmark = pytest.mark.anyio
 
 
-async def test_chat_ws_full_conversation_flow() -> None:
-    with (
-        TestClient(app) as client,
-        client.websocket_connect("/v1/chat/ws") as websocket,
-    ):
+@pytest.fixture()
+async def chat_client(
+    monkeypatch: pytest.MonkeyPatch,
+) -> AsyncIterator[AsyncClient]:
+    class RedisStub:
+        async def aclose(self) -> None:
+            pass
+
+    async def get_redis_stub() -> RedisStub:
+        return RedisStub()
+
+    @asynccontextmanager
+    async def scheduler_lifespan_stub() -> AsyncIterator[None]:
+        yield
+
+    monkeypatch.setattr("app.main.get_redis", get_redis_stub)
+    monkeypatch.setattr("app.main.scheduler_lifespan", scheduler_lifespan_stub)
+
+    async with app.router.lifespan_context(app):
+        async with AsyncClient(
+            transport=ASGIWebSocketTransport(app),
+            base_url="http://testserver",
+        ) as client:
+            yield client
+
+
+async def test_chat_ws_full_conversation_flow(chat_client: AsyncClient) -> None:
+    assert await app.state.chat_sessions.count() == 0
+
+    async with chat_client.websocket("ws://testserver/v1/chat/ws") as websocket:
         # 1. connect
-        connected = websocket.receive_json()
+        connected = await websocket.receive_json()
         assert connected["type"] == "connected"
         session_id = connected["session_id"]
+        assert await app.state.chat_sessions.count() == 1
 
         # 2. start conversation with a system prompt
-        websocket.send_json(
+        await websocket.send_json(
             {"type": "start", "system_prompt": "You are a helpful assistant."}
         )
-        started = websocket.receive_json()
-        print(f"{started=}")
+        started = await websocket.receive_json()
         assert started == {"type": "started", "session_id": session_id}
 
         # 3. send a user message and aggregate the streamed assistant chunks
-        websocket.send_json({"type": "user_message", "content": "hello"})
+        await websocket.send_json({"type": "user_message", "content": "hello"})
 
-        event = websocket.receive_json()
+        event = await websocket.receive_json()
         aggregated = ""
         while event["type"] == "assistant_chunk":
             assert event["index"] >= 0
             aggregated += event["content"]
-            event = websocket.receive_json()
+            event = await websocket.receive_json()
 
         assert event["type"] == "assistant_message"
         assert event["message"]["role"] == "assistant"
         assert event["message"]["content"] == aggregated
-        print(f"{event=}")
-        print(f"{aggregated=}")
         assert "hello" in aggregated.lower()
 
         # 4. request full history (system + user + assistant)
-        websocket.send_json({"type": "history_request"})
-        history = websocket.receive_json()
+        await websocket.send_json({"type": "history_request"})
+        history = await websocket.receive_json()
         assert history["type"] == "history"
         roles = [m["role"] for m in history["messages"]]
         assert roles == ["system", "user", "assistant"]
 
         # 5. gracefully end the conversation
-        websocket.send_json({"type": "end"})
+        await websocket.send_json({"type": "end"})
+
+    assert await app.state.chat_sessions.count() == 0
 
 
-async def test_chat_ws_invalid_event_returns_error() -> None:
-    with (
-        TestClient(app) as client,
-        client.websocket_connect("/v1/chat/ws") as websocket,
-    ):
-        websocket.receive_json()  # connected
+async def test_chat_ws_invalid_event_returns_error(chat_client: AsyncClient) -> None:
+    async with chat_client.websocket("ws://testserver/v1/chat/ws") as websocket:
+        await websocket.receive_json()  # connected
 
-        websocket.send_json({"type": "not_a_real_event"})
-        error = websocket.receive_json()
+        await websocket.send_json({"type": "not_a_real_event"})
+        error = await websocket.receive_json()
         assert error["type"] == "error"
 
-        websocket.send_json({"type": "end"})
+        await websocket.send_json({"type": "end"})
 
 
-async def test_chat_ws_ask_question_and_print_model_answer() -> None:
+async def test_chat_ws_ask_question_and_print_model_answer(
+    chat_client: AsyncClient,
+) -> None:
     """Ask a real question over the websocket chat and print the model's answer.
 
     Runs against whatever ``ChatAgent`` is wired up via ``app.state.chat_agent``
@@ -85,30 +110,27 @@ async def test_chat_ws_ask_question_and_print_model_answer() -> None:
     """
     question = "What is the capital of France?"
 
-    with (
-        TestClient(app) as client,
-        client.websocket_connect("/v1/chat/ws") as websocket,
-    ):
+    async with chat_client.websocket("ws://testserver/v1/chat/ws") as websocket:
         # 1. connect
-        connected = websocket.receive_json()
+        connected = await websocket.receive_json()
         assert connected["type"] == "connected"
         session_id = connected["session_id"]
 
         # 2. start conversation
-        websocket.send_json(
+        await websocket.send_json(
             {"type": "start", "system_prompt": "You are a helpful assistant."}
         )
-        started = websocket.receive_json()
+        started = await websocket.receive_json()
         assert started == {"type": "started", "session_id": session_id}
 
         # 3. ask the question and aggregate the streamed answer
-        websocket.send_json({"type": "user_message", "content": question})
+        await websocket.send_json({"type": "user_message", "content": question})
 
-        event = websocket.receive_json()
+        event = await websocket.receive_json()
         answer = ""
         while event["type"] == "assistant_chunk":
             answer += event["content"]
-            event = websocket.receive_json()
+            event = await websocket.receive_json()
 
         assert event["type"] == "assistant_message"
         answer = event["message"]["content"]
@@ -120,5 +142,4 @@ async def test_chat_ws_ask_question_and_print_model_answer() -> None:
         assert answer.strip() != ""
 
         # 4. gracefully end the conversation
-        websocket.send_json({"type": "end"})
-
+        await websocket.send_json({"type": "end"})