feat: add snapshot receiver

src/exo/routing/event_router.py

@@ -80,6 +80,9 @@ class EventRouter:
     def shutdown(self) -> None:
         self._tg.cancel_tasks()
 
+    def set_buffer_start(self, idx: int) -> None:
+        self.event_buffer.fast_forward_to(idx)
+
     async def _ingest(self, system_id: SystemId, recv: Receiver[Event]):
         idx = 0
         with recv as events:
@@ -95,7 +98,6 @@ class EventRouter:
                 self.out_for_delivery[event.event_id] = (anyio.current_time(), f_ev)
 
     async def _run_ext_in(self):
-        buf = OrderedBuffer[Event]()
         with self.external_inbound as events:
             async for event in events:
                 if event.session != self.session_id:
@@ -103,12 +105,12 @@ class EventRouter:
                 if event.origin != self.session_id.master_node_id:
                     continue
 
-                buf.ingest(event.origin_idx, event.event)
+                self.event_buffer.ingest(event.origin_idx, event.event)
                 event_id = event.event.event_id
                 if event_id in self.out_for_delivery:
                     self.out_for_delivery.pop(event_id)
 
-                drained = buf.drain_indexed()
+                drained = self.event_buffer.drain_indexed()
                 if drained:
                     self._nack_attempts = 0
                     if self._nack_cancel_scope:
@@ -119,7 +121,9 @@ class EventRouter:
                     or self._nack_cancel_scope.cancel_called
                 ):
                     # Request the next index.
-                    self._tg.start_soon(self._nack_request, buf.next_idx_to_release)
+                    self._tg.start_soon(
+                        self._nack_request, self.event_buffer.next_idx_to_release
+                    )
                     continue
 
                 for idx, event in drained:

src/exo/routing/snapshot_receiver.py

@@ -0,0 +1,109 @@
+"""Reassembles a snapshot from a stream of `SnapshotChunk`s.
+
+A receiver belongs to one node; it ignores chunks addressed to other
+requesters and chunks from prior sessions. Once a transfer's chunks have
+all been collected and the SHA-256 checks out, the snapshot is decoded into
+a `State`. Concurrent transfers (for the same requester) are tolerated:
+each is keyed by `transfer_id`.
+"""
+
+from __future__ import annotations
+
+import hashlib
+from dataclasses import dataclass, field
+from typing import final
+
+import zstandard
+from loguru import logger
+
+from exo.shared.types.common import NodeId, SessionId
+from exo.shared.types.snapshots import SnapshotChunk, SnapshotTransferId
+from exo.shared.types.state import State
+
+
+@final
+@dataclass
+class _Assembly:
+    """Partial state for one in-flight snapshot transfer."""
+
+    total_chunks: int
+    sha256_hex: str
+    schema_version: int
+    last_event_applied_idx: int
+    chunks: dict[int, bytes] = field(default_factory=dict)
+
+    def is_complete(self) -> bool:
+        return len(self.chunks) == self.total_chunks
+
+    def assemble(self) -> bytes:
+        return b"".join(self.chunks[i] for i in range(self.total_chunks))
+
+
+@dataclass
+class ReceivedSnapshot:
+    last_event_applied_idx: int
+    state: State
+
+
+class SnapshotReceiver:
+    """Filters and reassembles inbound chunks into a `ReceivedSnapshot`.
+
+    Stateless w.r.t. delivery: callers feed `SnapshotChunk`s in via `ingest`
+    and check the return value for completion.
+    """
+
+    def __init__(self, my_node_id: NodeId, session_id: SessionId) -> None:
+        self._my_node_id = my_node_id
+        self._session_id = session_id
+        self._assemblies: dict[SnapshotTransferId, _Assembly] = {}
+
+    def ingest(self, chunk: SnapshotChunk) -> ReceivedSnapshot | None:
+        """Absorb a chunk; return the snapshot once a transfer completes.
+
+        Returns None for partial transfers, mismatched recipients, stale
+        sessions, version mismatches, or corrupt payloads.
+        """
+        if chunk.requester_node_id != self._my_node_id:
+            return None
+        if chunk.session_id != self._session_id:
+            return None
+
+        existing = self._assemblies.get(chunk.transfer_id)
+        if existing is None:
+            existing = _Assembly(
+                total_chunks=chunk.total_chunks,
+                sha256_hex=chunk.sha256_hex,
+                schema_version=chunk.schema_version,
+                last_event_applied_idx=chunk.last_event_applied_idx,
+            )
+            self._assemblies[chunk.transfer_id] = existing
+        existing.chunks[chunk.chunk_index] = chunk.data
+
+        if not existing.is_complete():
+            return None
+
+        # Transfer complete — finalise and remove from the in-flight map.
+        del self._assemblies[chunk.transfer_id]
+        body = existing.assemble()
+        if hashlib.sha256(body).hexdigest() != existing.sha256_hex:
+            logger.warning(f"Snapshot {chunk.transfer_id} failed checksum; discarding")
+            return None
+        try:
+            decompressed = zstandard.ZstdDecompressor().decompress(body)
+            state = State.model_validate_json(decompressed.decode("utf-8"))
+        except (zstandard.ZstdError, ValueError) as e:
+            logger.opt(exception=e).warning(
+                f"Snapshot {chunk.transfer_id} could not be decoded; discarding"
+            )
+            return None
+        if state.schema_version != existing.schema_version:
+            # Should not happen — the master writes schema_version into both
+            # the chunk meta and the State payload — but treat it as corrupt.
+            logger.warning(
+                f"Snapshot {chunk.transfer_id} schema version mismatch "
+                f"(chunk={existing.schema_version}, state={state.schema_version})"
+            )
+            return None
+        return ReceivedSnapshot(
+            last_event_applied_idx=existing.last_event_applied_idx, state=state
+        )

