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