aris-server 0.1.0__py3-none-any.whl

aris_server/__init__.py

@@ -0,0 +1,3 @@
+"""Aris Server: open-weight models behind the Anthropic Messages API."""
+
+__version__ = "0.1.0"

aris_server/budget.py

@@ -0,0 +1,82 @@
+"""Think budget: close a model's reasoning on a complete thought once it has used its allowance."""
+
+import mlx.core as mx
+
+
+class ThinkBudget:
+    """Logits processor. Counts generated tokens while thinking. Past `budget`, outside a code fence, a paragraph
+    break becomes a single newline and the end-of-thinking token is forced after the next newline; `overrun` tokens
+    later it closes regardless (an open fence, or a line that never ends: models can loop in place). Inert once
+    thinking has ended.
+
+    Ending on a single newline matters: chat templates re-render past reasoning as `text\\n</think>`, so a turn that
+    closed after a paragraph break (`text\\n\\n</think>`) no longer matches its own prompt cache, and the next turn
+    re-reads everything from that point."""
+
+    def __init__(self, budget: int, end_id: int, newline_id: int, paragraph_ids: frozenset[int],
+                 fence_counts: dict[int, int], overrun: int = 1024):
+        self.budget, self.end_id, self.newline_id, self.overrun = budget, end_id, newline_id, overrun
+        self.paragraph_ids, self.fence_counts = paragraph_ids, fence_counts
+        self.n = self.fences = 0
+        self.done = self.forced = False
+        self._paragraphs = None
+
+    def _only(self, token_id, logits):
+        mask = mx.arange(logits.shape[-1]) == token_id
+        return mx.broadcast_to(mx.where(mask, mx.array(0.0, logits.dtype), mx.array(-mx.inf, logits.dtype)),
+                               logits.shape)
+
+    def __call__(self, tokens, logits):
+        if self.done or tokens is None or len(tokens) == 0:
+            return logits
+        last = int(tokens[-1].item())
+        if last == self.end_id:
+            self.done = True
+            return logits
+        self.n += 1
+        self.fences += self.fence_counts.get(last, 0)
+        if self.n < self.budget:
+            return logits
+        hard = self.n >= self.budget + self.overrun
+        if self.fences % 2 and not hard:
+            return logits
+        if last == self.newline_id:
+            self.done = self.forced = True
+            return self._only(self.end_id, logits)
+        if hard:
+            return self._only(self.newline_id, logits)
+        # A paragraph break the model would take becomes a single newline, which then closes thinking.
+        if self._paragraphs is None:
+            ids = sorted(self.paragraph_ids)
+            self._paragraphs = mx.array(ids) if ids else None
+        if self._paragraphs is None:
+            return logits
+        flat = logits.reshape(-1, logits.shape[-1])
+        merged = mx.logaddexp(flat[:, self.newline_id], mx.logsumexp(flat[:, self._paragraphs], axis=-1))
+        flat[:, self._paragraphs] = -mx.inf
+        flat[:, self.newline_id] = merged
+        return flat.reshape(logits.shape)
+
+
+class BudgetTokens:
+    """Per-tokenizer lookups the processor needs, computed once per model."""
+
+    def __init__(self, tokenizer):
+        end = tokenizer.think_end_tokens if hasattr(tokenizer, "think_end_tokens") else None
+        self.end_id = end[0] if end and len(end) == 1 else None
+        newline = tokenizer.encode("\n", add_special_tokens=False) if hasattr(tokenizer, "encode") else []
+        self.newline_id = newline[0] if len(newline) == 1 else None
+        self.paragraph_ids, self.fence_counts = set(), {}
+        decode = tokenizer.decode
+        for i in range(len(tokenizer.vocab) if hasattr(tokenizer, "vocab") else tokenizer.vocab_size):
+            text = decode([i])
+            if text.endswith("\n\n"):
+                self.paragraph_ids.add(i)
+            if "```" in text:
+                self.fence_counts[i] = text.count("```")
+        self.paragraph_ids = frozenset(self.paragraph_ids)
+
+    def make(self, budget: int) -> ThinkBudget | None:
+        if budget <= 0 or self.end_id is None or self.newline_id is None:
+            return None
+        return ThinkBudget(budget, self.end_id, self.newline_id, self.paragraph_ids, self.fence_counts)

aris_server/cli.py

@@ -0,0 +1,141 @@
+"""aris-server command line."""
+
+import argparse
+import dataclasses
+import logging
+import signal
+import sys
+from pathlib import Path
+
+from . import __version__
+from .config import builtin_models, home, load_models, resolve
+
+
+def _wired_limit():
+    from mlx_lm.utils import maybe_set_recommended_wired_limit
+
+    maybe_set_recommended_wired_limit()
+
+
+def _quiet_dependencies(level):
+    """mlx-lm logs through the root logger and httpx logs every Hugging Face request; keep only their warnings."""
+    for name in ("httpx", "huggingface_hub"):
+        logging.getLogger(name).setLevel(level)
+
+    def keep(record):
+        noisy = "mlx_lm" in record.pathname or record.name.startswith(("httpx", "huggingface_hub"))
+        return record.levelno >= level or not noisy
+
+    for handler in logging.getLogger().handlers:
+        handler.addFilter(keep)
+
+
+def main(argv=None):
+    ap = argparse.ArgumentParser(prog="aris-server", description="Open-weight models behind the Anthropic Messages API.")
+    ap.add_argument("--version", action="version", version=f"aris-server {__version__}")
+    sub = ap.add_subparsers(dest="cmd", required=True)
+    s = sub.add_parser("serve", help="load a model and serve it")
+    s.add_argument("--model", help="configured model name, local MLX folder, or Hugging Face repo id")
+    s.add_argument("--adapter", help="LoRA adapter folder to apply on top of the model")
+    s.add_argument("--kv-bits", type=int, choices=[4, 8, 16],
+                   help="bits per KV cache value (16 = full precision); overrides the model's kv_bits")
+    s.add_argument("--host", default="127.0.0.1")
+    s.add_argument("--port", type=int, default=8484)
+    s.add_argument("--api-key", help="require this key in x-api-key")
+    s.add_argument("--concurrency", type=int, default=4, help="sequences decoded at once")
+    s.add_argument("--prompt-cache-gb", type=float, help="memory for saved prompt caches (default: fit to the GPU)")
+    s.add_argument("--disk-cache-gb", type=float, default=20.0,
+                   help="disk for saved prompt caches, all models together (0 turns it off)")
+    s.add_argument("--reserve-gb", type=float, default=6.0, help="GPU memory kept free for active sequences")
+    s.add_argument("--log-level", default="INFO")
+    sub.add_parser("models", help="list configured and built-in models")
+    p = sub.add_parser("pull", help="download the weights a model would load on this Mac, ahead of serving it")
+    p.add_argument("model", help="configured or built-in model name, or Hugging Face repo id")
+    p.add_argument("--json", action="store_true", help="report each step as a JSON line on stdout, for apps")
+    svc = sub.add_parser("service", help="run as a login service that restarts itself (macOS launchd)")
+    svc.add_argument("action", choices=["install", "uninstall", "status"])
+    svc.add_argument("serve_args", nargs=argparse.REMAINDER, help="arguments for serve, after --")
+    a = ap.parse_args(argv)
+
+    if a.cmd == "service":
+        from . import service
+
+        if a.action == "install":
+            return service.install([x for x in a.serve_args if x != "--"])
+        return service.uninstall() if a.action == "uninstall" else service.status()
+
+    models = load_models()
+    if a.cmd == "models":
+        if not models:
+            print(f"no models configured in {home() / 'models.toml'}")
+        for name, spec in models.items():
+            print(f"{name:24} {' or '.join(spec.choices) or spec.path}")
+        for name, spec in builtin_models().items():
+            if name not in models:
+                print(f"{name:24} {' or '.join(spec.choices) or spec.path}  (built in)")
+        return 0
+    if a.cmd == "pull":
+        logging.basicConfig(level=logging.WARNING, format="%(levelname)s %(message)s")  # stdout is the report
+        _quiet_dependencies(logging.WARNING)
+        import mlx.core as mx
+
+        from .pull import pull
+
+        return pull(resolve(a.model, models), mx.device_info()["max_recommended_working_set_size"], as_json=a.json)
+
+    logging.basicConfig(level=getattr(logging, a.log_level.upper(), logging.INFO),
+                        format="%(asctime)s %(levelname)s %(message)s")
+    _quiet_dependencies(logging.DEBUG if a.log_level.upper() == "DEBUG" else logging.WARNING)
+    name = a.model or next(iter(models), None)
+    if not name:
+        print("no model: pass --model or configure one in models.toml", file=sys.stderr)
+        return 2
+    if a.host not in ("127.0.0.1", "localhost", "::1") and not a.api_key:
+        logging.warning("serving on %s without --api-key: anyone who can reach this port can use the model", a.host)
+
+    import mlx.core as mx
+
+    from .engine import Engine, NotEnoughMemory, check_memory, choose_weights, weights_and_config
+    from .http import Server
+
+    _wired_limit()
+    spec = resolve(name, models)
+    if a.adapter:
+        spec = dataclasses.replace(spec, adapter=str(Path(a.adapter).expanduser().resolve()))
+    if a.kv_bits:
+        spec = dataclasses.replace(spec, kv_bits=None if a.kv_bits == 16 else a.kv_bits)
+    if spec.adapter and not (Path(spec.adapter) / "adapter_config.json").is_file():
+        print(f"no adapter_config.json in {spec.adapter}", file=sys.stderr)
+        return 2
+    try:
+        # Before loading, and for a Hugging Face repo before downloading: a model this Mac can't hold would swap it to
+        # a standstill, and under the service, again after every restart.
+        limit = mx.device_info()["max_recommended_working_set_size"]
+        spec = choose_weights(spec, limit)
+        if (sized := weights_and_config(spec.path, spec.adapter)) is not None:
+            check_memory(name, limit, sized[0])
+        engine = Engine(spec, decode_concurrency=a.concurrency, reserve_gb=a.reserve_gb,
+                        disk_cache_gb=a.disk_cache_gb,
+                        prompt_cache_bytes=int(a.prompt_cache_gb * (1 << 30)) if a.prompt_cache_gb else None)
+    except NotEnoughMemory as e:
+        logging.error("%s", e)
+        return 2
+    server = Server((a.host, a.port), engine, models, a.api_key)
+    logging.info("Aris Server %s serving %s on http://%s:%d", __version__, name, a.host, a.port)
+
+    def stop(*_):  # a service manager stops with SIGTERM: finish the current step (and any cache write) first
+        raise KeyboardInterrupt
+
+    signal.signal(signal.SIGTERM, stop)
+    try:
+        server.serve_forever()
+    except KeyboardInterrupt:
+        pass
+    finally:
+        server.server_close()
+        engine.generator.stop_and_join()  # the generation thread writes the prompt caches as it exits
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())

