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]