"""MuseHub coordination bus — server-side persistence for agent coordination. Provides three operations: ``coord_push`` Accepts a batch of coordination records from a Muse CLI agent and writes them to the ``musehub_coord_records`` table. Write-once semantics for all kinds except ``heartbeat`` (which is upserted — same run_id updates the existing row rather than inserting a duplicate). ``coord_pull`` Returns coordination records for a repo, optionally filtered by kind and cursor ``since_id``. Enables agents to bootstrap local state from the hub (e.g. "what reservations are active on this repo right now?"). ``coord_watch_stream`` Async generator that polls the database every 2 seconds and yields SSE events for new records. Heartbeat comments are injected every 15 seconds of silence to keep proxies alive. Performance ----------- - ``coord_pull`` uses the ``ix_coord_repo_id_cursor`` index for O(log n) cursor seeks and the ``ix_coord_repo_kind_id`` index for filtered pulls. - ``coord_watch_stream`` holds no database connection between polls — it opens a fresh session per poll cycle to avoid idle connection accumulation. - Batch push is O(1) database round-trips for the insert batch via SQLAlchemy bulk insert + a single SELECT for upsert pre-check. Security -------- - All ``record_id`` values are validated at the Pydantic layer before this module is called — no path traversal possible. - ``run_id`` is trusted as-is (agent identity, not a path) but capped at 255 chars by the Pydantic model. - The ``repo_id`` FK constraint prevents cross-repo record injection. """ import asyncio import logging from collections.abc import AsyncIterator from datetime import datetime, timezone from sqlalchemy import and_, select from sqlalchemy.exc import IntegrityError from sqlalchemy.ext.asyncio import AsyncSession from musehub.db import coord_models as db from musehub.types.json_types import JSONValue from musehub.models.coord import ( CoordPollRequest, CoordPullResponse, CoordPushRequest, CoordPushResponse, CoordRecordIn, CoordRecordOut, ) from musehub.mcp.sse import sse_event, sse_heartbeat logger = logging.getLogger(__name__) # Polling interval for the SSE watch stream. _POLL_INTERVAL_SECONDS: float = 2.0 # Maximum silence duration before a heartbeat comment is injected into the # SSE stream to keep HTTP proxies from closing the connection. _HEARTBEAT_INTERVAL_SECONDS: float = 15.0 # Maximum records returned per poll cycle in the watch stream (prevents # a slow consumer from being overwhelmed by a burst of pushes). _WATCH_BATCH_SIZE: int = 100 def _utc_now() -> datetime: return datetime.now(tz=timezone.utc) def _unwrap_payload(payload: JSONValue) -> JSONValue: """Return the raw Python value from a PydanticJson RootModel (or pass through). Calls model_dump() so nested PydanticJson values inside dicts/lists are also recursively unwrapped — plain Python types pass through unchanged. """ if hasattr(payload, "model_dump"): return payload.model_dump() return payload def _row_to_out(row: db.MusehubCoordRecord) -> CoordRecordOut: """Convert an ORM row to the API response model.""" return CoordRecordOut( id=row.id, repo_id=row.repo_id, kind=row.kind, record_id=row.record_id, run_id=row.run_id, payload=row.payload, created_at=row.created_at, expires_at=row.expires_at, ) # ── Push ─────────────────────────────────────────────────────────────────────── async def coord_push( session: AsyncSession, repo_id: str, request: CoordPushRequest, ) -> CoordPushResponse: """Write coordination records for *repo_id* to the database. Write-once semantics apply to all kinds except ``heartbeat``: - **First push** of a (repo_id, kind, record_id) triple → insert → counted as ``inserted``. - **Subsequent push** of the same triple → silently skipped → counted as ``skipped``. ``heartbeat`` records are upserted: re-pushing the same (repo_id, kind, record_id) updates ``payload`` and ``created_at`` in place. This allows agents to refresh their heartbeat without accumulating duplicate rows. Args: session: Async SQLAlchemy session. repo_id: Target repository ID. request: Validated push request with 1–500 records. Returns: :class:`~musehub.models.coord.CoordPushResponse` with counts. """ inserted = 0 skipped = 0 heartbeats = [r for r in request.records if r.kind == "heartbeat"] others = [r for r in request.records if r.kind != "heartbeat"] # ── Heartbeats: upsert ──────────────────────────────────────────────────── for rec in heartbeats: # Try to find an existing heartbeat for this (repo_id, kind, record_id). result = await session.execute( select(db.MusehubCoordRecord).where( and_( db.MusehubCoordRecord.repo_id == repo_id, db.MusehubCoordRecord.kind == rec.kind, db.MusehubCoordRecord.record_id == rec.record_id, ) ) ) existing = result.scalar_one_or_none() if existing is not None: existing.payload = _unwrap_payload(rec.payload) existing.created_at = _utc_now() if rec.expires_at is not None: existing.expires_at = rec.expires_at # No insert — the id stays the same, but a new SSE event won't fire. # Callers that need watch events on heartbeat refresh should use a # new record_id per heartbeat cycle. skipped += 1 else: row = db.MusehubCoordRecord( repo_id=repo_id, kind=rec.kind, record_id=rec.record_id, run_id=rec.run_id, payload=_unwrap_payload(rec.payload), created_at=_utc_now(), expires_at=rec.expires_at, ) session.add(row) inserted += 1 # ── Non-heartbeats: write-once ──────────────────────────────────────────── for rec in others: row = db.MusehubCoordRecord( repo_id=repo_id, kind=rec.kind, record_id=rec.record_id, run_id=rec.run_id, payload=_unwrap_payload(rec.payload), created_at=_utc_now(), expires_at=rec.expires_at, ) session.add(row) try: await session.flush() inserted += 1 except IntegrityError: await session.rollback() skipped += 1 logger.debug( "coord_push: skipped duplicate %s/%s for repo %s", rec.kind, rec.record_id, repo_id, ) await session.commit() # ── Materialize into structured tables ─────────────────────────────────── # Only materialize kinds that have structured tables; fire-and-forget # (failures are logged but do not affect the push response). _MATERIALIZE_KINDS = {"reservation", "release", "task", "claim"} if any(r.kind in _MATERIALIZE_KINDS for r in request.records): from musehub.services.musehub_coord_server import materialize_coord_record for rec in request.records: if rec.kind in _MATERIALIZE_KINDS: await materialize_coord_record( session, repo_id, rec.kind, rec.record_id, _unwrap_payload(rec.payload), rec.expires_at, ) return CoordPushResponse(inserted=inserted, skipped=skipped) # ── Pull ─────────────────────────────────────────────────────────────────────── async def coord_pull( session: AsyncSession, repo_id: str, request: CoordPollRequest, ) -> CoordPullResponse: """Return coordination records for *repo_id* since *request.since_id*. Records are returned in insertion order (oldest first) so that callers can replay history deterministically. The returned ``cursor`` is the ``id`` of the last record in the response — pass it as ``since_id`` in the next call. Args: session: Async SQLAlchemy session. repo_id: Target repository ID. request: Validated poll request with ``since_id``, ``kinds``, ``limit``. Returns: :class:`~musehub.models.coord.CoordPullResponse`. """ stmt = ( select(db.MusehubCoordRecord) .where( and_( db.MusehubCoordRecord.repo_id == repo_id, db.MusehubCoordRecord.id > request.since_id, ) ) .order_by(db.MusehubCoordRecord.id.asc()) .limit(request.limit) ) if request.kinds: stmt = stmt.where(db.MusehubCoordRecord.kind.in_(request.kinds)) result = await session.execute(stmt) rows = result.scalars().all() records = [_row_to_out(r) for r in rows] cursor = records[-1].id if records else 0 return CoordPullResponse(records=records, cursor=cursor) # ── Watch (SSE stream) ───────────────────────────────────────────────────────── async def coord_watch_stream( repo_id: str, since_id: int, kinds: list[str], get_session: "type[AsyncIterator[AsyncSession]]", # noqa: F821 ) -> AsyncIterator[str]: """Async generator that yields SSE events for new coordination records. Polls the database every :data:`_POLL_INTERVAL_SECONDS` seconds. Yields: - ``event: coord_record`` — for each new record (one event per record). - ``": heartbeat"`` — every :data:`_HEARTBEAT_INTERVAL_SECONDS` of silence, to keep HTTP proxies alive. The client should reconnect using ``Last-Event-ID`` (the ``id`` field of the last received event) as the ``?since_id`` query parameter. Args: repo_id: Target repository ID. since_id: Only yield records with ``id > since_id``. kinds: Filter by record kind. Empty list = all kinds. get_session: Async context manager factory for database sessions. The watch stream opens one session per poll cycle so that idle connections are not held between polls. Yields: SSE event strings (``"event: ...\\ndata: ...\\n\\n"`` or ``": heartbeat\\n\\n"``). """ cursor = since_id last_event_at = asyncio.get_event_loop().time() while True: now = asyncio.get_event_loop().time() # Inject heartbeat if silent for too long. if now - last_event_at >= _HEARTBEAT_INTERVAL_SECONDS: yield sse_heartbeat() last_event_at = now # Poll the database for new records. try: async for session in get_session(): stmt = ( select(db.MusehubCoordRecord) .where( and_( db.MusehubCoordRecord.repo_id == repo_id, db.MusehubCoordRecord.id > cursor, ) ) .order_by(db.MusehubCoordRecord.id.asc()) .limit(_WATCH_BATCH_SIZE) ) if kinds: stmt = stmt.where(db.MusehubCoordRecord.kind.in_(kinds)) result = await session.execute(stmt) rows = result.scalars().all() for row in rows: cursor = row.id payload = _row_to_out(row).model_dump(mode="json") yield sse_event(payload, event_id=str(cursor), event_type="coord_record") last_event_at = asyncio.get_event_loop().time() except Exception: # noqa: BLE001 logger.exception("coord_watch_stream: poll error for repo %s", repo_id) await asyncio.sleep(_POLL_INTERVAL_SECONDS)