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: