From 8208ff3e0768f40616f6da73adab210ebfdbda11 Mon Sep 17 00:00:00 2001 From: James Howe Date: Wed, 9 Sep 2026 12:20:28 +0100 Subject: [PATCH] Stop delivering a room to a client removed from it Removing an agent from a room deletes the membership row, but a running client holds its own per-room subscription and keeps receiving everything said in the room. Matrix used to enforce the kick at the homeserver, and that enforcement went away with it. Same leak on the pull side: poll_events returns everything buffered for the agent with no membership check, unlike its sibling poll_room_events. A removal is now rung over the invite bus the same way an invitation is, and the transport drops the room's subscription and its cursor. poll_events filters by the rooms the agent is in. --- .../bridges/agent/protocol/event_buffer.py | 15 +- .../bridges/agent/protocol/service.py | 19 +- core/switch_core/provisioning/postgres.py | 19 +- core/switch_core/transport/invites.py | 36 +++- core/switch_core/transport/postgres.py | 37 ++++ .../agent/protocol/test_event_buffer.py | 28 ++- .../protocol/test_poll_events_membership.py | 74 ++++++++ .../test_postgres_provisioning.py | 36 ++++ .../transport/test_postgres_transport.py | 172 +++++++++++++++++- 9 files changed, 422 insertions(+), 14 deletions(-) create mode 100644 core/tests/switch_core/bridges/agent/protocol/test_poll_events_membership.py diff --git a/core/switch_core/bridges/agent/protocol/event_buffer.py b/core/switch_core/bridges/agent/protocol/event_buffer.py index d3256896c..41861a0bb 100644 --- a/core/switch_core/bridges/agent/protocol/event_buffer.py +++ b/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 diff --git a/core/switch_core/bridges/agent/protocol/service.py b/core/switch_core/bridges/agent/protocol/service.py index 8e5390d33..287e4e0d4 100644 --- a/core/switch_core/bridges/agent/protocol/service.py +++ b/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 diff --git a/core/switch_core/provisioning/postgres.py b/core/switch_core/provisioning/postgres.py index f5ac7c271..9aa4c947a 100644 --- a/core/switch_core/provisioning/postgres.py +++ b/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. diff --git a/core/switch_core/transport/invites.py b/core/switch_core/transport/invites.py index 14a94c601..48d27a504 100644 --- a/core/switch_core/transport/invites.py +++ b/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,10 +38,16 @@ 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 @@ -43,6 +55,12 @@ def register(self, user_id: str, handler: InviteHandler) -> None: 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 diff --git a/core/switch_core/transport/postgres.py b/core/switch_core/transport/postgres.py index 822d1ef02..5ef732079 100644 --- a/core/switch_core/transport/postgres.py +++ b/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) diff --git a/core/tests/switch_core/bridges/agent/protocol/test_event_buffer.py b/core/tests/switch_core/bridges/agent/protocol/test_event_buffer.py index c68fd10ea..4cdbc3f64 100644 --- a/core/tests/switch_core/bridges/agent/protocol/test_event_buffer.py +++ b/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)) diff --git a/core/tests/switch_core/bridges/agent/protocol/test_poll_events_membership.py b/core/tests/switch_core/bridges/agent/protocol/test_poll_events_membership.py new file mode 100644 index 000000000..8160114f4 --- /dev/null +++ b/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"] diff --git a/core/tests/switch_core/provisioning/test_postgres_provisioning.py b/core/tests/switch_core/provisioning/test_postgres_provisioning.py index 6c943f146..4f4ece530 100644 --- a/core/tests/switch_core/provisioning/test_postgres_provisioning.py +++ b/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: diff --git a/core/tests/switch_core/transport/test_postgres_transport.py b/core/tests/switch_core/transport/test_postgres_transport.py index 2373674d9..c2a491086 100644 --- a/core/tests/switch_core/transport/test_postgres_transport.py +++ b/core/tests/switch_core/transport/test_postgres_transport.py @@ -20,9 +20,11 @@ from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker from switch_core.db.models import Client, Room +from switch_core.db.stores.client_store import ClientStore from switch_core.db.stores.media_store import MediaStore from switch_core.db.stores.message_store import MessageStore from switch_core.db.stores.room_store import RoomStore +from switch_core.provisioning.postgres import PostgresProvisioning from switch_core.transport import ( InboundCustomEvent, InboundMedia, @@ -56,6 +58,18 @@ async def _make_room(session: AsyncSession) -> tuple[str, str, str, str]: return room.id, room.matrix_room_id, client.id, client.matrix_user_id +async def _make_client(session: AsyncSession, label: str) -> tuple[str, str]: + """Insert one more Client. Returns (client id, client mxid).""" + client = Client( + matrix_user_id=f"@{label}-{uuid.uuid4().hex[:8]}:test", + display_name=label, + type="agent", + ) + session.add(client) + await session.flush() + return client.id, client.matrix_user_id + + async def _watched_room(transport: PostgresTransport) -> str: """Wait until the transport is watching a room, and name it. @@ -118,12 +132,13 @@ class _Received: def __init__(self) -> None: self.events: list[object] = [] - def handlers(self) -> TransportHandlers: + def handlers(self, *, on_invite: object | None = None) -> TransportHandlers: return TransportHandlers( on_message=self._take, on_media=self._take, on_member_event=self._take, on_custom_event=self._take, + on_invite=on_invite, # type: ignore[arg-type] ) async def _take(self, _room, event) -> None: @@ -137,6 +152,7 @@ def _transport( user_id: str, listener: _FakeListener | None = None, ephemeral: EphemeralBus | None = None, + invites: InviteBus | None = None, ) -> PostgresTransport: return PostgresTransport( user_id=user_id, @@ -147,7 +163,7 @@ def _transport( message_store=MessageStore(), media_store=MediaStore(), listener=listener or _FakeListener(), - invites=InviteBus(), + invites=invites or InviteBus(), ephemeral=ephemeral or EphemeralBus(), ) @@ -858,3 +874,155 @@ async def test_it_hears_the_join_once( kinds = [type(event) for event in received.events] assert kinds == [InboundMembership] + + +class TestBeingRemovedFromARoom: + """A removal has to reach the client, not only the table. + + The homeserver enforced a kick: a removed client's sync stopped returning + the room, so nothing downstream had to check membership a second time. + Here the subscription is the client's own and outlives the row it was + resolved from, so a removal is rung over the same bus an invitation is and + the transport drops the room. + """ + + @pytest.fixture(autouse=True) + def _cleanup(self) -> Iterator[None]: + self._tasks: list[asyncio.Task] = [] + yield + for task in self._tasks: + task.cancel() + + def _provisioning( + self, + session_factory: async_sessionmaker[AsyncSession], + invites: InviteBus, + ) -> PostgresProvisioning: + return PostgresProvisioning( + session_factory=session_factory, + room_store=RoomStore(), + client_store=ClientStore(), + message_store=MessageStore(), + invites=invites, + ) + + async def _two_in_a_room( + self, session_factory: async_sessionmaker[AsyncSession] + ) -> tuple[str, str, tuple[str, str], tuple[str, str]]: + """A room with two members already recorded. + + Written directly rather than joined so that no arrival event is in the + way of what the test is reading. + """ + async with session_factory() as session: + room_id, room, leaving_id, leaving_user = await _make_room(session) + staying_id, staying_user = await _make_client(session, "staying") + await RoomStore().add_client(session, leaving_id, room_id) + await RoomStore().add_client(session, staying_id, room_id) + await session.commit() + return room_id, room, (leaving_id, leaving_user), (staying_id, staying_user) + + async def test_it_stops_being_delivered_the_rooms_messages( + self, session_factory: async_sessionmaker[AsyncSession] + ) -> None: + room_id, room, leaving, staying = await self._two_in_a_room(session_factory) + listener = _FakeListener() + invites = InviteBus() + received = _Received() + + removed = _transport( + session_factory, + client_id=leaving[0], + user_id=leaving[1], + listener=listener, + invites=invites, + ) + removed.register_handlers(received.handlers()) + member = _transport( + session_factory, + client_id=staying[0], + user_id=staying[1], + listener=listener, + invites=invites, + ) + member.register_handlers(_Received().handlers()) + self._tasks.append(asyncio.create_task(removed.receive_forever())) + self._tasks.append(asyncio.create_task(member.receive_forever())) + await _watched_room(removed) + await _watched_room(member) + + # While a member, it hears the room. + await member.send_message(room, "while a member", sender_name="agent one") + await listener.announce(room_id) + assert [event.body for event in received.events] == ["while a member"] + + await self._provisioning(session_factory, invites).kick_user(room, leaving[1]) + + async with session_factory() as session: + members = await RoomStore().get_client_ids(session, room_id) + assert leaving[0] not in members + assert staying[0] in members + + await member.send_message(room, "after removal", sender_name="agent one") + await listener.announce(room_id) + + assert [event.body for event in received.events] == ["while a member"] + assert room_id not in removed._watching + + async def test_being_put_back_does_not_hand_over_what_was_said_while_out( + self, session_factory: async_sessionmaker[AsyncSession] + ) -> None: + """The cursor goes with the subscription. + + Keeping it would mean a client added back resumed from where it left + off and was handed everything said while it was out, which is the same + leak one restart later. + """ + room_id, room, leaving, staying = await self._two_in_a_room(session_factory) + listener = _FakeListener() + invites = InviteBus() + received = _Received() + + removed = _transport( + session_factory, + client_id=leaving[0], + user_id=leaving[1], + listener=listener, + invites=invites, + ) + + async def _accept(room_ref, _event) -> None: + """What the client does with an invitation, so a put-back is one.""" + await removed.join_room(room_ref.room_id) + + removed.register_handlers(received.handlers(on_invite=_accept)) + member = _transport( + session_factory, + client_id=staying[0], + user_id=staying[1], + listener=listener, + invites=invites, + ) + member.register_handlers(_Received().handlers()) + self._tasks.append(asyncio.create_task(removed.receive_forever())) + self._tasks.append(asyncio.create_task(member.receive_forever())) + await _watched_room(removed) + await _watched_room(member) + + provisioning = self._provisioning(session_factory, invites) + await provisioning.kick_user(room, leaving[1]) + await member.send_message( + room, "said while it was out", sender_name="agent one" + ) + await listener.announce(room_id) + + await provisioning.invite_to_room(room, leaving[1]) + await _settled(removed) + + bodies = [ + event.body for event in received.events if isinstance(event, InboundMessage) + ] + assert bodies == [] + # It is back in, and told so, which is how it was told the first time. + assert room_id in removed._watching + assert [type(event) for event in received.events] == [InboundMembership]