Update state store after creating room

Author: tulir

mautrix/client/api/rooms.py

@@ -5,7 +5,7 @@
 # file, You can obtain one at http://mozilla.org/MPL/2.0/.
 from __future__ import annotations
 
-from typing import Any, Awaitable, Callable, Iterable
+from typing import Any, Awaitable, Callable
 import asyncio
 
 from multidict import CIMultiDict
@@ -64,7 +64,7 @@ async def create_room(
         topic: str | None = None,
         is_direct: bool = False,
         invitees: list[UserID] | None = None,
-        initial_state: Iterable[StateEvent | StrippedStateEvent | dict[str, JSON]] | None = None,
+        initial_state: list[StateEvent | StrippedStateEvent | dict[str, JSON]] | None = None,
         room_version: str = None,
         creation_content: RoomCreateStateEventContent | dict[str, JSON] | None = None,
         power_level_override: PowerLevelStateEventContent | dict[str, JSON] | None = None,

mautrix/client/store_updater.py

@@ -149,6 +149,28 @@ async def get_state(self, room_id: RoomID) -> list[StateEvent]:
             )
         return state
 
+    async def create_room(self, *args, **kwargs) -> RoomID:
+        room_id = await super().create_room(*args, **kwargs)
+        if self.state_store:
+            invitee_membership = Membership.INVITE
+            if kwargs.get("beeper_auto_join_invites"):
+                invitee_membership = Membership.JOIN
+            for user_id in kwargs.get("invitees", []):
+                await self.state_store.set_membership(room_id, user_id, invitee_membership)
+            for evt in kwargs.get("initial_state", []):
+                await self.state_store.update_state(
+                    StateEvent(
+                        type=EventType.find(evt["type"], t_class=EventType.Class.STATE),
+                        room_id=room_id,
+                        event_id=EventID("$fake-create-id"),
+                        sender=self.mxid,
+                        state_key=evt.get("state_key", ""),
+                        timestamp=0,
+                        content=evt["content"],
+                    )
+                )
+        return room_id
+
     async def send_state_event(
         self,
         room_id: RoomID,