From 996b9d7e5cc3313033bb83c818da0a6b44985e2b Mon Sep 17 00:00:00 2001 From: Vineeth Voruganti <13438633+VVoruganti@users.noreply.github.com> Date: Tue, 25 Aug 2026 10:47:39 -0400 Subject: [PATCH] fix: preserve session observer limit --- src/crud/session.py | 132 ++++++++++++++++++++----------- tests/crud/test_session.py | 155 ++++++++++++++++++++++++++++++++++++- 2 files changed, 238 insertions(+), 49 deletions(-) diff --git a/src/crud/session.py b/src/crud/session.py index 707e5ceb..51ca1f13 100644 --- a/src/crud/session.py +++ b/src/crud/session.py @@ -1052,8 +1052,10 @@ async def set_peers_for_session( """ Replace a session's ordinary peer set with ``peer_names``. - Active members keep their joined_at and configuration. Departed members - rejoin with the incoming configuration. Scope memberships are preserved. + Active members keep their joined_at but take the incoming configuration: + this is a replace, so the caller's map is the desired end state. Departed + members rejoin with the incoming configuration. Scope memberships are + preserved. Args: db: Database session @@ -1066,11 +1068,13 @@ async def set_peers_for_session( Raises: ResourceNotFoundException: If the session does not exist + ObserverException: If the resulting peer set would exceed the observer + limit """ - # Validate observer limit before making any changes - observer_count = count_observers_in_config(peer_names) - if observer_count > settings.SESSION_OBSERVERS_LIMIT: - raise ObserverException(session_name, observer_count) + # No observer pre-check here: an already-active membership keeps its stored + # configuration, so the incoming map is not what lands. Counting it would + # reject a request that lowers the observer count as often as one that raises + # it. _get_or_add_peers_to_session enforces the limit on the resulting rows. # Verify session exists stmt = ( @@ -1121,12 +1125,14 @@ async def set_peers_for_session( ) _reject_resolved_scope_peers(peers_result.resource) - # Add new peers to session + # Add new peers to session. This route replaces the session's peer set, so the + # incoming configuration is authoritative even for an already-active member. peers = await _get_or_add_peers_to_session( db, workspace_name=workspace_name, session_name=session_name, peer_names=peer_names, + replace_config=True, ) await db.commit() @@ -1165,13 +1171,22 @@ async def _get_or_add_peers_to_session( peer_names: dict[str, schemas.SessionPeerConfig], *, fetch_after_upsert: bool = True, + replace_config: bool = False, ) -> list[models.SessionPeer]: """ Upsert session-peer memberships for a session and optionally fetch the active memberships afterward. New peers are inserted, peers that previously left the session are rejoined, - and already-active peers keep their existing session-level configuration. + and already-active peers keep their existing joined_at. + + An already-active peer also keeps its stored configuration unless + ``replace_config`` is set: an add must not overwrite configuration it was + never asked about, while a replace states the desired end state. + + The observer limit is checked against the rows the upsert actually produced, + not against the incoming map, since under the add semantics the incoming map + is not necessarily what lands. Args: db: Database session @@ -1180,13 +1195,17 @@ async def _get_or_add_peers_to_session( peer_names: Mapping of peer names to session-level configuration fetch_after_upsert: If True, query and return the active session peers after the upsert. If False, skip that read and return an empty list. + replace_config: If True, an already-active membership takes the incoming + configuration instead of keeping its stored one. Set by replace-style + callers; leave False for add-style callers. Returns: Active SessionPeer objects after the upsert, or an empty list when the post-upsert fetch is skipped Raises: - ObserverException: If adding peers would exceed the observer limit + ObserverException: If the resulting active peer set would exceed the + observer limit """ # If no peers to add, skip the insert and just return existing active session peers if not peer_names: @@ -1205,43 +1224,10 @@ async def _get_or_add_peers_to_session( # costs document rows, not LLM calls, and counting them would # cap scopes-per-session at SESSION_OBSERVERS_LIMIT and surface as an # observer-shaped 400 through a facade that hides observers entirely. + # Resolved up front because the limit check below gates on whether this + # request asks for a *non-scope* observer. scopes_being_added = await scope_peer_names(db, workspace_name, peer_names.keys()) - # Only validate observer limit if we're adding non-scope peers with observe_others=True - new_observer_count = count_observers_in_config( - {n: c for n, c in peer_names.items() if n not in scopes_being_added} - ) - - if new_observer_count > 0: - # Use a single efficient query to count existing observers not being updated - # This uses PostgreSQL's JSONB operators to check the observe_others field directly - existing_observers_stmt = select(func.count()).where( - models.SessionPeer.session_name == session_name, - models.SessionPeer.workspace_name == workspace_name, - models.SessionPeer.left_at.is_(None), # Only active peers - models.SessionPeer.peer_name.notin_( - peer_names.keys() - ), # Exclude peers being updated - models.SessionPeer.configuration["observe_others"].astext.cast( - Boolean - ), # Only observers - # Existing scope memberships are excluded for the same reason as above. - ~exists( - select(models.Peer.id) - .where(models.Peer.workspace_name == workspace_name) - .where(models.Peer.name == models.SessionPeer.peer_name) - .where(scope_peer_clause()) - .correlate(models.SessionPeer) - ), - ) - result = await db.execute(existing_observers_stmt) - existing_observer_count = result.scalar() or 0 - - total_observers = existing_observer_count + new_observer_count - - if total_observers > settings.SESSION_OBSERVERS_LIMIT: - raise ObserverException(session_name, total_observers) - # Use upsert to handle both new peers and rejoining peers stmt = pg_insert(models.SessionPeer).values( [ @@ -1257,7 +1243,16 @@ async def _get_or_add_peers_to_session( ] ) - # On conflict, rejoin departed peers and leave active memberships unchanged. + # On conflict, rejoin departed peers. joined_at always survives on an active + # membership -- advancing it would move the peer_perspective search window + # past messages the peer was present for (issue #940). + # + # Configuration depends on the caller's semantics. An add ("ensure this peer + # is here") must not silently overwrite a config it never asked about, so an + # active membership keeps its stored one. A replace ("these are the session's + # peers, configured thus") states a desired end state, so the incoming config + # wins -- otherwise PUT /peers could never change the configuration of a peer + # already in the session. stmt = stmt.on_conflict_do_update( index_elements=["session_name", "peer_name", "workspace_name"], set_={ @@ -1266,7 +1261,9 @@ async def _get_or_add_peers_to_session( else_=models.SessionPeer.joined_at, ), "left_at": None, - "configuration": case( + "configuration": stmt.excluded.configuration + if replace_config + else case( (models.SessionPeer.left_at.is_not(None), stmt.excluded.configuration), else_=models.SessionPeer.configuration, ), @@ -1274,6 +1271,49 @@ async def _get_or_add_peers_to_session( ) await db.execute(stmt) + # Enforce the observer limit on the resulting rows rather than predicting them. + # Under add semantics an already-active membership keeps its stored + # configuration (see the CASE above), so the incoming config is not what lands + # and cannot be counted: predicting from it silently undercounts preserved + # observers and lets a session grow past the limit indefinitely by re-sending + # its current observers at a lower config alongside new ones. Counting after + # the upsert is correct under both configuration semantics and cannot desync + # from those branches. Raising here rolls the upsert back: ObserverException + # is never caught, and both get_db and tracked_db roll back on exception. + # + # Gated on the request actually asking for a non-scope observer so that a + # session already over the limit keeps behaving as it does today: it can + # still take non-observers and scope attachments, and only a request that + # would make it worse is rejected. + if any( + config.observe_others + for peer_name, config in peer_names.items() + if peer_name not in scopes_being_added + ): + observer_count = ( + await db.scalar( + select(func.count()).where( + models.SessionPeer.session_name == session_name, + models.SessionPeer.workspace_name == workspace_name, + models.SessionPeer.left_at.is_(None), # Only active peers + models.SessionPeer.configuration["observe_others"].astext.cast( + Boolean + ), # Only observers + # Scope memberships are excluded for the reason given above. + ~exists( + select(models.Peer.id) + .where(models.Peer.workspace_name == workspace_name) + .where(models.Peer.name == models.SessionPeer.peer_name) + .where(scope_peer_clause()) + .correlate(models.SessionPeer) + ), + ) + ) + or 0 + ) + if observer_count > settings.SESSION_OBSERVERS_LIMIT: + raise ObserverException(session_name, observer_count) + if not fetch_after_upsert: return [] diff --git a/tests/crud/test_session.py b/tests/crud/test_session.py index bcc3f089..2eec1f32 100644 --- a/tests/crud/test_session.py +++ b/tests/crud/test_session.py @@ -2,11 +2,12 @@ from datetime import datetime, timezone import pytest from nanoid import generate as generate_nanoid -from sqlalchemy import select +from sqlalchemy import Boolean, func, select from sqlalchemy.ext.asyncio import AsyncSession from src import crud, models, schemas -from src.exceptions import ResourceNotFoundException +from src.config import settings +from src.exceptions import ObserverException, ResourceNotFoundException class TestSessionCRUD: @@ -165,7 +166,9 @@ class TestSessionCRUD: ).one() assert active_joined_at == datetime(2020, 1, 1, tzinfo=timezone.utc) assert active_left_at is None - assert active_config == original_config.model_dump() + # A replace states the desired end state, so the incoming config lands even + # though the membership window is untouched. + assert active_config == updated_config.model_dump() await crud.set_peers_for_session( db_session, @@ -186,6 +189,152 @@ class TestSessionCRUD: assert rejoined_left_at is None assert rejoined_config == updated_config.model_dump() + @pytest.mark.asyncio + async def test_observer_limit_counts_preserved_config( + self, + db_session: AsyncSession, + sample_data: tuple[models.Workspace, models.Peer], + monkeypatch: pytest.MonkeyPatch, + ): + """On the add path an already-active observer keeps its stored config, so + it still counts against the limit when re-sent as a non-observer.""" + monkeypatch.setattr(settings, "SESSION_OBSERVERS_LIMIT", 2) + test_workspace, _ = sample_data + # Bound to a local: rollback below expires the ORM instance, and reloading + # it would lazy-load outside the greenlet context. + workspace_name = test_workspace.name + session_name = str(generate_nanoid()) + observer = schemas.SessionPeerConfig(observe_others=True, observe_me=False) + bystander = schemas.SessionPeerConfig(observe_others=False, observe_me=True) + existing = [str(generate_nanoid()) for _ in range(2)] + + await crud.get_or_create_session( + db_session, + schemas.SessionCreate( + name=session_name, peers=dict.fromkeys(existing, observer) + ), + workspace_name, + ) + + # Adding cannot demote an active member, so re-sending the two observers as + # non-observers leaves them observing and the third peer makes three. + with pytest.raises(ObserverException): + await crud.get_or_create_session( + db_session, + schemas.SessionCreate( + name=session_name, + peers={ + **dict.fromkeys(existing, bystander), + str(generate_nanoid()): observer, + }, + ), + workspace_name, + ) + + # The rejected request left nothing behind. + await db_session.rollback() + observer_count = await db_session.scalar( + select(func.count()).where( + models.SessionPeer.session_name == session_name, + models.SessionPeer.workspace_name == workspace_name, + models.SessionPeer.left_at.is_(None), + models.SessionPeer.configuration["observe_others"].astext.cast(Boolean), + ) + ) + assert observer_count == 2 + + @pytest.mark.asyncio + async def test_set_peers_observer_limit_counts_replaced_config( + self, + db_session: AsyncSession, + sample_data: tuple[models.Workspace, models.Peer], + monkeypatch: pytest.MonkeyPatch, + ): + """The replace path applies the incoming config, so demoting active + observers frees room under the limit in the same request.""" + monkeypatch.setattr(settings, "SESSION_OBSERVERS_LIMIT", 2) + test_workspace, _ = sample_data + workspace_name = test_workspace.name + session_name = str(generate_nanoid()) + db_session.add(models.Session(name=session_name, workspace_name=workspace_name)) + await db_session.flush() + observer = schemas.SessionPeerConfig(observe_others=True, observe_me=False) + bystander = schemas.SessionPeerConfig(observe_others=False, observe_me=True) + existing = [str(generate_nanoid()) for _ in range(2)] + + await crud.set_peers_for_session( + db_session, + workspace_name=workspace_name, + session_name=session_name, + peer_names=dict.fromkeys(existing, observer), + ) + + # Demoting both active observers while adding a new one leaves exactly one. + await crud.set_peers_for_session( + db_session, + workspace_name=workspace_name, + session_name=session_name, + peer_names={ + **dict.fromkeys(existing, bystander), + str(generate_nanoid()): observer, + }, + ) + observer_count = await db_session.scalar( + select(func.count()).where( + models.SessionPeer.session_name == session_name, + models.SessionPeer.workspace_name == workspace_name, + models.SessionPeer.left_at.is_(None), + models.SessionPeer.configuration["observe_others"].astext.cast(Boolean), + ) + ) + assert observer_count == 1 + + @pytest.mark.asyncio + async def test_observer_limit_lets_over_limit_session_take_non_observers( + self, + db_session: AsyncSession, + sample_data: tuple[models.Workspace, models.Peer], + monkeypatch: pytest.MonkeyPatch, + ): + """A session already past the limit still accepts non-observers, so + sessions that grew over it before enforcement do not become unusable.""" + monkeypatch.setattr(settings, "SESSION_OBSERVERS_LIMIT", 2) + test_workspace, _ = sample_data + workspace_name = test_workspace.name + session_name = str(generate_nanoid()) + observer = schemas.SessionPeerConfig(observe_others=True, observe_me=False) + bystander = schemas.SessionPeerConfig(observe_others=False, observe_me=True) + + await crud.get_or_create_session( + db_session, + schemas.SessionCreate( + name=session_name, + peers=dict.fromkeys( + [str(generate_nanoid()) for _ in range(2)], observer + ), + ), + workspace_name, + ) + + # Now the limit is below what the session already holds. + monkeypatch.setattr(settings, "SESSION_OBSERVERS_LIMIT", 1) + await crud.get_or_create_session( + db_session, + schemas.SessionCreate( + name=session_name, peers={str(generate_nanoid()): bystander} + ), + workspace_name, + ) + + active_count = await db_session.scalar( + select(func.count()).where( + models.SessionPeer.session_name == session_name, + models.SessionPeer.workspace_name == workspace_name, + models.SessionPeer.left_at.is_(None), + ) + ) + assert active_count == 3 + @pytest.mark.asyncio async def test_get_session_peer_configuration( self,