Auto-join space rooms when accepting space invites

src/app/pages/app-shell/CommunityList.tsx

@@ -6,6 +6,7 @@ import { AsyncStatus, useAsyncCallback } from '../../hooks/useAsyncCallback';
 import { getMemberDisplayName, getRoomAvatarUrl } from '../../utils/room';
 import { getMxIdLocalPart } from '../../utils/matrix';
 import { nameInitials } from '../../utils/common';
+import { joinSpaceRooms } from '../../utils/joinSpaceRooms';
 import { PlusIcon } from './AppIcons';
 import * as css from './CommunitiesDrawer.css';
 
@@ -70,6 +71,8 @@ function CommunityInviteRow({ room, onJoined }: CommunityInviteRowProps) {
   const [joinState, join] = useAsyncCallback<void, MatrixError, []>(
     useCallback(async () => {
       await mx.joinRoom(room.roomId);
+      // Join the community's rooms in the background; don't hold up navigation.
+      joinSpaceRooms(mx, room.roomId);
       onJoined(room.roomId);
     }, [mx, room, onJoined])
   );

src/app/pages/client/inbox/Invites.tsx

@@ -67,6 +67,7 @@ import { useIgnoredUsers } from '../../../hooks/useIgnoredUsers';
 import { useReportRoomSupported } from '../../../hooks/useReportRoomSupported';
 import { useSetting } from '../../../state/hooks/settings';
 import { settingsAtom } from '../../../state/settings';
+import { joinSpaceRooms } from '../../../utils/joinSpaceRooms';
 
 const COMPACT_CARD_WIDTH = 548;
 
@@ -176,6 +177,10 @@ function InviteCard({
       if (dmUserId) {
         await addRoomIdToMDirect(mx, invite.roomId, dmUserId);
       }
+      if (invite.isSpace) {
+        // Join the space's rooms in the background; don't hold up navigation.
+        joinSpaceRooms(mx, invite.roomId);
+      }
       onNavigate(invite.roomId, invite.isSpace);
     }, [mx, invite, userId, onNavigate])
   );

src/app/utils/joinSpaceRooms.test.ts

@@ -0,0 +1,71 @@
+import { describe, expect, it } from 'vitest';
+import { IHierarchyRoom } from 'matrix-js-sdk/lib/@types/spaces';
+import { getSpaceRoomsToJoin } from './joinSpaceRooms';
+
+const hierarchyRoom = (
+  roomId: string,
+  extra: Omit<Partial<IHierarchyRoom>, 'join_rule'> & { join_rule?: string } = {}
+): IHierarchyRoom =>
+  ({
+    room_id: roomId,
+    world_readable: false,
+    guest_can_join: false,
+    num_joined_members: 1,
+    children_state: [],
+    ...extra,
+  } as IHierarchyRoom);
+
+const child = (stateKey: string, via?: unknown) => ({
+  type: 'm.space.child',
+  state_key: stateKey,
+  sender: '@alice:example.org',
+  origin_server_ts: 0,
+  content: via === undefined ? {} : { via },
+});
+
+describe('getSpaceRoomsToJoin', () => {
+  const spaceId = '!space:example.org';
+
+  it('returns unjoined rooms and sub-spaces with their via servers, in hierarchy order', () => {
+    const rooms = [
+      hierarchyRoom(spaceId, {
+        room_type: 'm.space',
+        children_state: [
+          child('!sub:example.org', ['example.org']),
+          child('!a:example.org', ['a.org', 5]),
+        ],
+      } as Partial<IHierarchyRoom>),
+      hierarchyRoom('!sub:example.org', {
+        room_type: 'm.space',
+        children_state: [child('!b:example.org', ['b.org'])],
+      } as Partial<IHierarchyRoom>),
+      hierarchyRoom('!b:example.org'),
+      hierarchyRoom('!a:example.org'),
+    ];
+
+    expect(getSpaceRoomsToJoin(spaceId, rooms, () => false)).toEqual([
+      { roomId: '!sub:example.org', via: ['example.org'] },
+      { roomId: '!b:example.org', via: ['b.org'] },
+      { roomId: '!a:example.org', via: ['a.org'] },
+    ]);
+  });
+
+  it('skips the space itself, joined rooms, duplicates and invite-only rooms', () => {
+    const rooms = [
+      hierarchyRoom(spaceId, { room_type: 'm.space' } as Partial<IHierarchyRoom>),
+      hierarchyRoom('!joined:example.org'),
+      hierarchyRoom('!invite:example.org', { join_rule: 'invite' }),
+      hierarchyRoom('!knock:example.org', { join_rule: 'knock' }),
+      hierarchyRoom('!public:example.org', { join_rule: 'public' }),
+      hierarchyRoom('!restricted:example.org', { join_rule: 'restricted' }),
+      hierarchyRoom('!public:example.org', { join_rule: 'public' }),
+    ];
+
+    expect(
+      getSpaceRoomsToJoin(spaceId, rooms, (roomId) => roomId === '!joined:example.org')
+    ).toEqual([
+      { roomId: '!public:example.org', via: [] },
+      { roomId: '!restricted:example.org', via: [] },
+    ]);
+  });
+});

