Add per-chat real client routing

sdk/agent_api_wrapper.py

@@ -3,6 +3,8 @@ from __future__ import annotations
 import asyncio
 import logging
 import sys
+import re
+from urllib.parse import urlsplit, urlunsplit
 from pathlib import Path
 
 import aiohttp
@@ -26,10 +28,46 @@ logger = logging.getLogger(__name__)
 class AgentApiWrapper(AgentApi):
     """Capture tokens_used from MsgEventEnd without patching upstream code."""
 
-    def __init__(self, agent_id: str, url: str, **kwargs) -> None:
-        super().__init__(agent_id=agent_id, url=url, **kwargs)
+    def __init__(
+        self,
+        agent_id: str,
+        base_url: str | None = None,
+        *,
+        chat_id: int | str = 0,
+        url: str | None = None,
+        **kwargs,
+    ) -> None:
+        if base_url is None and url is None:
+            raise TypeError("AgentApiWrapper requires base_url or url")
+
+        self._base_url = self._normalize_base_url(base_url or url or "")
+        self._init_kwargs = dict(kwargs)
+        self.chat_id = chat_id
+        super().__init__(
+            agent_id=agent_id,
+            url=self._build_ws_url(self._base_url, chat_id),
+            **kwargs,
+        )
         self.last_tokens_used = 0
 
+    @staticmethod
+    def _normalize_base_url(base_url: str) -> str:
+        parsed = urlsplit(base_url)
+        path = re.sub(r"(?:/v1)?/agent_ws(?:/[^/]+)?$", "", parsed.path.rstrip("/"))
+        return urlunsplit((parsed.scheme, parsed.netloc, path, "", ""))
+
+    @staticmethod
+    def _build_ws_url(base_url: str, chat_id: int | str) -> str:
+        return base_url.rstrip("/") + f"/v1/agent_ws/{chat_id}/"
+
+    def for_chat(self, chat_id: int | str) -> "AgentApiWrapper":
+        return type(self)(
+            agent_id=self.id,
+            base_url=self._base_url,
+            chat_id=chat_id,
+            **self._init_kwargs,
+        )
+
     async def _listen(self):
         try:
             async for msg in self._ws:

sdk/real.py

@@ -17,11 +17,21 @@ class RealPlatformClient(PlatformClient):
         self._agent_api = agent_api
         self._prototype_state = prototype_state
         self._platform = platform
+        self._chat_apis: dict[str, AgentApiWrapper] = {}
 
     @property
     def agent_api(self) -> AgentApiWrapper:
         return self._agent_api
 
+    async def _get_chat_api(self, chat_id: str) -> AgentApiWrapper:
+        chat_key = str(chat_id)
+        chat_api = self._chat_apis.get(chat_key)
+        if chat_api is None:
+            chat_api = self._agent_api.for_chat(chat_key)
+            await chat_api.connect()
+            self._chat_apis[chat_key] = chat_api
+        return chat_api
+
     async def get_or_create_user(
         self,
         external_id: str,
@@ -66,8 +76,9 @@ class RealPlatformClient(PlatformClient):
         text: str,
         attachments: list[Attachment] | None = None,
     ) -> AsyncIterator[MessageChunk]:
-        self._agent_api.last_tokens_used = 0
-        async for event in self._agent_api.send_message(text):
+        chat_api = await self._get_chat_api(chat_id)
+        chat_api.last_tokens_used = 0
+        async for event in chat_api.send_message(text):
             yield MessageChunk(
                 message_id=user_id,
                 delta=event.text,
@@ -77,7 +88,7 @@ class RealPlatformClient(PlatformClient):
             message_id=user_id,
             delta="",
             finished=True,
-            tokens_used=self._agent_api.last_tokens_used,
+            tokens_used=chat_api.last_tokens_used,
         )
 
     async def get_settings(self, user_id: str) -> UserSettings:
@@ -85,3 +96,8 @@ class RealPlatformClient(PlatformClient):
 
     async def update_settings(self, user_id: str, action) -> None:
         await self._prototype_state.update_settings(user_id, action)
+
+    async def close(self) -> None:
+        for chat_api in list(self._chat_apis.values()):
+            await chat_api.close()
+        self._chat_apis.clear()