Use sliding window

flake.nix

@@ -110,7 +110,7 @@
               nixpkgs-fmt.enable = true;
               ruff-format = {
                 enable = true;
-                excludes = [ "rust/exo_rs/python/**/*.pyi" ];
+                excludes = [ "rust/exo_rs/exo_rs.pyi" ];
               };
               rustfmt = {
                 enable = true;

python/parts.nix

@@ -46,6 +46,7 @@ let
       exoOverlay = final: prev: {
         # Replace workspace exo_rs with Nix-built wheel.
         # Preserve passthru so mkVirtualEnv can resolve dependency groups.
+        # Copy .pyi stub + py.typed marker so basedpyright can find the types.
         exo-rs = pkgs.stdenv.mkDerivation {
           pname = "exo-rs";
           version = "0.1.0";
@@ -54,6 +55,12 @@ let
           nativeBuildInputs = [ final.pyprojectWheelHook ];
           dontStrip = true;
           passthru = prev.exo-rs.passthru or { };
+
+          postInstall = ''
+            local siteDir=$out/${final.python.sitePackages}/exo_rs
+            cp ${inputs.self}/rust/exo_rs/exo_rs.pyi $siteDir/
+            touch $siteDir/py.typed
+          '';
         };
       };
       buildSystemsOverlay = final: prev:

rust/exo_rs/exo_rs.pyi

@@ -5,7 +5,6 @@ import builtins
 import os
 import pathlib
 import typing
-
 __all__ = [
     "AllQueuesFullError",
     "FromSwarm",
@@ -26,35 +25,24 @@ class AllQueuesFullError(builtins.Exception):
 class FromSwarm:
     @typing.final
     class Connection(FromSwarm):
-        __match_args__ = (
-            "peer_id",
-            "connected",
-        )
+        __match_args__ = ("peer_id", "connected",)
         @property
         def peer_id(self) -> builtins.str: ...
         @property
         def connected(self) -> builtins.bool: ...
-        def __new__(
-            cls, peer_id: builtins.str, connected: builtins.bool
-        ) -> FromSwarm.Connection: ...
-
+        def __new__(cls, peer_id: builtins.str, connected: builtins.bool) -> FromSwarm.Connection: ...
+    
     @typing.final
     class Message(FromSwarm):
-        __match_args__ = (
-            "origin",
-            "topic",
-            "data",
-        )
+        __match_args__ = ("origin", "topic", "data",)
         @property
         def origin(self) -> builtins.str: ...
         @property
         def topic(self) -> builtins.str: ...
         @property
         def data(self) -> bytes: ...
-        def __new__(
-            cls, origin: builtins.str, topic: builtins.str, data: bytes
-        ) -> FromSwarm.Message: ...
-
+        def __new__(cls, origin: builtins.str, topic: builtins.str, data: bytes) -> FromSwarm.Message: ...
+    
     ...
 
 @typing.final
@@ -89,29 +77,24 @@ class MessageTooLargeError(builtins.Exception):
 
 @typing.final
 class NetworkingHandle:
-    def __new__(
-        cls,
-        identity: Keypair,
-        bootstrap_peers: typing.Sequence[builtins.str],
-        listen_port: builtins.int,
-    ) -> NetworkingHandle: ...
+    def __new__(cls, identity: Keypair, bootstrap_peers: typing.Sequence[builtins.str], listen_port: builtins.int) -> NetworkingHandle: ...
     def recv(self) -> typing.Awaitable[FromSwarm]: ...
     async def gossipsub_subscribe(self, topic: builtins.str) -> builtins.bool:
         r"""
         Subscribe to a `GossipSub` topic.
-
+        
         Returns `True` if the subscription worked. Returns `False` if we were already subscribed.
         """
     async def gossipsub_unsubscribe(self, topic: builtins.str) -> builtins.bool:
         r"""
         Unsubscribes from a `GossipSub` topic.
-
+        
         Returns `True` if we were subscribed to this topic. Returns `False` if we were not subscribed.
         """
     async def gossipsub_publish(self, topic: builtins.str, data: bytes) -> None:
         r"""
         Publishes a message with multiple topics to the `GossipSub` network.
-
+        
         If no peers are found that subscribe to this topic, throws `NoPeersSubscribedToTopicError` exception.
         """
 
@@ -125,30 +108,28 @@ class NoPeersSubscribedToTopicError(builtins.Exception):
 class Pidfile:
     r"""
     A PID file protected with a lock.
-
+    
     An instance of `Pidfile` can be used to manage a PID file: create it,
     lock it, detect already running daemons. It is backed by [`pidfile`]
     functions of `libbsd`/`libutil` which use `flopen` to lock the PID
     file.
-
+    
     When a PID file is created, the process ID of the current process is
     *not* written there, making it possible to lock the PID file before
     forking and only write the ID of the forked process when it is ready.
-
+    
     The PID file is deleted automatically when the `Pidfile` comes out of
     the scope. To close the PID file without deleting it, for example, in
     the parent process of a forked daemon, call `close()`.
-
+    
     [`exit`]: https://doc.rust-lang.org/std/process/fn.exit.html
     [`pidfile`]: https://linux.die.net/man/3/pidfile
     [`daemon`(3)]: https://linux.die.net/man/3/daemon
     """
-    def __new__(
-        cls, path: builtins.str | os.PathLike | pathlib.Path, mode: builtins.int
-    ) -> Pidfile:
+    def __new__(cls, path: builtins.str | os.PathLike | pathlib.Path, mode: builtins.int) -> Pidfile:
         r"""
         Creates a new PID file and locks it.
-
+        
         If the PID file cannot be locked, returns `PidfileError::AlreadyRunning` with
         a PID of the already running process, or `None` if no PID has been written to
         the PID file yet.
@@ -156,13 +137,13 @@ class Pidfile:
     def write(self) -> None:
         r"""
         Writes the current process ID to the PID file.
-
+        
         The file is truncated before writing.
         """
     def as_raw_fd(self) -> builtins.int:
         r"""
         Extracts the raw file descriptor.
-
+        
         This function is typically used to **borrow** an owned file descriptor.
         When used in this way, this method does **not** pass ownership of the
         raw file descriptor to the caller, and the file descriptor is only
@@ -178,3 +159,4 @@ class Pidfile:
 class PidfileError(builtins.Exception):
     def __repr__(self) -> builtins.str: ...
     def __str__(self) -> builtins.str: ...
+

rust/exo_rs/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "maturin"
 
 [project]
 name = "exo_rs"
-version = "0.2.18"
+version = "0.2.16"
 description = "Add your description here"
 readme = "README.md"
 authors = [
@@ -19,8 +19,8 @@ dev = ["exo_rs", "pytest>=8.4.0", "pytest-asyncio>=1.0.0"]
 
 [tool.maturin]
 #purelib = true
-python-source = "python"
-module-name = "exo_rs._core"
+#python-source = "python"
+module-name = "exo_rs"
 features = ["pyo3/extension-module", "pyo3/experimental-async"]
 
 [tool.pyo3-stub-gen]
@@ -30,6 +30,3 @@ generate-init-py = true
 log_cli = true
 log_cli_level = "INFO"
 asyncio_mode = "auto"
-
-[tool.uv]
-cache-keys = [{ file = "src/**/*.rs" }]

rust/exo_rs/python/.gitignore

@@ -1,4 +0,0 @@
-*.so
-*.pyd
-*.dll
-*.dylib

rust/exo_rs/python/exo_rs/__init__.py

@@ -1,4 +0,0 @@
-# This file is automatically generated by pyo3_stub_gen
-# ruff: noqa: F401
-
-__all__ = []

rust/exo_rs/python/exo_rs/ident/__init__.py

@@ -1,8 +0,0 @@
-# This file is automatically generated by pyo3_stub_gen
-# ruff: noqa: F401
-
-from exo_rs._core import Keypair
-
-__all__ = [
-    "Keypair",
-]

rust/exo_rs/python/exo_rs/networking/__init__.py

@@ -1,18 +0,0 @@
-# This file is automatically generated by pyo3_stub_gen
-# ruff: noqa: F401
-
-from exo_rs._core import (
-    AllQueuesFullError,
-    FromSwarm,
-    MessageTooLargeError,
-    NetworkingHandle,
-    NoPeersSubscribedToTopicError,
-)
-
-__all__ = [
-    "AllQueuesFullError",
-    "FromSwarm",
-    "MessageTooLargeError",
-    "NetworkingHandle",
-    "NoPeersSubscribedToTopicError",
-]

rust/exo_rs/python/exo_rs/pidfile/__init__.py

@@ -1,9 +0,0 @@
-# This file is automatically generated by pyo3_stub_gen
-# ruff: noqa: F401
-
-from exo_rs._core import Pidfile, PidfileError
-
-__all__ = [
-    "Pidfile",
-    "PidfileError",
-]

rust/exo_rs/src/lib.rs

@@ -1,14 +1,21 @@
-//! Python package for EXO Rust bindings.
-
-module_doc!("exo_rs", "Python package for EXO Rust bindings.");
+//! TODO: crate documentation
+//!
+//! this is here as a placeholder documentation
+//!
+//!
 
 mod allow_threading;
-pub mod ident;
-pub mod networking;
-pub mod pidfile;
+mod ident;
+mod networking;
+mod pidfile;
 
-use pyo3::{pyclass, pymodule};
-use pyo3_stub_gen::{define_stub_info_gatherer, module_doc, reexport_module_members};
+use crate::ident::PyKeypair;
+use crate::networking::networking_submodule;
+use crate::pidfile::pidfile_submodule;
+use pyo3::prelude::PyModule;
+use pyo3::types::PyModuleMethods;
+use pyo3::{Bound, PyResult, pyclass, pymodule};
+use pyo3_stub_gen::define_stub_info_gatherer;
 
 /// Namespace for all the constants used by this crate.
 pub(crate) mod r#const {
@@ -143,43 +150,28 @@ pub(crate) mod ext {
     }
 }
 
-#[pymodule(name = "_core", gil_used = true)]
-mod py_exo_rs {
-    #[pymodule_export]
-    use super::ident::PyKeypair;
-    #[pymodule_export]
-    use super::networking::{
-        PyAllQueuesFullError, PyFromSwarm, PyMessageTooLargeError, PyNetworkingHandle,
-        PyNoPeersSubscribedToTopicError,
-    };
-    #[pymodule_export]
-    use super::pidfile::{PyPidfile, PyPidfileError};
-    use pyo3::{
-        PyResult,
-        prelude::{Bound, PyModule},
-    };
+/// A Python module implemented in Rust. The name of this function must match
+/// the `lib.name` setting in the `Cargo.toml`, else Python will not be able to
+/// import the module.
+#[pymodule(name = "exo_rs", gil_used = true)]
+fn main_module(m: &Bound<'_, PyModule>) -> PyResult<()> {
+    // install logger
+    pyo3_log::init();
+    let mut builder = tokio::runtime::Builder::new_multi_thread();
+    builder.enable_all();
+    pyo3_async_runtimes::tokio::init(builder);
 
-    #[pymodule_init]
-    fn init(_m: &Bound<'_, PyModule>) -> PyResult<()> {
-        // install logger (TODO: change to tracing)
-        pyo3_log::init();
+    // TODO: for now this is all NOT a submodule, but figure out how to make the submodule system
+    //       work with maturin, where the types generate correctly, in the right folder, without
+    //       too many importing issues...
+    m.add_class::<PyKeypair>()?;
+    networking_submodule(m)?;
+    pidfile_submodule(m)?;
 
-        // create pyo3_async_runtimes
-        let mut builder = tokio::runtime::Builder::new_multi_thread();
-        builder.enable_all();
-        pyo3_async_runtimes::tokio::init(builder);
+    // top-level constructs
+    // TODO: ...
 
-        Ok(())
-    }
+    Ok(())
 }
 
-// make sure these re-exports match the #[pymodule_export] from above
-reexport_module_members!("exo_rs.ident" from "exo_rs._core";
-    "Keypair");
-reexport_module_members!("exo_rs.networking" from "exo_rs._core";
-    "AllQueuesFullError", "FromSwarm", "MessageTooLargeError", "NetworkingHandle",
-    "NoPeersSubscribedToTopicError");
-reexport_module_members!("exo_rs.pidfile" from "exo_rs._core";
-    "Pidfile", "PidfileError");
-
 define_stub_info_gatherer!(stub_info);

rust/exo_rs/src/networking.rs

@@ -5,20 +5,20 @@ use crate::r#const::MPSC_CHANNEL_SIZE;
 use crate::ext::{ByteArrayExt as _, FutureExt, PyErrExt as _};
 use crate::ext::{ResultExt as _, TokioMpscSenderExt as _};
 use crate::ident::PyKeypair;
+use crate::networking::exception::{
+    PyAllQueuesFullError, PyMessageTooLargeError, PyNoPeersSubscribedToTopicError,
+};
 use crate::pyclass;
 use futures_lite::{Stream, StreamExt as _};
 use libp2p::gossipsub::PublishError;
 use networking::swarm::{FromSwarm, ToSwarm, create_swarm};
 use pyo3::exceptions::PyRuntimeError;
+use pyo3::prelude::{PyModule, PyModuleMethods as _};
 use pyo3::types::PyBytes;
 use pyo3::{Bound, Py, PyAny, PyErr, PyResult, Python, pymethods};
 use pyo3_stub_gen::derive::{gen_stub_pyclass, gen_stub_pyclass_complex_enum, gen_stub_pymethods};
 use tokio::sync::{Mutex, mpsc, oneshot};
 
-pub use exception::{
-    PyAllQueuesFullError, PyMessageTooLargeError, PyNoPeersSubscribedToTopicError,
-};
-
 mod exception {
     use pyo3::types::PyTuple;
     use pyo3::{exceptions::PyException, prelude::*};
@@ -129,7 +129,7 @@ mod exception {
 
 #[gen_stub_pyclass]
 #[pyclass(name = "NetworkingHandle")]
-pub struct PyNetworkingHandle {
+struct PyNetworkingHandle {
     // channels
     pub to_swarm: mpsc::Sender<ToSwarm>,
     pub swarm: Arc<Mutex<Pin<Box<dyn Stream<Item = FromSwarm> + Send>>>>,
@@ -137,7 +137,7 @@ pub struct PyNetworkingHandle {
 
 #[gen_stub_pyclass_complex_enum]
 #[pyclass(name = "FromSwarm")]
-pub enum PyFromSwarm {
+enum PyFromSwarm {
     Connection {
         peer_id: String,
         connected: bool,
@@ -148,7 +148,6 @@ pub enum PyFromSwarm {
         data: Py<PyBytes>,
     },
 }
-
 impl From<FromSwarm> for PyFromSwarm {
     fn from(value: FromSwarm) -> Self {
         match value {
@@ -297,3 +296,14 @@ impl PyNetworkingHandle {
         Ok(())
     }
 }
+
+pub fn networking_submodule(m: &Bound<'_, PyModule>) -> PyResult<()> {
+    m.add_class::<exception::PyNoPeersSubscribedToTopicError>()?;
+    m.add_class::<exception::PyAllQueuesFullError>()?;
+    m.add_class::<exception::PyMessageTooLargeError>()?;
+
+    m.add_class::<PyNetworkingHandle>()?;
+    m.add_class::<PyFromSwarm>()?;
+
+    Ok(())
+}

rust/exo_rs/src/pidfile.rs

@@ -1,5 +1,6 @@
 use pidfile_rs::{Pidfile, PidfileError};
 use pyo3::exceptions::PyException;
+use pyo3::prelude::{PyModule, PyModuleMethods};
 use pyo3::{Bound, PyErr, PyResult, Python, pyclass, pymethods};
 use pyo3_stub_gen::derive::{gen_stub_pyclass, gen_stub_pymethods};
 use std::fs;
@@ -118,3 +119,10 @@ impl PyPidfile {
         self.0 = None;
     }
 }
+
+pub fn pidfile_submodule(m: &Bound<PyModule>) -> PyResult<()> {
+    m.add_class::<PyPidfileError>()?;
+    m.add_class::<PyPidfile>()?;
+
+    Ok(())
+}

rust/exo_rs/tests/test_python.py

@@ -1,15 +1,14 @@
 import asyncio
-from contextlib import suppress
 
 import pytest
 from _pytest.capture import CaptureFixture
-from exo_rs.ident import Keypair
-from exo_rs.networking import (
-    FromSwarm,
+from exo_rs import (
+    Keypair,
     NetworkingHandle,
     NoPeersSubscribedToTopicError,
+    Pidfile,
+    FromSwarm,
 )
-from exo_rs.pidfile import Pidfile
 
 
 @pytest.mark.asyncio
@@ -17,21 +16,16 @@ async def test_sleep_on_multiple_items() -> None:
     print("PYTHON: starting handle")
     h = NetworkingHandle(Keypair.generate(), [], 0)
 
-    recv_task = asyncio.create_task(_await_recv(h))
+    rt = asyncio.create_task(_await_recv(h))
 
-    try:
-        # sleep for 4 ticks
-        for _ in range(4):
-            await asyncio.sleep(1)
+    # sleep for 4 ticks
+    for i in range(4):
+        await asyncio.sleep(1)
 
-            try:
-                await h.gossipsub_publish("topic", b"somehting or other")
-            except NoPeersSubscribedToTopicError as e:
-                print("caught it", e)
-    finally:
-        recv_task.cancel()
-        with suppress(asyncio.CancelledError):
-            await recv_task
+        try:
+            await h.gossipsub_publish("topic", b"somehting or other")
+        except NoPeersSubscribedToTopicError as e:
+            print("caught it", e)
 
 
 def test_pidfile(capsys: CaptureFixture[str]):
@@ -52,5 +46,4 @@ async def _await_recv(h: NetworkingHandle):
 
 
 def scoped_lock_file():
-    lock_file = Pidfile("/tmp/lock.pid", 0o0600)
-    lock_file.close()
+    a = Pidfile("/tmp/lock.pid", 0o0600)

rust/parts.nix

@@ -32,17 +32,8 @@
             isRootCargoFile =
               (baseName == "Cargo.toml" || baseName == "Cargo.lock")
               && (builtins.dirOf path == toString inputs.self);
-            isExoRsPythonSource =
-              (lib.hasInfix "/rust/exo_rs/python/" path || lib.hasSuffix "/rust/exo_rs/python" path)
-              && (
-                type == "directory"
-                || lib.hasSuffix ".py" path
-                || lib.hasSuffix ".pyi" path
-                || baseName == "py.typed"
-              );
           in
           isRootCargoFile
-          || isExoRsPythonSource
           || (inRustDir && (craneLib.filterCargoSources path type || lib.hasSuffix ".toml" path || lib.hasSuffix ".md" path));
       };
 

src/exo/main.py

@@ -10,7 +10,7 @@ from typing import Self
 import anyio
 from anyio.lowlevel import checkpoint as anyio_checkpoint
 from daemon import DaemonContext  # pyright: ignore[reportMissingTypeStubs]
-from exo_rs.pidfile import Pidfile, PidfileError
+from exo_rs import Pidfile, PidfileError
 from loguru import logger
 from pydantic import PositiveInt
 

src/exo/routing/connection_message.py

@@ -1,4 +1,4 @@
-from exo_rs.networking import FromSwarm
+from exo_rs import FromSwarm
 
 from exo.shared.types.common import NodeId
 from exo.utils.pydantic_ext import FrozenModel

src/exo/routing/router.py

@@ -12,10 +12,10 @@ from anyio import (
     move_on_after,
     sleep_forever,
 )
-from exo_rs.ident import Keypair
-from exo_rs.networking import (
+from exo_rs import (
     AllQueuesFullError,
     FromSwarm,
+    Keypair,
     MessageTooLargeError,
     NetworkingHandle,
     NoPeersSubscribedToTopicError,

src/exo/utils/tests/test_pidfile.py

@@ -8,14 +8,14 @@ import textwrap
 from pathlib import Path
 from typing import Final
 
-from exo_rs.pidfile import Pidfile
+from exo_rs import Pidfile
 
 _CHILD_ACQUIRE_PIDFILE_SCRIPT: Final = textwrap.dedent(
     """
     import sys
     from pathlib import Path
 
-    from exo_rs.pidfile import Pidfile, PidfileError
+    from exo_rs import Pidfile, PidfileError
 
     path = Path(sys.argv[1])
     try:

src/exo/worker/engines/mlx/cache.py

@@ -229,6 +229,47 @@ def has_non_kv_caches(cache: KVCacheType) -> bool:
     return any(is_non_trimmable_cache_entry(c) for c in cache)
 
 
+# Max snapshots retained per cache entry. Each CacheSnapshot pins detached GPU
+# copies of every non-trimmable (SSM/ArraysCache, RotatingKVCache) layer, so
+# retaining one per ~4096-token prefill chunk makes snapshot memory grow linearly
+# with context — the dominant residual cost when a single entry is grown to long
+# contexts on hybrid models (~56 MB/snapshot on Qwen3.5-122B, so a full 256K
+# context = 64 snapshots ≈ 3.6 GB). A sliding window of the most-recent N caps
+# this at N×per-snapshot (~0.9 GB here) while preserving the restore points
+# in-place grows actually use (they always extend from the tip).
+_MAX_RETAINED_SNAPSHOTS = 16
+
+
+def _bounded_snapshots(snapshots: list[CacheSnapshot]) -> list[CacheSnapshot]:
+    """Deduplicate snapshots by token position and bound the retained count.
+
+    Returned list is sorted ascending by ``token_count``.
+    """
+    # Deduplicate by position, keeping the most-recently-appended snapshot per
+    # position. Repeated in-place grows re-snapshot positions the kept old
+    # snapshots already cover, which would otherwise grow `_snapshots`
+    # unbounded even at constant context.
+    # TODO: keying on token_count alone is safe only while a position uniquely
+    # identifies the prefix within an entry (grows are strict prefix-extensions).
+    # If edit-and-regenerate, sliding-window/prefix trimming, cross-entry
+    # snapshot sharing, per-request adapter/LoRA swap, or branchy decoding
+    # (beam/parallel/speculative) is added, enrich the key to
+    # (token_count, prefix_hash[, media/adapter id]) — else a stale snapshot
+    # could be restored for a different prefix (silent wrong output).
+    by_position: dict[int, CacheSnapshot] = {}
+    for snapshot in snapshots:
+        by_position[snapshot.token_count] = snapshot
+    deduped = [by_position[pos] for pos in sorted(by_position)]
+
+    # Sliding window: keep only the most-recent N positions. In-place grows
+    # always extend from the tip, so the newest snapshots are the ones future
+    # grows restore from — dropping the oldest is never incorrect: a later hit on
+    # a prefix older than the window finds no snapshot <= target, so get_kv_cache
+    # returns a fresh cache (matched_index=None) and the request takes a full cold
+    # prefill — correct, just slower than a partial-hit reuse for that one request.
+    return deduped[-_MAX_RETAINED_SNAPSHOTS:]
+
+
 class KVPrefixCache:
     def __init__(self, group: mx.distributed.Group | None):
         self.prompts: list[mx.array] = []  # mx array of tokens (ints)
@@ -261,7 +302,9 @@ class KVPrefixCache:
         self._evict_if_needed()
         self.prompts.append(prompt_tokens)
         self.caches.append(deepcopy(cache))
-        self._snapshots.append(ssm_snapshots)
+        self._snapshots.append(
+            _bounded_snapshots(ssm_snapshots) if ssm_snapshots else None
+        )
         self._media_regions.append(media_regions or [])
         self.prefill_tps.append(prefill_tps)
         self._access_counter += 1
@@ -288,7 +331,7 @@ class KVPrefixCache:
 
         self.prompts[index] = prompt_tokens
         self.caches[index] = deepcopy(cache)
-        self._snapshots[index] = merged or None
+        self._snapshots[index] = _bounded_snapshots(merged) or None
         self._media_regions[index] = media_regions or []
         self.prefill_tps[index] = prefill_tps
         self._access_counter += 1

src/exo/worker/tests/unittests/test_mlx/test_kv_prefix_cache.py

@@ -11,6 +11,7 @@ from mlx_lm.sample_utils import make_sampler
 from exo.shared.types.common import ModelId
 from exo.shared.types.text_generation import InputMessage, TextGenerationTaskParams
 from exo.worker.engines.mlx.cache import (
+    CacheSnapshot,
     KVPrefixCache,
     cache_length,
     encode_prompt,
@@ -77,6 +78,74 @@ class TestGetPrefixLength:
         assert get_prefix_length(a, b) == 0
 
 
+class TestSnapshotAccumulation:
+    """Locks in the fix for the actual per-grow Metal leak on hybrid (SSM)
+    models: `update_kv_cache` must not let `_snapshots` grow without bound when
+    the same entry is grown in place many times."""
+
+    def test_repeated_update_does_not_accumulate_snapshots(self):
+        with patch(
+            "exo.worker.engines.mlx.cache.get_memory_used_percentage",
+            return_value=0.0,
+        ):
+            kv_prefix_cache = KVPrefixCache(None)
+            initial = [
+                CacheSnapshot(states=[None], token_count=4096),
+                CacheSnapshot(states=[None], token_count=8192),
+            ]
+            kv_prefix_cache.add_kv_cache(
+                mx.arange(10000), [KVCache()], ssm_snapshots=initial
+            )
+
+            # Each in-place grow re-prefills from restore_pos and produces a
+            # fresh snapshot at a position the retained old snapshots already
+            # cover. Pre-fix this appended one snapshot per grow forever.
+            for _ in range(50):
+                fresh = [CacheSnapshot(states=[None], token_count=8192)]
+                kv_prefix_cache.update_kv_cache(
+                    0, mx.arange(10000), [KVCache()], fresh, restore_pos=8192
+                )
+
+            stored = kv_prefix_cache._snapshots[0]
+            assert stored is not None
+            # Bounded by the number of distinct snapshot positions (here 2),
+            # not by the 50 grows.
+            assert len(stored) == 2
+            assert sorted(s.token_count for s in stored) == [4096, 8192]
+            # The kept 8192 snapshot must be the most recently supplied one.
+            assert stored[1] is fresh[0]
+
+    def test_extension_caps_snapshots_to_sliding_window(self):
+        """Extending a single entry to a long context (one snapshot per ~4096
+        tokens) must cap retained snapshots to a sliding window of the most-recent
+        N, not keep all of them — that linear-in-context retention was the
+        residual OOM cause."""
+        from exo.worker.engines.mlx.cache import _MAX_RETAINED_SNAPSHOTS
+
+        with patch(
+            "exo.worker.engines.mlx.cache.get_memory_used_percentage",
+            return_value=0.0,
+        ):
+            kv_prefix_cache = KVPrefixCache(None)
+            # 64 distinct positions = a 262144-token context at 4096/chunk.
+            num_positions = 64
+            snaps = [
+                CacheSnapshot(states=[None], token_count=4096 * (i + 1))
+                for i in range(num_positions)
+            ]
+            kv_prefix_cache.add_kv_cache(
+                mx.arange(10), [KVCache()], ssm_snapshots=snaps
+            )
+
+            stored = kv_prefix_cache._snapshots[0]
+            assert stored is not None
+            # Capped at the window; the most-recent N positions are retained
+            # (in-place grows extend from the tip, so these are what get used).
+            assert len(stored) == _MAX_RETAINED_SNAPSHOTS
+            assert stored == snaps[-_MAX_RETAINED_SNAPSHOTS:]
+            assert stored[-1] is snaps[-1]  # tip always kept
+
+
 class TestKVPrefix:
     @pytest.fixture
     def mock_tokenizer(self):

uv.lock

@@ -623,7 +623,7 @@ requires-dist = [
 
 [[package]]
 name = "exo-rs"
-version = "0.2.18"
+version = "0.2.16"
 source = { editable = "rust/exo_rs" }
 
 [package.dev-dependencies]