feat: version state snapshots

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/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/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/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]:
     """