musehub_coord.py
python
sha256:d035733f21ccff27735fddebfbbe0ed24565a32a22db8de5885402262671ecd2
chore: bump version to 0.2.0rc15 for musehub#113 fix release
Sonnet 4.6
patch
15 days ago
| 1 | """MuseHub coordination bus — server-side persistence for agent coordination. |
| 2 | |
| 3 | Provides three operations: |
| 4 | |
| 5 | ``coord_push`` |
| 6 | Accepts a batch of coordination records from a Muse CLI agent and writes |
| 7 | them to the ``musehub_coord_records`` table. Write-once semantics for |
| 8 | all kinds except ``heartbeat`` (which is upserted — same run_id updates |
| 9 | the existing row rather than inserting a duplicate). |
| 10 | |
| 11 | ``coord_pull`` |
| 12 | Returns coordination records for a repo, optionally filtered by kind and |
| 13 | cursor ``since_id``. Enables agents to bootstrap local state from the hub |
| 14 | (e.g. "what reservations are active on this repo right now?"). |
| 15 | |
| 16 | ``coord_watch_stream`` |
| 17 | Async generator that polls the database every 2 seconds and yields SSE |
| 18 | events for new records. Heartbeat comments are injected every 15 seconds |
| 19 | of silence to keep proxies alive. |
| 20 | |
| 21 | Performance |
| 22 | ----------- |
| 23 | - ``coord_pull`` uses the ``ix_coord_repo_id_cursor`` index for O(log n) cursor |
| 24 | seeks and the ``ix_coord_repo_kind_id`` index for filtered pulls. |
| 25 | - ``coord_watch_stream`` holds no database connection between polls — it opens |
| 26 | a fresh session per poll cycle to avoid idle connection accumulation. |
| 27 | - Batch push is O(1) database round-trips for the insert batch via |
| 28 | SQLAlchemy bulk insert + a single SELECT for upsert pre-check. |
| 29 | |
| 30 | Security |
| 31 | -------- |
| 32 | - All ``record_id`` values are validated at the Pydantic layer before this |
| 33 | module is called — no path traversal possible. |
| 34 | - ``run_id`` is trusted as-is (agent identity, not a path) but capped at 255 |
| 35 | chars by the Pydantic model. |
| 36 | - The ``repo_id`` FK constraint prevents cross-repo record injection. |
| 37 | """ |
| 38 | |
| 39 | import asyncio |
| 40 | import logging |
| 41 | from collections.abc import AsyncIterator |
| 42 | from datetime import datetime, timezone |
| 43 | |
| 44 | from sqlalchemy import and_, select |
| 45 | from sqlalchemy.exc import IntegrityError |
| 46 | from sqlalchemy.ext.asyncio import AsyncSession |
| 47 | |
| 48 | from musehub.db import coord_models as db |
| 49 | from musehub.types.json_types import JSONValue |
| 50 | from musehub.models.coord import ( |
| 51 | CoordPollRequest, |
| 52 | CoordPullResponse, |
| 53 | CoordPushRequest, |
| 54 | CoordPushResponse, |
| 55 | CoordRecordIn, |
| 56 | CoordRecordOut, |
| 57 | ) |
| 58 | from musehub.mcp.sse import sse_event, sse_heartbeat |
| 59 | |
| 60 | logger = logging.getLogger(__name__) |
| 61 | |
| 62 | # Polling interval for the SSE watch stream. |
| 63 | _POLL_INTERVAL_SECONDS: float = 2.0 |
| 64 | |
| 65 | # Maximum silence duration before a heartbeat comment is injected into the |
| 66 | # SSE stream to keep HTTP proxies from closing the connection. |
| 67 | _HEARTBEAT_INTERVAL_SECONDS: float = 15.0 |
| 68 | |
| 69 | # Maximum records returned per poll cycle in the watch stream (prevents |
| 70 | # a slow consumer from being overwhelmed by a burst of pushes). |
| 71 | _WATCH_BATCH_SIZE: int = 100 |
| 72 | |
| 73 | def _utc_now() -> datetime: |
| 74 | return datetime.now(tz=timezone.utc) |
| 75 | |
| 76 | def _unwrap_payload(payload: JSONValue) -> JSONValue: |
| 77 | """Return the raw Python value from a PydanticJson RootModel (or pass through). |
| 78 | |
| 79 | Calls model_dump() so nested PydanticJson values inside dicts/lists are |
| 80 | also recursively unwrapped — plain Python types pass through unchanged. |
| 81 | """ |
| 82 | if hasattr(payload, "model_dump"): |
| 83 | return payload.model_dump() |
| 84 | return payload |
| 85 | |
| 86 | def _row_to_out(row: db.MusehubCoordRecord) -> CoordRecordOut: |
| 87 | """Convert an ORM row to the API response model.""" |
| 88 | return CoordRecordOut( |
| 89 | id=row.id, |
| 90 | repo_id=row.repo_id, |
| 91 | kind=row.kind, |
| 92 | record_id=row.record_id, |
| 93 | run_id=row.run_id, |
| 94 | payload=row.payload, |
| 95 | created_at=row.created_at, |
| 96 | expires_at=row.expires_at, |
| 97 | ) |
| 98 | |
| 99 | # ── Push ─────────────────────────────────────────────────────────────────────── |
| 100 | |
| 101 | async def coord_push( |
| 102 | session: AsyncSession, |
| 103 | repo_id: str, |
| 104 | request: CoordPushRequest, |
| 105 | ) -> CoordPushResponse: |
| 106 | """Write coordination records for *repo_id* to the database. |
| 107 | |
| 108 | Write-once semantics apply to all kinds except ``heartbeat``: |
| 109 | - **First push** of a (repo_id, kind, record_id) triple → insert → counted as ``inserted``. |
| 110 | - **Subsequent push** of the same triple → silently skipped → counted as ``skipped``. |
| 111 | |
| 112 | ``heartbeat`` records are upserted: re-pushing the same |
| 113 | (repo_id, kind, record_id) updates ``payload`` and ``created_at`` |
| 114 | in place. This allows agents to refresh their heartbeat without |
| 115 | accumulating duplicate rows. |
| 116 | |
| 117 | Args: |
| 118 | session: Async SQLAlchemy session. |
| 119 | repo_id: Target repository ID. |
| 120 | request: Validated push request with 1–500 records. |
| 121 | |
| 122 | Returns: |
| 123 | :class:`~musehub.models.coord.CoordPushResponse` with counts. |
| 124 | """ |
| 125 | inserted = 0 |
| 126 | skipped = 0 |
| 127 | |
| 128 | heartbeats = [r for r in request.records if r.kind == "heartbeat"] |
| 129 | others = [r for r in request.records if r.kind != "heartbeat"] |
| 130 | |
| 131 | # ── Heartbeats: upsert ──────────────────────────────────────────────────── |
| 132 | for rec in heartbeats: |
| 133 | # Try to find an existing heartbeat for this (repo_id, kind, record_id). |
| 134 | result = await session.execute( |
| 135 | select(db.MusehubCoordRecord).where( |
| 136 | and_( |
| 137 | db.MusehubCoordRecord.repo_id == repo_id, |
| 138 | db.MusehubCoordRecord.kind == rec.kind, |
| 139 | db.MusehubCoordRecord.record_id == rec.record_id, |
| 140 | ) |
| 141 | ) |
| 142 | ) |
| 143 | existing = result.scalar_one_or_none() |
| 144 | if existing is not None: |
| 145 | existing.payload = _unwrap_payload(rec.payload) |
| 146 | existing.created_at = _utc_now() |
| 147 | if rec.expires_at is not None: |
| 148 | existing.expires_at = rec.expires_at |
| 149 | # No insert — the id stays the same, but a new SSE event won't fire. |
| 150 | # Callers that need watch events on heartbeat refresh should use a |
| 151 | # new record_id per heartbeat cycle. |
| 152 | skipped += 1 |
| 153 | else: |
| 154 | row = db.MusehubCoordRecord( |
| 155 | repo_id=repo_id, |
| 156 | kind=rec.kind, |
| 157 | record_id=rec.record_id, |
| 158 | run_id=rec.run_id, |
| 159 | payload=_unwrap_payload(rec.payload), |
| 160 | created_at=_utc_now(), |
| 161 | expires_at=rec.expires_at, |
| 162 | ) |
| 163 | session.add(row) |
| 164 | inserted += 1 |
| 165 | |
| 166 | # ── Non-heartbeats: write-once ──────────────────────────────────────────── |
| 167 | for rec in others: |
| 168 | row = db.MusehubCoordRecord( |
| 169 | repo_id=repo_id, |
| 170 | kind=rec.kind, |
| 171 | record_id=rec.record_id, |
| 172 | run_id=rec.run_id, |
| 173 | payload=_unwrap_payload(rec.payload), |
| 174 | created_at=_utc_now(), |
| 175 | expires_at=rec.expires_at, |
| 176 | ) |
| 177 | session.add(row) |
| 178 | try: |
| 179 | await session.flush() |
| 180 | inserted += 1 |
| 181 | except IntegrityError: |
| 182 | await session.rollback() |
| 183 | skipped += 1 |
| 184 | logger.debug( |
| 185 | "coord_push: skipped duplicate %s/%s for repo %s", |
| 186 | rec.kind, |
| 187 | rec.record_id, |
| 188 | repo_id, |
| 189 | ) |
| 190 | |
| 191 | await session.commit() |
| 192 | |
| 193 | # ── Materialize into structured tables ─────────────────────────────────── |
| 194 | # Only materialize kinds that have structured tables; fire-and-forget |
| 195 | # (failures are logged but do not affect the push response). |
| 196 | _MATERIALIZE_KINDS = {"reservation", "release", "task", "claim"} |
| 197 | if any(r.kind in _MATERIALIZE_KINDS for r in request.records): |
| 198 | from musehub.services.musehub_coord_server import materialize_coord_record |
| 199 | for rec in request.records: |
| 200 | if rec.kind in _MATERIALIZE_KINDS: |
| 201 | await materialize_coord_record( |
| 202 | session, repo_id, |
| 203 | rec.kind, rec.record_id, _unwrap_payload(rec.payload), rec.expires_at, |
| 204 | ) |
| 205 | |
| 206 | return CoordPushResponse(inserted=inserted, skipped=skipped) |
| 207 | |
| 208 | # ── Pull ─────────────────────────────────────────────────────────────────────── |
| 209 | |
| 210 | async def coord_pull( |
| 211 | session: AsyncSession, |
| 212 | repo_id: str, |
| 213 | request: CoordPollRequest, |
| 214 | ) -> CoordPullResponse: |
| 215 | """Return coordination records for *repo_id* since *request.since_id*. |
| 216 | |
| 217 | Records are returned in insertion order (oldest first) so that callers can |
| 218 | replay history deterministically. The returned ``cursor`` is the ``id`` of |
| 219 | the last record in the response — pass it as ``since_id`` in the next call. |
| 220 | |
| 221 | Args: |
| 222 | session: Async SQLAlchemy session. |
| 223 | repo_id: Target repository ID. |
| 224 | request: Validated poll request with ``since_id``, ``kinds``, ``limit``. |
| 225 | |
| 226 | Returns: |
| 227 | :class:`~musehub.models.coord.CoordPullResponse`. |
| 228 | """ |
| 229 | stmt = ( |
| 230 | select(db.MusehubCoordRecord) |
| 231 | .where( |
| 232 | and_( |
| 233 | db.MusehubCoordRecord.repo_id == repo_id, |
| 234 | db.MusehubCoordRecord.id > request.since_id, |
| 235 | ) |
| 236 | ) |
| 237 | .order_by(db.MusehubCoordRecord.id.asc()) |
| 238 | .limit(request.limit) |
| 239 | ) |
| 240 | |
| 241 | if request.kinds: |
| 242 | stmt = stmt.where(db.MusehubCoordRecord.kind.in_(request.kinds)) |
| 243 | |
| 244 | result = await session.execute(stmt) |
| 245 | rows = result.scalars().all() |
| 246 | |
| 247 | records = [_row_to_out(r) for r in rows] |
| 248 | cursor = records[-1].id if records else 0 |
| 249 | |
| 250 | return CoordPullResponse(records=records, cursor=cursor) |
| 251 | |
| 252 | # ── Watch (SSE stream) ───────────────────────────────────────────────────────── |
| 253 | |
| 254 | async def coord_watch_stream( |
| 255 | repo_id: str, |
| 256 | since_id: int, |
| 257 | kinds: list[str], |
| 258 | get_session: "type[AsyncIterator[AsyncSession]]", # noqa: F821 |
| 259 | ) -> AsyncIterator[str]: |
| 260 | """Async generator that yields SSE events for new coordination records. |
| 261 | |
| 262 | Polls the database every :data:`_POLL_INTERVAL_SECONDS` seconds. Yields: |
| 263 | |
| 264 | - ``event: coord_record`` — for each new record (one event per record). |
| 265 | - ``": heartbeat"`` — every :data:`_HEARTBEAT_INTERVAL_SECONDS` of silence, |
| 266 | to keep HTTP proxies alive. |
| 267 | |
| 268 | The client should reconnect using ``Last-Event-ID`` (the ``id`` field of |
| 269 | the last received event) as the ``?since_id`` query parameter. |
| 270 | |
| 271 | Args: |
| 272 | repo_id: Target repository ID. |
| 273 | since_id: Only yield records with ``id > since_id``. |
| 274 | kinds: Filter by record kind. Empty list = all kinds. |
| 275 | get_session: Async context manager factory for database sessions. |
| 276 | The watch stream opens one session per poll cycle so that |
| 277 | idle connections are not held between polls. |
| 278 | |
| 279 | Yields: |
| 280 | SSE event strings (``"event: ...\\ndata: ...\\n\\n"`` or ``": heartbeat\\n\\n"``). |
| 281 | """ |
| 282 | cursor = since_id |
| 283 | last_event_at = asyncio.get_event_loop().time() |
| 284 | |
| 285 | while True: |
| 286 | now = asyncio.get_event_loop().time() |
| 287 | |
| 288 | # Inject heartbeat if silent for too long. |
| 289 | if now - last_event_at >= _HEARTBEAT_INTERVAL_SECONDS: |
| 290 | yield sse_heartbeat() |
| 291 | last_event_at = now |
| 292 | |
| 293 | # Poll the database for new records. |
| 294 | try: |
| 295 | async for session in get_session(): |
| 296 | stmt = ( |
| 297 | select(db.MusehubCoordRecord) |
| 298 | .where( |
| 299 | and_( |
| 300 | db.MusehubCoordRecord.repo_id == repo_id, |
| 301 | db.MusehubCoordRecord.id > cursor, |
| 302 | ) |
| 303 | ) |
| 304 | .order_by(db.MusehubCoordRecord.id.asc()) |
| 305 | .limit(_WATCH_BATCH_SIZE) |
| 306 | ) |
| 307 | if kinds: |
| 308 | stmt = stmt.where(db.MusehubCoordRecord.kind.in_(kinds)) |
| 309 | |
| 310 | result = await session.execute(stmt) |
| 311 | rows = result.scalars().all() |
| 312 | |
| 313 | for row in rows: |
| 314 | cursor = row.id |
| 315 | payload = _row_to_out(row).model_dump(mode="json") |
| 316 | yield sse_event(payload, event_id=str(cursor), event_type="coord_record") |
| 317 | last_event_at = asyncio.get_event_loop().time() |
| 318 | |
| 319 | except Exception: # noqa: BLE001 |
| 320 | logger.exception("coord_watch_stream: poll error for repo %s", repo_id) |
| 321 | |
| 322 | await asyncio.sleep(_POLL_INTERVAL_SECONDS) |
File History
11 commits
sha256:d035733f21ccff27735fddebfbbe0ed24565a32a22db8de5885402262671ecd2
chore: bump version to 0.2.0rc15 for musehub#113 fix release
Sonnet 4.6
patch
15 days ago
sha256:0032d6cfa33bc3c8367436ad768e7dd0e339b4332153160247da8266cb5fa352
Merge branch 'task/version-tags-phase3-server' into dev
Human
17 days ago
sha256:4669620efda9ff41c55bdefd1f7bfe1c239d468428744c84ead9957e5a003a53
merge: rescue snapshot-recovery hardening (c00aa21d) into d…
Opus 4.8
minor
⚠
30 days ago
sha256:a59da49c4611b970fc4b6ae48678ce4943261c213a07ddbd73ce9201df869b4a
fix: remove false-positive proposal_comments index drop fro…
Sonnet 4.6
patch
34 days ago
sha256:0a240d6dbff234f07d98a28a4a9a68db702f3f9ff9260196f24219bdb1c0b6f3
feat: render markdown mists as HTML with heading anchor links
Sonnet 4.6
patch
35 days ago
sha256:24a7d47486ebc4ebd1832830580e177ec6f877b48dced8c000e198cdec4ce9d6
Merge 'task/bump-version-rc12' into 'dev' — proposal: Bump …
Human
36 days ago
sha256:b9ff931d147e0114a1f17060f415b89ed551c170a91ff226c70437aa5c85f9ee
Merge 'task/bump-version-rc12' into 'dev' — proposal: Bump …
Human
36 days ago
sha256:d1122d21e73471879b460037b22c0b50fded7c423444a176f248428f75dac39c
Merge 'task/fix-issue-pagination-cursor' into 'dev' — propo…
Human
36 days ago
sha256:01e18975e73d2b3cd5b6db7929c895bef9aa6e0d4391dc5b2adfc548b41318dd
Merge 'feat/adding-debug-logs-to-staging' into 'dev' — prop…
Human
36 days ago
sha256:6b1949fc2797ca4c1936a637a4cbfec828ef56cf52398a2e74ca3c4f494e728f
fix: use wire_bytes not mpack_bytes_raw in compute_object_b…
Sonnet 4.6
patch
48 days ago
sha256:b99f2455dc346966d040133f5203297e6e3ef5803a93728a2c30568d0a0f7583
rename: delta_add → delta_upsert across wire format, models…
Sonnet 4.6
patch
50 days ago