Serialize Matrix chat sends

sdk/real.py

@@ -20,6 +20,7 @@ class RealPlatformClient(PlatformClient):
         self._platform = platform
         self._chat_apis: dict[str, AgentApiWrapper] = {}
         self._chat_api_lock = asyncio.Lock()
+        self._chat_send_locks: dict[str, asyncio.Lock] = {}
 
     @property
     def agent_api(self) -> AgentApiWrapper:
@@ -41,6 +42,14 @@ class RealPlatformClient(PlatformClient):
                     self._chat_apis[chat_key] = chat_api
         return chat_api
 
+    def _get_chat_send_lock(self, chat_id: str) -> asyncio.Lock:
+        chat_key = str(chat_id)
+        lock = self._chat_send_locks.get(chat_key)
+        if lock is None:
+            lock = asyncio.Lock()
+            self._chat_send_locks[chat_key] = lock
+        return lock
+
     async def get_or_create_user(
         self,
         external_id: str,
@@ -85,21 +94,23 @@ class RealPlatformClient(PlatformClient):
         text: str,
         attachments: list[Attachment] | None = None,
     ) -> AsyncIterator[MessageChunk]:
-        chat_api = await self._get_chat_api(chat_id)
-        if hasattr(chat_api, "last_tokens_used"):
-            chat_api.last_tokens_used = 0
-        async for event in chat_api.send_message(text):
+        lock = self._get_chat_send_lock(chat_id)
+        async with lock:
+            chat_api = await self._get_chat_api(chat_id)
+            if hasattr(chat_api, "last_tokens_used"):
+                chat_api.last_tokens_used = 0
+            async for event in chat_api.send_message(text):
+                yield MessageChunk(
+                    message_id=user_id,
+                    delta=event.text,
+                    finished=False,
+                )
             yield MessageChunk(
                 message_id=user_id,
-                delta=event.text,
-                finished=False,
+                delta="",
+                finished=True,
+                tokens_used=getattr(chat_api, "last_tokens_used", 0),
             )
-        yield MessageChunk(
-            message_id=user_id,
-            delta="",
-            finished=True,
-            tokens_used=getattr(chat_api, "last_tokens_used", 0),
-        )
 
     async def get_settings(self, user_id: str) -> UserSettings:
         return await self._prototype_state.get_settings(user_id)
@@ -113,6 +124,7 @@ class RealPlatformClient(PlatformClient):
             if callable(close):
                 await close()
         self._chat_apis.clear()
+        self._chat_send_locks.clear()
         if not callable(getattr(self._agent_api, "for_chat", None)):
             close = getattr(self._agent_api, "close", None)
             if callable(close):