gabriel / musehub public
test_wire_multibatch_push.py python
537 lines 19.8 KB
Raw
sha256:a34090cc4a394a78bd72cbbe34b08cc59525141e19135b6c0ab154f10611b9ef debug(push/stream): instrument O-frame decode path with INF… Sonnet 4.6 patch 121 days ago
1 """TDD — multi-batch push: objects sent across sequential push/stream requests.
2
3 The muse CLI sends objects in batches of CHUNK_OBJECTS (500). Each batch is a
4 separate HTTP POST to /push/stream. Only the final batch carries commits and
5 snapshots. The server's referential integrity check on the final batch must
6 find ALL objects referenced by ALL snapshot manifests — not just the ones in
7 the last batch's MWP stream.
8
9 This means objects from earlier batches must be:
10 (a) stored in R2 and
11 (b) committed to musehub_objects in the DB
12
13 before the final batch's integrity check runs.
14
15 Failure mode
16 ------------
17 If any earlier batch rolls back its DB transaction (e.g. due to a connection
18 error mid-stream), the objects from that batch never land in the DB. The final
19 batch's integrity check queries the DB, finds them missing, and returns 422.
20
21 Invariants encoded here
22 -----------------------
23
24 MB-1 Objects sent in a non-final batch (no commits) are stored in the DB
25 after that batch's request completes.
26
27 MB-2 A final batch with commits succeeds (RESULT.ok=True) when snapshot
28 manifests reference objects that were sent in earlier batches.
29
30 MB-3 A final batch with commits fails (RESULT.ok=False, 422-style message)
31 when snapshot manifests reference objects that were NEVER sent in any
32 batch and are not pre-registered.
33
34 MB-5 If a non-final batch's DB transaction is rolled back (simulated), the
35 final batch fails — proving that DB commit of earlier batches is load-
36 bearing, not optional.
37 """
38 from __future__ import annotations
39
40 from datetime import datetime, timezone
41
42 from collections.abc import Mapping
43 from muse.core.types import blob_id, now_utc_iso
44 from musehub.db.musehub_models import MusehubRepo
45 from musehub.types.json_types import JSONObject, JSONValue
46 from unittest.mock import AsyncMock, patch
47
48 import msgpack
49 import pytest
50 from sqlalchemy import select
51 from sqlalchemy.ext.asyncio import AsyncSession
52
53 from muse.core.mpack import MuseWireFrameWriter
54 from musehub.models.wire import (
55 SFRAME_COMMIT_PACK,
56 SFRAME_END,
57 SFRAME_ERROR,
58 SFRAME_HEADER,
59 SFRAME_OBJECT,
60 SFRAME_RESULT,
61 )
62
63 _fw = MuseWireFrameWriter()
64
65
66 # ---------------------------------------------------------------------------
67 # Helpers
68 # ---------------------------------------------------------------------------
69
70 def _pack(data: JSONValue) -> bytes:
71 return msgpack.packb(data, use_bin_type=True)
72
73
74 def _wrap(ft: str, data: JSONValue) -> bytes:
75 return _fw.wrap(frame_type=ft, payload=_pack(data))
76
77
78 def _header_frame(n_objects: int = 0, n_commits: int = 0, branch: str = "main") -> bytes:
79 return _wrap(SFRAME_HEADER, {
80 "t": SFRAME_HEADER,
81 "branch": branch,
82 "force": False,
83 "have": [],
84 "head": blob_id(b"head"),
85 "n_objects": n_objects,
86 "n_commits": n_commits,
87 })
88
89
90 def _object_frame(oid: str, content: bytes) -> bytes:
91 return _wrap(SFRAME_OBJECT, {
92 "t": SFRAME_OBJECT,
93 "id": oid,
94 "content": content,
95 "path": "file.bin",
96 "enc": "raw",
97 })
98
99
100 def _commit_pack_frame(commits: list[JSONObject], snapshots: list[JSONObject] | None = None) -> bytes:
101 return _wrap(SFRAME_COMMIT_PACK, {
102 "t": SFRAME_COMMIT_PACK,
103 "commits": commits,
104 "snapshots": snapshots or [],
105 })
106
107
108 def _end_frame(n_objects: int = 0, n_commits: int = 0) -> bytes:
109 return _wrap(SFRAME_END, {
110 "t": SFRAME_END,
111 "n_objects": n_objects,
112 "n_commits": n_commits,
113 })
114
115
116 async def _collect_frames(gen: AsyncIterator[bytes]) -> list[JSONObject]:
117 unpacker = msgpack.Unpacker(raw=False)
118 async for chunk in gen:
119 unpacker.feed(chunk)
120 return list(unpacker)
121
122
123 def _make_commit(snapshot_id: str, branch: str = "main") -> JSONObject:
124 cid = blob_id(f"commit-{now_utc_iso()}".encode())
125 return {
126 "commit_id": cid,
127 "parent_commit_id": None,
128 "parent2_commit_id": None,
129 "snapshot_id": snapshot_id,
130 "branch": branch,
131 "message": "multibatch test commit",
132 "author": "gabriel",
133 "committed_at": now_utc_iso(),
134 "signature": "",
135 "signer_key_id": "",
136 "agent_id": "",
137 "model_id": "",
138 "metadata": {},
139 }
140
141
142 def _make_snapshot(snapshot_id: str, manifest: JSONObject) -> JSONObject:
143 return {
144 "snapshot_id": snapshot_id,
145 "manifest": manifest,
146 "committed_at": now_utc_iso(),
147 }
148
149
150 async def _make_repo(db_session: AsyncSession, name: str) -> MusehubRepo:
151 import secrets as _secrets
152 from musehub.db.musehub_models import MusehubRepo, MusehubBranch
153 from musehub.core.genesis import compute_repo_id, compute_branch_id
154 owner_user_id = _secrets.token_hex(16)
155 slug = name.lower().replace(" ", "-")
156 created_at = datetime.now(tz=timezone.utc)
157 repo_id = compute_repo_id(owner_user_id, slug, "", created_at.isoformat())
158 repo = MusehubRepo(
159 repo_id=repo_id, name=name, owner="gabriel", slug=slug,
160 visibility="public", owner_user_id=owner_user_id,
161 description="", tags=[], created_at=created_at,
162 )
163 db_session.add(repo)
164 await db_session.commit()
165 branch = MusehubBranch(
166 branch_id=compute_branch_id(repo_id, "main"),
167 repo_id=repo_id, name="main",
168 )
169 db_session.add(branch)
170 await db_session.commit()
171 await db_session.refresh(repo)
172 return repo
173
174
175 def _stub_r2(monkeypatch: pytest.MonkeyPatch) -> Mapping[str, bytes]:
176 """Patch R2 with an in-memory store. Returns the store dict for inspection."""
177 _store: dict[str, bytes] = {}
178
179 async def _put(oid: str, data: bytes, **_: JSONValue) -> str:
180 _store[oid] = data
181 return f"https://r2.fake/{oid}"
182
183 async def _get(oid: str, **_: JSONValue) -> bytes | None:
184 return _store.get(oid)
185
186 async def _exists(oid: str, **_: JSONValue) -> bool:
187 return oid in _store
188
189 backend = AsyncMock()
190 backend.put = _put
191 backend.get = _get
192 backend.exists = _exists
193 monkeypatch.setattr("musehub.services.musehub_wire.get_backend", lambda: backend)
194 return _store
195
196
197 def _make_objects(n: int, seed: str = "") -> list[tuple[str, bytes]]:
198 """Return n (oid, content) pairs."""
199 return [
200 (lambda c: (blob_id(c), c))(f"{seed}object-{i}".encode())
201 for i in range(n)
202 ]
203
204
205 async def _push_batch(
206 session: AsyncSession,
207 repo_id: str,
208 objects: list[tuple[str, bytes]],
209 commits: list[dict] | None = None,
210 snapshots: list[dict] | None = None,
211 branch: str = "main",
212 ) -> list[dict]:
213 """Send one push/stream batch to wire_push_stream. Returns decoded frames."""
214 from musehub.services.musehub_wire import wire_push_stream
215
216 n_objects = len(objects)
217 n_commits = len(commits or [])
218
219 async def body() -> None:
220 frames = _header_frame(n_objects=n_objects, n_commits=n_commits, branch=branch)
221 for oid, content in objects:
222 frames += _object_frame(oid, content)
223 frames += _commit_pack_frame(commits or [], snapshots or [])
224 frames += _end_frame(n_objects=n_objects, n_commits=n_commits)
225 yield frames
226
227 return await _collect_frames(
228 wire_push_stream(session, repo_id, body(), "gabriel")
229 )
230
231
232 # ---------------------------------------------------------------------------
233 # MB-1 — objects from a non-final batch are stored in the DB
234 # ---------------------------------------------------------------------------
235
236 @pytest.mark.asyncio
237 async def test_mb1_nonfinal_batch_objects_land_in_db(
238 db_session: AsyncSession, monkeypatch: pytest.MonkeyPatch
239 ) -> None:
240 """Objects sent in a batch with no commits must be stored in musehub_objects
241 after that batch's push/stream request completes."""
242 from musehub.db import musehub_models as db
243
244 _stub_r2(monkeypatch)
245 repo = await _make_repo(db_session, "MB-1 Repo")
246
247 objects = _make_objects(3, seed="mb1-")
248 oids = [oid for oid, _ in objects]
249
250 # Non-final batch: objects only, no commits
251 frames = await _push_batch(db_session, str(repo.repo_id), objects)
252 result = next((f for f in frames if f.get("t") == SFRAME_RESULT), None)
253 assert result is not None and result["ok"] is True, (
254 f"non-final batch must return ok=True; frames: {frames}"
255 )
256 await db_session.commit()
257
258 stored = set(
259 (await db_session.execute(
260 select(db.MusehubObject.object_id).where(
261 db.MusehubObject.object_id.in_(oids)
262 )
263 )).scalars().all()
264 )
265 assert stored == set(oids), (
266 f"objects from non-final batch must be in DB;\n"
267 f" expected: {set(oids)}\n"
268 f" found: {stored}\n"
269 f" missing: {set(oids) - stored}"
270 )
271
272
273 # ---------------------------------------------------------------------------
274 # MB-2 — final batch succeeds when earlier batches stored their objects
275 # ---------------------------------------------------------------------------
276
277 @pytest.mark.asyncio
278 async def test_mb2_final_batch_succeeds_with_objects_from_earlier_batches(
279 db_session: AsyncSession, monkeypatch: pytest.MonkeyPatch
280 ) -> None:
281 """push/stream final batch must return RESULT.ok=True when snapshot manifests
282 reference objects that were sent in earlier (non-final) batches."""
283 _stub_r2(monkeypatch)
284 repo = await _make_repo(db_session, "MB-2 Repo")
285
286 # Batch 0 and 1: objects only
287 batch0 = _make_objects(5, seed="mb2-b0-")
288 batch1 = _make_objects(5, seed="mb2-b1-")
289
290 for batch in (batch0, batch1):
291 frames = await _push_batch(db_session, str(repo.repo_id), batch)
292 result = next((f for f in frames if f.get("t") == SFRAME_RESULT), None)
293 assert result is not None and result["ok"] is True, (
294 f"intermediate batch must succeed; frames: {frames}"
295 )
296 await db_session.commit()
297
298 # Final batch: one more object + commit referencing ALL objects
299 batch2 = _make_objects(2, seed="mb2-b2-")
300 all_objects = batch0 + batch1 + batch2
301 manifest = {f"file_{i}.bin": oid for i, (oid, _) in enumerate(all_objects)}
302
303 snap_id = blob_id(b"mb2-snap")
304 commit = _make_commit(snap_id)
305 snap = _make_snapshot(snap_id, manifest)
306
307 frames = await _push_batch(
308 db_session, str(repo.repo_id),
309 objects=batch2,
310 commits=[commit],
311 snapshots=[snap],
312 )
313 result_frames = [f for f in frames if f.get("t") == SFRAME_RESULT]
314 assert result_frames, f"no RESULT frame; got: {[f.get('t') for f in frames]}"
315 assert result_frames[0]["ok"] is True, (
316 f"final batch must succeed when earlier-batch objects are in DB; "
317 f"got: {result_frames[0]}"
318 )
319
320
321 # ---------------------------------------------------------------------------
322 # MB-3 — final batch fails when referenced objects were never sent
323 # ---------------------------------------------------------------------------
324
325 @pytest.mark.asyncio
326 async def test_mb3_final_batch_fails_when_objects_never_sent(
327 db_session: AsyncSession, monkeypatch: pytest.MonkeyPatch
328 ) -> None:
329 """push/stream must reject a commit whose snapshot references an object
330 that was never sent in any batch and is not in the DB (not pre-registered).
331 This is the baseline that confirms the integrity check is working."""
332 _stub_r2(monkeypatch)
333 repo = await _make_repo(db_session, "MB-3 Repo")
334
335 ghost_oid = blob_id(b"ghost-object-never-sent")
336 snap_id = blob_id(b"mb3-snap")
337 commit = _make_commit(snap_id)
338 snap = _make_snapshot(snap_id, {"ghost.bin": ghost_oid})
339
340 frames = await _push_batch(
341 db_session, str(repo.repo_id),
342 objects=[],
343 commits=[commit],
344 snapshots=[snap],
345 )
346 result_frames = [f for f in frames if f.get("t") == SFRAME_RESULT]
347 error_frames = [f for f in frames if f.get("t") == SFRAME_ERROR]
348 assert result_frames or error_frames, (
349 f"push must be rejected when snapshot references unsent object; "
350 f"got frame types: {[f.get('t') for f in frames]}"
351 )
352 if result_frames:
353 assert result_frames[0]["ok"] is False, (
354 f"push must fail for ghost object; got: {result_frames[0]}"
355 )
356
357
358 # ---------------------------------------------------------------------------
359 # MB-5 — rolled-back earlier batch causes final batch to fail (causal proof)
360 # ---------------------------------------------------------------------------
361
362 @pytest.mark.asyncio
363 async def test_mb5_rolled_back_earlier_batch_causes_final_failure(
364 db_session: AsyncSession, monkeypatch: pytest.MonkeyPatch
365 ) -> None:
366 """If an earlier batch's DB transaction is rolled back (simulating a
367 connection error mid-stream), its objects are absent from the DB. The
368 final batch's integrity check must detect this and fail.
369
370 This test is the causal proof: DB commit of earlier batches is load-bearing.
371 """
372 _stub_r2(monkeypatch)
373 repo = await _make_repo(db_session, "MB-5 Repo")
374
375 # Simulate a batch whose objects were stored in R2 but rolled back in the DB.
376 # We do this by inserting objects directly into the R2 stub but NOT calling
377 # push_batch (so no DB rows are written).
378 dropped_objects = _make_objects(3, seed="mb5-dropped-")
379 # These objects exist in R2 (simulated) but have no DB rows.
380 # (In production: first attempt connected, stored to R2, then SSL error
381 # caused DB rollback. Retry re-sent them and they DID land in DB.
382 # Here we test the intermediate failure state.)
383
384 # Final batch references the dropped (DB-absent) objects
385 manifest = {f"dropped_{i}.bin": oid for i, (oid, _) in enumerate(dropped_objects)}
386 snap_id = blob_id(b"mb5-snap")
387 commit = _make_commit(snap_id)
388 snap = _make_snapshot(snap_id, manifest)
389
390 frames = await _push_batch(
391 db_session, str(repo.repo_id),
392 objects=[], # not re-sending them in this batch either
393 commits=[commit],
394 snapshots=[snap],
395 )
396 result_frames = [f for f in frames if f.get("t") == SFRAME_RESULT]
397 if result_frames:
398 assert result_frames[0]["ok"] is False, (
399 "final batch must fail when earlier batch was rolled back "
400 f"and objects are absent from DB; got: {result_frames[0]}"
401 )
402 else:
403 # An error frame (not RESULT) is also acceptable
404 assert any(
405 f.get("t") not in (SFRAME_RESULT,) for f in frames
406 ), "expected failure response for missing objects"
407
408
409 # ---------------------------------------------------------------------------
410 # MB-6 — have-excluded objects absent from DB causes final batch to fail
411 # ---------------------------------------------------------------------------
412
413 @pytest.mark.asyncio
414 async def test_mb6_have_excluded_objects_absent_from_db_rejects_push(
415 db_session: AsyncSession, monkeypatch: pytest.MonkeyPatch
416 ) -> None:
417 """The root cause of the staging 824-missing-objects 422.
418
419 The push CLI computes a `have` set from the remote's branch heads. Objects
420 reachable from those heads are excluded from the wire bundle — the server is
421 assumed to already have them. If the server does NOT have those objects in
422 its DB (e.g. the branch was created server-side but its objects were never
423 pushed), the final batch's integrity check must reject the push.
424
425 Invariant MB-6: 4789 objects walked − 3965 loaded = 824 excluded by `have`
426 = 824 missing on server. The 824 are deterministic because they are always
427 the same objects from `main`'s history that were never pushed to staging.
428
429 This test proves the server-side invariant is sound. The fix is to push
430 the have-anchor branch first (see MB-7).
431 """
432 _stub_r2(monkeypatch)
433 repo = await _make_repo(db_session, "MB-6 Repo")
434
435 # Objects that the CLI would EXCLUDE from the wire bundle because the
436 # remote has a branch head (e.g. `main`) that covers them. On staging,
437 # those objects were never actually pushed — so the server DB has no rows.
438 have_excluded = _make_objects(5, seed="mb6-have-excluded-")
439 excluded_oids = [oid for oid, _ in have_excluded]
440 # Do NOT send them in any batch — simulating the have-exclusion.
441
442 snap_id = blob_id(b"mb6-snap")
443 commit = _make_commit(snap_id)
444 snap = _make_snapshot(snap_id, {f"shared_{i}.bin": oid for i, oid in enumerate(excluded_oids)})
445
446 frames = await _push_batch(
447 db_session, str(repo.repo_id),
448 objects=[], # excluded_oids not in wire bundle
449 commits=[commit],
450 snapshots=[snap],
451 )
452 result_frames = [f for f in frames if f.get("t") == SFRAME_RESULT]
453 error_frames = [f for f in frames if f.get("t") == SFRAME_ERROR]
454 assert result_frames or error_frames, (
455 f"expected rejection; got frame types: {[f.get('t') for f in frames]}"
456 )
457 if result_frames:
458 assert result_frames[0]["ok"] is False, (
459 f"push must fail when have-excluded objects absent from DB; "
460 f"got: {result_frames[0]}"
461 )
462
463
464 # ---------------------------------------------------------------------------
465 # MB-7 — push have-anchor branch first, then dependent branch succeeds
466 # ---------------------------------------------------------------------------
467
468 @pytest.mark.asyncio
469 async def test_mb7_push_have_anchor_branch_first_then_dev_succeeds(
470 db_session: AsyncSession, monkeypatch: pytest.MonkeyPatch
471 ) -> None:
472 """The fix for the staging 422: push `main` before pushing `dev`.
473
474 When `dev`'s snapshot manifests reference objects from `main`'s history
475 and the CLI's `have` computation excludes those objects from the wire
476 bundle, the server must already have them. Pushing `main` first ensures
477 that.
478
479 Invariant MB-7: push branch A (puts shared objects in DB) → push branch B
480 with have-exclusion of those objects → integrity check passes.
481
482 This is the minimal server-side proof that the two-push sequence fixes the
483 824-missing-objects staging 422.
484 """
485 _stub_r2(monkeypatch)
486 repo = await _make_repo(db_session, "MB-7 Repo")
487
488 # Shared objects — would be excluded by have=[main_head] when pushing dev.
489 shared_objects = _make_objects(5, seed="mb7-shared-")
490 shared_oids = [oid for oid, _ in shared_objects]
491
492 # ── Step 1: push "main" (the have-anchor branch).
493 # This puts shared_oids into the DB.
494 main_snap_id = blob_id(b"mb7-main-snap")
495 main_commit = _make_commit(main_snap_id, branch="main")
496 main_snap = _make_snapshot(main_snap_id, {f"shared_{i}.bin": oid for i, oid in enumerate(shared_oids)})
497
498 main_frames = await _push_batch(
499 db_session, str(repo.repo_id),
500 objects=shared_objects,
501 commits=[main_commit],
502 snapshots=[main_snap],
503 branch="main",
504 )
505 await db_session.commit()
506
507 main_result = [f for f in main_frames if f.get("t") == SFRAME_RESULT]
508 assert main_result and main_result[0]["ok"] is True, (
509 f"main push must succeed; got: {main_result}"
510 )
511
512 # ── Step 2: push "dev" with have-excluded shared objects.
513 # shared_oids are NOT sent in the wire bundle (the CLI excluded them via
514 # have=[main_head]). The server must find them in the DB from step 1.
515 dev_snap_id = blob_id(b"mb7-dev-snap")
516 dev_only_objects = _make_objects(3, seed="mb7-dev-only-")
517 dev_commit = _make_commit(dev_snap_id, branch="dev")
518 dev_snap = _make_snapshot(dev_snap_id, {
519 **{f"shared_{i}.bin": oid for i, oid in enumerate(shared_oids)},
520 **{f"dev_{i}.bin": oid for i, (oid, _) in enumerate(dev_only_objects)},
521 })
522
523 dev_frames = await _push_batch(
524 db_session, str(repo.repo_id),
525 objects=dev_only_objects, # shared_oids intentionally excluded
526 commits=[dev_commit],
527 snapshots=[dev_snap],
528 branch="dev",
529 )
530 dev_result = [f for f in dev_frames if f.get("t") == SFRAME_RESULT]
531 assert dev_result, (
532 f"expected RESULT frame from dev push; got: {[f.get('t') for f in dev_frames]}"
533 )
534 assert dev_result[0]["ok"] is True, (
535 f"dev push must succeed when shared objects are in DB from main push; "
536 f"got: {dev_result[0]}"
537 )
File History 1 commit
sha256:a34090cc4a394a78bd72cbbe34b08cc59525141e19135b6c0ab154f10611b9ef debug(push/stream): instrument O-frame decode path with INF… Sonnet 4.6 patch 121 days ago