aris_server/config.py

@@ -0,0 +1,119 @@
+"""Model settings: which weights a model name means and how effort and thinking map onto it."""
+
+import os
+import tomllib
+from dataclasses import dataclass, field
+from pathlib import Path
+
+EFFORTS = ("low", "medium", "high", "xhigh", "max")
+KV_BITS_CHOICES = (4, 8)
+DEFAULT_THINK_BUDGET = {"low": 1024, "medium": 4096, "high": 8192, "xhigh": 16384, "max": 0}
+
+
+def home() -> Path:
+    return Path(os.environ.get("ARIS_SERVER_HOME", Path.home() / ".aris-server"))
+
+
+# A model given several weights (`path` as a list, best first) loads the first that leaves this much context on the
+# Mac it runs on: a coding agent's conversation outgrows less, and compacting it costs a long re-read each time.
+DEFAULT_MIN_CONTEXT = 65536
+
+
+@dataclass
+class ModelSpec:
+    name: str
+    path: str
+    context: int | None = None
+    # Weights to choose from by memory, best first (`path` is the first): for example the same model at 8 and 4 bits.
+    choices: tuple[str, ...] = ()
+    min_context: int = DEFAULT_MIN_CONTEXT
+    adapter: str | None = None  # a LoRA adapter folder (adapter_config.json + adapters.safetensors) on top of path
+    # Bits per value of the KV cache (8 or 4), None for full precision: 8 bits holds about twice the context in the
+    # same memory.
+    kv_bits: int | None = None
+    # Template variable that carries effort (e.g. "reasoning_effort"), and how request efforts map to its values.
+    effort_kwarg: str | None = None
+    effort_map: dict[str, str] = field(default_factory=dict)
+    # Thinking tokens allowed per effort before thinking is closed at the next paragraph break; 0 = unlimited.
+    think_budget: dict[str, int] = field(default_factory=lambda: dict(DEFAULT_THINK_BUDGET))
+    template_kwargs: dict = field(default_factory=dict)
+    # Share of the think budget for turns that continue the model's own tool loop: the plan exists by then, and
+    # overthinking models re-plan from scratch on every tool result.
+    followup_budget: float = 0.5
+    # Sampling when the request doesn't say (Claude Code sends none with thinking on). Qwen-family thinking defaults.
+    temperature: float = 0.6
+    top_p: float = 0.95
+    top_k: int = 20
+
+    def budget_for(self, effort: str) -> int:
+        return int(self.think_budget.get(effort, DEFAULT_THINK_BUDGET.get(effort, 0)))
+
+
+# Models the server knows by name with no models.toml entry; an entry there with the same name replaces one. Their
+# weights are on the Hugging Face Hub, best first: the server picks what each Mac holds before downloading anything.
+BUILTIN_MODELS = {
+    # Aris 3.8-27B: Arete-27B (Nous Research, on Qwen3.8-27B) trained on its own passing agent runs. The settings are
+    # the ones its evals ran with: Claude Code's effort onto Arete's trained effort levels, and a think budget each.
+    "aris-3.8-27b": {
+        "path": ["YourAris/Aris-3.8-27B-MLX-8bit", "YourAris/Aris-3.8-27B-MLX-4bit"],
+        "kv_bits": 8,
+        "effort_kwarg": "reasoning_effort",
+        "effort_map": {"low": "low", "medium": "medium", "high": "medium", "xhigh": "xhigh", "max": "xhigh"},
+        "think_budget": {"low": 1024, "medium": 3072, "high": 4096, "xhigh": 12288, "max": 0},
+    },
+}
+
+
+def load_models(path: Path | None = None) -> dict[str, ModelSpec]:
+    """The models configured in models.toml (not the built-in ones: see builtin_models)."""
+    path = path or home() / "models.toml"
+    if not path.is_file():
+        return {}
+    data = tomllib.loads(path.read_text())
+    return {name: _spec(name, raw, f"model {name!r} in {path}") for name, raw in (data.get("models") or {}).items()}
+
+
+def builtin_models() -> dict[str, ModelSpec]:
+    return {name: _spec(name, raw, f"built-in model {name!r}") for name, raw in BUILTIN_MODELS.items()}
+
+
+def _spec(name: str, raw: dict, where: str) -> ModelSpec:
+    if "path" not in raw:
+        raise ValueError(f"{where} has no path")
+    paths = raw["path"] if isinstance(raw["path"], list) else [raw["path"]]
+    if not paths or not all(isinstance(p, str) and p for p in paths):
+        raise ValueError(f"{where}: path must be a folder or repo id, or a list of them")
+    paths = [_weights_path(p) for p in paths]
+    budget = dict(DEFAULT_THINK_BUDGET)
+    budget.update(raw.get("think_budget") or {})
+    if raw.get("kv_bits") not in (None, *KV_BITS_CHOICES):
+        raise ValueError(f"{where}: kv_bits must be one of {KV_BITS_CHOICES}")
+    return ModelSpec(
+        name=name,
+        path=paths[0],
+        choices=tuple(paths) if len(paths) > 1 else (),
+        min_context=int(raw.get("min_context", DEFAULT_MIN_CONTEXT)),
+        context=raw.get("context"),
+        adapter=str(Path(raw["adapter"]).expanduser()) if raw.get("adapter") else None,
+        kv_bits=raw.get("kv_bits"),
+        effort_kwarg=raw.get("effort_kwarg"),
+        effort_map=dict(raw.get("effort_map") or {}),
+        think_budget=budget,
+        template_kwargs=dict(raw.get("template_kwargs") or {}),
+        **{k: raw[k] for k in ("temperature", "top_p", "top_k", "followup_budget") if k in raw},
+    )
+
+
+def _weights_path(p: str) -> str:
+    """A local folder with its home expanded, or a Hugging Face repo id as given."""
+    return str(Path(p).expanduser()) if p.startswith(("~", "/", ".")) else p
+
+
+def resolve(name: str, models: dict[str, ModelSpec]) -> ModelSpec:
+    """A configured model by name, then a built-in one, or else an unconfigured one by local path or Hugging Face repo
+    id."""
+    if name in models:
+        return models[name]
+    if name in BUILTIN_MODELS:
+        return builtin_models()[name]
+    return ModelSpec(name=name, path=_weights_path(name))

