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

README.md

@@ -185,7 +185,7 @@ ollama run llama3.2
 <p align="right">(<a href="#readme-top">back to top</a>)</p>
 
 ### Websocket AI chat service
-A general-purpose, websocket-based chat service lives at `/v1/chat/ws`, built with FastAPI 0.139.0
+A general-purpose, websocket-based chat service lives at `/v1/chat/ws`, built with FastAPI 0.159.1
 and Pydantic models for every message shape exchanged over the wire. It follows the
 [Pydantic AI chat-app](https://github.com/pydantic/pydantic-ai/blob/main/docs/web.md) interaction
 pattern - a session id, role-labeled messages, and streamed assistant chunks - so a UI can be wired
@@ -203,10 +203,15 @@ up later without changing the backend contract.
   chat flow.
 
 **Lifespan wiring**: `app/main.py`'s `lifespan` builds the configured `ChatAgent` via
-`build_chat_agent(settings.chat)` and a fresh `ChatSessionManager` on startup (`app.chat_agent`
+`build_chat_agent(settings.chat)` and a fresh `ChatSessionManager` on startup. The agent is stored
-/ `app.state.chat_sessions`), and calls `chat_agent.aclose()` on shutdown to release any held
+on `app.chat_agent`, and the session manager on `app.state.chat_sessions`; the websocket route
+uses that manager to track active sessions. On shutdown, `chat_agent.aclose()` releases any held
 connections.
 
+The WebSocket integration tests use HTTPX2's async `ASGIWebSocketTransport`
+(`httpx2[ws]`) to exercise the application without starting a server. They manage
+the app lifespan explicitly because the transport does not start it.
+
 **Configuration** (`ChatConfig` in `app/config.py`, all overridable via environment variables):
 
 | Env var | Default | Description |

app/api/chat.py

@@ -100,7 +100,7 @@ async def chat_websocket(websocket: WebSocket) -> None:
     await websocket.accept()
 
     agent: ChatAgent = websocket.app.chat_agent
-    session_manager: ChatSessionManager = ChatSessionManager()
+    session_manager: ChatSessionManager = websocket.app.state.chat_sessions
     session = await session_manager.create()
 
     await websocket.send_json(Connected(session_id=session.id).model_dump(mode="json"))

app/main.py

@@ -22,6 +22,7 @@ from app.middleware.profiler import ProfilingMiddleware
 from app.redis import get_redis
 from app.services.auth import AuthBearer
 from app.services.chat_agent import build_chat_agent
+from app.services.chat_session import ChatSessionManager
 from app.services.scheduler import scheduler_lifespan
 
 templates = Jinja2Templates(directory=Path(__file__).parent.parent / "templates")
@@ -32,11 +33,9 @@ async def lifespan(app: FastAPI):
     app.logger = get_logger()
     app.redis = await get_redis()
     global_settings.postgres_url.unicode_string()
-    # Chat service: initialize the pluggable model-client adapter. The session
+    # Initialize the chat agent and per-application session registry.
-    # manager is now a singleton and will be auto-instantiated on first access.
-    # See app/services/chat_agent.py to swap the local stub for a real model
-    # client (OpenAI, Ollama, etc.).
     app.chat_agent = build_chat_agent(global_settings.chat)
+    app.state.chat_sessions = ChatSessionManager()
     try:
         # app.postgres_pool = await asyncpg.create_pool(
         #     dsn=postgres_dsn,

app/services/chat_session.py

@@ -13,7 +13,6 @@ from uuid import UUID, uuid4
 import attrs
 
 from app.schemas.chat import ChatMessage
-from app.utils.singleton import SingletonMetaNoArgs
 
 
 @attrs.define(slots=True)
@@ -28,11 +27,11 @@ class ChatSession:
         return list(self.messages)
 
 
-class ChatSessionManager(metaclass=SingletonMetaNoArgs):
+class ChatSessionManager:
     """Tracks active chat sessions keyed by their opaque session id.
 
-    Implemented as a singleton to ensure exactly one instance per application,
+    Each application lifespan owns its own manager, maintaining a consistent
-    maintaining a consistent registry of all active websocket chat sessions.
+    registry of active websocket chat sessions for that application.
     """
 
     def __init__(self) -> None:

pyproject.toml

@@ -47,6 +47,7 @@ test = [
     "inline-snapshot==0.36.1",
     "dirty-equals==0.11",
     "polyfactory==3.5.0",
+    "httpx2[ws]==2.16.0",
 ]
 
 [tool.pyrefly.errors]

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
+Uses HTTPX2's async ASGI WebSocket transport. The app lifespan is managed
-the ``httpx.AsyncClient`` fixture used elsewhere) since it triggers the
+explicitly because the transport does not trigger lifespan events itself.
-FastAPI lifespan on ``__enter__``/``__exit__``, wiring up
-``app.chat_agent`` / ``app.state.chat_sessions`` exactly like a real
-server run.
 """
 
+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:
+@pytest.fixture()
-    with (
+async def chat_client(
-        TestClient(app) as client,
+    monkeypatch: pytest.MonkeyPatch,
-        client.websocket_connect("/v1/chat/ws") as websocket,
+) -> 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()
+        started = await websocket.receive_json()
-        print(f"{started=}")
         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"})
+        await websocket.send_json({"type": "history_request"})
-        history = websocket.receive_json()
+        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:
+async def test_chat_ws_invalid_event_returns_error(chat_client: AsyncClient) -> None:
-    with (
+    async with chat_client.websocket("ws://testserver/v1/chat/ws") as websocket:
-        TestClient(app) as client,
+        await websocket.receive_json()  # connected
-        client.websocket_connect("/v1/chat/ws") as websocket,
-    ):
-        websocket.receive_json()  # connected
 
-        websocket.send_json({"type": "not_a_real_event"})
+        await websocket.send_json({"type": "not_a_real_event"})
-        error = websocket.receive_json()
+        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 (
+    async with chat_client.websocket("ws://testserver/v1/chat/ws") as websocket:
-        TestClient(app) as client,
-        client.websocket_connect("/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"})
-

uv.lock

@@ -359,6 +359,7 @@ dev = [
 ]
 test = [
     { name = "dirty-equals" },
+    { name = "httpx2", extra = ["ws"] },
     { name = "inline-snapshot" },
     { name = "polyfactory" },
     { name = "pytest" },
@@ -404,6 +405,7 @@ dev = [
 ]
 test = [
     { name = "dirty-equals", specifier = "==0.11" },
+    { name = "httpx2", extras = ["ws"], specifier = "==2.16.0" },
     { name = "inline-snapshot", specifier = "==0.36.1" },
     { name = "polyfactory", specifier = "==3.5.0" },
     { name = "pytest", specifier = "==9.1.1" },
@@ -513,6 +515,19 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/7e/f5/f66802a942d491edb555dd61e3a9961140fd64c90bce1eafd741609d334d/httpcore-1.0.9-py3-none-any.whl", hash = "sha256:2d400746a40668fc9dec9810239072b40b4484b640a8c38fd654a024c7a1bf55", size = 78784, upload-time = "2025-04-24T22:06:20.566Z" },
 ]
 
+[[package]]
+name = "httpcore2"
+version = "2.16.0"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "h11" },
+    { name = "truststore" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/c0/61/afa9d6b7452f9eca73f73d45324eb0bb52f098fd5e6049bf6b8263be9b4e/httpcore2-2.16.0.tar.gz", hash = "sha256:c72cd0413563b17e1cd8354677dff887466f8ec2551679759e93126534428643", size = 69886, upload-time = "2026-10-21T07:09:45.31Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/af/c2/ce417c397afa5181f0e3c6ff01a4d01639414e1adef08730d97571297151/httpcore2-2.16.0-py3-none-any.whl", hash = "sha256:5c4358319f2602d984ad93182c9fee5160f22e9da7babdda870a5e7c4e2dc34f", size = 85070, upload-time = "2026-10-21T07:09:42.802Z" },
+]
+
 [[package]]
 name = "httptools"
 version = "0.9.0"
@@ -562,6 +577,36 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/2a/39/e50c7c3a983047577ee07d2a9e53faf5a69493943ec3f6a384bdc792deb2/httpx-0.28.1-py3-none-any.whl", hash = "sha256:d909fcccc110f8c7faf814ca82a9a4d816bc5a6dbfea25d6591d6985b8ba59ad", size = 73517, upload-time = "2024-12-06T15:37:21.509Z" },
 ]
 
+[[package]]
+name = "httpx2"
+version = "2.16.0"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "anyio", marker = "sys_platform != 'emscripten'" },
+    { name = "httpcore2", marker = "sys_platform != 'emscripten'" },
+    { name = "httpx2-jsfetch", marker = "sys_platform == 'emscripten'" },
+    { name = "idna" },
+    { name = "truststore", marker = "sys_platform != 'emscripten'" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/43/73/5ca2cedf5e4e50423d1c258877865a6b62793633b561c58499f8aeb33b7e/httpx2-2.16.0.tar.gz", hash = "sha256:0d54c98128e8f2ff7fbbde0fca8c359e5f032edeaaaac443cdb7f31f685bcd7f", size = 103116, upload-time = "2026-10-21T07:09:46.348Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/7b/3d/2e4bc4fd1ef9770871669adb348d612a6e4786b303a23f9935b81c48dd60/httpx2-2.16.0-py3-none-any.whl", hash = "sha256:fd2c5d7909e8d86fe92d53020c310ca0e4e431b838c207c5754b9bbf0308d49a", size = 97880, upload-time = "2026-10-21T07:09:44.01Z" },
+]
+
+[package.optional-dependencies]
+ws = [
+    { name = "wsproto" },
+]
+
+[[package]]
+name = "httpx2-jsfetch"
+version = "1.0"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/cd/c4/0e5636363151a2a1795e0a77617168b9ca438e1748ec05fc9b5687f93d64/httpx2_jsfetch-1.0.tar.gz", hash = "sha256:70a0e3eabfef7cce5ad9c629f7d01ca05e418f586646f4ddf14782e4c1454c60", size = 6872, upload-time = "2026-08-07T00:13:07.492Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/9b/43/832f631d32e4f1211caa2ba368317739fe71f0b8530e4c9d15dc454bac2a/httpx2_jsfetch-1.0-py3-none-any.whl", hash = "sha256:cb916b707601e69a07721aabc8f3f6659be3a6893bc1ff5c6f9e02241df2da32", size = 6382, upload-time = "2026-08-07T00:13:06.567Z" },
+]
+
 [[package]]
 name = "idna"
 version = "3.20"
@@ -1287,6 +1332,15 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/ad/66/0d785f0bc5e4315a96c989bb476d0fc07ea4f85132550c7b156ca2035d52/traitlets-5.16.1-py3-none-any.whl", hash = "sha256:f775618166caa0396c8e337099240f2bd3e5e917d203b2e6fbe21a58d3cb1f6b", size = 86211, upload-time = "2026-08-03T08:32:34.48Z" },
 ]
 
+[[package]]
+name = "truststore"
+version = "0.10.5"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/ee/9f/c5201d42a484c061e528825fc8e2d565f5abd50a4ced6fb7d29c4ec99b2b/truststore-0.10.5.tar.gz", hash = "sha256:30d36967ccaded5cbb38d602c433f53600036c79d502f4533a49b60a03bbefcd", size = 28091, upload-time = "2026-10-12T22:27:31.808Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/51/e9/3a7820be2bb0fe53b6bc9c3be26d3d1158004e4c3ab953aa6840b955b1e9/truststore-0.10.5-py3-none-any.whl", hash = "sha256:9aaaedaefaf06d8b206278cf8b5012bc897f485a874503501e12d776df78951c", size = 19017, upload-time = "2026-10-12T22:27:30.377Z" },
+]
+
 [[package]]
 name = "tryceratops"
 version = "2.4.1"
@@ -1541,3 +1595,15 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/31/06/c969504aa9887010bac12484df2f03aa0e7ca5ada615a989da04e8f33046/whenever-0.10.5-cp314-cp314t-win_amd64.whl", hash = "sha256:06c32f8931bbc52b6b3f513e831aeb936e231f038b2a66cd11d4a00c131b82e6", size = 544821, upload-time = "2026-08-07T15:02:39.579Z" },
     { url = "https://files.pythonhosted.org/packages/f2/84/c937c7ec9de2aa59fbcdd60b52a8da60069636d2252c7461e43d389c04a4/whenever-0.10.5-py3-none-any.whl", hash = "sha256:193091633fc1dece4c86e7303cbd4039320b4f45c69dc5ae69b8f711bca82603", size = 129302, upload-time = "2026-08-07T15:03:29.285Z" },
 ]
+
+[[package]]
+name = "wsproto"
+version = "1.3.2"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "h11" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/c7/79/12135bdf8b9c9367b8701c2c19a14c913c120b882d50b014ca0d38083c2c/wsproto-1.3.2.tar.gz", hash = "sha256:b86885dcf294e15204919950f666e06ffc6c7c114ca900b060d6e16293528294", size = 50116, upload-time = "2025-11-20T18:18:01.871Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/a4/f5/10b68b7b1544245097b2a1b8238f66f2fc6dcaeb24ba5d917f52bd2eed4f/wsproto-1.3.2-py3-none-any.whl", hash = "sha256:61eea322cdf56e8cc904bd3ad7573359a242ba65688716b0710a5eb12beab584", size = 24405, upload-time = "2025-11-20T18:18:00.454Z" },
+]