"""Wire verb atomic tests — push T0–T6, fetch F0–F4. Usage: python3 scripts/wire_test.py # push + fetch vs localhost python3 scripts/wire_test.py --target staging # vs staging python3 scripts/wire_test.py --push-only # push tests only python3 scripts/wire_test.py --fetch-only # fetch tests only Cert verification uses the mkcert CA root for localhost (proper TLS, no skip). """ from __future__ import annotations import asyncio import hashlib import os import sys import time from pathlib import Path import httpx import msgpack sys.path.insert(0, str(Path.home() / "ecosystem" / "muse")) sys.path.insert(0, str(Path.home() / "ecosystem" / "musehub")) MKCERT_CA = str(Path.home() / "Library/Application Support/mkcert/rootCA.pem") TARGETS = { "local": ("https://localhost:1337", MKCERT_CA), "staging": ("https://staging.musehub.ai", True), # True = system CAs } target_name = "local" run_push = True run_fetch = True for arg in sys.argv[1:]: if arg in TARGETS: target_name = arg elif arg == "--push-only": run_fetch = False elif arg == "--fetch-only": run_push = False BASE_URL, SSL_VERIFY = TARGETS[target_name] OWNER = "gabriel" REPO = "timing-test" PUSH_ROUTE = f"/{OWNER}/{REPO}/push/stream" FETCH_ROUTE = f"/{OWNER}/{REPO}/fetch/stream" REFS_ROUTE = f"/{OWNER}/{REPO}/refs" WIRE_CONTENT_TYPE = "application/x-muse-wire" def _oid(data: bytes) -> str: return "sha256:" + hashlib.sha256(data).hexdigest() def _pack_frame(frame_type: str, payload: dict) -> bytes: from muse.core.mpack import MuseWireFrameWriter fw = MuseWireFrameWriter() return fw.wrap(frame_type=frame_type, payload=msgpack.packb(payload, use_bin_type=True)) def _header_frame(n_objects: int, n_commits: int) -> bytes: return _pack_frame("H", { "t": "H", "branch": "main", "force": True, "have": [], "head": "", "n_objects": n_objects, "n_commits": n_commits, }) def _object_frame(size_bytes: int, index: int) -> tuple[str, bytes]: raw = os.urandom(max(64, size_bytes)) # Ensure content never starts with b'#!' to avoid the polyglot shebang guard. if raw[:2] == b"#!": raw = b"\x00" + raw[1:] oid = _oid(raw) frame = _pack_frame("O", { "t": "O", "id": oid, "path": f"file_{index:04d}.dat", "enc": "raw", "content": raw, }) return oid, frame def _content_snapshot_id(manifest: dict) -> str: """Compute a content-addressed snapshot ID from a manifest dict. Mirrors muse.core.snapshot.compute_snapshot_id: sha256 of NUL-joined sorted "path NUL hex_oid" pairs, so the same file tree always yields the same ID regardless of insertion order. """ _SEP = "\x00" parts = sorted( f"{path}{_SEP}{oid.split(':', 1)[1]}" for path, oid in manifest.items() ) payload = _SEP.join(parts).encode() return _oid(payload) def _commit_frame(n_commits: int, object_ids: list[str]) -> tuple[bytes, str, list[str]]: """Build a C frame with n_commits chained commits. Returns (frame_bytes, tip_commit_id, all_commit_ids). Wire format uses ``parent_commit_id`` (not ``parent_ids``) — WireCommit on the server has separate parent_commit_id / parent2_commit_id fields. Commit IDs are random so repeated test runs never hit on_conflict_do_nothing on already-stored rows. Snapshot IDs are content-addressed from the manifest so muse pull's hash-verification passes. """ commits = [] prev_id = "" commit_ids = [_oid(os.urandom(32)) for _ in range(n_commits)] manifests = [] for i in range(n_commits): manifest: dict[str, str] = {} if object_ids: obj_idx = i % len(object_ids) manifest[f"file_{obj_idx:04d}.dat"] = object_ids[obj_idx] manifests.append(manifest) snap_id_list = [_content_snapshot_id(m) for m in manifests] for i in range(n_commits): snap_id = snap_id_list[i] cid = commit_ids[i] commits.append({ "commit_id": cid, "parent_commit_id": prev_id if prev_id else None, "snapshot_id": snap_id, "branch": "main", "message": f"wire test commit {i}", "author": OWNER, "committed_at": "2026-04-28T00:00:00+00:00", "signature": "", "signer_key_id": "", "agent_id": "wire-test", "model_id": "n/a", "metadata": {}, }) prev_id = cid snapshots = [ {"snapshot_id": snap_id_list[i], "manifest": manifests[i]} for i in range(n_commits) ] tip = commit_ids[-1] if commit_ids else "" return _pack_frame("C", {"t": "C", "commits": commits, "snapshots": snapshots}), tip, commit_ids def _end_frame(n_objects: int, n_commits: int) -> bytes: return _pack_frame("E", {"t": "E", "n_objects": n_objects, "n_commits": n_commits}) def _auth_header(method: str, route: str) -> str: from muse.cli.config import get_signing_identity from muse.core.msign import build_msign_header s = get_signing_identity(remote_url=BASE_URL) if s is None: raise RuntimeError(f"No signing identity for {BASE_URL}") return build_msign_header(s, method, f"{BASE_URL}{route}", b"") def _auth_header_body(method: str, route: str, body: bytes) -> str: from muse.cli.config import get_signing_identity from muse.core.msign import build_msign_header s = get_signing_identity(remote_url=BASE_URL) if s is None: raise RuntimeError(f"No signing identity for {BASE_URL}") return build_msign_header(s, method, f"{BASE_URL}{route}", body) # ── Push ────────────────────────────────────────────────────────────────────── async def do_push(n_objects: int, n_commits: int, label: str = "") -> dict: object_ids, obj_frames = [], [] for i in range(n_objects): oid, frame = _object_frame(512, i) object_ids.append(oid) obj_frames.append(frame) c_frame, tip_commit_id, all_commit_ids = _commit_frame(n_commits, object_ids) body = ( _header_frame(n_objects, n_commits) + b"".join(obj_frames) + c_frame + _end_frame(n_objects, n_commits) ) headers = { "Content-Type": WIRE_CONTENT_TYPE, "Accept": WIRE_CONTENT_TYPE, "Authorization": _auth_header("POST", PUSH_ROUTE), } result: dict = { "label": label or f"{n_objects}obj+{n_commits}commits", "body_kb": round(len(body) / 1024, 1), "t_connect": None, "t_first_byte": None, "t_total": None, "status": None, "ok": None, "error": None, "progress": [], "tip_commit_id": tip_commit_id, "all_commit_ids": all_commit_ids, } try: t0 = time.perf_counter() t_first = None chunks = [] async with httpx.AsyncClient(http2=False, timeout=300.0, verify=SSL_VERIFY) as client: async with client.stream("POST", f"{BASE_URL}{PUSH_ROUTE}", content=body, headers=headers) as resp: result["t_connect"] = round((time.perf_counter() - t0) * 1000, 1) result["status"] = resp.status_code async for chunk in resp.aiter_bytes(): if t_first is None: t_first = time.perf_counter() - t0 result["t_first_byte"] = round(t_first * 1000, 1) chunks.append(chunk) result["t_total"] = round((time.perf_counter() - t0) * 1000, 1) unpacker = msgpack.Unpacker(raw=False) unpacker.feed(b"".join(chunks)) for frame in unpacker: if not isinstance(frame, dict): continue ft = frame.get("t") if ft == "P": result["progress"].append(frame.get("msg", "")) elif ft == "R": result["ok"] = frame.get("ok") elif ft == "X": result["ok"] = False result["error"] = frame.get("msg", "error frame") except httpx.ReadTimeout: result["error"] = "client timeout (>300s)" result["t_total"] = 300_000 except Exception as exc: result["error"] = str(exc)[:120] return result # ── Fetch ───────────────────────────────────────────────────────────────────── async def get_refs() -> dict: """GET /refs → {"branch_heads": {"main": "sha256:..."}, ...}""" async with httpx.AsyncClient(http2=False, timeout=30.0, verify=SSL_VERIFY) as client: resp = await client.get( f"{BASE_URL}{REFS_ROUTE}", headers={"Accept": "application/x-msgpack"}, ) if resp.status_code != 200: raise RuntimeError(f"GET /refs returned HTTP {resp.status_code}: {resp.text[:200]}") data = msgpack.unpackb(resp.content, raw=False) return data async def do_fetch( want: list[str], have: list[str], label: str = "", ) -> dict: body_bytes = msgpack.packb({"want": want, "have": have}, use_bin_type=True) headers = { "Content-Type": "application/x-msgpack", "Accept": WIRE_CONTENT_TYPE, "Authorization": _auth_header_body("POST", FETCH_ROUTE, body_bytes), } result: dict = { "label": label or f"fetch want={len(want)} have={len(have)}", "body_kb": round(len(body_bytes) / 1024, 1), "t_connect": None, "t_first_byte": None, "t_total": None, "status": None, "ok": None, "error": None, "n_objects": 0, "n_commits": 0, } try: t0 = time.perf_counter() t_first = None chunks = [] async with httpx.AsyncClient(http2=False, timeout=300.0, verify=SSL_VERIFY) as client: async with client.stream( "POST", f"{BASE_URL}{FETCH_ROUTE}", content=body_bytes, headers=headers, ) as resp: result["t_connect"] = round((time.perf_counter() - t0) * 1000, 1) result["status"] = resp.status_code async for chunk in resp.aiter_bytes(): if t_first is None: t_first = time.perf_counter() - t0 result["t_first_byte"] = round(t_first * 1000, 1) chunks.append(chunk) result["t_total"] = round((time.perf_counter() - t0) * 1000, 1) unpacker = msgpack.Unpacker(raw=False) unpacker.feed(b"".join(chunks)) for frame in unpacker: if not isinstance(frame, dict): continue ft = frame.get("t") if ft == "H": result["ok"] = True elif ft == "O": result["n_objects"] += 1 elif ft == "C": commits = frame.get("commits") or [] result["n_commits"] = len(commits) elif ft == "E": result["ok"] = True elif ft == "X": result["ok"] = False result["error"] = frame.get("msg", "X frame") if result["ok"] is None: result["ok"] = False result["error"] = "no E frame received" except httpx.ReadTimeout: result["error"] = "client timeout (>300s)" result["t_total"] = 300_000 except Exception as exc: result["error"] = str(exc)[:120] return result def _fmt_push(r: dict) -> str: ok = "✅" if r["ok"] else ("❌" if r["ok"] is False else "?") err = f" ERROR: {r['error']}" if r["error"] else "" p = f" [{len(r['progress'])} P-frames]" if r["progress"] else "" return ( f" {r['label']:35s} {r['body_kb']:7.1f} KB" f" connect={str(r['t_connect'] or '?'):>7}ms" f" first={str(r['t_first_byte'] or '?'):>7}ms" f" total={str(r['t_total'] or '?'):>8}ms" f" {ok}{p}{err}" ) def _fmt_fetch(r: dict) -> str: ok = "✅" if r["ok"] else ("❌" if r["ok"] is False else "?") err = f" ERROR: {r['error']}" if r["error"] else "" obj_info = f" [{r['n_objects']}o {r['n_commits']}c]" if r["ok"] else "" return ( f" {r['label']:35s} {r['body_kb']:7.1f} KB" f" connect={str(r['t_connect'] or '?'):>7}ms" f" first={str(r['t_first_byte'] or '?'):>7}ms" f" total={str(r['t_total'] or '?'):>8}ms" f" {ok}{obj_info}{err}" ) PUSH_TESTS = [ (1, 0, "T0: 1obj+0commits (XS)"), (0, 1, "T1: 0obj+1commit (XS)"), (1, 1, "T2: 1obj+1commit (XS)"), (10, 5, "T3: 10obj+5commits (XS)"), (500, 0, "T4: 500obj+0commits (S)"), (500, 100, "T5: 500obj+100commits(M)"), (500, 830, "T6: 500obj+830commits(L final batch)"), ] FETCH_TESTS = [ ("F0: 1obj+1commit (XS)", 1, 1, 0.0), ("F1: 10obj+5commits (XS)", 10, 5, 0.0), ("F2: 500obj+100commits (M)", 500, 100, 0.0), ("F3: 500obj+830commits (L)", 500, 830, 0.0), ("F4: incremental 830c (L)", 500, 830, 0.5), ] async def main() -> None: if run_push: print(f"\n── Push atomic tests → {BASE_URL}{PUSH_ROUTE} ──") print(f" {'label':35s} {'body':>8} {'connect':>11} {'first':>10} {'total':>12} status") print(" " + "-" * 100) for n_obj, n_com, label in PUSH_TESTS: r = await do_push(n_obj, n_com, label=label) print(_fmt_push(r)) sys.stdout.flush() if r["error"] or (r["t_total"] and r["t_total"] > 90_000): print(f"\n❌ WALL HIT — stopping here.") break if r["ok"] is False and not r["error"]: print(f"\n❌ Server returned ok=False — stopping here.") break print() if run_fetch: print(f"\n── Fetch atomic tests → {BASE_URL}{FETCH_ROUTE} ──") print(f" {'label':35s} {'body':>8} {'connect':>11} {'first':>10} {'total':>12} status") print(" " + "-" * 100) for label, n_obj, n_com, have_frac in FETCH_TESTS: pr = await do_push(n_obj, n_com) if not pr["ok"]: print(f" {label:35s} push failed: {pr['error']}") continue tip = pr["tip_commit_id"] all_ids = pr["all_commit_ids"] if have_frac > 0 and all_ids: have_cutoff = max(1, int(len(all_ids) * have_frac)) have = [all_ids[have_cutoff - 1]] else: have = [] r = await do_fetch([tip], have, label=label) print(_fmt_fetch(r)) sys.stdout.flush() if r["error"] or (r["t_total"] and r["t_total"] > 90_000): print(f"\n❌ WALL HIT — stopping here.") break print() if __name__ == "__main__": asyncio.run(main())