gabriel / musehub public
wire_test.py python
425 lines 14.7 KB
Raw
sha256:a34090cc4a394a78bd72cbbe34b08cc59525141e19135b6c0ab154f10611b9ef debug(push/stream): instrument O-frame decode path with INF… Sonnet 4.6 patch 123 days ago
1 """Wire verb atomic tests — push T0–T6, fetch F0–F4.
2
3 Usage:
4 python3 scripts/wire_test.py # push + fetch vs localhost
5 python3 scripts/wire_test.py --target staging # vs staging
6 python3 scripts/wire_test.py --push-only # push tests only
7 python3 scripts/wire_test.py --fetch-only # fetch tests only
8
9 Cert verification uses the mkcert CA root for localhost (proper TLS, no skip).
10 """
11 from __future__ import annotations
12
13 import asyncio
14 import hashlib
15 import os
16 import sys
17 import time
18 from pathlib import Path
19
20 import httpx
21 import msgpack
22
23 sys.path.insert(0, str(Path.home() / "ecosystem" / "muse"))
24 sys.path.insert(0, str(Path.home() / "ecosystem" / "musehub"))
25
26 MKCERT_CA = str(Path.home() / "Library/Application Support/mkcert/rootCA.pem")
27
28 TARGETS = {
29 "local": ("https://localhost:1337", MKCERT_CA),
30 "staging": ("https://staging.musehub.ai", True), # True = system CAs
31 }
32
33 target_name = "local"
34 run_push = True
35 run_fetch = True
36 for arg in sys.argv[1:]:
37 if arg in TARGETS:
38 target_name = arg
39 elif arg == "--push-only":
40 run_fetch = False
41 elif arg == "--fetch-only":
42 run_push = False
43
44 BASE_URL, SSL_VERIFY = TARGETS[target_name]
45 OWNER = "gabriel"
46 REPO = "timing-test"
47 PUSH_ROUTE = f"/{OWNER}/{REPO}/push/stream"
48 FETCH_ROUTE = f"/{OWNER}/{REPO}/fetch/stream"
49 REFS_ROUTE = f"/{OWNER}/{REPO}/refs"
50 WIRE_CONTENT_TYPE = "application/x-muse-wire"
51
52
53 def _oid(data: bytes) -> str:
54 return "sha256:" + hashlib.sha256(data).hexdigest()
55
56
57 def _pack_frame(frame_type: str, payload: dict) -> bytes:
58 from muse.core.mpack import MuseWireFrameWriter
59 fw = MuseWireFrameWriter()
60 return fw.wrap(frame_type=frame_type, payload=msgpack.packb(payload, use_bin_type=True))
61
62
63 def _header_frame(n_objects: int, n_commits: int) -> bytes:
64 return _pack_frame("H", {
65 "t": "H", "branch": "main", "force": True,
66 "have": [], "head": "", "n_objects": n_objects, "n_commits": n_commits,
67 })
68
69
70 def _object_frame(size_bytes: int, index: int) -> tuple[str, bytes]:
71 raw = os.urandom(max(64, size_bytes))
72 # Ensure content never starts with b'#!' to avoid the polyglot shebang guard.
73 if raw[:2] == b"#!":
74 raw = b"\x00" + raw[1:]
75 oid = _oid(raw)
76 frame = _pack_frame("O", {
77 "t": "O", "id": oid, "path": f"file_{index:04d}.dat",
78 "enc": "raw", "content": raw,
79 })
80 return oid, frame
81
82
83 def _content_snapshot_id(manifest: dict) -> str:
84 """Compute a content-addressed snapshot ID from a manifest dict.
85
86 Mirrors muse.core.snapshot.compute_snapshot_id: sha256 of NUL-joined
87 sorted "path NUL hex_oid" pairs, so the same file tree always yields
88 the same ID regardless of insertion order.
89 """
90 _SEP = "\x00"
91 parts = sorted(
92 f"{path}{_SEP}{oid.split(':', 1)[1]}" for path, oid in manifest.items()
93 )
94 payload = _SEP.join(parts).encode()
95 return _oid(payload)
96
97
98 def _commit_frame(n_commits: int, object_ids: list[str]) -> tuple[bytes, str, list[str]]:
99 """Build a C frame with n_commits chained commits.
100
101 Returns (frame_bytes, tip_commit_id, all_commit_ids).
102
103 Wire format uses ``parent_commit_id`` (not ``parent_ids``) — WireCommit
104 on the server has separate parent_commit_id / parent2_commit_id fields.
105
106 Commit IDs are random so repeated test runs never hit on_conflict_do_nothing
107 on already-stored rows. Snapshot IDs are content-addressed from the manifest
108 so muse pull's hash-verification passes.
109 """
110 commits = []
111 prev_id = ""
112 commit_ids = [_oid(os.urandom(32)) for _ in range(n_commits)]
113
114 manifests = []
115 for i in range(n_commits):
116 manifest: dict[str, str] = {}
117 if object_ids:
118 obj_idx = i % len(object_ids)
119 manifest[f"file_{obj_idx:04d}.dat"] = object_ids[obj_idx]
120 manifests.append(manifest)
121
122 snap_id_list = [_content_snapshot_id(m) for m in manifests]
123
124 for i in range(n_commits):
125 snap_id = snap_id_list[i]
126 cid = commit_ids[i]
127 commits.append({
128 "commit_id": cid,
129 "parent_commit_id": prev_id if prev_id else None,
130 "snapshot_id": snap_id, "branch": "main",
131 "message": f"wire test commit {i}", "author": OWNER,
132 "committed_at": "2026-04-28T00:00:00+00:00",
133 "signature": "", "signer_key_id": "",
134 "agent_id": "wire-test", "model_id": "n/a", "metadata": {},
135 })
136 prev_id = cid
137
138 snapshots = [
139 {"snapshot_id": snap_id_list[i], "manifest": manifests[i]}
140 for i in range(n_commits)
141 ]
142 tip = commit_ids[-1] if commit_ids else ""
143 return _pack_frame("C", {"t": "C", "commits": commits, "snapshots": snapshots}), tip, commit_ids
144
145
146 def _end_frame(n_objects: int, n_commits: int) -> bytes:
147 return _pack_frame("E", {"t": "E", "n_objects": n_objects, "n_commits": n_commits})
148
149
150 def _auth_header(method: str, route: str) -> str:
151 from muse.cli.config import get_signing_identity
152 from muse.core.msign import build_msign_header
153 s = get_signing_identity(remote_url=BASE_URL)
154 if s is None:
155 raise RuntimeError(f"No signing identity for {BASE_URL}")
156 return build_msign_header(s, method, f"{BASE_URL}{route}", b"")
157
158
159 def _auth_header_body(method: str, route: str, body: bytes) -> str:
160 from muse.cli.config import get_signing_identity
161 from muse.core.msign import build_msign_header
162 s = get_signing_identity(remote_url=BASE_URL)
163 if s is None:
164 raise RuntimeError(f"No signing identity for {BASE_URL}")
165 return build_msign_header(s, method, f"{BASE_URL}{route}", body)
166
167
168 # ── Push ──────────────────────────────────────────────────────────────────────
169
170
171 async def do_push(n_objects: int, n_commits: int, label: str = "") -> dict:
172 object_ids, obj_frames = [], []
173 for i in range(n_objects):
174 oid, frame = _object_frame(512, i)
175 object_ids.append(oid)
176 obj_frames.append(frame)
177
178 c_frame, tip_commit_id, all_commit_ids = _commit_frame(n_commits, object_ids)
179 body = (
180 _header_frame(n_objects, n_commits)
181 + b"".join(obj_frames)
182 + c_frame
183 + _end_frame(n_objects, n_commits)
184 )
185
186 headers = {
187 "Content-Type": WIRE_CONTENT_TYPE,
188 "Accept": WIRE_CONTENT_TYPE,
189 "Authorization": _auth_header("POST", PUSH_ROUTE),
190 }
191
192 result: dict = {
193 "label": label or f"{n_objects}obj+{n_commits}commits",
194 "body_kb": round(len(body) / 1024, 1),
195 "t_connect": None, "t_first_byte": None, "t_total": None,
196 "status": None, "ok": None, "error": None, "progress": [],
197 "tip_commit_id": tip_commit_id,
198 "all_commit_ids": all_commit_ids,
199 }
200
201 try:
202 t0 = time.perf_counter()
203 t_first = None
204 chunks = []
205
206 async with httpx.AsyncClient(http2=False, timeout=300.0, verify=SSL_VERIFY) as client:
207 async with client.stream("POST", f"{BASE_URL}{PUSH_ROUTE}", content=body, headers=headers) as resp:
208 result["t_connect"] = round((time.perf_counter() - t0) * 1000, 1)
209 result["status"] = resp.status_code
210 async for chunk in resp.aiter_bytes():
211 if t_first is None:
212 t_first = time.perf_counter() - t0
213 result["t_first_byte"] = round(t_first * 1000, 1)
214 chunks.append(chunk)
215
216 result["t_total"] = round((time.perf_counter() - t0) * 1000, 1)
217
218 unpacker = msgpack.Unpacker(raw=False)
219 unpacker.feed(b"".join(chunks))
220 for frame in unpacker:
221 if not isinstance(frame, dict):
222 continue
223 ft = frame.get("t")
224 if ft == "P":
225 result["progress"].append(frame.get("msg", ""))
226 elif ft == "R":
227 result["ok"] = frame.get("ok")
228 elif ft == "X":
229 result["ok"] = False
230 result["error"] = frame.get("msg", "error frame")
231
232 except httpx.ReadTimeout:
233 result["error"] = "client timeout (>300s)"
234 result["t_total"] = 300_000
235 except Exception as exc:
236 result["error"] = str(exc)[:120]
237
238 return result
239
240
241 # ── Fetch ─────────────────────────────────────────────────────────────────────
242
243
244 async def get_refs() -> dict:
245 """GET /refs → {"branch_heads": {"main": "sha256:..."}, ...}"""
246 async with httpx.AsyncClient(http2=False, timeout=30.0, verify=SSL_VERIFY) as client:
247 resp = await client.get(
248 f"{BASE_URL}{REFS_ROUTE}",
249 headers={"Accept": "application/x-msgpack"},
250 )
251 if resp.status_code != 200:
252 raise RuntimeError(f"GET /refs returned HTTP {resp.status_code}: {resp.text[:200]}")
253 data = msgpack.unpackb(resp.content, raw=False)
254 return data
255
256
257 async def do_fetch(
258 want: list[str],
259 have: list[str],
260 label: str = "",
261 ) -> dict:
262 body_bytes = msgpack.packb({"want": want, "have": have}, use_bin_type=True)
263
264 headers = {
265 "Content-Type": "application/x-msgpack",
266 "Accept": WIRE_CONTENT_TYPE,
267 "Authorization": _auth_header_body("POST", FETCH_ROUTE, body_bytes),
268 }
269
270 result: dict = {
271 "label": label or f"fetch want={len(want)} have={len(have)}",
272 "body_kb": round(len(body_bytes) / 1024, 1),
273 "t_connect": None, "t_first_byte": None, "t_total": None,
274 "status": None, "ok": None, "error": None,
275 "n_objects": 0, "n_commits": 0,
276 }
277
278 try:
279 t0 = time.perf_counter()
280 t_first = None
281 chunks = []
282
283 async with httpx.AsyncClient(http2=False, timeout=300.0, verify=SSL_VERIFY) as client:
284 async with client.stream(
285 "POST", f"{BASE_URL}{FETCH_ROUTE}",
286 content=body_bytes, headers=headers,
287 ) as resp:
288 result["t_connect"] = round((time.perf_counter() - t0) * 1000, 1)
289 result["status"] = resp.status_code
290 async for chunk in resp.aiter_bytes():
291 if t_first is None:
292 t_first = time.perf_counter() - t0
293 result["t_first_byte"] = round(t_first * 1000, 1)
294 chunks.append(chunk)
295
296 result["t_total"] = round((time.perf_counter() - t0) * 1000, 1)
297
298 unpacker = msgpack.Unpacker(raw=False)
299 unpacker.feed(b"".join(chunks))
300 for frame in unpacker:
301 if not isinstance(frame, dict):
302 continue
303 ft = frame.get("t")
304 if ft == "H":
305 result["ok"] = True
306 elif ft == "O":
307 result["n_objects"] += 1
308 elif ft == "C":
309 commits = frame.get("commits") or []
310 result["n_commits"] = len(commits)
311 elif ft == "E":
312 result["ok"] = True
313 elif ft == "X":
314 result["ok"] = False
315 result["error"] = frame.get("msg", "X frame")
316
317 if result["ok"] is None:
318 result["ok"] = False
319 result["error"] = "no E frame received"
320
321 except httpx.ReadTimeout:
322 result["error"] = "client timeout (>300s)"
323 result["t_total"] = 300_000
324 except Exception as exc:
325 result["error"] = str(exc)[:120]
326
327 return result
328
329
330 def _fmt_push(r: dict) -> str:
331 ok = "✅" if r["ok"] else ("❌" if r["ok"] is False else "?")
332 err = f" ERROR: {r['error']}" if r["error"] else ""
333 p = f" [{len(r['progress'])} P-frames]" if r["progress"] else ""
334 return (
335 f" {r['label']:35s} {r['body_kb']:7.1f} KB"
336 f" connect={str(r['t_connect'] or '?'):>7}ms"
337 f" first={str(r['t_first_byte'] or '?'):>7}ms"
338 f" total={str(r['t_total'] or '?'):>8}ms"
339 f" {ok}{p}{err}"
340 )
341
342
343 def _fmt_fetch(r: dict) -> str:
344 ok = "✅" if r["ok"] else ("❌" if r["ok"] is False else "?")
345 err = f" ERROR: {r['error']}" if r["error"] else ""
346 obj_info = f" [{r['n_objects']}o {r['n_commits']}c]" if r["ok"] else ""
347 return (
348 f" {r['label']:35s} {r['body_kb']:7.1f} KB"
349 f" connect={str(r['t_connect'] or '?'):>7}ms"
350 f" first={str(r['t_first_byte'] or '?'):>7}ms"
351 f" total={str(r['t_total'] or '?'):>8}ms"
352 f" {ok}{obj_info}{err}"
353 )
354
355
356 PUSH_TESTS = [
357 (1, 0, "T0: 1obj+0commits (XS)"),
358 (0, 1, "T1: 0obj+1commit (XS)"),
359 (1, 1, "T2: 1obj+1commit (XS)"),
360 (10, 5, "T3: 10obj+5commits (XS)"),
361 (500, 0, "T4: 500obj+0commits (S)"),
362 (500, 100, "T5: 500obj+100commits(M)"),
363 (500, 830, "T6: 500obj+830commits(L final batch)"),
364 ]
365
366 FETCH_TESTS = [
367 ("F0: 1obj+1commit (XS)", 1, 1, 0.0),
368 ("F1: 10obj+5commits (XS)", 10, 5, 0.0),
369 ("F2: 500obj+100commits (M)", 500, 100, 0.0),
370 ("F3: 500obj+830commits (L)", 500, 830, 0.0),
371 ("F4: incremental 830c (L)", 500, 830, 0.5),
372 ]
373
374
375 async def main() -> None:
376 if run_push:
377 print(f"\n── Push atomic tests → {BASE_URL}{PUSH_ROUTE} ──")
378 print(f" {'label':35s} {'body':>8} {'connect':>11} {'first':>10} {'total':>12} status")
379 print(" " + "-" * 100)
380
381 for n_obj, n_com, label in PUSH_TESTS:
382 r = await do_push(n_obj, n_com, label=label)
383 print(_fmt_push(r))
384 sys.stdout.flush()
385 if r["error"] or (r["t_total"] and r["t_total"] > 90_000):
386 print(f"\n❌ WALL HIT — stopping here.")
387 break
388 if r["ok"] is False and not r["error"]:
389 print(f"\n❌ Server returned ok=False — stopping here.")
390 break
391
392 print()
393
394 if run_fetch:
395 print(f"\n── Fetch atomic tests → {BASE_URL}{FETCH_ROUTE} ──")
396 print(f" {'label':35s} {'body':>8} {'connect':>11} {'first':>10} {'total':>12} status")
397 print(" " + "-" * 100)
398
399 for label, n_obj, n_com, have_frac in FETCH_TESTS:
400 pr = await do_push(n_obj, n_com)
401 if not pr["ok"]:
402 print(f" {label:35s} push failed: {pr['error']}")
403 continue
404
405 tip = pr["tip_commit_id"]
406 all_ids = pr["all_commit_ids"]
407
408 if have_frac > 0 and all_ids:
409 have_cutoff = max(1, int(len(all_ids) * have_frac))
410 have = [all_ids[have_cutoff - 1]]
411 else:
412 have = []
413
414 r = await do_fetch([tip], have, label=label)
415 print(_fmt_fetch(r))
416 sys.stdout.flush()
417 if r["error"] or (r["t_total"] and r["t_total"] > 90_000):
418 print(f"\n❌ WALL HIT — stopping here.")
419 break
420
421 print()
422
423
424 if __name__ == "__main__":
425 asyncio.run(main())
File History 1 commit
sha256:a34090cc4a394a78bd72cbbe34b08cc59525141e19135b6c0ab154f10611b9ef debug(push/stream): instrument O-frame decode path with INF… Sonnet 4.6 patch 123 days ago