From 4492f66bca7515e265ab38286858aeb6db0cfcc3 Mon Sep 17 00:00:00 2001 From: Aakash Kattelu Date: Tue, 25 Aug 2026 10:51:59 -0400 Subject: [PATCH] fix(crud): preserve joined_at for active session peers (#1059) * fix(crud): preserve joined_at for active session peers Re-adding an already-active peer no longer advances the membership window, so peer_perspective search keeps messages from the original join. Genuine rejoins still start a new window. * docs: document set_peers membership window and wrap test docstrings * fix: preserve session observer limit --------- Co-authored-by: Vineeth Voruganti <13438633+VVoruganti@users.noreply.github.com> --- src/crud/session.py | 160 +++++++++++------- tests/crud/test_session.py | 328 ++++++++++++++++++++++++++++++++++++- tests/test_search.py | 71 +++++++- 3 files changed, 499 insertions(+), 60 deletions(-) diff --git a/src/crud/session.py b/src/crud/session.py index efa3f664..51ca1f13 100644 --- a/src/crud/session.py +++ b/src/crud/session.py @@ -1050,25 +1050,31 @@ async def set_peers_for_session( peer_names: dict[str, schemas.SessionPeerConfig], ) -> list[models.SessionPeer]: """ - Set peers for a session, overwriting any existing peers. - If peers don't exist, they will be created. + Replace a session's ordinary peer set with ``peer_names``. + + 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 workspace_name: Name of the workspace session_name: Name of the session - peer_names: Set of peer names to set for the session + peer_names: Mapping of peer names to session-level configuration Returns: List of SessionPeer objects for all peers in the 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 = ( @@ -1084,20 +1090,21 @@ async def set_peers_for_session( f"Session {session_name} not found in workspace {workspace_name}" ) - # Soft delete every *ordinary* active membership. Scope memberships are - # deliberately preserved: this route replaces the peers the caller names, and a - # caller detaches a scope by simply *omitting* it from an otherwise valid - # replacement map — never naming it, so no request-level guard can see it. - # Without the exclusion a plain replacement silently bypasses the facade that - # owns scope membership and its removal reconciliation. Being part of the - # UPDATE, this holds regardless of the request body or concurrent scope - # creation. + # Soft delete every *ordinary* active membership not in the incoming map. + # Scope memberships are deliberately preserved: this route replaces the peers + # the caller names, and a caller detaches a scope by simply *omitting* it from + # an otherwise valid replacement map — never naming it, so no request-level + # guard can see it. Without the exclusion a plain replacement silently + # bypasses the facade that owns scope membership and its removal + # reconciliation. Being part of the UPDATE, this holds regardless of the + # request body or concurrent scope creation. update_stmt = ( update(models.SessionPeer) .where( models.SessionPeer.session_name == session_name, models.SessionPeer.workspace_name == workspace_name, models.SessionPeer.left_at.is_(None), # Only update active peers + models.SessionPeer.peer_name.notin_(peer_names.keys()), ~exists( select(models.Peer.id) .where(models.Peer.workspace_name == workspace_name) @@ -1118,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() @@ -1162,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 @@ -1177,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: @@ -1202,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( [ @@ -1254,15 +1243,27 @@ async def _get_or_add_peers_to_session( ] ) - # On conflict, update joined_at and clear left_at (rejoin scenario) - # If left_at is not None (peer has left the session): Use the new configuration (stmt.excluded.configuration) - # If left_at is None (peer is still active): Keep the existing configuration (models.SessionPeer.configuration) + # 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_={ - "joined_at": func.now(), + "joined_at": case( + (models.SessionPeer.left_at.is_not(None), func.now()), + 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, ), @@ -1270,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 8b4b5795..2eec1f32 100644 --- a/tests/crud/test_session.py +++ b/tests/crud/test_session.py @@ -1,14 +1,340 @@ +from datetime import datetime, timezone + import pytest from nanoid import generate as generate_nanoid +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: """Test suite for session CRUD operations""" + @pytest.mark.asyncio + async def test_get_or_create_session_preserves_active_joined_at( + self, + db_session: AsyncSession, + sample_data: tuple[models.Workspace, models.Peer], + ): + """Active re-adds keep joined_at and config; a genuine rejoin starts a + new window.""" + test_workspace, test_peer = sample_data + session_name = str(generate_nanoid()) + original_config = schemas.SessionPeerConfig( + observe_others=True, observe_me=False + ) + updated_config = schemas.SessionPeerConfig( + observe_others=False, observe_me=True + ) + session_peer_stmt = select( + models.SessionPeer.joined_at, + models.SessionPeer.left_at, + models.SessionPeer.configuration, + ).where( + models.SessionPeer.session_name == session_name, + models.SessionPeer.peer_name == test_peer.name, + models.SessionPeer.workspace_name == test_workspace.name, + ) + + await crud.get_or_create_session( + db_session, + schemas.SessionCreate( + name=session_name, peers={test_peer.name: original_config} + ), + test_workspace.name, + ) + first_joined_at, first_left_at, first_config = ( + await db_session.execute(session_peer_stmt) + ).one() + assert first_left_at is None + assert first_config == original_config.model_dump() + + await crud.get_or_create_session( + db_session, + schemas.SessionCreate( + name=session_name, peers={test_peer.name: updated_config} + ), + test_workspace.name, + ) + second_joined_at, second_left_at, second_config = ( + await db_session.execute(session_peer_stmt) + ).one() + assert second_joined_at == first_joined_at + assert second_left_at is None + assert second_config == original_config.model_dump() + + session_peer = ( + await db_session.execute( + select(models.SessionPeer).where( + models.SessionPeer.session_name == session_name, + models.SessionPeer.peer_name == test_peer.name, + models.SessionPeer.workspace_name == test_workspace.name, + ) + ) + ).scalar_one() + session_peer.left_at = datetime.now(timezone.utc) + await db_session.commit() + + await crud.get_or_create_session( + db_session, + schemas.SessionCreate( + name=session_name, peers={test_peer.name: updated_config} + ), + test_workspace.name, + ) + rejoined_joined_at, rejoined_left_at, rejoined_config = ( + await db_session.execute(session_peer_stmt) + ).one() + assert rejoined_joined_at > second_joined_at + assert rejoined_left_at is None + assert rejoined_config == updated_config.model_dump() + + @pytest.mark.asyncio + async def test_set_peers_preserves_active_joined_at( + self, + db_session: AsyncSession, + sample_data: tuple[models.Workspace, models.Peer], + ): + """PUT-peers keeps active membership windows and refreshes real rejoins.""" + test_workspace, test_peer = sample_data + session_name = str(generate_nanoid()) + original_config = schemas.SessionPeerConfig( + observe_others=True, observe_me=False + ) + updated_config = schemas.SessionPeerConfig( + observe_others=False, observe_me=True + ) + db_session.add( + models.Session(name=session_name, workspace_name=test_workspace.name) + ) + await db_session.flush() + + session_peer_stmt = select( + models.SessionPeer.joined_at, + models.SessionPeer.left_at, + models.SessionPeer.configuration, + ).where( + models.SessionPeer.session_name == session_name, + models.SessionPeer.peer_name == test_peer.name, + models.SessionPeer.workspace_name == test_workspace.name, + ) + + await crud.set_peers_for_session( + db_session, + workspace_name=test_workspace.name, + session_name=session_name, + peer_names={test_peer.name: original_config}, + ) + first_left_at, first_config = ( + await db_session.execute( + select( + models.SessionPeer.left_at, + models.SessionPeer.configuration, + ).where( + models.SessionPeer.session_name == session_name, + models.SessionPeer.peer_name == test_peer.name, + models.SessionPeer.workspace_name == test_workspace.name, + ) + ) + ).one() + assert first_left_at is None + assert first_config == original_config.model_dump() + + session_peer = ( + await db_session.execute( + select(models.SessionPeer).where( + models.SessionPeer.session_name == session_name, + models.SessionPeer.peer_name == test_peer.name, + models.SessionPeer.workspace_name == test_workspace.name, + ) + ) + ).scalar_one() + session_peer.joined_at = datetime(2020, 1, 1, tzinfo=timezone.utc) + await db_session.commit() + + await crud.set_peers_for_session( + db_session, + workspace_name=test_workspace.name, + session_name=session_name, + peer_names={test_peer.name: updated_config}, + ) + active_joined_at, active_left_at, active_config = ( + await db_session.execute(session_peer_stmt) + ).one() + assert active_joined_at == datetime(2020, 1, 1, tzinfo=timezone.utc) + assert active_left_at is None + # 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, + workspace_name=test_workspace.name, + session_name=session_name, + peer_names={}, + ) + await crud.set_peers_for_session( + db_session, + workspace_name=test_workspace.name, + session_name=session_name, + peer_names={test_peer.name: updated_config}, + ) + rejoined_joined_at, rejoined_left_at, rejoined_config = ( + await db_session.execute(session_peer_stmt) + ).one() + assert rejoined_joined_at > active_joined_at + 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, diff --git a/tests/test_search.py b/tests/test_search.py index 84f3ffa3..b4b69c43 100644 --- a/tests/test_search.py +++ b/tests/test_search.py @@ -4,9 +4,10 @@ import datetime import pytest from nanoid import generate as generate_nanoid +from sqlalchemy import update from sqlalchemy.ext.asyncio import AsyncSession -from src import crud, models +from src import crud, models, schemas from src.utils.search import search @@ -704,3 +705,71 @@ async def test_grep_messages_observer_scoping_left_session_still_visible( matched_ids = [m.public_id for matches, _ in results for m in matches] assert msg_during.public_id in matched_ids assert msg_after.public_id in matched_ids + + +@pytest.mark.asyncio +async def test_peer_perspective_search_after_active_readd( + db_session: AsyncSession, +): + """Active re-add keeps existing messages visible; a genuine rejoin starts a + new window.""" + workspace = models.Workspace(name=generate_nanoid()) + peer1 = models.Peer(name="peer1", workspace_name=workspace.name) + peer2 = models.Peer(name="peer2", workspace_name=workspace.name) + session = models.Session(name="session1", workspace_name=workspace.name) + db_session.add_all([workspace, peer1, peer2, session]) + await db_session.flush() + + past_time = datetime.datetime.now(datetime.timezone.utc) - datetime.timedelta( + hours=1 + ) + await db_session.execute( + models.session_peers_table.insert().values( + workspace_name=workspace.name, + session_name=session.name, + peer_name=peer1.name, + joined_at=past_time, + left_at=None, + ) + ) + msg_old = models.Message( + content="old persistent message", + session_name=session.name, + peer_name=peer2.name, + workspace_name=workspace.name, + seq_in_session=1, + created_at=past_time + datetime.timedelta(minutes=1), + ) + db_session.add(msg_old) + await db_session.commit() + + session_create = schemas.SessionCreate( + name=session.name, + peers={peer1.name: schemas.SessionPeerConfig()}, + ) + await crud.get_or_create_session(db_session, session_create, workspace.name) + results = await search( + "persistent", + filters={"peer_perspective": peer1.name, "workspace_id": workspace.name}, + limit=10, + ) + assert msg_old.public_id in [m.public_id for m in results] + + await db_session.execute( + update(models.SessionPeer) + .where( + models.SessionPeer.session_name == session.name, + models.SessionPeer.peer_name == peer1.name, + models.SessionPeer.workspace_name == workspace.name, + ) + .values(left_at=datetime.datetime.now(datetime.timezone.utc)) + ) + await db_session.commit() + + await crud.get_or_create_session(db_session, session_create, workspace.name) + results = await search( + "persistent", + filters={"peer_perspective": peer1.name, "workspace_id": workspace.name}, + limit=10, + ) + assert msg_old.public_id not in [m.public_id for m in results]