aris_server/disk_cache.py

@@ -0,0 +1,245 @@
+"""Prompt caches on disk, so a restart, a crash or a reopened thread resumes in seconds instead of re-reading a long
+conversation (minutes for a 27B model at 40k tokens). A conversation evicted from memory is also loaded back from
+here. Each finished turn's cache and the shared system-prompt snapshot
+are written as they enter the in-memory cache; a conversation's older snapshot is deleted when its next turn is saved.
+
+Each setup (weights, adapter, KV precision) has a folder of its own, and one cap covers them all. A server reads only
+its own folder, so the caches of setups no running server holds go first, then its own oldest. A running server holds
+a shared lock on its folder; another server sharing the directory clears only folders nobody holds.
+"""
+
+import fcntl
+import hashlib
+import json
+import logging
+import os
+import re
+import time
+from pathlib import Path
+
+import mlx.core as mx
+from mlx_lm.models.cache import load_prompt_cache, save_prompt_cache
+
+MIN_TOKENS = 1024  # shorter prompts are quick to re-read
+LOCK = "lock"      # shared by each server using the folder, exclusive for one clearing it
+
+
+def model_id(model_path: str, adapter: str | None = None, kv_bits: int | None = None) -> str:
+    """Caches are only valid for the exact weights that made them (the base model and any adapter on it) and for the
+    precision of the KV cache they hold."""
+    parts = []
+    for folder in (model_path, adapter):
+        if folder is None:
+            continue
+        path = Path(folder)
+        parts.append(str(path))
+        if path.is_dir():
+            parts += [f"{p.name}:{p.stat().st_size}:{int(p.stat().st_mtime)}" for p in sorted(path.glob("*.safetensors"))]
+    if kv_bits:
+        parts.append(f"kv_bits={kv_bits}")
+    return hashlib.sha1("|".join(parts).encode()).hexdigest()[:16]
+
+
+def _stat(data: Path) -> tuple[float, int]:
+    """A cache file's (mtime, size), or (0, 0) once it is gone."""
+    try:
+        s = data.stat()
+    except FileNotFoundError:
+        return 0.0, 0
+    return s.st_mtime, s.st_size
+
+
+def _same_file(path: Path, f) -> bool:
+    """Whether `path` is still the file `f` has open: not once the file, or its folder, is deleted."""
+    try:
+        return os.path.samestat(os.stat(path), os.fstat(f.fileno()))
+    except FileNotFoundError:
+        return False
+
+
+def _claim(folder: Path):
+    """The folder's lock, held exclusively, or None while a server uses the folder or another is clearing it."""
+    try:
+        lock = open(folder / LOCK, "a")
+    except OSError:
+        return None
+    try:
+        fcntl.flock(lock, fcntl.LOCK_EX | fcntl.LOCK_NB)
+    except OSError:
+        lock.close()
+        return None
+    return lock
+
+
+def _remove(data: Path):
+    """Delete a cache and its metadata. One that can't be deleted (in a folder made read-only, say) stays, with a
+    warning: the disk cache must never stop the server."""
+    try:
+        data.unlink(missing_ok=True)
+        data.with_suffix(".json").unlink(missing_ok=True)
+    except OSError as e:
+        logging.warning("could not delete prompt cache %s: %s", data.name, e)
+
+
+class DiskCache:
+    def __init__(self, root: Path, model_path: str, max_bytes: int, adapter: str | None = None,
+                 kv_bits: int | None = None):
+        self.root = root
+        self.dir = root / model_id(model_path, adapter, kv_bits)
+        self.max_bytes = max_bytes
+        self._lock = self._hold()
+        self.index: dict[Path, tuple[tuple[int, ...], str]] = {}
+        for data in self.dir.glob("*.safetensors"):
+            if not data.with_suffix(".json").exists():  # a write cut off by a crash or a kill
+                _remove(data)
+        for f in self.dir.glob("*.json"):
+            data = f.with_suffix(".safetensors")
+            if not data.exists():
+                _remove(data)
+                continue
+            try:
+                meta = json.loads(f.read_text())
+                self.index[data] = (tuple(meta["tokens"]), meta["type"])
+            except (OSError, ValueError, KeyError, TypeError) as e:  # cut off mid-write: no turn can match it
+                # Raised here, it stopped the server before it served, and under the service at every restart.
+                logging.warning("dropping prompt cache %s: its metadata can't be read: %s", data.name, e)
+                _remove(data)
+        self._trim()  # started on another setup: what the last one left over the cap goes now
+
+    def _hold(self):
+        """A shared lock on this folder for as long as the server runs."""
+        while True:
+            self.dir.mkdir(parents=True, exist_ok=True)
+            try:
+                lock = open(self.dir / LOCK, "a")
+            except FileNotFoundError:  # another server removed the folder just now
+                continue
+            fcntl.flock(lock, fcntl.LOCK_SH)
+            if _same_file(self.dir / LOCK, lock):
+                return lock
+            lock.close()  # the folder was cleared while this waited for it: make it again
+
+    def __del__(self):
+        lock = getattr(self, "_lock", None)
+        if lock is not None:
+            lock.close()
+
+    def has(self, tokens) -> bool:
+        """Whether these tokens' cache is saved. One whose file is gone (deleted by hand, say) isn't: it is saved again."""
+        key = tuple(tokens)
+        return any(t == key and data.exists() for data, (t, _) in self.index.items())
+
+    def save(self, tokens, prompt_cache, cache_type: str):
+        if len(tokens) < MIN_TOKENS or cache_type not in ("system", "assistant"):
+            return
+        if not _same_file(self.dir / LOCK, self._lock):  # the folder was deleted while serving: make it, hold it again
+            logging.warning("prompt cache folder %s was deleted while serving; making it again", self.dir.name)
+            lock = self._hold()
+            self._lock.close()
+            self._lock = lock
+        t0 = time.time()
+        key = tuple(tokens)
+        data = self.dir / f"{hashlib.sha1(json.dumps(key).encode()).hexdigest()[:20]}.safetensors"
+        save_prompt_cache(str(data), prompt_cache)
+        # The metadata file goes last: a file without it is an interrupted write and is ignored.
+        data.with_suffix(".json").write_text(json.dumps({"tokens": list(key), "type": cache_type}))
+        # Each turn saves the read prompt (what a retry of the same request needs) and then the finished answer (what
+        # the next turn continues). Of the same conversation's earlier snapshots, keep only the longest.
+        earlier = sorted((p for p, (t, kind) in self.index.items()
+                          if kind != "system" and p != data and key[:len(t)] == t), key=lambda p: len(self.index[p][0]))
+        for other in earlier[:-1]:
+            self._delete(other)
+        self.index[data] = (key, cache_type)
+        size = data.stat().st_size
+        self._trim()
+        logging.info("saved %d-token %s cache to disk (%.0f MB, %.1fs)", len(key), cache_type, size / 1e6,
+                     time.time() - t0)
+
+    def _delete(self, data: Path):
+        self.index.pop(data, None)
+        _remove(data)
+
+    def _trim(self):
+        """Keep this setup's caches, with those of setups no running server holds, under the cap: the others' go first,
+        oldest first, since this server can't read them; then this setup's oldest."""
+        claimed = [(folder, lock) for folder in self.root.iterdir()
+                   if folder != self.dir and folder.is_dir() and re.fullmatch(r"[0-9a-f]{16}", folder.name)
+                   and (lock := _claim(folder)) is not None]
+        try:
+            spare = sorted((p for folder, _ in claimed for p in folder.glob("*.safetensors")), key=_stat)
+            own = sorted(self.index, key=_stat)
+            total = sum(_stat(p)[1] for p in spare + own)
+            freed = removed = 0
+            for data in spare:
+                if total <= self.max_bytes:
+                    break
+                size = _stat(data)[1]
+                _remove(data)
+                total -= size
+                freed += size
+                removed += 1
+            for data in own:
+                if total <= self.max_bytes:
+                    break
+                total -= _stat(data)[1]
+                self._delete(data)
+            if removed:
+                logging.info("disk cache: removed %d caches (%.1f GB) that other models or settings left", removed,
+                             freed / (1 << 30))
+            for folder, _ in claimed:
+                if any(folder.glob("*.safetensors")):
+                    continue
+                try:
+                    for leftover in [*folder.glob("*.json"), folder / LOCK]:
+                        leftover.unlink(missing_ok=True)
+                    folder.rmdir()
+                except OSError:  # something else is in it, or it can't be changed: leave the folder
+                    pass
+        finally:
+            for _, lock in claimed:
+                lock.close()
+
+    def nearest(self, tokens, longer_than: int = 0) -> tuple[Path, int] | None:
+        """The saved cache covering the longest proper prefix of `tokens`, if it covers more than `longer_than`. One
+        whose file is gone from disk (deleted by hand, say) is forgotten here, before the caller lets go of what memory
+        holds for the turn."""
+        key = tuple(tokens)
+        while True:
+            best = None
+            for data, (t, _) in self.index.items():
+                if longer_than < len(t) < len(key) and key[:len(t)] == t and (best is None or len(t) > best[1]):
+                    best = (data, len(t))
+            if best is None or best[0].exists():
+                return best
+            logging.warning("prompt cache %s is gone from disk; forgetting it", best[0].name)
+            self._delete(best[0])
+
+    def load(self, data: Path):
+        """The cache saved in `data`, or None if it can't be read (damaged, or gone): then it is dropped, so no later
+        turn tries it again."""
+        try:
+            cache = load_prompt_cache(str(data))
+            # Loading is lazy and bound to the loading thread's stream; generation may run on another thread.
+            mx.eval([c.state for c in cache])
+        except Exception as e:  # noqa: BLE001  MLX raises RuntimeError for a damaged or missing file
+            logging.warning("dropping unreadable prompt cache %s: %s", data.name, e)
+            self._delete(data)
+            return None
+        return cache
+
+    def restore(self, insert, budget: int) -> int:
+        """Load the most recent caches that fit in `budget` bytes through insert(tokens, cache, type)."""
+        loaded = 0
+        newest = sorted(self.index, key=lambda p: p.stat().st_mtime, reverse=True)
+        systems = [p for p in newest if self.index[p][1] == "system"][:1]
+        for data in systems + [p for p in newest if self.index[p][1] != "system"]:
+            size = data.stat().st_size
+            if loaded + size > budget:
+                continue
+            cache = self.load(data)
+            if cache is None:  # damaged: dropped
+                continue
+            tokens, cache_type = self.index[data]
+            insert(list(tokens), cache, cache_type)
+            loaded += size
+        return loaded