fix: preserve session observer limit
This commit is contained in:
Vineeth Voruganti
parent
0ad0e42336
commit
996b9d7e5c
|
|
@ -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 []
|
||||
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
Loading…
Reference in New Issue