src/exo/routing/tests/test_event_buffer.py

@@ -141,3 +141,28 @@ async def test_drain_and_ingest_with_new_sequence(buffer: OrderedBuffer[Event]):
     assert [e[0] for e in drained] == [2]
     assert buffer.next_idx_to_release == 3
     assert 4 in buffer.store
+
+
+@pytest.mark.asyncio
+async def test_fast_forward_discards_buffered_stale_events(
+    buffer: OrderedBuffer[Event],
+):
+    buffer.ingest(*make_indexed_event(0))
+    buffer.ingest(*make_indexed_event(2))
+    buffer.ingest(*make_indexed_event(4))
+
+    buffer.fast_forward_to(3)
+
+    assert buffer.next_idx_to_release == 3
+    assert set(buffer.store) == {4}
+
+
+@pytest.mark.asyncio
+async def test_fast_forward_only_moves_forward(buffer: OrderedBuffer[Event]):
+    buffer.ingest(*make_indexed_event(0))
+    buffer.ingest(*make_indexed_event(1))
+    buffer.drain()
+
+    buffer.fast_forward_to(1)
+
+    assert buffer.next_idx_to_release == 2

src/exo/routing/tests/test_snapshot_receiver.py

@@ -0,0 +1,151 @@
+import hashlib
+
+import pytest
+import zstandard
+
+from exo.routing.snapshot_receiver import SnapshotReceiver
+from exo.shared.types.common import NodeId, SessionId
+from exo.shared.types.snapshots import SnapshotChunk, SnapshotTransferId
+from exo.shared.types.state import State
+
+
+@pytest.fixture
+def session_id() -> SessionId:
+    return SessionId(master_node_id=NodeId("master"), election_clock=0)
+
+
+@pytest.fixture
+def my_node() -> NodeId:
+    return NodeId("worker-1")
+
+
+def _encode(state: State) -> bytes:
+    return zstandard.ZstdCompressor().compress(state.model_dump_json().encode("utf-8"))
+
+
+def _make_chunks(
+    body: bytes,
+    *,
+    chunk_size: int,
+    requester_node_id: NodeId,
+    session_id: SessionId,
+    state: State,
+    transfer_id: SnapshotTransferId | None = None,
+) -> list[SnapshotChunk]:
+    sha256 = hashlib.sha256(body).hexdigest()
+    transfer_id = transfer_id or SnapshotTransferId()
+    pieces = [body[i : i + chunk_size] for i in range(0, len(body), chunk_size)] or [
+        b""
+    ]
+    return [
+        SnapshotChunk.from_data(
+            data=piece,
+            transfer_id=transfer_id,
+            requester_node_id=requester_node_id,
+            session_id=session_id,
+            schema_version=state.schema_version,
+            last_event_applied_idx=state.last_event_applied_idx,
+            chunk_index=i,
+            total_chunks=len(pieces),
+            sha256_hex=sha256,
+        )
+        for i, piece in enumerate(pieces)
+    ]
+
+
+def test_completes_on_full_transfer(my_node: NodeId, session_id: SessionId):
+    state = State(last_event_applied_idx=42)
+    chunks = _make_chunks(
+        _encode(state),
+        chunk_size=64,
+        requester_node_id=my_node,
+        session_id=session_id,
+        state=state,
+    )
+    receiver = SnapshotReceiver(my_node, session_id)
+
+    received = None
+    for chunk in chunks:
+        received = receiver.ingest(chunk)
+    assert received is not None
+    assert received.last_event_applied_idx == 42
+    assert received.state.last_event_applied_idx == 42
+
+
+def test_handles_out_of_order_chunks(my_node: NodeId, session_id: SessionId):
+    state = State(last_event_applied_idx=99)
+    chunks = _make_chunks(
+        _encode(state),
+        chunk_size=32,
+        requester_node_id=my_node,
+        session_id=session_id,
+        state=state,
+    )
+    receiver = SnapshotReceiver(my_node, session_id)
+
+    # Reverse them.
+    received = None
+    for chunk in reversed(chunks):
+        received = receiver.ingest(chunk)
+    assert received is not None
+    assert received.last_event_applied_idx == 99
+
+
+def test_ignores_chunks_for_other_recipients(my_node: NodeId, session_id: SessionId):
+    state = State(last_event_applied_idx=1)
+    other = NodeId("worker-2")
+    chunks = _make_chunks(
+        _encode(state),
+        chunk_size=64,
+        requester_node_id=other,
+        session_id=session_id,
+        state=state,
+    )
+    receiver = SnapshotReceiver(my_node, session_id)
+    for chunk in chunks:
+        assert receiver.ingest(chunk) is None
+
+
+def test_ignores_chunks_from_stale_session(my_node: NodeId, session_id: SessionId):
+    state = State(last_event_applied_idx=1)
+    other_session = SessionId(master_node_id=NodeId("other-master"), election_clock=99)
+    chunks = _make_chunks(
+        _encode(state),
+        chunk_size=64,
+        requester_node_id=my_node,
+        session_id=other_session,
+        state=state,
+    )
+    receiver = SnapshotReceiver(my_node, session_id)
+    for chunk in chunks:
+        assert receiver.ingest(chunk) is None
+
+
+def test_discards_on_checksum_mismatch(my_node: NodeId, session_id: SessionId):
+    state = State(last_event_applied_idx=1)
+    chunks = _make_chunks(
+        _encode(state),
+        chunk_size=64,
+        requester_node_id=my_node,
+        session_id=session_id,
+        state=state,
+    )
+    # Corrupt the last byte of the last chunk.
+    original = chunks[-1]
+    chunks[-1] = SnapshotChunk.from_data(
+        data=original.data + b"\x00garbage",
+        transfer_id=original.transfer_id,
+        requester_node_id=original.requester_node_id,
+        session_id=original.session_id,
+        schema_version=original.schema_version,
+        last_event_applied_idx=original.last_event_applied_idx,
+        chunk_index=original.chunk_index,
+        total_chunks=original.total_chunks,
+        sha256_hex=original.sha256_hex,
+    )
+
+    receiver = SnapshotReceiver(my_node, session_id)
+    received = None
+    for chunk in chunks:
+        received = receiver.ingest(chunk)
+    assert received is None

