diff --git a/workspace/backend/alembic/versions/083_event_poll_indexes.py b/workspace/backend/alembic/versions/083_event_poll_indexes.py new file mode 100644 index 000000000..ef1cd49b7 --- /dev/null +++ b/workspace/backend/alembic/versions/083_event_poll_indexes.py @@ -0,0 +1,31 @@ +"""Selective ordered indexes for hot event-poll prefixes.""" +from alembic import op +import sqlalchemy as sa + +revision = '083' +down_revision = '082' +branch_labels = None +depends_on = None + +INDEXES = ( + ('idx_events_poll_message', 'workspace.message.posted%'), + ('idx_events_poll_control', 'workspace.agent.control%'), + ('idx_events_poll_tool_result', 'workspace.tool_result%'), +) + + +def upgrade(): + with op.get_context().autocommit_block(): + for name, pattern in INDEXES: + op.execute(f"CREATE INDEX CONCURRENTLY IF NOT EXISTS {name} " + f"ON events (network_id, timestamp DESC, id DESC) WHERE type LIKE '{pattern}'") + valid = op.get_bind().execute(sa.text( + 'SELECT indisvalid FROM pg_index WHERE indexrelid = to_regclass(:name)'), {'name': name}).scalar_one() + if not valid: + raise RuntimeError(f'Invalid {name}: drop concurrently and retry') + + +def downgrade(): + with op.get_context().autocommit_block(): + for name, _ in reversed(INDEXES): + op.execute(f'DROP INDEX CONCURRENTLY IF EXISTS {name}') diff --git a/workspace/backend/app/cache.py b/workspace/backend/app/cache.py index e7fe2f032..2678aa37f 100644 --- a/workspace/backend/app/cache.py +++ b/workspace/backend/app/cache.py @@ -118,6 +118,58 @@ def json_read_through( return value +def json_singleflight(key, ttl_seconds, compute, cache_if=None): + """Share identical cache fills across workers, with a bounded wait. + + Redis/lock failures fall through to the database. Never swallow compute + errors or cache failed computations. The lock's ownership token prevents + an expired leader from releasing a newer leader's lock. + """ + if key is None: + return compute() + missing = object() + + def read(): + raw = get_bytes(key) + if raw is not None: + try: + return json.loads(raw) + except (TypeError, ValueError): + pass + return missing + + value = read() + if value is not missing: + return value + client = _lazy_client() + lock = None + acquired = False + if client is not None: + try: + lock = client.lock(key + ':fill', timeout=5, blocking_timeout=0.1, sleep=0.01) + acquired = lock.acquire() + except Exception: + pass + try: + # A preceding leader may have filled the cache while we waited. + value = read() if client is not None else missing + if value is not missing: + return value + value = compute() + if cache_if is None or cache_if(value): + try: + set_bytes(key, json.dumps(value, default=str, separators=(',', ':')).encode(), ttl_seconds) + except (TypeError, ValueError): + pass + return value + finally: + if acquired: + try: + lock.release() + except Exception: + pass + + # --------------------------------------------------------------------------- # Pub/Sub — used by SSE streaming # --------------------------------------------------------------------------- diff --git a/workspace/backend/app/models.py b/workspace/backend/app/models.py index 15c7de7be..1660c41f3 100644 --- a/workspace/backend/app/models.py +++ b/workspace/backend/app/models.py @@ -70,6 +70,12 @@ class EventRecord(Base): Index("idx_events_network_type_target_ts", "network_id", "type", "target", "timestamp"), Index("idx_events_channel_latest", "network_id", "target", timestamp.desc(), id.desc(), postgresql_where=text("target LIKE 'channel/%'")), + Index("idx_events_poll_message", "network_id", timestamp.desc(), id.desc(), + postgresql_where=text("type LIKE 'workspace.message.posted%'")), + Index("idx_events_poll_control", "network_id", timestamp.desc(), id.desc(), + postgresql_where=text("type LIKE 'workspace.agent.control%'")), + Index("idx_events_poll_tool_result", "network_id", timestamp.desc(), id.desc(), + postgresql_where=text("type LIKE 'workspace.tool_result%'")), ) diff --git a/workspace/backend/app/routers/events.py b/workspace/backend/app/routers/events.py index f1781c6f2..ad5e71bfc 100644 --- a/workspace/backend/app/routers/events.py +++ b/workspace/backend/app/routers/events.py @@ -26,6 +26,7 @@ from app.response import ResponseCode, json_response, success_response from app.routers.network import _verify_workspace_access, _workspace_filter from app.services.unread_window import unread_cutoff_ms +from app.services.poll_queries import event_type_filter, resume_cursor, stream_head from openagents.core.onm_events import Event from openagents.core.onm_mods import EventRejected, PipelineContext @@ -840,251 +841,245 @@ def poll_events( except Exception: pass - query = select(EventRecord).where(EventRecord.network_id == workspace_id) - if hidden_targets: - query = query.where(EventRecord.target.notin_(list(hidden_targets))) - if dm_scope: - query = query.where(_dm_privacy_clause(_dm_addresses(viewer))) - - # Filter events to only channels where the agent is a member - if member: - member_channel_names = db.execute( - select(Channel.name).where( - Channel.workspace_id == workspace_id, - Channel.id.in_( - select(ChannelMember.channel_id).where(ChannelMember.agent_name == member) + def load_response(): + query = select(EventRecord).where(EventRecord.network_id == workspace_id) + if hidden_targets: + query = query.where(EventRecord.target.notin_(list(hidden_targets))) + if dm_scope: + query = query.where(_dm_privacy_clause(_dm_addresses(viewer))) + + # Filter events to only channels where the agent is a member + if member: + member_channel_names = db.execute( + select(Channel.name).where( + Channel.workspace_id == workspace_id, + Channel.id.in_( + select(ChannelMember.channel_id).where(ChannelMember.agent_name == member) + ), + ) + ).scalars().all() + channel_targets = [f"channel/{name}" for name in member_channel_names] + if channel_targets: + query = query.where(EventRecord.target.in_(channel_targets)) + else: + # Agent is not a member of any channel — return empty + return success_response({"events": [], "has_more": False}) + + if conversation: + parts = [p.strip() for p in conversation.split(",", 1)] + if len(parts) != 2 or not parts[0] or not parts[1]: + return json_response(ResponseCode.BAD_REQUEST, "conversation must be two comma-separated addresses") + a, b = parts + query = query.where( + EventRecord.visibility == "direct", + ~EventRecord.target.startswith("channel/"), + or_( + and_(EventRecord.source == a, EventRecord.target == b), + and_(EventRecord.source == b, EventRecord.target == a), ), ) - ).scalars().all() - channel_targets = [f"channel/{name}" for name in member_channel_names] - if channel_targets: - query = query.where(EventRecord.target.in_(channel_targets)) - else: - # Agent is not a member of any channel — return empty - return success_response({"events": [], "has_more": False}) - - if conversation: - parts = [p.strip() for p in conversation.split(",", 1)] - if len(parts) != 2 or not parts[0] or not parts[1]: - return json_response(ResponseCode.BAD_REQUEST, "conversation must be two comma-separated addresses") - a, b = parts - query = query.where( - EventRecord.visibility == "direct", - ~EventRecord.target.startswith("channel/"), - or_( - and_(EventRecord.source == a, EventRecord.target == b), - and_(EventRecord.source == b, EventRecord.target == a), - ), - ) - if after: - cursor_row = db.execute( - select(EventRecord.timestamp, EventRecord.id).where(EventRecord.id == after) - ).one_or_none() - if cursor_row is not None: - # Use (timestamp, id) tuple to avoid skipping/duplicating events with the same timestamp - query = query.where( - or_( - EventRecord.timestamp > cursor_row.timestamp, - and_(EventRecord.timestamp == cursor_row.timestamp, EventRecord.id > cursor_row.id), + after_position = None + if after: + cursor_row = db.execute( + select(EventRecord.timestamp, EventRecord.id).where(EventRecord.id == after) + ).one_or_none() + if cursor_row is not None: + after_position = {"id": cursor_row.id, "timestamp": cursor_row.timestamp} + # Use (timestamp, id) tuple to avoid skipping/duplicating events with the same timestamp + query = query.where( + or_( + EventRecord.timestamp > cursor_row.timestamp, + and_(EventRecord.timestamp == cursor_row.timestamp, EventRecord.id > cursor_row.id), + ) ) - ) - if before: - cursor_row = db.execute( - select(EventRecord.timestamp, EventRecord.id).where(EventRecord.id == before) - ).one_or_none() - if cursor_row is not None: - query = query.where( - or_( - EventRecord.timestamp < cursor_row.timestamp, - and_(EventRecord.timestamp == cursor_row.timestamp, EventRecord.id < cursor_row.id), + if before: + cursor_row = db.execute( + select(EventRecord.timestamp, EventRecord.id).where(EventRecord.id == before) + ).one_or_none() + if cursor_row is not None: + query = query.where( + or_( + EventRecord.timestamp < cursor_row.timestamp, + and_(EventRecord.timestamp == cursor_row.timestamp, EventRecord.id < cursor_row.id), + ) ) - ) - if target: - query = query.where(EventRecord.target == target) + if target: + query = query.where(EventRecord.target == target) - if channel: - query = query.where(EventRecord.target == f"channel/{channel}") + if channel: + query = query.where(EventRecord.target == f"channel/{channel}") - if type: - query = query.where(EventRecord.type.startswith(type)) + if type: + query = query.where(event_type_filter(type)) - if exclude_message_types: - excluded = [t.strip() for t in exclude_message_types.split(",") if t.strip()] - if excluded: - # payload->>'message_type' — NULL (missing key / non-message event) - # must be kept, so the NOT IN check alone is not enough. - message_type = EventRecord.payload["message_type"].as_string() - query = query.where( - or_(message_type.is_(None), message_type.notin_(excluded)) - ) + if exclude_message_types: + excluded = [t.strip() for t in exclude_message_types.split(",") if t.strip()] + if excluded: + # payload->>'message_type' — NULL (missing key / non-message event) + # must be kept, so the NOT IN check alone is not enough. + message_type = EventRecord.payload["message_type"].as_string() + query = query.where( + or_(message_type.is_(None), message_type.notin_(excluded)) + ) - if target_agents: - # Return only events routed to this agent — mirrors the adapter's - # client-side `target_agents` filter so agents stop pulling the whole - # network's traffic and discarding ~(N-1)/N of it. - # - # The set matches exactly what the adapter keeps (verified against live - # traffic — this filter == the client filter for every agent): - # • events whose metadata.target_agents contains this agent, PLUS - # • *untargeted human* messages (no target_agents key) — the adapter - # broadcasts those. - # Untargeted *agent/system* messages are the bulk of a busy workspace's - # traffic (agent-to-agent chatter) and the adapter always discards them - # when they carry no target list, so returning them was the difference - # between "filter is correct" and "filter actually reduces load". They - # are excluded here. Events targeted at *other* agents and the - # `["__no_response__"]` sentinel are excluded too. - if db.bind.dialect.name == "postgresql": - # `@>` array containment (GIN-friendly); `? 'target_agents'` detects - # the untargeted case, scoped to human senders. - query = query.where( - or_( - EventRecord.metadata_.contains({"target_agents": [target_agents]}), - and_( - EventRecord.source.startswith("human:"), - ~EventRecord.metadata_.has_key("target_agents"), - ), + if target_agents: + # Return only events routed to this agent — mirrors the adapter's + # client-side `target_agents` filter so agents stop pulling the whole + # network's traffic and discarding ~(N-1)/N of it. + # + # The set matches exactly what the adapter keeps (verified against live + # traffic — this filter == the client filter for every agent): + # • events whose metadata.target_agents contains this agent, PLUS + # • *untargeted human* messages (no target_agents key) — the adapter + # broadcasts those. + # Untargeted *agent/system* messages are the bulk of a busy workspace's + # traffic (agent-to-agent chatter) and the adapter always discards them + # when they carry no target list, so returning them was the difference + # between "filter is correct" and "filter actually reduces load". They + # are excluded here. Events targeted at *other* agents and the + # `["__no_response__"]` sentinel are excluded too. + if db.bind.dialect.name == "postgresql": + # `@>` array containment (GIN-friendly); `? 'target_agents'` detects + # the untargeted case, scoped to human senders. + query = query.where( + or_( + EventRecord.metadata_.contains({"target_agents": [target_agents]}), + and_( + EventRecord.source.startswith("human:"), + ~EventRecord.metadata_.has_key("target_agents"), + ), + ) ) - ) - else: - # SQLite (tests): no JSONB operators — fall back to a text match. - # Over-inclusive on the containment side but safe; the adapter - # re-filters client-side. - meta_text = cast(EventRecord.metadata_, Text) - query = query.where( - or_( - meta_text.like(f'%"{target_agents}"%'), - and_( - EventRecord.source.startswith("human:"), - meta_text.notlike('%target_agents%'), - ), + else: + # SQLite (tests): no JSONB operators — fall back to a text match. + # Over-inclusive on the containment side but safe; the adapter + # re-filters client-side. + meta_text = cast(EventRecord.metadata_, Text) + query = query.where( + or_( + meta_text.like(f'%"{target_agents}"%'), + and_( + EventRecord.source.startswith("human:"), + meta_text.notlike('%target_agents%'), + ), + ) ) + + if search: + # Search within payload JSON for content field (works with both JSONB and JSON) + query = query.where( + cast(EventRecord.payload, Text).ilike(f"%{search}%") ) - if search: - # Search within payload JSON for content field (works with both JSONB and JSON) - query = query.where( - cast(EventRecord.payload, Text).ilike(f"%{search}%") - ) + # Head-cursor snapshot for target_agents polls. When we filter to one + # agent's events, the last *returned* event can lag far behind the stream + # tip (most traffic is for other agents), so a naive "cursor = last event" + # would make the agent re-scan the same range every poll — reintroducing + # the O(traffic) scan we're trying to remove. Capture the tip of the + # *unfiltered* (network+type[+channel/target]) stream BEFORE running the + # main query so the client can skip past other agents' events once it has + # drained its own. Snapshotting first guarantees we never advance past an + # event the main query didn't get to see (at worst a harmless re-scan). + head_snapshot = stream_head(db, workspace_id, type, channel, target) if target_agents else None + + if sort == "desc": + query = query.order_by(EventRecord.timestamp.desc(), EventRecord.id.desc()).limit(limit + 1) + else: + query = query.order_by(EventRecord.timestamp.asc(), EventRecord.id.asc()).limit(limit + 1) + rows = db.execute(query).scalars().all() + + has_more = len(rows) > limit + events = rows[:limit] + + composing = False + if not search and not conversation: + from app.composing import has_any_composing + composing = has_any_composing(workspace_id) + + response_data = { + "events": [ + { + "id": e.id, + "type": e.type, + "source": e.source, + "target": e.target, + "payload": e.payload, + "metadata": e.metadata_, + "timestamp": e.timestamp, + "visibility": e.visibility, + } + for e in events + ], + "has_more": has_more, + "oldest_id": (events[-1].id if sort == "desc" else events[0].id) if events else None, + "newest_id": (events[0].id if sort == "desc" else events[-1].id) if events else None, + } + if target_agents: + # Cursor the agent should resume from. If more of *its own* events + # remain (has_more), continue from the last one returned; otherwise + # jump to the stream tip so it skips past other agents' traffic + # instead of re-scanning it next poll. Additive field — older clients + # ignore it and fall back to newest_id. + next_cursor = response_data["newest_id"] + if not has_more: + newest = (events[0] if sort == "desc" else events[-1]) if events else None + next_cursor = resume_cursor(head_snapshot, + {"id": newest.id, "timestamp": newest.timestamp} if newest else None, + after_position) + response_data["next_cursor"] = next_cursor + if composing: + response_data["composing"] = True + + response = success_response(response_data) + + # Populate cache for subsequent polls within the TTL window. + # success_response returns a dict; Redis stores the serialized JSON. + if cache_key is not None and isinstance(response, dict): + try: + import json as _json + serialized = _json.dumps( + response, default=str, separators=(",", ":") + ).encode("utf-8") + # The outer singleflight helper stores the exact response. + # Level 2 maintenance — track the head cursor for these + # filters, and cache the "empty" response when the client was + # already at head. + if head_tracker_key is not None: + newest_id = response.get("data", {}).get("newest_id") + if events and newest_id: + # Update head tracker — newest_id is the tip we just saw. + # Longer TTL because head updates are cheap and we want + # subsequent at-head checks to find it. + cache.set_bytes( + head_tracker_key, + str(newest_id).encode("utf-8"), + ttl_seconds=30.0, + ) + elif not events and incoming_after: + # DB returned empty AND the client had a cursor. This + # confirms "after = head" for this filter. Populate + # both the head tracker (so other clients can match) + # and the shared at-head empty response. + filter_hash = head_tracker_key.split(":")[-1] + cache.set_bytes( + head_tracker_key, + incoming_after.encode("utf-8"), + ttl_seconds=30.0, + ) + cache.set_bytes( + "v1events:athead:" + filter_hash, + serialized, + ttl_seconds=1.5, + ) + except Exception: + pass - # Head-cursor snapshot for target_agents polls. When we filter to one - # agent's events, the last *returned* event can lag far behind the stream - # tip (most traffic is for other agents), so a naive "cursor = last event" - # would make the agent re-scan the same range every poll — reintroducing - # the O(traffic) scan we're trying to remove. Capture the tip of the - # *unfiltered* (network+type[+channel/target]) stream BEFORE running the - # main query so the client can skip past other agents' events once it has - # drained its own. Snapshotting first guarantees we never advance past an - # event the main query didn't get to see (at worst a harmless re-scan). - head_id_snapshot = None - if target_agents: - head_q = select(EventRecord.id).where(EventRecord.network_id == workspace_id) - if type: - head_q = head_q.where(EventRecord.type.startswith(type)) - if channel: - head_q = head_q.where(EventRecord.target == f"channel/{channel}") - elif target: - head_q = head_q.where(EventRecord.target == target) - head_id_snapshot = db.execute( - head_q.order_by(EventRecord.timestamp.desc(), EventRecord.id.desc()).limit(1) - ).scalar() - - if sort == "desc": - query = query.order_by(EventRecord.timestamp.desc(), EventRecord.id.desc()).limit(limit + 1) - else: - query = query.order_by(EventRecord.timestamp.asc(), EventRecord.id.asc()).limit(limit + 1) - rows = db.execute(query).scalars().all() - - has_more = len(rows) > limit - events = rows[:limit] - - composing = False - if not search and not conversation: - from app.composing import has_any_composing - composing = has_any_composing(workspace_id) - - response_data = { - "events": [ - { - "id": e.id, - "type": e.type, - "source": e.source, - "target": e.target, - "payload": e.payload, - "metadata": e.metadata_, - "timestamp": e.timestamp, - "visibility": e.visibility, - } - for e in events - ], - "has_more": has_more, - "oldest_id": (events[-1].id if sort == "desc" else events[0].id) if events else None, - "newest_id": (events[0].id if sort == "desc" else events[-1].id) if events else None, - } - if target_agents: - # Cursor the agent should resume from. If more of *its own* events - # remain (has_more), continue from the last one returned; otherwise - # jump to the stream tip so it skips past other agents' traffic - # instead of re-scanning it next poll. Additive field — older clients - # ignore it and fall back to newest_id. - next_cursor = response_data["newest_id"] - if not has_more and head_id_snapshot: - next_cursor = head_id_snapshot - response_data["next_cursor"] = next_cursor - if composing: - response_data["composing"] = True - - response = success_response(response_data) - - # Populate cache for subsequent polls within the TTL window. - # success_response returns a dict; Redis stores the serialized JSON. - if cache_key is not None and isinstance(response, dict): - try: - import json as _json - serialized = _json.dumps( - response, default=str, separators=(",", ":") - ).encode("utf-8") - # Level 1: exact-match cache (includes cursor). Slightly - # longer TTL helps dedup adjacent polls from the same agent. - cache.set_bytes(cache_key, serialized, ttl_seconds=1.5) - - # Level 2 maintenance — track the head cursor for these - # filters, and cache the "empty" response when the client was - # already at head. - if head_tracker_key is not None: - newest_id = response.get("data", {}).get("newest_id") - if events and newest_id: - # Update head tracker — newest_id is the tip we just saw. - # Longer TTL because head updates are cheap and we want - # subsequent at-head checks to find it. - cache.set_bytes( - head_tracker_key, - str(newest_id).encode("utf-8"), - ttl_seconds=30.0, - ) - elif not events and incoming_after: - # DB returned empty AND the client had a cursor. This - # confirms "after = head" for this filter. Populate - # both the head tracker (so other clients can match) - # and the shared at-head empty response. - filter_hash = head_tracker_key.split(":")[-1] - cache.set_bytes( - head_tracker_key, - incoming_after.encode("utf-8"), - ttl_seconds=30.0, - ) - cache.set_bytes( - "v1events:athead:" + filter_hash, - serialized, - ttl_seconds=1.5, - ) - except Exception: - pass + return response - return response + return cache.json_singleflight(cache_key, 1.5, load_response, cache_if=lambda value: isinstance(value, dict)) # --------------------------------------------------------------------------- diff --git a/workspace/backend/app/services/poll_queries.py b/workspace/backend/app/services/poll_queries.py new file mode 100644 index 000000000..2af0f041b --- /dev/null +++ b/workspace/backend/app/services/poll_queries.py @@ -0,0 +1,81 @@ +"""Indexed event polling and short-lived, commit-invalidated stream heads.""" +import hashlib +import json +import uuid + +from sqlalchemy import event, literal_column, select +from sqlalchemy.orm import Session + +from app import cache +from app.models import EventRecord + +POLL_PREFIXES = ('workspace.message.posted', 'workspace.agent.control', 'workspace.tool_result') + + +def event_type_filter(prefix): + # Constants let PostgreSQL prove the partial-index predicate even with a + # generic prepared plan. Preserve the API's existing LIKE-prefix semantics + # (including the underscore in tool_result); never interpolate user input. + if prefix in POLL_PREFIXES: + return EventRecord.type.like(literal_column("'" + prefix + "%'")) + return EventRecord.type.startswith(prefix) + + +def _generation_key(workspace_id): + return 'v1events:head-generation:' + str(workspace_id) + + +@event.listens_for(Session, 'before_flush') +def _remember_event_writes(session, flush_context, instances): + for row in session.new.union(session.dirty).union(session.deleted): + if isinstance(row, EventRecord): + session.info.setdefault('poll_head_workspaces', set()).add(str(row.network_id)) + + +@event.listens_for(Session, 'after_commit') +def _invalidate_committed_heads(session): + if session.in_nested_transaction(): + return + for workspace_id in session.info.pop('poll_head_workspaces', set()): + # A generation prevents a query started before the commit from later + # repopulating the current cache with an obsolete head. All ORM event + # writers are covered, including those publishing before commit. + cache.set_bytes(_generation_key(workspace_id), uuid.uuid4().hex.encode(), 60) + + +@event.listens_for(Session, 'after_rollback') +def _discard_rolled_back_heads(session): + if not session.in_nested_transaction(): + session.info.pop('poll_head_workspaces', None) + + +def stream_head(db, workspace_id, prefix, channel=None, target=None): + generation = cache.get_bytes(_generation_key(workspace_id)) or b'0' + scope = json.dumps([str(workspace_id), prefix, channel, target], separators=(',', ':')) + key = 'v1events:stream-head:' + hashlib.sha256(scope.encode()).hexdigest() + ':' + generation.hex() + + def fetch(): + query = select(EventRecord.id, EventRecord.timestamp).where(EventRecord.network_id == workspace_id) + if prefix: + query = query.where(event_type_filter(prefix)) + if channel: + query = query.where(EventRecord.target == 'channel/' + channel) + elif target: + query = query.where(EventRecord.target == target) + row = db.execute(query.order_by(EventRecord.timestamp.desc(), EventRecord.id.desc()).limit(1)).first() + return {'id': row.id, 'timestamp': row.timestamp} if row else None + + # Null heads are cached too: streams that have never had an event are a + # common source of repeated cursor-free polls. External SQL writers are + # bounded by this TTL even if they bypass the ORM commit hook. + return cache.json_singleflight(key, 1.0, fetch) + + +def resume_cursor(head, newest, incoming): + # A cached head may predate the message query. Never move backward from a + # returned event or valid incoming cursor. At equal timestamps prefer the + # DB-ordered position, avoiding a Python-vs-PostgreSQL collation mismatch. + position = newest or incoming + if head and (not position or head['timestamp'] > position['timestamp']): + position = head + return position['id'] if position else None diff --git a/workspace/backend/tests/test_poll_queries.py b/workspace/backend/tests/test_poll_queries.py new file mode 100644 index 000000000..01f2ede85 --- /dev/null +++ b/workspace/backend/tests/test_poll_queries.py @@ -0,0 +1,209 @@ +import concurrent.futures +import os +import threading +import time +from types import SimpleNamespace + +import pytest +from sqlalchemy import select +from sqlalchemy.dialects import postgresql + +from app import cache +from app.models import EventRecord +from app.services.poll_queries import event_type_filter, resume_cursor, stream_head, _generation_key + +WID = '11111111-1111-4111-8111-111111111111' + + +@pytest.fixture +def redis_memory(monkeypatch): + store = {} + locks = {} + guard = threading.Lock() + + def get(key): + entry = store.get(key) + return entry[0] if entry and entry[1] > time.monotonic() else None + + def put(key, value, ttl_seconds): + store[key] = (value, time.monotonic() + ttl_seconds) + + class Client: + def lock(self, key, **kw): + with guard: + lock = locks.setdefault(key, threading.Lock()) + return SimpleNamespace(acquire=lambda: lock.acquire(timeout=kw['blocking_timeout']), release=lock.release) + + monkeypatch.setattr(cache, 'get_bytes', get) + monkeypatch.setattr(cache, 'set_bytes', put) + monkeypatch.setattr(cache, '_lazy_client', lambda: Client()) + return store + + +def add(db, id='a', timestamp=10, **kw): + row = EventRecord(id=id, timestamp=timestamp, network_id=kw.pop('network_id', WID), + source='openagents:bot', target=kw.pop('target', 'channel/a'), + type=kw.pop('type', 'workspace.message.posted'), payload={}, **kw) + db.add(row) + db.commit() + return row + + +def test_concurrent_identical_fills_compute_once(redis_memory): + calls = [] + start = threading.Barrier(6) + def compute(): + calls.append(True) + time.sleep(.03) + return {'events': ['one']} + def request(_): + start.wait() + return cache.json_singleflight('same', 1, compute) + with concurrent.futures.ThreadPoolExecutor(6) as pool: + assert list(pool.map(request, range(6))) == [{'events': ['one']}] * 6 + assert len(calls) == 1 + + +def test_errors_not_cached_and_lock_released(redis_memory): + def fail(): + raise ValueError('failed query') + with pytest.raises(ValueError): + cache.json_singleflight('key', 1, fail) + assert cache.json_singleflight('key', 1, lambda: 42) == 42 + assert cache.json_singleflight('key', 1, fail) == 42 + + +def test_null_cached_and_http_errors_not_cached(redis_memory): + assert cache.json_singleflight('empty', 1, lambda: None) is None + assert cache.json_singleflight('empty', 1, lambda: 42) is None + marker = object() + assert cache.json_singleflight('error', 1, lambda: marker, lambda v: isinstance(v, dict)) is marker + assert 'error' not in redis_memory + + +def test_redis_down_and_lock_timeout_fall_through(monkeypatch): + monkeypatch.setattr(cache, 'get_bytes', lambda k: None) + monkeypatch.setattr(cache, 'set_bytes', lambda *a: None) + monkeypatch.setattr(cache, '_lazy_client', lambda: None) + assert cache.json_singleflight('x', 1, lambda: 42) == 42 + monkeypatch.setattr(cache, '_lazy_client', lambda: SimpleNamespace(lock=lambda *a, **kw: + SimpleNamespace(acquire=lambda: False))) + assert cache.json_singleflight('x', 1, lambda: 43) == 43 + + +def test_cached_head_invalidated_on_insert_edit_delete_but_not_flush(db, redis_memory): + add(db) + assert stream_head(db, WID, 'workspace.message.posted')['id'] == 'a' + db.add(EventRecord(id='b', network_id=WID, source='x', target='channel/a', + type='workspace.message.posted', timestamp=20)) + db.flush() + assert stream_head(db, WID, 'workspace.message.posted')['id'] == 'a' + db.commit() + assert stream_head(db, WID, 'workspace.message.posted')['id'] == 'b' + b = db.get(EventRecord, 'b'); b.type = 'other'; db.commit() + assert stream_head(db, WID, 'workspace.message.posted')['id'] == 'a' + db.delete(db.get(EventRecord, 'a')); db.commit() + assert stream_head(db, WID, 'workspace.message.posted') is None + + +def test_rollback_and_nested_commit(db, redis_memory): + add(db) + generation = cache.get_bytes(_generation_key(WID)) + row = db.get(EventRecord, 'a'); row.timestamp = 20; db.flush(); db.rollback() + assert cache.get_bytes(_generation_key(WID)) == generation + with db.begin_nested(): + row = db.get(EventRecord, 'a'); row.timestamp = 30; db.flush() + assert cache.get_bytes(_generation_key(WID)) == generation + db.commit() + assert cache.get_bytes(_generation_key(WID)) != generation + + +def test_old_fill_cannot_poison_generation_after_commit(redis_memory): + calls = [] + def execute(query): + calls.append(True) + if len(calls) == 1: + cache.set_bytes(_generation_key(WID), b'new-generation', 60) + row = SimpleNamespace(id='old' if len(calls) == 1 else 'new', timestamp=len(calls)) + return SimpleNamespace(first=lambda: row) + db = SimpleNamespace(execute=execute) + assert stream_head(db, WID, 'workspace.message.posted')['id'] == 'old' + assert stream_head(db, WID, 'workspace.message.posted')['id'] == 'new' + assert len(calls) == 2 + + +def test_head_scope_and_legacy_subtypes(db, redis_memory): + add(db, 'a', 10) + add(db, 'b', 20, target='channel/b', type='workspace.message.posted.extra') + add(db, 'c', 30, type='workspace.agent.control') + assert stream_head(db, WID, 'workspace.message.posted')['id'] == 'b' + assert stream_head(db, WID, 'workspace.message.posted', channel='a')['id'] == 'a' + assert stream_head(db, WID, 'workspace.agent.control')['id'] == 'c' + assert stream_head(db, WID, 'workspace.tool_result') is None + assert stream_head(db, '22222222-2222-4222-8222-222222222222', None) is None + + +@pytest.mark.parametrize('head,newest,incoming,expected', [ + ({'id':'old','timestamp':1}, {'id':'new','timestamp':2}, None, 'new'), + ({'id':'old','timestamp':1}, None, {'id':'incoming','timestamp':2}, 'incoming'), + ({'id':'head','timestamp':3}, {'id':'new','timestamp':2}, None, 'head'), + ({'id':'head','timestamp':2}, {'id':'new','timestamp':2}, None, 'new'), + (None, None, {'id':'incoming','timestamp':2}, 'incoming'), + (None, None, None, None), +]) +def test_cursor_never_regresses(head, newest, incoming, expected): + assert resume_cursor(head, newest, incoming) == expected + + +@pytest.mark.parametrize('prefix', ['workspace.message.posted', 'workspace.agent.control', + 'workspace.tool_result', 'custom.prefix', "x' OR TRUE --"]) +def test_prefix_filter_preserves_semantics_and_parameter_safety(db, prefix): + for i, typ in enumerate(['workspace.message.posted', 'workspace.message.posted.extra', + 'workspace.agent.control', 'workspace.tool_result', 'workspace.toolXresult', 'custom.prefix.extra']): + add(db, str(i), i, type=typ) + old = select(EventRecord.id).where(EventRecord.type.startswith(prefix)).order_by(EventRecord.id) + new = select(EventRecord.id).where(event_type_filter(prefix)).order_by(EventRecord.id) + assert db.execute(old).all() == db.execute(new).all() + compiled = new.compile(dialect=postgresql.dialect()) + if prefix.startswith("x'"): + assert prefix not in str(compiled) + + +def test_arrival_after_head_snapshot_does_not_regress_cursor(client, workspace, db, monkeypatch): + add(db, 'old', 10, network_id=workspace['id']) + def old_head(*args): + # Simulate an event committed after the captured stream head but before + # the message SELECT. It must be returned and become the resume cursor. + add(db, 'arrived', 20, network_id=workspace['id'], metadata_={'target_agents':['alpha']}) + return {'id':'old', 'timestamp':10} + monkeypatch.setattr('app.routers.events.stream_head', old_head) + response = client.get('/v1/events', params={'network':workspace['id'], 'after':'old', + 'target_agents':'alpha', 'type':'workspace.message.posted'}, + headers={'X-Workspace-Token':workspace['token']}) + data = response.json()['data'] + assert [e['id'] for e in data['events']] == ['arrived'] + assert data['next_cursor'] == 'arrived' + + +@pytest.mark.skipif(os.getenv('OPENAGENTS_TEST_ISOLATED_REDIS') != '1', reason='requires disposable Redis') +def test_real_redis_lock_and_expiration(): + import uuid + key = 'test:poll-singleflight:' + uuid.uuid4().hex + calls = [] + barrier = threading.Barrier(8) + def compute(): + calls.append(True) + time.sleep(.03) + return {'events': ['one']} + def request(_): + barrier.wait() + return cache.json_singleflight(key, .15, compute) + try: + with concurrent.futures.ThreadPoolExecutor(8) as pool: + assert list(pool.map(request, range(8))) == [{'events':['one']}] * 8 + assert len(calls) == 1 + time.sleep(.2) + assert cache.json_singleflight(key, .15, compute) == {'events':['one']} + assert len(calls) == 2 + finally: + cache.delete_key(key)