fix(prefill): inject immutable retained-token cap

docs/ops/distributed-prefill-kv-network.md

@@ -417,6 +417,10 @@ hash boundary. The default segment is 256 tokens, keeping each allens step below
 the five-minute research budget at the measured Prefill rate. Worker job status
 updates `tokens_computed` after every segment; Terminal heartbeats display
 tokens, percentage, and ETA.
+The worker injects `sink + window` as an explicit immutable retained-token cap
+into `PrefillJobStore`. Snapshot capacity checks therefore remain window-aware
+even if an RPC compatibility object omits its window field; long prompts cannot
+fall back to full-length estimates after deployment or process restart.
 
 Karpathy-style optimization lives in `autoresearch/prefill/`. Humans edit
 `program.md`, the research agent edits only `candidate.py`, and immutable

inference_engine/distributed/prefill_worker.py

@@ -60,15 +60,22 @@ def estimate_final_snapshot_bytes(
     token_count: int,
     compatibility: CacheCompatibility,
     bytes_per_token: int,
+    *,
+    max_retained_tokens: int = 0,
 ) -> int:
     """Estimate retained KV, capped by sink + sliding-window capacity."""
     if token_count <= 0 or bytes_per_token <= 0:
         raise ValueError("token_count and bytes_per_token must be > 0")
     retained_tokens = int(token_count)
-    if compatibility.window_size > 0:
+    retained_limit = int(max_retained_tokens)
+    if retained_limit <= 0 and compatibility.window_size > 0:
+        retained_limit = (
+            compatibility.sink_size + compatibility.window_size
+        )
+    if retained_limit > 0:
         retained_tokens = min(
             retained_tokens,
-            compatibility.sink_size + compatibility.window_size,
+            retained_limit,
         )
     return retained_tokens * int(bytes_per_token)
 
@@ -111,6 +118,7 @@ def __init__(
         completed_ttl_s: float = 600.0,
         max_prompt_tokens: int = 131_072,
         estimated_snapshot_bytes_per_token: int = 16,
+        max_retained_tokens: int = 0,
     ) -> None:
         if min(
             max_concurrent_jobs,
@@ -120,6 +128,8 @@ def __init__(
             estimated_snapshot_bytes_per_token,
         ) <= 0:
             raise ValueError("worker limits must be > 0")
+        if max_retained_tokens < 0:
+            raise ValueError("max_retained_tokens must be >= 0")
         if (engine is None) == (engine_factory is None):
             raise ValueError("provide exactly one of engine or engine_factory")
         self.engine = engine
@@ -132,6 +142,7 @@ def __init__(
         self.estimated_snapshot_bytes_per_token = int(
             estimated_snapshot_bytes_per_token,
         )
+        self.max_retained_tokens = int(max_retained_tokens)
         self._jobs: dict[str, PrefillJob] = {}
         self._requests: dict[tuple[str, str], str] = {}
         self._lock = threading.RLock()
@@ -216,6 +227,7 @@ def submit(
                 len(job.token_ids),
                 self.cache_store.compatibility,
                 self.estimated_snapshot_bytes_per_token,
+                max_retained_tokens=self.max_retained_tokens,
             )
             if estimated_bytes > self.cache_store.max_bytes:
                 raise ValueError(
@@ -296,6 +308,7 @@ def _run(self, job_id: str) -> None:
                     len(job.token_ids),
                     self.cache_store.compatibility,
                     self.estimated_snapshot_bytes_per_token,
+                    max_retained_tokens=self.max_retained_tokens,
                 ),
             )
             engine = self._engine_for_current_thread()

scripts/start_prefill_worker_node.py

@@ -143,6 +143,7 @@ def engine_factory() -> MLXPrefillComputeEngine:
         estimated_snapshot_bytes_per_token=(
             args.estimated_snapshot_bytes_per_token
         ),
+        max_retained_tokens=args.sink + args.window,
     )
     jobs.warmup()
 

tests/inference_engine/bridge/test_prefill_worker_launchd.py

@@ -85,3 +85,10 @@ def test_two_mac_deployment_uses_allens_as_prefill_only():
     assert "<string>--window</string><string>2048</string>" in worker
     assert "<string>--prefill-tps</string><string>1</string>" in worker
     assert "scripts/start_prefill_cache_node.py" in PEER_PLIST.read_text()
+
+
+def test_worker_injects_explicit_retained_token_cap():
+    source = (
+        ROOT / "scripts" / "start_prefill_worker_node.py"
+    ).read_text()
+    assert "max_retained_tokens=args.sink + args.window" in source

tests/inference_engine/distributed/test_prefill_worker.py

@@ -55,6 +55,12 @@ def test_snapshot_estimate_caps_at_sink_plus_sliding_window():
     )
     no_window = CacheCompatibility(model_id="m", window_size=0)
     assert estimate_final_snapshot_bytes(2731, no_window, 10) == 27_310
+    assert estimate_final_snapshot_bytes(
+        2731,
+        no_window,
+        400_000,
+        max_retained_tokens=2052,
+    ) == 820_800_000
     with pytest.raises(ValueError, match="must be > 0"):
         estimate_final_snapshot_bytes(0, compatibility, 400_000)
 
@@ -246,6 +252,8 @@ def test_job_store_validation_queue_stats_and_gc():
         PrefillJobStore(None, cache)
     with pytest.raises(ValueError, match="exactly one"):
         PrefillJobStore(_Engine(), cache, engine_factory=_Engine)
+    with pytest.raises(ValueError, match="max_retained_tokens"):
+        PrefillJobStore(_Engine(), cache, max_retained_tokens=-1)
     blocking = _Engine()
     blocking.block.set()
     jobs = PrefillJobStore(blocking, cache, max_jobs=1, max_prompt_tokens=4)