src/app/utils/joinSpaceRooms.ts

@@ -0,0 +1,94 @@
+import { MatrixClient } from 'matrix-js-sdk';
+import { IHierarchyRoom } from 'matrix-js-sdk/lib/@types/spaces';
+import { rateLimitedActions } from './matrix';
+
+export type SpaceRoomToJoin = {
+  roomId: string;
+  via: string[];
+};
+
+// Rooms with these join rules can't be joined without an invite, so there's
+// no point attempting them.
+const UNJOINABLE_JOIN_RULES = new Set<string>(['invite', 'knock', 'private']);
+
+/**
+ * From a space hierarchy response, picks every room (including sub-spaces)
+ * below the space that the user hasn't joined yet and could join, along with
+ * the via servers advertised by the parent's m.space.child event. The
+ * hierarchy's order is kept, so a sub-space comes before its own children
+ * (restricted rooms may require membership of their parent sub-space).
+ */
+export const getSpaceRoomsToJoin = (
+  spaceId: string,
+  hierarchyRooms: IHierarchyRoom[],
+  isJoined: (roomId: string) => boolean
+): SpaceRoomToJoin[] => {
+  const viaMap = new Map<string, string[]>();
+  hierarchyRooms.forEach((room) => {
+    room.children_state?.forEach((childState) => {
+      const via: unknown = childState.content?.via;
+      if (Array.isArray(via) && !viaMap.has(childState.state_key)) {
+        viaMap.set(
+          childState.state_key,
+          via.filter((server): server is string => typeof server === 'string')
+        );
+      }
+    });
+  });
+
+  const seen = new Set<string>();
+  return hierarchyRooms
+    .filter((room) => {
+      if (room.room_id === spaceId || seen.has(room.room_id)) return false;
+      seen.add(room.room_id);
+      if (isJoined(room.room_id)) return false;
+      const joinRule: string | undefined = room.join_rule;
+      return !joinRule || !UNJOINABLE_JOIN_RULES.has(joinRule);
+    })
+    .map((room) => ({
+      roomId: room.room_id,
+      via: viaMap.get(room.room_id) ?? [],
+    }));
+};
+
+const HIERARCHY_PAGE_LIMIT = 100;
+const HIERARCHY_MAX_PAGES = 10;
+
+const fetchSpaceHierarchy = async (
+  mx: MatrixClient,
+  spaceId: string
+): Promise<IHierarchyRoom[]> => {
+  const rooms: IHierarchyRoom[] = [];
+  let from: string | undefined;
+  for (let page = 0; page < HIERARCHY_MAX_PAGES; page += 1) {
+    // eslint-disable-next-line no-await-in-loop
+    const res = await mx.getRoomHierarchy(spaceId, HIERARCHY_PAGE_LIMIT, undefined, false, from);
+    rooms.push(...res.rooms);
+    from = res.next_batch;
+    if (!from) break;
+  }
+  return rooms;
+};
+
+/**
+ * Joins every room in a space the user has just joined. Rooms that fail to
+ * join (e.g. no longer joinable) are skipped; the space join itself has
+ * already succeeded, so this never throws.
+ */
+export const joinSpaceRooms = async (mx: MatrixClient, spaceId: string): Promise<void> => {
+  try {
+    const hierarchy = await fetchSpaceHierarchy(mx, spaceId);
+    const isJoined = (roomId: string) => mx.getRoom(roomId)?.getMyMembership() === 'join';
+    const toJoin = getSpaceRoomsToJoin(spaceId, hierarchy, isJoined);
+
+    // Failed joins are skipped by rateLimitedActions; rate limits are retried.
+    await rateLimitedActions(
+      toJoin,
+      (room) => mx.joinRoom(room.roomId, { viaServers: room.via }),
+      3
+    );
+  } catch (err) {
+    // eslint-disable-next-line no-console
+    console.warn(`Failed to join rooms of space ${spaceId}`, err);
+  }
+};