musehub_wire_shared.py
python
sha256:f99af7b1a7f36c4d537d1c630d4b71fc39222b1255f82e930929e2fc89015e11
fix: relax browse_repo perf budget to 500ms — 200ms was too…
Sonnet 4.6
102 days ago
| 1 | """Shared types, helpers, and error classes for the wire protocol service.""" |
| 2 | |
| 3 | import asyncio |
| 4 | import collections |
| 5 | import hashlib |
| 6 | import logging |
| 7 | import msgpack as _msgpack |
| 8 | import time as _time_module |
| 9 | from datetime import datetime, timezone |
| 10 | |
| 11 | _BLOB_PUT_SEM: asyncio.Semaphore | None = None |
| 12 | |
| 13 | def _get_blob_put_sem() -> asyncio.Semaphore: |
| 14 | global _BLOB_PUT_SEM |
| 15 | if _BLOB_PUT_SEM is None: |
| 16 | _BLOB_PUT_SEM = asyncio.Semaphore(100) |
| 17 | return _BLOB_PUT_SEM |
| 18 | |
| 19 | from sqlalchemy import func, select, text as _sa_text |
| 20 | from sqlalchemy.dialects.postgresql import insert as _pg_insert |
| 21 | from sqlalchemy.ext.asyncio import AsyncSession |
| 22 | |
| 23 | from musehub.db.musehub_abuse_models import MusehubBlockedHash, MusehubDailyPushBytes, MusehubPushAnomaly |
| 24 | from musehub.db.musehub_collaborator_models import MusehubCollaborator |
| 25 | from musehub.db.musehub_jobs_models import MusehubBackgroundJob |
| 26 | from musehub.db.musehub_repo_models import ( |
| 27 | MusehubBranch, |
| 28 | MusehubCommit, |
| 29 | MusehubCommitGraph, |
| 30 | MusehubCommitRef, |
| 31 | MusehubObject, |
| 32 | MusehubObjectRef, |
| 33 | MusehubMPackIndex, |
| 34 | MusehubPackIndex, |
| 35 | MusehubRepo, |
| 36 | MusehubSnapshot, |
| 37 | MusehubSnapshotRef, |
| 38 | ) |
| 39 | from musehub.models.wire import ( |
| 40 | WireCommit, |
| 41 | WireMPack, |
| 42 | WireFetchRequest, |
| 43 | WireObject, |
| 44 | WireRefsResponse, |
| 45 | WireSnapshot, |
| 46 | ) |
| 47 | from muse.core.types import blob_id, decode_pubkey, split_id |
| 48 | from muse.core.ids import commit_identity_bytes as _muse_commit_identity_bytes |
| 49 | from muse.core.provenance import provenance_payload, verify_commit_ed25519 |
| 50 | from muse.core.compression import ZSTD_AVAILABLE, compress_and_encode, compute_delta |
| 51 | from musehub.core.genesis import compute_branch_id |
| 52 | from musehub.types.json_types import IntDict, JSONObject, JSONValue, StrDict |
| 53 | from musehub.types.pydantic_types import PydanticJson, unwrap |
| 54 | from musehub.config import settings |
| 55 | from musehub.storage import get_backend |
| 56 | from musehub.storage.backends import read_object_bytes |
| 57 | from musehub.services.musehub_jobs import enqueue_push_intel, enqueue_profile_snapshot |
| 58 | from collections.abc import AsyncIterator |
| 59 | from typing import TypedDict |
| 60 | |
| 61 | _monotonic = _time_module.monotonic |
| 62 | |
| 63 | class RepairResult(TypedDict): |
| 64 | repaired: bool |
| 65 | |
| 66 | class FetchPresignResult(TypedDict): |
| 67 | presign: bool |
| 68 | blob_urls: dict[str, str] |
| 69 | commits: list[dict] |
| 70 | snapshots: list[dict] |
| 71 | branch_heads: dict[str, str] |
| 72 | repo_id: str |
| 73 | domain: str |
| 74 | default_branch: str |
| 75 | expires_at: str | None |
| 76 | commit_count: int |
| 77 | blob_count: int |
| 78 | |
| 79 | class FetchMPackResult(TypedDict): |
| 80 | mpack_url: str | None |
| 81 | mpack_id: str | None |
| 82 | commit_count: int |
| 83 | blob_count: int |
| 84 | |
| 85 | class FetchNotIndexedError(Exception): |
| 86 | def __init__(self, missing_count: int) -> None: |
| 87 | super().__init__(f"{missing_count} object(s) not yet indexed") |
| 88 | self.missing_count = missing_count |
| 89 | |
| 90 | type _ObjFetchMap = dict[str, tuple[str, str, str | None, bytes | None]] |
| 91 | type _ObjectMap = dict[str, MusehubObject] |
| 92 | type _BytesMap = dict[str, bytes] |
| 93 | type _ChildMap = dict[str, list[str]] |
| 94 | |
| 95 | logger = logging.getLogger(__name__) |
| 96 | |
| 97 | def _commit_identity_bytes(wire_commit: "WireCommit") -> bytes: |
| 98 | parent_ids: list[str] = [] |
| 99 | if wire_commit.parent_commit_id: |
| 100 | parent_ids.append(wire_commit.parent_commit_id) |
| 101 | if wire_commit.parent2_commit_id: |
| 102 | parent_ids.append(wire_commit.parent2_commit_id) |
| 103 | return _muse_commit_identity_bytes( |
| 104 | parent_ids=parent_ids, |
| 105 | snapshot_id=wire_commit.snapshot_id or "", |
| 106 | message=wire_commit.message, |
| 107 | committed_at_iso=wire_commit.committed_at, |
| 108 | author=wire_commit.author, |
| 109 | signer_public_key=wire_commit.signer_public_key, |
| 110 | ) |
| 111 | |
| 112 | |
| 113 | async def _reconstruct_manifest(session: AsyncSession, snapshot_id: str) -> StrDict: |
| 114 | import msgpack as _mp |
| 115 | |
| 116 | chain: list[tuple[bytes | None, bytes | None]] = [] |
| 117 | current_id: str | None = snapshot_id |
| 118 | |
| 119 | while current_id: |
| 120 | row = (await session.execute( |
| 121 | select(MusehubSnapshot.snapshot_id, |
| 122 | MusehubSnapshot.manifest_blob, |
| 123 | MusehubSnapshot.delta_blob, |
| 124 | MusehubSnapshot.parent_snapshot_id) |
| 125 | .where(MusehubSnapshot.snapshot_id == current_id) |
| 126 | )).one_or_none() |
| 127 | if row is None: |
| 128 | break |
| 129 | _, manifest_blob, delta_blob, parent_id = row |
| 130 | if manifest_blob is not None: |
| 131 | base: dict[str, str] = dict(_mp.unpackb(manifest_blob, raw=False)) |
| 132 | for _, d_blob in reversed(chain): |
| 133 | if d_blob: |
| 134 | d = _mp.unpackb(d_blob, raw=False) |
| 135 | if isinstance(d, dict) and "add" in d: |
| 136 | base.update(d.get("add") or {}) |
| 137 | for p in d.get("rm") or []: |
| 138 | base.pop(p, None) |
| 139 | else: |
| 140 | base.update(d) |
| 141 | return base |
| 142 | chain.append((manifest_blob, delta_blob)) |
| 143 | current_id = parent_id |
| 144 | |
| 145 | base = {} |
| 146 | for _, d_blob in reversed(chain): |
| 147 | if d_blob: |
| 148 | d = _mp.unpackb(d_blob, raw=False) |
| 149 | if isinstance(d, dict) and "add" in d: |
| 150 | base.update(d.get("add") or {}) |
| 151 | for p in d.get("rm") or []: |
| 152 | base.pop(p, None) |
| 153 | else: |
| 154 | base.update(d) |
| 155 | return base |
| 156 | |
| 157 | |
| 158 | async def _upsert_object_refs( |
| 159 | session: AsyncSession, |
| 160 | repo_id: str, |
| 161 | object_ids: list[str], |
| 162 | ) -> None: |
| 163 | if not object_ids: |
| 164 | return |
| 165 | _CHUNK = 5000 |
| 166 | for i in range(0, len(object_ids), _CHUNK): |
| 167 | chunk = object_ids[i : i + _CHUNK] |
| 168 | await session.execute( |
| 169 | _pg_insert(MusehubObjectRef) |
| 170 | .values([{"repo_id": repo_id, "object_id": oid} for oid in chunk]) |
| 171 | .on_conflict_do_nothing(index_elements=["repo_id", "object_id"]) |
| 172 | ) |
| 173 | |
| 174 | class MPackValidationError(ValueError): |
| 175 | pass |
| 176 | |
| 177 | class ObjectHashMismatch(ValueError): |
| 178 | pass |
| 179 | |
| 180 | class NonFastForwardError(Exception): |
| 181 | pass |
| 182 | |
| 183 | |
| 184 | def _is_fast_forward(incoming_head: str, current_head: str, wire_bytes: bytes) -> bool: |
| 185 | if not current_head or incoming_head == current_head: |
| 186 | return True |
| 187 | |
| 188 | if wire_bytes[:4] != b"MUSE": |
| 189 | raise ValueError(f"mpack is not MUSE binary format (got {wire_bytes[:4]!r})") |
| 190 | from muse.core.mpack import parse_wire_mpack as _parse_wire_mpack |
| 191 | mpack: JSONObject = _parse_wire_mpack(wire_bytes) |
| 192 | commits: list[JSONValue] = mpack.get("commits") or [] |
| 193 | |
| 194 | parent_map: dict[str, list[str]] = {} |
| 195 | for c in commits: |
| 196 | cid = c.get("commit_id", "") |
| 197 | parents: list[str] = [] |
| 198 | if c.get("parent_commit_id"): |
| 199 | parents.append(c["parent_commit_id"]) |
| 200 | if c.get("parent2_commit_id"): |
| 201 | parents.append(c["parent2_commit_id"]) |
| 202 | parent_map[cid] = parents |
| 203 | |
| 204 | seen: set[str] = set() |
| 205 | queue: list[str] = [incoming_head] |
| 206 | while queue: |
| 207 | cid = queue.pop() |
| 208 | if cid in seen: |
| 209 | continue |
| 210 | seen.add(cid) |
| 211 | if cid == current_head: |
| 212 | return True |
| 213 | for pid in parent_map.get(cid, []): |
| 214 | queue.append(pid) |
| 215 | |
| 216 | return False |
| 217 | |
| 218 | |
| 219 | async def _is_ancestor_db( |
| 220 | session: "AsyncSession", |
| 221 | ancestor_id: str, |
| 222 | descendant_id: str, |
| 223 | repo_id: str, |
| 224 | max_hops: int = 1000, |
| 225 | ) -> bool: |
| 226 | from musehub.graph.walk import walk_dag_async |
| 227 | |
| 228 | async def _adj(cid: str) -> list[str]: |
| 229 | row = (await session.execute( |
| 230 | select(MusehubCommitGraph).where(MusehubCommitGraph.commit_id == cid) |
| 231 | )).scalar_one_or_none() |
| 232 | if row is None: |
| 233 | return [] |
| 234 | return list(row.parent_ids or []) |
| 235 | |
| 236 | async for cid in walk_dag_async(descendant_id, _adj, max_nodes=max_hops): |
| 237 | if cid == ancestor_id: |
| 238 | return True |
| 239 | return False |
| 240 | |
| 241 | |
| 242 | class PolyglotObjectError(ValueError): |
| 243 | def __init__(self, object_id: str, detail: str) -> None: |
| 244 | super().__init__(f"polyglot object rejected ({object_id!r}): {detail}") |
| 245 | self.object_id = object_id |
| 246 | |
| 247 | def _verify_object_hash(object_id: str, content: bytes) -> None: |
| 248 | if not object_id.startswith("sha256:"): |
| 249 | raise ObjectHashMismatch( |
| 250 | f"object_id {object_id!r} is not sha256-prefixed — " |
| 251 | "all content-addressed IDs must use the sha256: prefix" |
| 252 | ) |
| 253 | actual_id = blob_id(content) |
| 254 | if actual_id != object_id: |
| 255 | _, declared = split_id(object_id) |
| 256 | _, actual = split_id(actual_id) |
| 257 | raise ObjectHashMismatch( |
| 258 | f"object_id hash mismatch for {object_id!r}: " |
| 259 | f"declared={declared[:16]}… actual={actual[:16]}…" |
| 260 | ) |
| 261 | |
| 262 | def _utc_now() -> datetime: |
| 263 | return datetime.now(tz=timezone.utc) |
| 264 | |
| 265 | def _parse_iso(s: str) -> datetime: |
| 266 | try: |
| 267 | return datetime.fromisoformat(s.replace("Z", "+00:00")) |
| 268 | except (ValueError, AttributeError): |
| 269 | return _utc_now() |
| 270 | |
| 271 | def _str_values(d: JSONValue) -> StrDict: |
| 272 | if not isinstance(d, dict): |
| 273 | return {} |
| 274 | return {str(k): str(v) for k, v in d.items()} |
| 275 | |
| 276 | def _str_list(v: JSONValue) -> list[str]: |
| 277 | if not isinstance(v, list): |
| 278 | return [] |
| 279 | return [str(x) for x in v] |
| 280 | |
| 281 | def _int_safe(v: JSONValue, default: int = 0) -> int: |
| 282 | return int(v) if isinstance(v, (int, float)) else default |
| 283 | |
| 284 | def _to_wire_commit(row: MusehubCommit) -> WireCommit: |
| 285 | parent_ids: list[str] = row.parent_ids if isinstance(row.parent_ids, list) else [] |
| 286 | return WireCommit( |
| 287 | commit_id=row.commit_id, |
| 288 | branch=row.branch or "", |
| 289 | snapshot_id=row.snapshot_id, |
| 290 | message=row.message or "", |
| 291 | committed_at=row.timestamp.isoformat() if row.timestamp else "", |
| 292 | parent_commit_id=parent_ids[0] if len(parent_ids) >= 1 else None, |
| 293 | parent2_commit_id=parent_ids[1] if len(parent_ids) >= 2 else None, |
| 294 | author=row.author or "", |
| 295 | metadata={}, |
| 296 | structured_delta=row.structured_delta if isinstance(row.structured_delta, dict) else None, |
| 297 | sem_ver_bump=str(row.sem_ver_bump or "none"), |
| 298 | breaking_changes=_str_list(row.breaking_changes), |
| 299 | agent_id=str(row.agent_id or ""), |
| 300 | model_id=str(row.model_id or ""), |
| 301 | toolchain_id=str(row.toolchain_id or ""), |
| 302 | prompt_hash=str(row.prompt_hash or ""), |
| 303 | signature=str(row.signature or ""), |
| 304 | signer_public_key=str(row.signer_public_key or ""), |
| 305 | signer_key_id=str(row.signer_key_id or ""), |
| 306 | format_version=1, |
| 307 | reviewed_by=_str_list(row.reviewed_by), |
| 308 | test_runs=_int_safe(row.test_runs), |
| 309 | ) |
| 310 | |
| 311 | def _snap_row_to_wire(row: MusehubSnapshot) -> JSONObject: |
| 312 | import msgpack as _mp |
| 313 | delta_upsert: dict[str, str] = {} |
| 314 | delta_remove: list[str] = [] |
| 315 | if row.delta_blob: |
| 316 | d = _mp.unpackb(row.delta_blob, raw=False) |
| 317 | if isinstance(d, dict) and "add" in d: |
| 318 | delta_upsert = d.get("add") or {} |
| 319 | delta_remove = d.get("rm") or [] |
| 320 | else: |
| 321 | delta_upsert = d |
| 322 | elif row.manifest_blob: |
| 323 | delta_upsert = dict(_mp.unpackb(row.manifest_blob, raw=False)) |
| 324 | return { |
| 325 | "snapshot_id": row.snapshot_id, |
| 326 | "parent_snapshot_id": row.parent_snapshot_id, |
| 327 | "delta_upsert": delta_upsert, |
| 328 | "delta_remove": delta_remove, |
| 329 | "directories": list(row.directories) if row.directories else [], |
| 330 | "created_at": row.created_at.isoformat() if row.created_at else "", |
| 331 | } |
File History
1 commit
sha256:f99af7b1a7f36c4d537d1c630d4b71fc39222b1255f82e930929e2fc89015e11
fix: relax browse_repo perf budget to 500ms — 200ms was too…
Sonnet 4.6
102 days ago