src/exo/shared/tests/test_state_serialization.py

@@ -25,6 +25,7 @@ def test_state_serialization_roundtrip() -> None:
     json_repr = state.model_dump_json()
     restored_state = State.model_validate_json(json_repr)
 
+    assert restored_state.schema_version == state.schema_version
     assert (
         state.topology.to_snapshot().nodes
         == restored_state.topology.to_snapshot().nodes

src/exo/shared/types/snapshots.py

@@ -0,0 +1,54 @@
+"""Wire types for snapshot transfer between master and a joining node.
+
+Snapshots can be tens of MB; the gossipsub message ceiling is around 1 MB.
+We slice the compressed snapshot body into chunks and publish each chunk on
+the SNAPSHOT_RESPONSES topic. The receiver collects chunks for its own
+`requester_node_id`, validates the SHA-256 of the reassembled body, and
+materialises the State.
+"""
+
+import base64
+
+from exo.shared.types.common import Id, NodeId, SessionId
+from exo.utils.pydantic_ext import FrozenModel
+
+
+class SnapshotTransferId(Id):
+    """Identifies a single snapshot transfer (one master response to one
+    `RequestSnapshot`). Distinct transfers may interleave; the id lets
+    receivers keep them apart."""
+
+
+class SnapshotChunk(FrozenModel):
+    """One slice of a snapshot in flight.
+
+    `data_b64` carries a base64-encoded slice of the zstd-compressed JSON
+    dump of State. Concatenating the *decoded* bytes of all chunks for a
+    `transfer_id` in order of `chunk_index` yields the full compressed
+    body; `sha256_hex` is the SHA-256 of that decoded blob.
+
+    We use base64 explicitly because the topic layer JSON-encodes messages,
+    and JSON can't carry raw binary. Helpers `from_data` / `data` keep the
+    base64 detail at the boundaries.
+    """
+
+    transfer_id: SnapshotTransferId
+    requester_node_id: NodeId
+    session_id: SessionId
+    schema_version: int
+    last_event_applied_idx: int
+    chunk_index: int
+    total_chunks: int
+    sha256_hex: str
+    data_b64: str
+
+    @classmethod
+    def from_data(cls, *, data: bytes, **kwargs: object) -> "SnapshotChunk":
+        return cls(data_b64=base64.b64encode(data).decode("ascii"), **kwargs)  # pyright: ignore[reportArgumentType]
+
+    @property
+    def data(self) -> bytes:
+        return base64.b64decode(self.data_b64)
+
+
+__all__ = ["SnapshotChunk", "SnapshotTransferId"]

src/exo/shared/types/state.py

@@ -42,6 +42,9 @@ class State(FrozenModel):
         strict=True,
         arbitrary_types_allowed=True,
     )
+    # Bump when a State change makes older snapshots unsafe to restore.
+    schema_version: int = Field(default=1, ge=1)
+
     instances: Mapping[InstanceId, Instance] = {}
     runners: Mapping[RunnerId, RunnerStatus] = {}
     downloads: Mapping[NodeId, Sequence[DownloadProgress]] = {}

src/exo/utils/event_buffer.py

@@ -47,6 +47,18 @@ class OrderedBuffer[T]:
         logger.trace(f"Releasing event {ret}")
         return ret
 
+    def fast_forward_to(self, idx: int) -> None:
+        """Skip every event before idx.
+
+        Snapshot restore uses this after applying state that already includes
+        events before idx. Any buffered or future event below idx is stale.
+        """
+        if idx <= self.next_idx_to_release:
+            return
+        self.next_idx_to_release = idx
+        for stale_idx in [i for i in self.store if i < idx]:
+            del self.store[stale_idx]
+
 
 class MultiSourceBuffer[SourceId, T]:
     """