"""TDD — multi-batch push: objects sent across sequential push/stream requests. The muse CLI sends objects in batches of CHUNK_OBJECTS (500). Each batch is a separate HTTP POST to /push/stream. Only the final batch carries commits and snapshots. The server's referential integrity check on the final batch must find ALL objects referenced by ALL snapshot manifests — not just the ones in the last batch's MWP stream. This means objects from earlier batches must be: (a) stored in R2 and (b) committed to musehub_objects in the DB before the final batch's integrity check runs. Failure mode ------------ If any earlier batch rolls back its DB transaction (e.g. due to a connection error mid-stream), the objects from that batch never land in the DB. The final batch's integrity check queries the DB, finds them missing, and returns 422. Invariants encoded here ----------------------- MB-1 Objects sent in a non-final batch (no commits) are stored in the DB after that batch's request completes. MB-2 A final batch with commits succeeds (RESULT.ok=True) when snapshot manifests reference objects that were sent in earlier batches. MB-3 A final batch with commits fails (RESULT.ok=False, 422-style message) when snapshot manifests reference objects that were NEVER sent in any batch and are not pre-registered. MB-5 If a non-final batch's DB transaction is rolled back (simulated), the final batch fails — proving that DB commit of earlier batches is load- bearing, not optional. """ from __future__ import annotations from datetime import datetime, timezone from collections.abc import Mapping from muse.core.types import blob_id, now_utc_iso from musehub.db.musehub_models import MusehubRepo from musehub.types.json_types import JSONObject, JSONValue from unittest.mock import AsyncMock, patch import msgpack import pytest from sqlalchemy import select from sqlalchemy.ext.asyncio import AsyncSession from muse.core.mpack import MuseWireFrameWriter from musehub.models.wire import ( SFRAME_COMMIT_PACK, SFRAME_END, SFRAME_ERROR, SFRAME_HEADER, SFRAME_OBJECT, SFRAME_RESULT, ) _fw = MuseWireFrameWriter() # --------------------------------------------------------------------------- # Helpers # --------------------------------------------------------------------------- def _pack(data: JSONValue) -> bytes: return msgpack.packb(data, use_bin_type=True) def _wrap(ft: str, data: JSONValue) -> bytes: return _fw.wrap(frame_type=ft, payload=_pack(data)) def _header_frame(n_objects: int = 0, n_commits: int = 0, branch: str = "main") -> bytes: return _wrap(SFRAME_HEADER, { "t": SFRAME_HEADER, "branch": branch, "force": False, "have": [], "head": blob_id(b"head"), "n_objects": n_objects, "n_commits": n_commits, }) def _object_frame(oid: str, content: bytes) -> bytes: return _wrap(SFRAME_OBJECT, { "t": SFRAME_OBJECT, "id": oid, "content": content, "path": "file.bin", "enc": "raw", }) def _commit_pack_frame(commits: list[JSONObject], snapshots: list[JSONObject] | None = None) -> bytes: return _wrap(SFRAME_COMMIT_PACK, { "t": SFRAME_COMMIT_PACK, "commits": commits, "snapshots": snapshots or [], }) def _end_frame(n_objects: int = 0, n_commits: int = 0) -> bytes: return _wrap(SFRAME_END, { "t": SFRAME_END, "n_objects": n_objects, "n_commits": n_commits, }) async def _collect_frames(gen: AsyncIterator[bytes]) -> list[JSONObject]: unpacker = msgpack.Unpacker(raw=False) async for chunk in gen: unpacker.feed(chunk) return list(unpacker) def _make_commit(snapshot_id: str, branch: str = "main") -> JSONObject: cid = blob_id(f"commit-{now_utc_iso()}".encode()) return { "commit_id": cid, "parent_commit_id": None, "parent2_commit_id": None, "snapshot_id": snapshot_id, "branch": branch, "message": "multibatch test commit", "author": "gabriel", "committed_at": now_utc_iso(), "signature": "", "signer_key_id": "", "agent_id": "", "model_id": "", "metadata": {}, } def _make_snapshot(snapshot_id: str, manifest: JSONObject) -> JSONObject: return { "snapshot_id": snapshot_id, "manifest": manifest, "committed_at": now_utc_iso(), } async def _make_repo(db_session: AsyncSession, name: str) -> MusehubRepo: import secrets as _secrets from musehub.db.musehub_models import MusehubRepo, MusehubBranch from musehub.core.genesis import compute_repo_id, compute_branch_id owner_user_id = _secrets.token_hex(16) slug = name.lower().replace(" ", "-") created_at = datetime.now(tz=timezone.utc) repo_id = compute_repo_id(owner_user_id, slug, "", created_at.isoformat()) repo = MusehubRepo( repo_id=repo_id, name=name, owner="gabriel", slug=slug, visibility="public", owner_user_id=owner_user_id, description="", tags=[], created_at=created_at, ) db_session.add(repo) await db_session.commit() branch = MusehubBranch( branch_id=compute_branch_id(repo_id, "main"), repo_id=repo_id, name="main", ) db_session.add(branch) await db_session.commit() await db_session.refresh(repo) return repo def _stub_r2(monkeypatch: pytest.MonkeyPatch) -> Mapping[str, bytes]: """Patch R2 with an in-memory store. Returns the store dict for inspection.""" _store: dict[str, bytes] = {} async def _put(oid: str, data: bytes, **_: JSONValue) -> str: _store[oid] = data return f"https://r2.fake/{oid}" async def _get(oid: str, **_: JSONValue) -> bytes | None: return _store.get(oid) async def _exists(oid: str, **_: JSONValue) -> bool: return oid in _store backend = AsyncMock() backend.put = _put backend.get = _get backend.exists = _exists monkeypatch.setattr("musehub.services.musehub_wire.get_backend", lambda: backend) return _store def _make_objects(n: int, seed: str = "") -> list[tuple[str, bytes]]: """Return n (oid, content) pairs.""" return [ (lambda c: (blob_id(c), c))(f"{seed}object-{i}".encode()) for i in range(n) ] async def _push_batch( session: AsyncSession, repo_id: str, objects: list[tuple[str, bytes]], commits: list[dict] | None = None, snapshots: list[dict] | None = None, branch: str = "main", ) -> list[dict]: """Send one push/stream batch to wire_push_stream. Returns decoded frames.""" from musehub.services.musehub_wire import wire_push_stream n_objects = len(objects) n_commits = len(commits or []) async def body() -> None: frames = _header_frame(n_objects=n_objects, n_commits=n_commits, branch=branch) for oid, content in objects: frames += _object_frame(oid, content) frames += _commit_pack_frame(commits or [], snapshots or []) frames += _end_frame(n_objects=n_objects, n_commits=n_commits) yield frames return await _collect_frames( wire_push_stream(session, repo_id, body(), "gabriel") ) # --------------------------------------------------------------------------- # MB-1 — objects from a non-final batch are stored in the DB # --------------------------------------------------------------------------- @pytest.mark.asyncio async def test_mb1_nonfinal_batch_objects_land_in_db( db_session: AsyncSession, monkeypatch: pytest.MonkeyPatch ) -> None: """Objects sent in a batch with no commits must be stored in musehub_objects after that batch's push/stream request completes.""" from musehub.db import musehub_models as db _stub_r2(monkeypatch) repo = await _make_repo(db_session, "MB-1 Repo") objects = _make_objects(3, seed="mb1-") oids = [oid for oid, _ in objects] # Non-final batch: objects only, no commits frames = await _push_batch(db_session, str(repo.repo_id), objects) result = next((f for f in frames if f.get("t") == SFRAME_RESULT), None) assert result is not None and result["ok"] is True, ( f"non-final batch must return ok=True; frames: {frames}" ) await db_session.commit() stored = set( (await db_session.execute( select(db.MusehubObject.object_id).where( db.MusehubObject.object_id.in_(oids) ) )).scalars().all() ) assert stored == set(oids), ( f"objects from non-final batch must be in DB;\n" f" expected: {set(oids)}\n" f" found: {stored}\n" f" missing: {set(oids) - stored}" ) # --------------------------------------------------------------------------- # MB-2 — final batch succeeds when earlier batches stored their objects # --------------------------------------------------------------------------- @pytest.mark.asyncio async def test_mb2_final_batch_succeeds_with_objects_from_earlier_batches( db_session: AsyncSession, monkeypatch: pytest.MonkeyPatch ) -> None: """push/stream final batch must return RESULT.ok=True when snapshot manifests reference objects that were sent in earlier (non-final) batches.""" _stub_r2(monkeypatch) repo = await _make_repo(db_session, "MB-2 Repo") # Batch 0 and 1: objects only batch0 = _make_objects(5, seed="mb2-b0-") batch1 = _make_objects(5, seed="mb2-b1-") for batch in (batch0, batch1): frames = await _push_batch(db_session, str(repo.repo_id), batch) result = next((f for f in frames if f.get("t") == SFRAME_RESULT), None) assert result is not None and result["ok"] is True, ( f"intermediate batch must succeed; frames: {frames}" ) await db_session.commit() # Final batch: one more object + commit referencing ALL objects batch2 = _make_objects(2, seed="mb2-b2-") all_objects = batch0 + batch1 + batch2 manifest = {f"file_{i}.bin": oid for i, (oid, _) in enumerate(all_objects)} snap_id = blob_id(b"mb2-snap") commit = _make_commit(snap_id) snap = _make_snapshot(snap_id, manifest) frames = await _push_batch( db_session, str(repo.repo_id), objects=batch2, commits=[commit], snapshots=[snap], ) result_frames = [f for f in frames if f.get("t") == SFRAME_RESULT] assert result_frames, f"no RESULT frame; got: {[f.get('t') for f in frames]}" assert result_frames[0]["ok"] is True, ( f"final batch must succeed when earlier-batch objects are in DB; " f"got: {result_frames[0]}" ) # --------------------------------------------------------------------------- # MB-3 — final batch fails when referenced objects were never sent # --------------------------------------------------------------------------- @pytest.mark.asyncio async def test_mb3_final_batch_fails_when_objects_never_sent( db_session: AsyncSession, monkeypatch: pytest.MonkeyPatch ) -> None: """push/stream must reject a commit whose snapshot references an object that was never sent in any batch and is not in the DB (not pre-registered). This is the baseline that confirms the integrity check is working.""" _stub_r2(monkeypatch) repo = await _make_repo(db_session, "MB-3 Repo") ghost_oid = blob_id(b"ghost-object-never-sent") snap_id = blob_id(b"mb3-snap") commit = _make_commit(snap_id) snap = _make_snapshot(snap_id, {"ghost.bin": ghost_oid}) frames = await _push_batch( db_session, str(repo.repo_id), objects=[], commits=[commit], snapshots=[snap], ) result_frames = [f for f in frames if f.get("t") == SFRAME_RESULT] error_frames = [f for f in frames if f.get("t") == SFRAME_ERROR] assert result_frames or error_frames, ( f"push must be rejected when snapshot references unsent object; " f"got frame types: {[f.get('t') for f in frames]}" ) if result_frames: assert result_frames[0]["ok"] is False, ( f"push must fail for ghost object; got: {result_frames[0]}" ) # --------------------------------------------------------------------------- # MB-5 — rolled-back earlier batch causes final batch to fail (causal proof) # --------------------------------------------------------------------------- @pytest.mark.asyncio async def test_mb5_rolled_back_earlier_batch_causes_final_failure( db_session: AsyncSession, monkeypatch: pytest.MonkeyPatch ) -> None: """If an earlier batch's DB transaction is rolled back (simulating a connection error mid-stream), its objects are absent from the DB. The final batch's integrity check must detect this and fail. This test is the causal proof: DB commit of earlier batches is load-bearing. """ _stub_r2(monkeypatch) repo = await _make_repo(db_session, "MB-5 Repo") # Simulate a batch whose objects were stored in R2 but rolled back in the DB. # We do this by inserting objects directly into the R2 stub but NOT calling # push_batch (so no DB rows are written). dropped_objects = _make_objects(3, seed="mb5-dropped-") # These objects exist in R2 (simulated) but have no DB rows. # (In production: first attempt connected, stored to R2, then SSL error # caused DB rollback. Retry re-sent them and they DID land in DB. # Here we test the intermediate failure state.) # Final batch references the dropped (DB-absent) objects manifest = {f"dropped_{i}.bin": oid for i, (oid, _) in enumerate(dropped_objects)} snap_id = blob_id(b"mb5-snap") commit = _make_commit(snap_id) snap = _make_snapshot(snap_id, manifest) frames = await _push_batch( db_session, str(repo.repo_id), objects=[], # not re-sending them in this batch either commits=[commit], snapshots=[snap], ) result_frames = [f for f in frames if f.get("t") == SFRAME_RESULT] if result_frames: assert result_frames[0]["ok"] is False, ( "final batch must fail when earlier batch was rolled back " f"and objects are absent from DB; got: {result_frames[0]}" ) else: # An error frame (not RESULT) is also acceptable assert any( f.get("t") not in (SFRAME_RESULT,) for f in frames ), "expected failure response for missing objects" # --------------------------------------------------------------------------- # MB-6 — have-excluded objects absent from DB causes final batch to fail # --------------------------------------------------------------------------- @pytest.mark.asyncio async def test_mb6_have_excluded_objects_absent_from_db_rejects_push( db_session: AsyncSession, monkeypatch: pytest.MonkeyPatch ) -> None: """The root cause of the staging 824-missing-objects 422. The push CLI computes a `have` set from the remote's branch heads. Objects reachable from those heads are excluded from the wire bundle — the server is assumed to already have them. If the server does NOT have those objects in its DB (e.g. the branch was created server-side but its objects were never pushed), the final batch's integrity check must reject the push. Invariant MB-6: 4789 objects walked − 3965 loaded = 824 excluded by `have` = 824 missing on server. The 824 are deterministic because they are always the same objects from `main`'s history that were never pushed to staging. This test proves the server-side invariant is sound. The fix is to push the have-anchor branch first (see MB-7). """ _stub_r2(monkeypatch) repo = await _make_repo(db_session, "MB-6 Repo") # Objects that the CLI would EXCLUDE from the wire bundle because the # remote has a branch head (e.g. `main`) that covers them. On staging, # those objects were never actually pushed — so the server DB has no rows. have_excluded = _make_objects(5, seed="mb6-have-excluded-") excluded_oids = [oid for oid, _ in have_excluded] # Do NOT send them in any batch — simulating the have-exclusion. snap_id = blob_id(b"mb6-snap") commit = _make_commit(snap_id) snap = _make_snapshot(snap_id, {f"shared_{i}.bin": oid for i, oid in enumerate(excluded_oids)}) frames = await _push_batch( db_session, str(repo.repo_id), objects=[], # excluded_oids not in wire bundle commits=[commit], snapshots=[snap], ) result_frames = [f for f in frames if f.get("t") == SFRAME_RESULT] error_frames = [f for f in frames if f.get("t") == SFRAME_ERROR] assert result_frames or error_frames, ( f"expected rejection; got frame types: {[f.get('t') for f in frames]}" ) if result_frames: assert result_frames[0]["ok"] is False, ( f"push must fail when have-excluded objects absent from DB; " f"got: {result_frames[0]}" ) # --------------------------------------------------------------------------- # MB-7 — push have-anchor branch first, then dependent branch succeeds # --------------------------------------------------------------------------- @pytest.mark.asyncio async def test_mb7_push_have_anchor_branch_first_then_dev_succeeds( db_session: AsyncSession, monkeypatch: pytest.MonkeyPatch ) -> None: """The fix for the staging 422: push `main` before pushing `dev`. When `dev`'s snapshot manifests reference objects from `main`'s history and the CLI's `have` computation excludes those objects from the wire bundle, the server must already have them. Pushing `main` first ensures that. Invariant MB-7: push branch A (puts shared objects in DB) → push branch B with have-exclusion of those objects → integrity check passes. This is the minimal server-side proof that the two-push sequence fixes the 824-missing-objects staging 422. """ _stub_r2(monkeypatch) repo = await _make_repo(db_session, "MB-7 Repo") # Shared objects — would be excluded by have=[main_head] when pushing dev. shared_objects = _make_objects(5, seed="mb7-shared-") shared_oids = [oid for oid, _ in shared_objects] # ── Step 1: push "main" (the have-anchor branch). # This puts shared_oids into the DB. main_snap_id = blob_id(b"mb7-main-snap") main_commit = _make_commit(main_snap_id, branch="main") main_snap = _make_snapshot(main_snap_id, {f"shared_{i}.bin": oid for i, oid in enumerate(shared_oids)}) main_frames = await _push_batch( db_session, str(repo.repo_id), objects=shared_objects, commits=[main_commit], snapshots=[main_snap], branch="main", ) await db_session.commit() main_result = [f for f in main_frames if f.get("t") == SFRAME_RESULT] assert main_result and main_result[0]["ok"] is True, ( f"main push must succeed; got: {main_result}" ) # ── Step 2: push "dev" with have-excluded shared objects. # shared_oids are NOT sent in the wire bundle (the CLI excluded them via # have=[main_head]). The server must find them in the DB from step 1. dev_snap_id = blob_id(b"mb7-dev-snap") dev_only_objects = _make_objects(3, seed="mb7-dev-only-") dev_commit = _make_commit(dev_snap_id, branch="dev") dev_snap = _make_snapshot(dev_snap_id, { **{f"shared_{i}.bin": oid for i, oid in enumerate(shared_oids)}, **{f"dev_{i}.bin": oid for i, (oid, _) in enumerate(dev_only_objects)}, }) dev_frames = await _push_batch( db_session, str(repo.repo_id), objects=dev_only_objects, # shared_oids intentionally excluded commits=[dev_commit], snapshots=[dev_snap], branch="dev", ) dev_result = [f for f in dev_frames if f.get("t") == SFRAME_RESULT] assert dev_result, ( f"expected RESULT frame from dev push; got: {[f.get('t') for f in dev_frames]}" ) assert dev_result[0]["ok"] is True, ( f"dev push must succeed when shared objects are in DB from main push; " f"got: {dev_result[0]}" )