Removing an agent from a room does not stop message delivery

core/switch_core/bridges/agent/protocol/event_buffer.py

@@ -255,8 +255,19 @@ def remove(self, agent_id: str) -> None:
     # Confirming clients get at-least-once instead.
     # ------------------------------------------------------------------
 
-    async def poll(self, agent_id: str, timeout: float = 30) -> list[AgentEvent]:
-        return await self._legacy_poll(agent_id, "legacy:all", timeout=timeout)
+    async def poll(
+        self, agent_id: str, timeout: float = 30, rooms: set[str] | None = None
+    ) -> list[AgentEvent]:
+        """Everything queued for this agent, optionally limited to `rooms`.
+
+        `rooms` is how the caller applies membership. The buffer is keyed by
+        agent and knows nothing about who is in what, so an event queued while
+        the agent was a member stays queued after it is removed; passing the
+        rooms it is in now is what keeps that event from being handed over.
+        """
+        return await self._legacy_poll(
+            agent_id, "legacy:all", timeout=timeout, rooms=rooms
+        )
 
     async def poll_room(
         self, agent_id: str, room_id: str, timeout: float = 30

core/switch_core/bridges/agent/protocol/service.py

@@ -1779,11 +1779,26 @@ async def download_media(
     # ── Events ───────────────────────────────────────────────────────────────
 
     async def poll_events(self, agent_id: str, timeout: float = 10) -> list[AgentEvent]:
-        """Poll for events across all rooms the agent is in."""
+        """Poll for events across all rooms the agent is in.
+
+        Membership is applied here rather than trusted from the buffer, which
+        is keyed by agent and knows nothing about who is in what. Its
+        room-scoped sibling `poll_room_events` calls `require_room_member`;
+        without the same check this call would hand over events from a room
+        the agent has since been removed from.
+        """
         async with self.session_factory() as session:
             await self.agent_session_store.touch_heartbeat(session, agent_id, None)
             await session.commit()
-        return await self.event_buffer.poll(agent_id, timeout=timeout)
+            # Archived rooms included: what is being applied here is
+            # membership, and an agent still in an archived room was not
+            # removed from it.
+            rooms = await self.room_store.get_rooms_for_agent(
+                session, agent_id, include_archived=True
+            )
+        return await self.event_buffer.poll(
+            agent_id, timeout=timeout, rooms={room.id for room in rooms}
+        )
 
     async def poll_notifications(
         self, agent_id: str, timeout: float = 10

core/switch_core/provisioning/postgres.py

@@ -8,9 +8,10 @@
   itself, which is what keeps its transport watching the room. When nobody is
   live the membership is written directly, because a client that is not
   running has nothing to wake and will find the room when it starts.
-- **Removal** deletes the membership. No leave event is written: a departure
-  is not something a reader of the room needs explained, and the timeline the
-  log serves is what was said.
+- **Removal** deletes the membership and wakes a live client so it drops the
+  room, mirroring the invitation. No leave event is written: a departure is
+  not something a reader of the room needs explained, and the timeline the log
+  serves is what was said.
 - **Discarding a room** is left to the caller's own delete, which cascades.
   Nothing here has a second copy to clean up.
 """
@@ -122,13 +123,23 @@ async def invite_to_room(self, room_id: str, user_id: str) -> None:
             await session.commit()
 
     async def kick_user(self, room_id: str, user_id: str) -> None:
-        """Remove a membership. Already out is success, per the port."""
+        """Remove a membership, stopping a live client's delivery with it.
+
+        Already out is success, per the port.
+
+        The wake-up is the other half of the removal, not a nicety. A running
+        client holds its own subscription to the room, so deleting the row on
+        its own leaves it reading a room it is no longer in. It is rung after
+        the commit so a client that reacts by re-reading its rooms cannot see
+        the membership it was just removed from.
+        """
         async with self._session_factory() as session:
             switch_room_id, client_id, _ = await self._resolve(
                 session, room_id, user_id
             )
             await self._room_store.remove_client(session, client_id, switch_room_id)
             await session.commit()
+        await self._invites.remove(user_id, room_id)
 
     async def delete_room(self, room_id: str) -> None:
         """Nothing of its own to discard.

core/switch_core/transport/invites.py

@@ -1,4 +1,4 @@
-"""How a client is told it has been added to a room.
+"""How a client is told its membership of a room changed.
 
 Over Matrix this was an event: the admin invited a user, the user's sync loop
 saw the invitation and `ClientBase.on_invite` joined. Nothing else had to know
@@ -12,6 +12,12 @@
 handle: an invitation, delivered to a live client, auto-accepted by the same
 code path as before.
 
+Removal is the same signal in reverse, and it is not optional. The homeserver
+used to enforce a kick: a removed client's sync simply stopped returning the
+room. Here the subscription is the client's own, so deleting the membership row
+leaves it reading a room it is no longer in. A removal must therefore be rung
+as well as written.
+
 **It is in-process, and that is a real limitation.** A client running in
 another replica of switch-core would not hear it. That is the same constraint
 Matrix sync sessions imposed and the reason switch-core is single-replica
@@ -32,17 +38,29 @@
 
 
 class InviteBus:
-    """Routes an invitation to the live transport for a user, if there is one."""
+    """Routes a membership change to the live transport for a user, if any.
+
+    Named for the signal it carried first. It now carries removals too, since
+    both are the same thing: telling a running client that the set of rooms it
+    should be reading has changed underneath it.
+    """
 
     def __init__(self) -> None:
         self._handlers: dict[str, InviteHandler] = {}
+        self._removal_handlers: dict[str, InviteHandler] = {}
 
     def register(self, user_id: str, handler: InviteHandler) -> None:
         self._handlers[user_id] = handler
 
     def unregister(self, user_id: str) -> None:
         self._handlers.pop(user_id, None)
 
+    def register_removal(self, user_id: str, handler: InviteHandler) -> None:
+        self._removal_handlers[user_id] = handler
+
+    def unregister_removal(self, user_id: str) -> None:
+        self._removal_handlers.pop(user_id, None)
+
     async def invite(self, user_id: str, transport_room_id: str) -> bool:
         """Tell `user_id` it is in `transport_room_id`.
 
@@ -55,3 +73,17 @@ async def invite(self, user_id: str, transport_room_id: str) -> bool:
             return False
         await handler(transport_room_id)
         return True
+
+    async def remove(self, user_id: str, transport_room_id: str) -> bool:
+        """Tell `user_id` it is no longer in `transport_room_id`.
+
+        Returns whether anyone was listening, for symmetry with `invite`, but
+        the caller has nothing to fall back to: a client that is not running
+        holds no subscription to drop, and one that starts later reads its
+        rooms from the table the removal already updated.
+        """
+        handler = self._removal_handlers.get(user_id)
+        if handler is None:
+            return False
+        await handler(transport_room_id)
+        return True

core/switch_core/transport/postgres.py

@@ -175,11 +175,13 @@ async def receive_forever(self) -> None:
             await self._watch(transport_room_id)
         self._receiving = True
         self._invites.register(self.user_id, self._on_invited)
+        self._invites.register_removal(self.user_id, self._on_removed)
         try:
             await self._closed.wait()
         finally:
             self._receiving = False
             self._invites.unregister(self.user_id)
+            self._invites.unregister_removal(self.user_id)
             self._unwatch_all()
             delivery.cancel()
 
@@ -206,6 +208,11 @@ async def _deliver_forever(self) -> None:
             self._delivering = True
             try:
                 for room_id in rooms:
+                    # Re-checked per room, not once up front: a removal that
+                    # lands while this batch is draining must stop the room it
+                    # has not reached yet.
+                    if room_id not in self._watching:
+                        continue
                     try:
                         await self._drain_room(room_id)
                     except asyncio.CancelledError:
@@ -245,6 +252,17 @@ async def _on_invited(self, transport_room_id: str) -> None:
             ),
         )
 
+    async def _on_removed(self, transport_room_id: str) -> None:
+        """This client was taken out of a room while it was running.
+
+        The subscription is dropped here rather than left to the next restart.
+        Under Matrix a kick ended the room's delivery at the homeserver, so
+        nothing downstream had to check membership again; here the
+        subscription is this client's own and outlives the row unless it is
+        taken back.
+        """
+        self._unwatch(transport_room_id)
+
     async def _watch(
         self, transport_room_id: str, *, from_seq: int | None = None
     ) -> None:
@@ -280,6 +298,25 @@ async def _watch(
             self._pending.add(room_id)
             self._wake.set()
 
+    def _unwatch(self, transport_room_id: str) -> None:
+        """Stop delivering one room. Not watching it is success.
+
+        The cursor goes with the subscription. Keeping it would mean a client
+        added back to the room resumed from where it left off and was handed
+        everything said while it was out, which is the leak this closes said a
+        different way.
+        """
+        room_id = self._room_ids.get(transport_room_id)
+        if room_id is None or room_id not in self._watching:
+            return
+        self._listener.unsubscribe(room_id, self._on_room_advanced)
+        self._ephemeral.unsubscribe(transport_room_id, self._on_ephemeral)
+        del self._watching[room_id]
+        self._cursors.pop(room_id, None)
+        # A wake-up already queued for this room would otherwise be drained
+        # after the subscription was dropped.
+        self._pending.discard(room_id)
+
     def _unwatch_all(self) -> None:
         for room_id, transport_room_id in self._watching.items():
             self._listener.unsubscribe(room_id, self._on_room_advanced)

core/tests/switch_core/bridges/agent/protocol/test_event_buffer.py

@@ -20,10 +20,10 @@
 ROOM = "room-1"
 
 
-def _message(addressed: bool) -> AgentEvent:
+def _message(addressed: bool, room_id: str = ROOM) -> AgentEvent:
     return AgentEvent(
         type="message",
-        room_id=ROOM,
+        room_id=room_id,
         payload=MessagePayload(
             addressed=addressed,
             sender="@u:s",
@@ -102,6 +102,30 @@ async def test_room_join_fans_out_only_when_listening() -> None:
     assert notifs[0].type == "room_join"
 
 
+async def test_polling_everything_can_be_limited_to_the_rooms_given() -> None:
+    """The buffer is keyed by agent and knows nothing about who is in what.
+
+    An event queued while the agent was a member stays queued after it is
+    removed, so the caller passes the rooms it is in now and that is what
+    keeps the event from being handed over.
+    """
+    q = EventBuffer()
+    q.enqueue(AGENT, ROOM, _message(addressed=False))
+    q.enqueue(AGENT, "room-2", _message(addressed=False, room_id="room-2"))
+
+    polled = await q.poll(AGENT, timeout=0, rooms={"room-2"})
+
+    assert [event.room_id for event in polled] == ["room-2"]
+
+
+async def test_polling_with_no_rooms_at_all_returns_nothing() -> None:
+    """An agent in no rooms is not a caller asking for every room."""
+    q = EventBuffer()
+    q.enqueue(AGENT, ROOM, _message(addressed=False))
+
+    assert await q.poll(AGENT, timeout=0, rooms=set()) == []
+
+
 async def test_remove_clears_notification_queue() -> None:
     q = EventBuffer()
     q.enqueue(AGENT, ROOM, _message(addressed=True))

core/tests/switch_core/bridges/agent/protocol/test_poll_events_membership.py

@@ -0,0 +1,74 @@
+"""Polling every room applies membership, as polling one room does.
+
+The buffer is keyed by agent, so an event queued while the agent was a member
+is still there after it is removed. `poll_room_events` asks
+`require_room_member` first; this is the same question asked by the call that
+has no room in it.
+"""
+
+from __future__ import annotations
+
+from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker
+
+from switch_core.bridges.agent.protocol.event_buffer import EventBuffer
+from switch_core.bridges.agent.protocol.types import AgentEvent, MessagePayload
+from switch_core.db.models import Room
+from switch_core.db.stores.agent_session_store import AgentSessionStore
+from switch_core.db.stores.room_store import RoomStore
+from tests.switch_core.bridges.agent.protocol.registration_harness import (
+    make_owner,
+    make_service,
+    register,
+)
+
+
+def _message(room_id: str, body: str) -> AgentEvent:
+    return AgentEvent(
+        type="message",
+        room_id=room_id,
+        payload=MessagePayload(
+            addressed=False,
+            sender="@someone:test",
+            sender_name="someone",
+            message_id="$m",
+            body=body,
+            timestamp=0,
+        ),
+    )
+
+
+async def _room(
+    session_factory: async_sessionmaker[AsyncSession], name: str, agent_id: str
+) -> str:
+    async with session_factory() as session:
+        room = Room(matrix_room_id=f"!{name}:test", name=name, description="")
+        session.add(room)
+        await session.flush()
+        await RoomStore().add_agents(session, room.id, [agent_id])
+        await session.commit()
+        return room.id
+
+
+async def test_events_from_a_room_the_agent_was_removed_from_are_not_returned(
+    session_factory: async_sessionmaker[AsyncSession],
+) -> None:
+    svc = make_service(session_factory)
+    svc.room_store = RoomStore()  # type: ignore[attr-defined]
+    svc.agent_session_store = AgentSessionStore()  # type: ignore[attr-defined]
+    svc.event_buffer = EventBuffer()  # type: ignore[attr-defined]
+
+    owner = await make_owner(session_factory)
+    agent_id = await register(svc, "poller", owner)
+    kept = await _room(session_factory, "kept", agent_id)
+    left = await _room(session_factory, "left", agent_id)
+
+    svc.event_buffer.enqueue(agent_id, left, _message(left, "said in the old room"))
+    svc.event_buffer.enqueue(agent_id, kept, _message(kept, "said in this one"))
+
+    async with session_factory() as session:
+        await RoomStore().remove_agents(session, left, [agent_id])
+        await session.commit()
+
+    events = await svc.poll_events(agent_id, timeout=0)
+
+    assert [event.payload.body for event in events] == ["said in this one"]

core/tests/switch_core/provisioning/test_postgres_provisioning.py

@@ -177,6 +177,42 @@ async def test_removing_a_member_takes_the_membership_away(
             )
         assert membership is None
 
+    async def test_removing_a_running_member_tells_it_to_stop_reading(
+        self, session_factory: async_sessionmaker[AsyncSession]
+    ) -> None:
+        """Deleting the row is only half of a removal.
+
+        A running client holds its own subscription to the room, so the row
+        going away leaves it reading a room it is no longer in. The homeserver
+        used to end the delivery itself; here it has to be rung.
+        """
+        async with session_factory() as session:
+            room = await _room(session)
+            client = await _client(session)
+            await session.commit()
+            transport_room_id, user_id = room.matrix_room_id, client.matrix_user_id
+            room_id, client_id = room.id, client.id
+
+        told: list[str] = []
+
+        async def _handler(removed_from: str) -> None:
+            told.append(removed_from)
+
+        invites = InviteBus()
+        provisioning = _provisioning(session_factory, invites)
+        await provisioning.invite_to_room(transport_room_id, user_id)
+        invites.register_removal(user_id, _handler)
+        await provisioning.kick_user(transport_room_id, user_id)
+
+        assert told == [transport_room_id]
+        # Rung after the row is gone, so a client that responds by re-reading
+        # its rooms cannot see the membership it was just removed from.
+        async with session_factory() as session:
+            membership = await session.get(
+                ClientRoom, {"client_id": client_id, "room_id": room_id}
+            )
+        assert membership is None
+
     async def test_removing_a_member_who_is_already_out_is_success(
         self, session_factory: async_sessionmaker[AsyncSession]
     ) -> None: