feat: add meta-instance dashboard UI components

Add MetaInstanceCard component and integrate meta-instance display into
both welcome and chat sidebars. Includes store types, state management,
CRUD operations, and status derivation (active/provisioning/error).

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

.mlx_typings/mlx/nn/layers/base.pyi

@@ -200,7 +200,7 @@ class Module(dict):
     ) -> mx.MX_ARRAY_TREE:  # -> dict[Any, Any | dict[Any, Any | dict[Any, Any] | list[Any]] | dict[Any, Any] | list[Any]]:
         """Return the submodules that do not contain other modules."""
 
-    def update(self, parameters: dict, strict: bool = ...) -> Module:
+    def update(self, parameters: dict[str, Any], strict: bool = ...) -> Module:
         """Replace the parameters of this Module with the provided ones in the
         dict of dicts and lists.
 

.mlx_typings/mlx/utils.pyi

@@ -7,7 +7,10 @@ from typing import Any, Callable, Dict, List, Optional, Tuple, Union
 from mlx.core import MX_ARRAY_TREE
 
 def tree_map(
-    fn: Callable, tree: Any, *rest: Any, is_leaf: Optional[Callable] = ...
+    fn: Callable[..., Any],
+    tree: Any,
+    *rest: Any,
+    is_leaf: Callable[..., bool] | None = ...,
 ) -> Any:
     """Applies ``fn`` to the leaves of the Python tree ``tree`` and
     returns a new collection with the results.
@@ -44,11 +47,11 @@ def tree_map(
     """
 
 def tree_map_with_path(
-    fn: Callable,
+    fn: Callable[..., Any],
     tree: Any,
     *rest: Any,
-    is_leaf: Optional[Callable] = ...,
-    path: Optional[Any] = ...,
+    is_leaf: Callable[..., bool] | None = ...,
+    path: str | None = ...,
 ) -> Any:
     """Applies ``fn`` to the path and leaves of the Python tree ``tree`` and
     returns a new collection with the results.
@@ -80,9 +83,9 @@ def tree_map_with_path(
 def tree_flatten(
     tree: Any,
     prefix: str = ...,
-    is_leaf: Optional[Callable] = ...,
-    destination: Optional[Union[List[Tuple[str, Any]], Dict[str, Any]]] = ...,
-) -> Union[List[Tuple[str, Any]], Dict[str, Any]]:
+    is_leaf: Callable[..., bool] | None = ...,
+    destination: list[tuple[str, Any]] | dict[str, Any] | None = ...,
+) -> list[tuple[str, Any]] | dict[str, Any]:
     """Flattens a Python tree to a list of key, value tuples.
 
     The keys are using the dot notation to define trees of arbitrary depth and
@@ -118,7 +121,7 @@ def tree_flatten(
             the Python tree.
     """
 
-def tree_unflatten(tree: Union[List[Tuple[str, Any]], Dict[str, Any]]) -> Any:
+def tree_unflatten(tree: list[tuple[str, Any]] | dict[str, Any]) -> Any:
     """Recreate a Python tree from its flat representation.
 
     .. code-block:: python

Cargo.lock

@@ -890,7 +890,7 @@ dependencies = [
  "delegate",
  "env_logger",
  "extend",
- "futures",
+ "futures-lite",
  "libp2p",
  "log",
  "networking",
@@ -914,6 +914,12 @@ dependencies = [
  "syn 2.0.111",
 ]
 
+[[package]]
+name = "fastrand"
+version = "2.3.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "37909eebbb50d72f9059c3b6d82c0463f2ff062c9e95845c43a6c9c0355411be"
+
 [[package]]
 name = "ff"
 version = "0.13.1"
@@ -1022,7 +1028,10 @@ version = "2.6.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "f78e10609fe0e0b3f4157ffab1876319b5b0db102a2c60dc4626306dc46b44ad"
 dependencies = [
+ "fastrand",
  "futures-core",
+ "futures-io",
+ "parking",
  "pin-project-lite",
 ]
 
@@ -2753,7 +2762,7 @@ dependencies = [
  "delegate",
  "either",
  "extend",
- "futures",
+ "futures-lite",
  "futures-timer",
  "keccak-const",
  "libp2p",

Cargo.toml

@@ -29,14 +29,13 @@ util = { path = "rust/util" }
 # Macro dependecies
 extend = "1.2"
 delegate = "0.13"
-pin-project = "1"
 
 # Utility dependencies
 keccak-const = "0.2"
 
 # Async dependencies
 tokio = "1.46"
-futures = "0.3"
+futures-lite = "2.6.1"
 futures-timer = "3.0"
 
 # Data structures

README.md

@@ -72,12 +72,19 @@ There are two ways to run exo:
 
 ### Run from Source (macOS)
 
-If you have [Nix](https://nixos.org/) installed, you can skip most of the steps below and run exo directly (after accepting the Cachix cache):
+If you have [Nix](https://nixos.org/) installed, you can skip most of the steps below and run exo directly:
 
 ```bash
 nix run .#exo
 ```
 
+**Note:** To accept the Cachix binary cache (and avoid the Xcode Metal ToolChain), add to `/etc/nix/nix.conf`:
+```
+trusted-users = root    (or your username)
+experimental-features = nix-command flakes
+```
+Then restart the Nix daemon: `sudo launchctl kickstart -k system/org.nixos.nix-daemon`
+
 **Prerequisites:**
 - [Xcode](https://developer.apple.com/xcode/) (provides the Metal ToolChain required for MLX compilation)
 - [brew](https://github.com/Homebrew/brew) (for simple package management on macOS)

bench/exo_bench.py

@@ -1,29 +1,48 @@
+# type: ignore
 #!/usr/bin/env python3
-# pyright: reportAny=false, reportUnknownMemberType=false, reportUnknownVariableType=false, reportUnknownArgumentType=false
+"""Tool-calling eval for exo's OpenAI-compatible API.
+
+Tests whether models correctly:
+- Trigger tool calls when appropriate
+- Return valid JSON arguments matching function schemas
+- Handle multi-turn tool use (call -> result -> final answer)
+- Avoid calling tools when unnecessary
+
+Start exo with a model first, then run:
+    uv run python tool_call_eval.py --model <model-id>
+    uv run python tool_call_eval.py --model <model-id> --host 10.0.0.5 --port 52415
+    uv run python tool_call_eval.py --model <model-id> --repeat 3
+    uv run python tool_call_eval.py --model <model-id> --scenarios weather_simple calculator_multi_turn
+"""
+
 from __future__ import annotations
 
 import argparse
 import contextlib
-import http.client
 import itertools
 import json
-import os
 import sys
 import time
 from collections.abc import Callable
 from pathlib import Path
 from statistics import mean
 from typing import Any
-from urllib.parse import urlencode
 
+from harness import (
+    ExoClient,
+    ExoHttpError,
+    add_common_instance_args,
+    instance_id_from_instance,
+    nodes_used_in_instance,
+    resolve_model_short_id,
+    run_planning_phase,
+    settle_and_fetch_placements,
+    wait_for_instance_gone,
+    wait_for_instance_ready,
+)
 from loguru import logger
 from transformers import AutoTokenizer
 
-# Backoff constants for cluster settling retry
-_SETTLE_INITIAL_BACKOFF_S = 1.0
-_SETTLE_MAX_BACKOFF_S = 60.0
-_SETTLE_BACKOFF_MULTIPLIER = 2.0
-
 # Monkey-patch for transformers 5.x compatibility
 # Kimi's tokenization_kimi.py imports bytes_to_unicode from the old location
 # which was moved in transformers 5.0.0rc2
@@ -103,154 +122,6 @@ def load_tokenizer_for_bench(model_id: str) -> Any:
     return AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
 
 
-class ExoHttpError(RuntimeError):
-    def __init__(self, status: int, reason: str, body_preview: str):
-        super().__init__(f"HTTP {status} {reason}: {body_preview}")
-        self.status = status
-
-
-class ExoClient:
-    def __init__(self, host: str, port: int, timeout_s: float = 7200.0):
-        self.host = host
-        self.port = port
-        self.timeout_s = timeout_s
-
-    def request_json(
-        self,
-        method: str,
-        path: str,
-        params: dict[str, Any] | None = None,
-        body: dict[str, Any] | None = None,
-        headers: dict[str, str] | None = None,
-    ) -> Any:
-        if not path.startswith("/"):
-            path = "/" + path
-        if params:
-            path = path + "?" + urlencode(params)
-
-        conn = http.client.HTTPConnection(self.host, self.port, timeout=self.timeout_s)
-        try:
-            payload: bytes | None = None
-            hdrs: dict[str, str] = {"Accept": "application/json"}
-
-            if body is not None:
-                payload = json.dumps(body).encode("utf-8")
-                hdrs["Content-Type"] = "application/json"
-            if headers:
-                hdrs.update(headers)
-
-            conn.request(method.upper(), path, body=payload, headers=hdrs)
-            resp = conn.getresponse()
-            raw = resp.read()
-            text = raw.decode("utf-8", errors="replace") if raw else ""
-
-            if resp.status >= 400:
-                raise ExoHttpError(resp.status, resp.reason, text[:300])
-
-            if not text:
-                return None
-            return json.loads(text)
-        finally:
-            conn.close()
-
-    def post_bench_chat_completions(self, payload: dict[str, Any]) -> dict[str, Any]:
-        return self.request_json("POST", "/bench/chat/completions", body=payload)
-
-
-def unwrap_instance(instance: dict[str, Any]) -> dict[str, Any]:
-    if len(instance) != 1:
-        raise KeyError(f"Expected 1 key, got keys={list(instance.keys())}")
-
-    tag = next(iter(instance))
-    inner = instance[tag]
-    if not isinstance(inner, dict):
-        raise TypeError(f"payload for {tag} must be dict, got {type(inner)}")
-    return inner
-
-
-def instance_id_from_instance(instance: dict[str, Any]) -> str:
-    inner = unwrap_instance(instance)
-    return str(inner["instanceId"])
-
-
-def nodes_used_in_instance(instance: dict[str, Any]) -> int:
-    inner = unwrap_instance(instance)
-    return len(inner["shardAssignments"]["nodeToRunner"])
-
-
-def runner_ids_from_instance(instance: dict[str, Any]) -> list[str]:
-    inner = unwrap_instance(instance)
-    runner_to_shard = inner["shardAssignments"]["runnerToShard"]
-    return list(runner_to_shard.keys())
-
-
-def runner_ready(runner: dict[str, Any]) -> bool:
-    return "RunnerReady" in runner
-
-
-def runner_failed(runner: dict[str, Any]) -> bool:
-    return "RunnerFailed" in runner
-
-
-def get_runner_failed_message(runner: dict[str, Any]) -> str | None:
-    if "RunnerFailed" in runner:
-        return runner["RunnerFailed"].get("errorMessage")
-    return None
-
-
-def wait_for_instance_ready(
-    client: ExoClient, instance_id: str, timeout: float = 24000.0
-) -> None:
-    start_time = time.time()
-    instance_existed = False
-    while time.time() - start_time < timeout:
-        state = client.request_json("GET", "/state")
-        instances = state.get("instances", {})
-
-        if instance_id not in instances:
-            if instance_existed:
-                # Instance was deleted after being created - likely due to runner failure
-                raise RuntimeError(
-                    f"Instance {instance_id} was deleted (runner may have failed)"
-                )
-            time.sleep(0.1)
-            continue
-
-        instance_existed = True
-        instance = instances[instance_id]
-        runner_ids = runner_ids_from_instance(instance)
-        runners = state.get("runners", {})
-
-        # Check for failed runners first
-        for rid in runner_ids:
-            runner = runners.get(rid, {})
-            if runner_failed(runner):
-                error_msg = get_runner_failed_message(runner) or "Unknown error"
-                raise RuntimeError(f"Runner {rid} failed: {error_msg}")
-
-        if all(runner_ready(runners.get(rid, {})) for rid in runner_ids):
-            return
-
-        time.sleep(0.1)
-
-    raise TimeoutError(f"Instance {instance_id} did not become ready within {timeout=}")
-
-
-def wait_for_instance_gone(
-    client: ExoClient, instance_id: str, timeout: float = 3.0
-) -> None:
-    start_time = time.time()
-    while time.time() - start_time < timeout:
-        try:
-            client.request_json("GET", f"/instance/{instance_id}")
-            time.sleep(0.4)
-        except ExoHttpError as e:
-            if e.status == 404:
-                return
-
-    raise TimeoutError(f"Instance {instance_id} did not get deleted within {timeout=}")
-
-
 def format_peak_memory(b: float) -> str:
     for unit in ["B", "KB", "MB", "GB", "TB"]:
         if b < 1024.0:
@@ -269,184 +140,6 @@ def parse_int_list(values: list[str]) -> list[int]:
     return items
 
 
-def resolve_model_short_id(client: ExoClient, model_arg: str) -> tuple[str, str]:
-    models = client.request_json("GET", "/models") or {}
-    data = models.get("data") or []
-
-    for m in data:
-        if m.get("name").lower() == model_arg.lower():
-            short_id = str(m["name"])
-            full_id = str(m.get("hugging_face_id") or m["name"])
-            return short_id, full_id
-
-    for m in data:
-        if m.get("hugging_face_id") == model_arg:
-            short_id = str(m["name"])
-            full_id = str(m["hugging_face_id"])
-            return short_id, full_id
-
-    raise ValueError(f"Model not found in /models: {model_arg}")
-
-
-def run_planning_phase(
-    client: ExoClient,
-    full_model_id: str,
-    preview: dict[str, Any],
-    danger_delete: bool,
-    timeout: float,
-    settle_deadline: float | None,
-) -> None:
-    """Check disk space and ensure model is downloaded before benchmarking."""
-    # Get model size from /models
-    models = client.request_json("GET", "/models") or {}
-    model_bytes = 0
-    for m in models.get("data", []):
-        if m.get("hugging_face_id") == full_model_id:
-            model_bytes = m.get("storage_size_megabytes", 0) * 1024 * 1024
-            break
-
-    if not model_bytes:
-        logger.warning(
-            f"Could not determine size for {full_model_id}, skipping disk check"
-        )
-        return
-
-    # Get nodes from preview
-    inner = unwrap_instance(preview["instance"])
-    node_ids = list(inner["shardAssignments"]["nodeToRunner"].keys())
-    runner_to_shard = inner["shardAssignments"]["runnerToShard"]
-
-    state = client.request_json("GET", "/state")
-    downloads = state.get("downloads", {})
-    node_disk = state.get("nodeDisk", {})
-
-    for node_id in node_ids:
-        node_downloads = downloads.get(node_id, [])
-
-        # Check if model already downloaded on this node
-        already_downloaded = any(
-            "DownloadCompleted" in p
-            and unwrap_instance(p["DownloadCompleted"]["shardMetadata"])["modelCard"][
-                "modelId"
-            ]
-            == full_model_id
-            for p in node_downloads
-        )
-        if already_downloaded:
-            continue
-
-        # Wait for disk info if settle_deadline is set
-        disk_info = node_disk.get(node_id, {})
-        backoff = _SETTLE_INITIAL_BACKOFF_S
-        while not disk_info and settle_deadline and time.monotonic() < settle_deadline:
-            remaining = settle_deadline - time.monotonic()
-            logger.info(
-                f"Waiting for disk info on {node_id} ({remaining:.0f}s remaining)..."
-            )
-            time.sleep(min(backoff, remaining))
-            backoff = min(backoff * _SETTLE_BACKOFF_MULTIPLIER, _SETTLE_MAX_BACKOFF_S)
-            state = client.request_json("GET", "/state")
-            node_disk = state.get("nodeDisk", {})
-            disk_info = node_disk.get(node_id, {})
-
-        if not disk_info:
-            logger.warning(f"No disk info for {node_id}, skipping space check")
-            continue
-
-        avail = disk_info.get("available", {}).get("inBytes", 0)
-        if avail >= model_bytes:
-            continue
-
-        if not danger_delete:
-            raise RuntimeError(
-                f"Insufficient disk on {node_id}: need {model_bytes // (1024**3)}GB, "
-                f"have {avail // (1024**3)}GB. Use --danger-delete-downloads to free space."
-            )
-
-        # Delete from smallest to largest
-        completed = [
-            (
-                unwrap_instance(p["DownloadCompleted"]["shardMetadata"])["modelCard"][
-                    "modelId"
-                ],
-                p["DownloadCompleted"]["totalBytes"]["inBytes"],
-            )
-            for p in node_downloads
-            if "DownloadCompleted" in p
-        ]
-        for del_model, size in sorted(completed, key=lambda x: x[1]):
-            logger.info(f"Deleting {del_model} from {node_id} ({size // (1024**2)}MB)")
-            client.request_json("DELETE", f"/download/{node_id}/{del_model}")
-            avail += size
-            if avail >= model_bytes:
-                break
-
-        if avail < model_bytes:
-            raise RuntimeError(f"Could not free enough space on {node_id}")
-
-    # Start downloads (idempotent)
-    for node_id in node_ids:
-        runner_id = inner["shardAssignments"]["nodeToRunner"][node_id]
-        shard = runner_to_shard[runner_id]
-        client.request_json(
-            "POST",
-            "/download/start",
-            body={
-                "targetNodeId": node_id,
-                "shardMetadata": shard,
-            },
-        )
-        logger.info(f"Started download on {node_id}")
-
-    # Wait for downloads
-    start = time.time()
-    while time.time() - start < timeout:
-        state = client.request_json("GET", "/state")
-        downloads = state.get("downloads", {})
-        all_done = True
-        for node_id in node_ids:
-            done = any(
-                "DownloadCompleted" in p
-                and unwrap_instance(p["DownloadCompleted"]["shardMetadata"])[
-                    "modelCard"
-                ]["modelId"]
-                == full_model_id
-                for p in downloads.get(node_id, [])
-            )
-            failed = [
-                p["DownloadFailed"]["errorMessage"]
-                for p in downloads.get(node_id, [])
-                if "DownloadFailed" in p
-                and unwrap_instance(p["DownloadFailed"]["shardMetadata"])["modelCard"][
-                    "modelId"
-                ]
-                == full_model_id
-            ]
-            if failed:
-                raise RuntimeError(f"Download failed on {node_id}: {failed[0]}")
-            if not done:
-                all_done = False
-        if all_done:
-            return
-        time.sleep(1)
-
-    raise TimeoutError("Downloads did not complete in time")
-
-
-def placement_filter(instance_meta: str, wanted: str) -> bool:
-    s = (instance_meta or "").lower()
-    if wanted == "both":
-        return ("ring" in s) or ("jaccl" in s)
-    return wanted in s
-
-
-def sharding_filter(sharding: str, wanted: str) -> bool:
-    s = (sharding or "").lower()
-    if wanted == "both":
-        return ("pipeline" in s) or ("tensor" in s)
-    return wanted in s
-
-
 def run_one_completion(
     client: ExoClient, model_id: str, pp_hint: int, tg: int, prompt_sizer: PromptSizer
 ) -> tuple[dict[str, Any], int]:
@@ -538,76 +231,12 @@ class PromptSizer:
         return content, tok
 
 
-def fetch_and_filter_placements(
-    client: ExoClient, full_model_id: str, args: argparse.Namespace
-) -> list[dict[str, Any]]:
-    previews_resp = client.request_json(
-        "GET", "/instance/previews", params={"model_id": full_model_id}
-    )
-    previews = previews_resp.get("previews") or []
-
-    selected: list[dict[str, Any]] = []
-    for p in previews:
-        if p.get("error") is not None:
-            continue
-        if not placement_filter(str(p.get("instance_meta", "")), args.instance_meta):
-            continue
-        if not sharding_filter(str(p.get("sharding", "")), args.sharding):
-            continue
-
-        instance = p.get("instance")
-        if not isinstance(instance, dict):
-            continue
-
-        n = nodes_used_in_instance(instance)
-        # Skip tensor ring single node as it is pointless when pipeline ring
-        if n == 1 and (
-            (args.sharding == "both" and "tensor" in p.get("sharding", "").lower())
-            or (
-                args.instance_meta == "both"
-                and "jaccl" in p.get("instance_meta", "").lower()
-            )
-        ):
-            continue
-
-        if (
-            args.skip_pipeline_jaccl
-            and (
-                args.instance_meta == "both"
-                and "jaccl" in p.get("instance_meta", "").lower()
-            )
-            and (
-                args.sharding == "both" and "pipeline" in p.get("sharding", "").lower()
-            )
-        ):
-            continue
-
-        if (
-            args.skip_tensor_ring
-            and (
-                args.instance_meta == "both"
-                and "ring" in p.get("instance_meta", "").lower()
-            )
-            and (args.sharding == "both" and "tensor" in p.get("sharding", "").lower())
-        ):
-            continue
-
-        if args.min_nodes <= n <= args.max_nodes:
-            selected.append(p)
-
-    return selected
-
-
 def main() -> int:
     ap = argparse.ArgumentParser(
         prog="exo-bench",
         description="Benchmark exo model throughput across placement previews.",
     )
-    ap.add_argument("--host", default=os.environ.get("EXO_HOST", "localhost"))
-    ap.add_argument(
-        "--port", type=int, default=int(os.environ.get("EXO_PORT", "52415"))
-    )
-    ap.add_argument("--model", required=True, help="Model short id or huggingface id")
+    add_common_instance_args(ap)
     ap.add_argument(
         "--pp",
         nargs="+",
@@ -620,34 +249,6 @@ def main() -> int:
         required=True,
         help="Generation lengths (ints). Accepts commas.",
     )
-    ap.add_argument(
-        "--max-nodes",
-        type=int,
-        default=4,
-        help="Only consider placements using <= this many nodes.",
-    )
-    ap.add_argument(
-        "--min-nodes",
-        type=int,
-        default=1,
-        help="Only consider placements using >= this many nodes.",
-    )
-    ap.add_argument(
-        "--instance-meta", choices=["ring", "jaccl", "both"], default="both"
-    )
-    ap.add_argument(
-        "--sharding", choices=["pipeline", "tensor", "both"], default="both"
-    )
-    ap.add_argument(
-        "--skip-pipeline-jaccl",
-        action="store_true",
-        help="Skip pipeline+jaccl placements, as it's often pointless.",
-    )
-    ap.add_argument(
-        "--skip-tensor-ring",
-        action="store_true",
-        help="Skip tensor+ring placements, as it's so slow.",
-    )
     ap.add_argument(
         "--repeat", type=int, default=1, help="Repetitions per (pp,tg) pair."
     )
@@ -657,9 +258,6 @@ def main() -> int:
         default=0,
         help="Warmup runs per placement (uses first pp/tg).",
     )
-    ap.add_argument(
-        "--timeout", type=float, default=7200.0, help="HTTP timeout (seconds)."
-    )
     ap.add_argument(
         "--json-out",
         default="bench/results.json",
@@ -674,17 +272,6 @@ def main() -> int:
         action="store_true",
         help="Force all pp×tg combinations (cartesian product) even when lists have equal length.",
     )
-    ap.add_argument(
-        "--settle-timeout",
-        type=float,
-        default=0,
-        help="Max seconds to wait for the cluster to produce valid placements (0 = try once).",
-    )
-    ap.add_argument(
-        "--danger-delete-downloads",
-        action="store_true",
-        help="Delete existing models from smallest to largest to make room for benchmark model.",
-    )
     args = ap.parse_args()
 
     pp_list = parse_int_list(args.pp)
@@ -719,24 +306,10 @@ def main() -> int:
         logger.error("[exo-bench] tokenizer usable but prompt sizing failed")
         raise
 
-    settle_deadline = (
-        time.monotonic() + args.settle_timeout if args.settle_timeout > 0 else None
+    selected = settle_and_fetch_placements(
+        client, full_model_id, args, settle_timeout=args.settle_timeout
     )
 
-    selected = fetch_and_filter_placements(client, full_model_id, args)
-
-    if not selected and settle_deadline:
-        backoff = _SETTLE_INITIAL_BACKOFF_S
-        while not selected and time.monotonic() < settle_deadline:
-            remaining = settle_deadline - time.monotonic()
-            logger.warning(
-                f"No valid placements yet (cluster may still be settling). "
-                f"Retrying in {backoff:.1f}s ({remaining:.0f}s remaining)..."
-            )
-            time.sleep(min(backoff, remaining))
-            backoff = min(backoff * _SETTLE_BACKOFF_MULTIPLIER, _SETTLE_MAX_BACKOFF_S)
-            selected = fetch_and_filter_placements(client, full_model_id, args)
-
     if not selected:
         logger.error("No valid placements matched your filters.")
         return 1
@@ -760,6 +333,10 @@ def main() -> int:
     if args.dry_run:
         return 0
 
+    settle_deadline = (
+        time.monotonic() + args.settle_timeout if args.settle_timeout > 0 else None
+    )
+
     logger.info("Planning phase: checking downloads...")
     run_planning_phase(
         client,

bench/harness.py

@@ -0,0 +1,477 @@
+# type: ignore
+from __future__ import annotations
+
+import argparse
+import http.client
+import json
+import os
+import time
+from typing import Any
+from urllib.parse import urlencode
+
+from loguru import logger
+
+_SETTLE_INITIAL_BACKOFF_S = 1.0
+_SETTLE_MAX_BACKOFF_S = 60.0
+_SETTLE_BACKOFF_MULTIPLIER = 2.0
+
+
+class ExoHttpError(RuntimeError):
+    def __init__(self, status: int, reason: str, body_preview: str):
+        super().__init__(f"HTTP {status} {reason}: {body_preview}")
+        self.status = status
+
+
+class ExoClient:
+    def __init__(self, host: str, port: int, timeout_s: float = 7200.0):
+        self.host = host
+        self.port = port
+        self.timeout_s = timeout_s
+
+    def request_json(
+        self,
+        method: str,
+        path: str,
+        params: dict[str, Any] | None = None,
+        body: dict[str, Any] | None = None,
+        headers: dict[str, str] | None = None,
+    ) -> Any:
+        if not path.startswith("/"):
+            path = "/" + path
+        if params:
+            path = path + "?" + urlencode(params)
+
+        conn = http.client.HTTPConnection(self.host, self.port, timeout=self.timeout_s)
+        try:
+            payload: bytes | None = None
+            hdrs: dict[str, str] = {"Accept": "application/json"}
+
+            if body is not None:
+                payload = json.dumps(body).encode("utf-8")
+                hdrs["Content-Type"] = "application/json"
+            if headers:
+                hdrs.update(headers)
+
+            conn.request(method.upper(), path, body=payload, headers=hdrs)
+            resp = conn.getresponse()
+            raw = resp.read()
+            text = raw.decode("utf-8", errors="replace") if raw else ""
+
+            if resp.status >= 400:
+                raise ExoHttpError(resp.status, resp.reason, text[:300])
+
+            if not text:
+                return None
+            return json.loads(text)
+        finally:
+            conn.close()
+
+    def post_bench_chat_completions(self, payload: dict[str, Any]) -> dict[str, Any]:
+        return self.request_json("POST", "/bench/chat/completions", body=payload)
+
+
+def unwrap_instance(instance: dict[str, Any]) -> dict[str, Any]:
+    if len(instance) != 1:
+        raise KeyError(f"Expected 1 key, got keys={list(instance.keys())}")
+
+    tag = next(iter(instance))
+    inner = instance[tag]
+    if not isinstance(inner, dict):
+        raise TypeError(f"payload for {tag} must be dict, got {type(inner)}")
+    return inner
+
+
+def instance_id_from_instance(instance: dict[str, Any]) -> str:
+    inner = unwrap_instance(instance)
+    return str(inner["instanceId"])
+
+
+def nodes_used_in_instance(instance: dict[str, Any]) -> int:
+    inner = unwrap_instance(instance)
+    return len(inner["shardAssignments"]["nodeToRunner"])
+
+
+def runner_ids_from_instance(instance: dict[str, Any]) -> list[str]:
+    inner = unwrap_instance(instance)
+    runner_to_shard = inner["shardAssignments"]["runnerToShard"]
+    return list(runner_to_shard.keys())
+
+
+def runner_ready(runner: dict[str, Any]) -> bool:
+    return "RunnerReady" in runner
+
+
+def runner_failed(runner: dict[str, Any]) -> bool:
+    return "RunnerFailed" in runner
+
+
+def get_runner_failed_message(runner: dict[str, Any]) -> str | None:
+    if "RunnerFailed" in runner:
+        return runner["RunnerFailed"].get("errorMessage")
+    return None
+
+
+def wait_for_instance_ready(
+    client: ExoClient, instance_id: str, timeout: float = 24000.0
+) -> None:
+    start_time = time.time()
+    instance_existed = False
+    while time.time() - start_time < timeout:
+        state = client.request_json("GET", "/state")
+        instances = state.get("instances", {})
+
+        if instance_id not in instances:
+            if instance_existed:
+                # Instance was deleted after being created - likely due to runner failure
+                raise RuntimeError(
+                    f"Instance {instance_id} was deleted (runner may have failed)"
+                )
+            time.sleep(0.1)
+            continue
+
+        instance_existed = True
+        instance = instances[instance_id]
+        runner_ids = runner_ids_from_instance(instance)
+        runners = state.get("runners", {})
+
+        # Check for failed runners first
+        for rid in runner_ids:
+            runner = runners.get(rid, {})
+            if runner_failed(runner):
+                error_msg = get_runner_failed_message(runner) or "Unknown error"
+                raise RuntimeError(f"Runner {rid} failed: {error_msg}")
+
+        if all(runner_ready(runners.get(rid, {})) for rid in runner_ids):
+            return
+
+        time.sleep(0.1)
+
+    raise TimeoutError(f"Instance {instance_id} did not become ready within {timeout=}")
+
+
+def wait_for_instance_gone(
+    client: ExoClient, instance_id: str, timeout: float = 3.0
+) -> None:
+    start_time = time.time()
+    while time.time() - start_time < timeout:
+        try:
+            client.request_json("GET", f"/instance/{instance_id}")
+            time.sleep(0.4)
+        except ExoHttpError as e:
+            if e.status == 404:
+                return
+            raise
+
+    raise TimeoutError(f"Instance {instance_id} did not get deleted within {timeout=}")
+
+
+def resolve_model_short_id(client: ExoClient, model_arg: str) -> tuple[str, str]:
+    models = client.request_json("GET", "/models") or {}
+    data = models.get("data") or []
+
+    for m in data:
+        if (m.get("name") or "").lower() == model_arg.lower():
+            short_id = str(m["name"])
+            full_id = str(m.get("hugging_face_id") or m["name"])
+            return short_id, full_id
+
+    for m in data:
+        if m.get("hugging_face_id") == model_arg:
+            short_id = str(m["name"])
+            full_id = str(m["hugging_face_id"])
+            return short_id, full_id
+
+    raise ValueError(f"Model not found in /models: {model_arg}")
+
+
+def placement_filter(instance_meta: str, wanted: str) -> bool:
+    s = (instance_meta or "").lower()
+    if wanted == "both":
+        return ("ring" in s) or ("jaccl" in s)
+    return wanted in s
+
+
+def sharding_filter(sharding: str, wanted: str) -> bool:
+    s = (sharding or "").lower()
+    if wanted == "both":
+        return ("pipeline" in s) or ("tensor" in s)
+    return wanted in s
+
+
+def fetch_and_filter_placements(
+    client: ExoClient, full_model_id: str, args: argparse.Namespace
+) -> list[dict[str, Any]]:
+    previews_resp = client.request_json(
+        "GET", "/instance/previews", params={"model_id": full_model_id}
+    )
+    previews = previews_resp.get("previews") or []
+
+    selected: list[dict[str, Any]] = []
+    for p in previews:
+        if p.get("error") is not None:
+            continue
+        if not placement_filter(str(p.get("instance_meta", "")), args.instance_meta):
+            continue
+        if not sharding_filter(str(p.get("sharding", "")), args.sharding):
+            continue
+
+        instance = p.get("instance")
+        if not isinstance(instance, dict):
+            continue
+
+        n = nodes_used_in_instance(instance)
+        # Skip tensor ring single node as it is pointless when pipeline ring
+        if n == 1 and (
+            (args.sharding == "both" and "tensor" in p.get("sharding", "").lower())
+            or (
+                args.instance_meta == "both"
+                and "jaccl" in p.get("instance_meta", "").lower()
+            )
+        ):
+            continue
+
+        if (
+            args.skip_pipeline_jaccl
+            and (
+                args.instance_meta == "both"
+                and "jaccl" in p.get("instance_meta", "").lower()
+            )
+            and (
+                args.sharding == "both" and "pipeline" in p.get("sharding", "").lower()
+            )
+        ):
+            continue
+
+        if (
+            args.skip_tensor_ring
+            and (
+                args.instance_meta == "both"
+                and "ring" in p.get("instance_meta", "").lower()
+            )
+            and (args.sharding == "both" and "tensor" in p.get("sharding", "").lower())
+        ):
+            continue
+
+        if args.min_nodes <= n <= args.max_nodes:
+            selected.append(p)
+
+    return selected
+
+
+def settle_and_fetch_placements(
+    client: ExoClient,
+    full_model_id: str,
+    args: argparse.Namespace,
+    settle_timeout: float = 0,
+) -> list[dict[str, Any]]:
+    selected = fetch_and_filter_placements(client, full_model_id, args)
+
+    if not selected and settle_timeout > 0:
+        backoff = _SETTLE_INITIAL_BACKOFF_S
+        deadline = time.monotonic() + settle_timeout
+        while not selected and time.monotonic() < deadline:
+            remaining = deadline - time.monotonic()
+            logger.warning(
+                f"No valid placements yet (cluster may still be settling). "
+                f"Retrying in {backoff:.1f}s ({remaining:.0f}s remaining)..."
+            )
+            time.sleep(min(backoff, remaining))
+            backoff = min(backoff * _SETTLE_BACKOFF_MULTIPLIER, _SETTLE_MAX_BACKOFF_S)
+            selected = fetch_and_filter_placements(client, full_model_id, args)
+
+    return selected
+
+
+def run_planning_phase(
+    client: ExoClient,
+    full_model_id: str,
+    preview: dict[str, Any],
+    danger_delete: bool,
+    timeout: float,
+    settle_deadline: float | None,
+) -> None:
+    """Check disk space and ensure model is downloaded before benchmarking."""
+    # Get model size from /models
+    models = client.request_json("GET", "/models") or {}
+    model_bytes = 0
+    for m in models.get("data", []):
+        if m.get("hugging_face_id") == full_model_id:
+            model_bytes = m.get("storage_size_megabytes", 0) * 1024 * 1024
+            break
+
+    if not model_bytes:
+        logger.warning(
+            f"Could not determine size for {full_model_id}, skipping disk check"
+        )
+        return
+
+    # Get nodes from preview
+    inner = unwrap_instance(preview["instance"])
+    node_ids = list(inner["shardAssignments"]["nodeToRunner"].keys())
+    runner_to_shard = inner["shardAssignments"]["runnerToShard"]
+
+    state = client.request_json("GET", "/state")
+    downloads = state.get("downloads", {})
+    node_disk = state.get("nodeDisk", {})
+
+    for node_id in node_ids:
+        node_downloads = downloads.get(node_id, [])
+
+        # Check if model already downloaded on this node
+        already_downloaded = any(
+            "DownloadCompleted" in p
+            and unwrap_instance(p["DownloadCompleted"]["shardMetadata"])["modelCard"][
+                "modelId"
+            ]
+            == full_model_id
+            for p in node_downloads
+        )
+        if already_downloaded:
+            continue
+
+        # Wait for disk info if settle_deadline is set
+        disk_info = node_disk.get(node_id, {})
+        backoff = _SETTLE_INITIAL_BACKOFF_S
+        while not disk_info and settle_deadline and time.monotonic() < settle_deadline:
+            remaining = settle_deadline - time.monotonic()
+            logger.info(
+                f"Waiting for disk info on {node_id} ({remaining:.0f}s remaining)..."
+            )
+            time.sleep(min(backoff, remaining))
+            backoff = min(backoff * _SETTLE_BACKOFF_MULTIPLIER, _SETTLE_MAX_BACKOFF_S)
+            state = client.request_json("GET", "/state")
+            node_disk = state.get("nodeDisk", {})
+            disk_info = node_disk.get(node_id, {})
+
+        if not disk_info:
+            logger.warning(f"No disk info for {node_id}, skipping space check")
+            continue
+
+        avail = disk_info.get("available", {}).get("inBytes", 0)
+        if avail >= model_bytes:
+            continue
+
+        if not danger_delete:
+            raise RuntimeError(
+                f"Insufficient disk on {node_id}: need {model_bytes // (1024**3)}GB, "
+                f"have {avail // (1024**3)}GB. Use --danger-delete-downloads to free space."
+            )
+
+        # Delete from smallest to largest
+        completed = [
+            (
+                unwrap_instance(p["DownloadCompleted"]["shardMetadata"])["modelCard"][
+                    "modelId"
+                ],
+                p["DownloadCompleted"]["totalBytes"]["inBytes"],
+            )
+            for p in node_downloads
+            if "DownloadCompleted" in p
+        ]
+        for del_model, size in sorted(completed, key=lambda x: x[1]):
+            logger.info(f"Deleting {del_model} from {node_id} ({size // (1024**2)}MB)")
+            client.request_json("DELETE", f"/download/{node_id}/{del_model}")
+            avail += size
+            if avail >= model_bytes:
+                break
+
+        if avail < model_bytes:
+            raise RuntimeError(f"Could not free enough space on {node_id}")
+
+    # Start downloads (idempotent)
+    for node_id in node_ids:
+        runner_id = inner["shardAssignments"]["nodeToRunner"][node_id]
+        shard = runner_to_shard[runner_id]
+        client.request_json(
+            "POST",
+            "/download/start",
+            body={
+                "targetNodeId": node_id,
+                "shardMetadata": shard,
+            },
+        )
+        logger.info(f"Started download on {node_id}")
+
+    # Wait for downloads
+    start = time.time()
+    while time.time() - start < timeout:
+        state = client.request_json("GET", "/state")
+        downloads = state.get("downloads", {})
+        all_done = True
+        for node_id in node_ids:
+            done = any(
+                "DownloadCompleted" in p
+                and unwrap_instance(p["DownloadCompleted"]["shardMetadata"])[
+                    "modelCard"
+                ]["modelId"]
+                == full_model_id
+                for p in downloads.get(node_id, [])
+            )
+            failed = [
+                p["DownloadFailed"]["errorMessage"]
+                for p in downloads.get(node_id, [])
+                if "DownloadFailed" in p
+                and unwrap_instance(p["DownloadFailed"]["shardMetadata"])["modelCard"][
+                    "modelId"
+                ]
+                == full_model_id
+            ]
+            if failed:
+                raise RuntimeError(f"Download failed on {node_id}: {failed[0]}")
+            if not done:
+                all_done = False
+        if all_done:
+            return
+        time.sleep(1)
+
+    raise TimeoutError("Downloads did not complete in time")
+
+
+def add_common_instance_args(ap: argparse.ArgumentParser) -> None:
+    ap.add_argument("--host", default=os.environ.get("EXO_HOST", "localhost"))
+    ap.add_argument(
+        "--port", type=int, default=int(os.environ.get("EXO_PORT", "52415"))
+    )
+    ap.add_argument("--model", required=True, help="Model short id or huggingface id")
+    ap.add_argument(
+        "--max-nodes",
+        type=int,
+        default=4,
+        help="Only consider placements using <= this many nodes.",
+    )
+    ap.add_argument(
+        "--min-nodes",
+        type=int,
+        default=1,
+        help="Only consider placements using >= this many nodes.",
+    )
+    ap.add_argument(
+        "--instance-meta", choices=["ring", "jaccl", "both"], default="both"
+    )
+    ap.add_argument(
+        "--sharding", choices=["pipeline", "tensor", "both"], default="both"
+    )
+    ap.add_argument(
+        "--skip-pipeline-jaccl",
+        action="store_true",
+        help="Skip pipeline+jaccl placements, as it's often pointless.",
+    )
+    ap.add_argument(
+        "--skip-tensor-ring",
+        action="store_true",
+        help="Skip tensor+ring placements, as it's so slow.",
+    )
+    ap.add_argument(
+        "--timeout", type=float, default=7200.0, help="HTTP timeout (seconds)."
+    )
+    ap.add_argument(
+        "--settle-timeout",
+        type=float,
+        default=0,
+        help="Max seconds to wait for the cluster to produce valid placements (0 = try once).",
+    )
+    ap.add_argument(
+        "--danger-delete-downloads",
+        action="store_true",
+        help="Delete existing models from smallest to largest to make room for benchmark model.",
+    )

bench/pyproject.toml

@@ -4,6 +4,7 @@ version = "0.1.0"
 description = "Benchmarking tool for exo distributed inference"
 requires-python = ">=3.13"
 dependencies = [
+    "httpx>=0.27.0",
     "loguru>=0.7.3",
     "transformers>=5.0.0",
     "huggingface-hub>=0.33.4",

bench/scenarios.toml

@@ -0,0 +1,306 @@
+# Tool definitions — each becomes an OpenAI function tool.
+# All scenarios get all tools unless they specify a `tools` list.
+
+[tools.get_current_weather]
+description = "Get the current weather in a given location"
+required = ["location"]
+
+[tools.get_current_weather.properties.location]
+type = "string"
+description = "City and state, e.g. San Francisco, CA"
+
+[tools.get_current_weather.properties.unit]
+type = "string"
+enum = ["celsius", "fahrenheit"]
+description = "Temperature unit"
+
+[tools.calculate]
+description = "Evaluate a mathematical expression and return the numeric result"
+required = ["expression"]
+
+[tools.calculate.properties.expression]
+type = "string"
+description = "The math expression to evaluate, e.g. '2 + 3 * 4'"
+
+[tools.search_products]
+description = "Search for products in a catalog by query, category, and price"
+required = ["query"]
+
+[tools.search_products.properties.query]
+type = "string"
+description = "Search query string"
+
+[tools.search_products.properties.category]
+type = "string"
+enum = ["electronics", "clothing", "food", "books"]
+description = "Product category to filter by"
+
+[tools.search_products.properties.max_price]
+type = "number"
+description = "Maximum price in USD"
+
+[tools.create_todos]
+description = "Create a structured todo list"
+required = ["todos"]
+
+[tools.create_todos.properties.todos]
+type = "array"
+description = "List of todo items"
+
+[tools.create_todos.properties.todos.items]
+type = "object"
+required = ["content", "status", "priority"]
+
+[tools.create_todos.properties.todos.items.properties.content]
+type = "string"
+description = "The todo item text"
+
+[tools.create_todos.properties.todos.items.properties.status]
+type = "string"
+description = "Status: pending, in_progress, or completed"
+
+[tools.create_todos.properties.todos.items.properties.priority]
+type = "string"
+description = "Priority: low, normal, or high"
+
+# -- Should call a tool --
+
+[[scenarios]]
+name = "weather_simple"
+description = "Basic weather query -> get_current_weather"
+expect_tool_call = true
+expected_function = "get_current_weather"
+required_arg_keys = ["location"]
+
+[[scenarios.messages]]
+role = "user"
+content = "What's the weather like in Tokyo right now?"
+
+[[scenarios]]
+name = "calculator_simple"
+description = "Math question -> calculate"
+expect_tool_call = true
+expected_function = "calculate"
+required_arg_keys = ["expression"]
+
+[[scenarios.messages]]
+role = "user"
+content = "Use the calculator to compute 3847 * 926 + 17293"
+
+[[scenarios]]
+name = "search_with_filters"
+description = "Product search with category and price filter"
+expect_tool_call = true
+expected_function = "search_products"
+required_arg_keys = ["query"]
+
+[[scenarios.messages]]
+role = "user"
+content = "Find me electronics under $50"
+
+# -- Multi-turn: tool call then follow-up --
+
+[[scenarios]]
+name = "weather_multi_turn"
+description = "Weather query -> tool result -> natural language summary"
+expect_tool_call = true
+expected_function = "get_current_weather"
+required_arg_keys = ["location"]
+
+[scenarios.tool_result]
+temperature = "18C"
+condition = "partly cloudy"
+humidity = "65%"
+wind = "12 km/h NW"
+
+[[scenarios.messages]]
+role = "user"
+content = "What's the weather in Paris?"
+
+[[scenarios]]
+name = "calculator_multi_turn"
+description = "Math query -> tool result -> model reports the answer"
+expect_tool_call = true
+expected_function = "calculate"
+required_arg_keys = ["expression"]
+
+[scenarios.tool_result]
+result = 491682
+
+[[scenarios.messages]]
+role = "user"
+content = "Use the calculator to compute 1847 * 263 + 5921"
+
+[[scenarios]]
+name = "search_multi_turn"
+description = "Search query -> tool result -> model summarizes products"
+expect_tool_call = true
+expected_function = "search_products"
+required_arg_keys = ["query"]
+
+[[scenarios.tool_result.results]]
+name = "Hands-On Machine Learning"
+price = 45.99
+rating = 4.8
+
+[[scenarios.tool_result.results]]
+name = "Deep Learning with Python"
+price = 39.99
+rating = 4.6
+
+[[scenarios.messages]]
+role = "user"
+content = "Search for books about machine learning"
+
+# -- Sequential tool calls --
+
+[[scenarios]]
+name = "chained_tool_calls_same"
+description = "Thinking + weather(Tokyo) -> result -> model must call weather(London)"
+expect_tool_call = true
+expected_function = "get_current_weather"
+required_arg_keys = ["location"]
+
+[[scenarios.messages]]
+role = "user"
+content = "Compare the weather in Tokyo and London."
+
+[[scenarios.messages]]
+role = "assistant"
+content = "I'll check both cities. Let me start with Tokyo."
+
+[[scenarios.messages.tool_calls]]
+id = "call_1"
+name = "get_current_weather"
+arguments = { location = "Tokyo" }
+
+[[scenarios.messages]]
+role = "tool"
+tool_call_id = "call_1"
+content = '{"temperature": "25C", "condition": "sunny"}'
+
+[[scenarios]]
+name = "chained_tool_calls_different"
+description = "Thinking + weather(Berlin) -> result -> model must call calculator"
+expect_tool_call = true
+expected_function = "calculate"
+required_arg_keys = ["expression"]
+
+[[scenarios.messages]]
+role = "user"
+content = "What's the weather in Berlin, and also use the calculator to compute 4819 * 37 + 291."
+
+[[scenarios.messages]]
+role = "assistant"
+content = "I'll handle both. Let me check Berlin's weather first."
+
+[[scenarios.messages.tool_calls]]
+id = "call_2"
+name = "get_current_weather"
+arguments = { location = "Berlin" }
+
+[[scenarios.messages]]
+role = "tool"
+tool_call_id = "call_2"
+content = '{"temperature": "12C", "condition": "rainy"}'
+
+[[scenarios]]
+name = "chained_tool_calls_three"
+description = "Two prior thinking+tool calls -> results -> model must make a third"
+expect_tool_call = true
+expected_function = "get_current_weather"
+required_arg_keys = ["location"]
+
+[[scenarios.messages]]
+role = "user"
+content = "Compare weather in Tokyo, Paris, and London."
+
+[[scenarios.messages]]
+role = "assistant"
+content = "I'll check all three cities. Starting with Tokyo."
+
+[[scenarios.messages.tool_calls]]
+id = "call_3"
+name = "get_current_weather"
+arguments = { location = "Tokyo" }
+
+[[scenarios.messages]]
+role = "tool"
+tool_call_id = "call_3"
+content = '{"temperature": "25C", "condition": "sunny"}'
+
+[[scenarios.messages]]
+role = "assistant"
+content = "Got Tokyo. Now checking Paris."
+
+[[scenarios.messages.tool_calls]]
+id = "call_4"
+name = "get_current_weather"
+arguments = { location = "Paris" }
+
+[[scenarios.messages]]
+role = "tool"
+tool_call_id = "call_4"
+content = '{"temperature": "18C", "condition": "cloudy"}'
+
+# -- Nested object schema (regression for lossy chat template rendering) --
+
+[[scenarios]]
+name = "nested_schema_tool_call"
+description = "Tool call with nested object array schema -> create_todos"
+expect_tool_call = true
+expected_function = "create_todos"
+required_arg_keys = ["todos"]
+nested_array_key = "todos"
+required_item_keys = ["content", "status", "priority"]
+tools = ["create_todos"]
+
+[[scenarios.messages]]
+role = "user"
+content = "Create a todo list with 3 items to learn Python"
+
+# -- Tool name integrity (regression for harmony token leaking into name) --
+
+[tools.glob]
+description = "Search for files matching a glob pattern in the codebase"
+required = ["pattern"]
+
+[tools.glob.properties.pattern]
+type = "string"
+description = "The glob pattern to match files against, e.g. '**/*.py'"
+
+[tools.glob.properties.path]
+type = "string"
+description = "The directory to search in"
+
+[[scenarios]]
+name = "tool_name_integrity"
+description = "Tool name must not contain harmony tokens like <|channel|>"
+expect_tool_call = true
+expected_function = "glob"
+required_arg_keys = ["pattern"]
+tools = ["glob"]
+
+[[scenarios.messages]]
+role = "user"
+content = "Find all Python files in the src directory"
+
+# -- Should NOT call a tool --
+
+[[scenarios]]
+name = "no_tool_joke"
+description = "Joke request should NOT trigger any tool"
+expect_tool_call = false
+
+[[scenarios.messages]]
+role = "user"
+content = "Tell me a funny joke about cats."
+
+[[scenarios]]
+name = "no_tool_factual"
+description = "Factual question answerable from training data"
+expect_tool_call = false
+
+[[scenarios.messages]]
+role = "user"
+content = "What is the capital of Japan?"

dashboard/src/lib/components/ChatForm.svelte

@@ -14,6 +14,7 @@
     totalTokens,
     thinkingEnabled as thinkingEnabledStore,
     setConversationThinking,
+    stopGeneration,
   } from "$lib/stores/app.svelte";
   import ChatAttachments from "./ChatAttachments.svelte";
   import ImageParamsPanel from "./ImageParamsPanel.svelte";
@@ -103,7 +104,7 @@
   const modelSupportsThinking = $derived(() => {
     if (!currentModel) return false;
     const caps = modelCapabilities[currentModel] || [];
-    return caps.includes("thinking") && caps.includes("text");
+    return caps.includes("thinking_toggle") && caps.includes("text");
   });
 
   const isEditOnlyWithoutImage = $derived(
@@ -653,86 +654,92 @@
         style="min-height: 28px; max-height: 150px;"
       ></textarea>
 
-      <button
-        type="submit"
-        disabled={!canSend || loading || isEditOnlyWithoutImage}
-        class="px-2.5 sm:px-4 py-1.5 sm:py-2 rounded text-xs sm:text-xs tracking-[0.1em] sm:tracking-[0.15em] uppercase font-medium transition-all duration-200 whitespace-nowrap
-					{!canSend || loading || isEditOnlyWithoutImage
-          ? 'bg-exo-medium-gray/50 text-exo-light-gray cursor-not-allowed'
-          : 'bg-exo-yellow text-exo-black hover:bg-exo-yellow-darker hover:shadow-[0_0_20px_rgba(255,215,0,0.3)]'}"
-        aria-label={shouldShowEditMode
-          ? "Edit image"
-          : isImageModel()
-            ? "Generate image"
-            : "Send message"}
-      >
-        {#if loading}
+      {#if loading}
+        <button
+          type="button"
+          onclick={() => stopGeneration()}
+          class="px-2.5 sm:px-4 py-1.5 sm:py-2 rounded text-xs sm:text-xs tracking-[0.1em] sm:tracking-[0.15em] font-medium transition-all duration-200 whitespace-nowrap bg-exo-medium-gray/70 text-exo-light-gray hover:bg-exo-medium-gray hover:text-white"
+          aria-label="Stop generation"
+        >
           <span class="inline-flex items-center gap-1 sm:gap-2">
-            <span
-              class="w-2.5 h-2.5 sm:w-3 sm:h-3 border-2 border-current border-t-transparent rounded-full animate-spin"
-            ></span>
-            <span class="hidden sm:inline"
-              >{shouldShowEditMode
-                ? "EDITING"
-                : isImageModel()
-                  ? "GENERATING"
-                  : "PROCESSING"}</span
-            >
-            <span class="sm:hidden">...</span>
-          </span>
-        {:else if shouldShowEditMode}
-          <span class="inline-flex items-center gap-1.5">
             <svg
-              class="w-3.5 h-3.5"
-              fill="none"
+              class="w-3 h-3 sm:w-3.5 sm:h-3.5"
+              fill="currentColor"
               viewBox="0 0 24 24"
-              stroke="currentColor"
-              stroke-width="2"
             >
-              <path
-                stroke-linecap="round"
-                stroke-linejoin="round"
-                d="M11 5H6a2 2 0 00-2 2v11a2 2 0 002 2h11a2 2 0 002-2v-5m-1.414-9.414a2 2 0 112.828 2.828L11.828 15H9v-2.828l8.586-8.586z"
-              />
+              <rect x="6" y="6" width="12" height="12" rx="1" />
             </svg>
-            <span>EDIT</span>
+            <span class="hidden sm:inline">Cancel</span>
           </span>
-        {:else if isEditOnlyWithoutImage}
-          <span class="inline-flex items-center gap-1.5">
-            <svg
-              class="w-3.5 h-3.5"
-              fill="none"
-              viewBox="0 0 24 24"
-              stroke="currentColor"
-              stroke-width="2"
-            >
-              <path
-                stroke-linecap="round"
-                stroke-linejoin="round"
-                d="M11 5H6a2 2 0 00-2 2v11a2 2 0 002 2h11a2 2 0 002-2v-5m-1.414-9.414a2 2 0 112.828 2.828L11.828 15H9v-2.828l8.586-8.586z"
-              />
-            </svg>
-            <span>EDIT</span>
-          </span>
-        {:else if isImageModel()}
-          <span class="inline-flex items-center gap-1.5">
-            <svg
-              class="w-3.5 h-3.5"
-              fill="none"
-              viewBox="0 0 24 24"
-              stroke="currentColor"
-              stroke-width="2"
-            >
-              <rect x="3" y="3" width="18" height="18" rx="2" ry="2" />
-              <circle cx="8.5" cy="8.5" r="1.5" />
-              <polyline points="21 15 16 10 5 21" />
-            </svg>
-            <span>GENERATE</span>
-          </span>
-        {:else}
-          SEND
-        {/if}
-      </button>
+        </button>
+      {:else}
+        <button
+          type="submit"
+          disabled={!canSend || isEditOnlyWithoutImage}
+          class="px-2.5 sm:px-4 py-1.5 sm:py-2 rounded text-xs sm:text-xs tracking-[0.1em] sm:tracking-[0.15em] uppercase font-medium transition-all duration-200 whitespace-nowrap
+					{!canSend || isEditOnlyWithoutImage
+            ? 'bg-exo-medium-gray/50 text-exo-light-gray cursor-not-allowed'
+            : 'bg-exo-yellow text-exo-black hover:bg-exo-yellow-darker hover:shadow-[0_0_20px_rgba(255,215,0,0.3)]'}"
+          aria-label={shouldShowEditMode
+            ? "Edit image"
+            : isImageModel()
+              ? "Generate image"
+              : "Send message"}
+        >
+          {#if shouldShowEditMode}
+            <span class="inline-flex items-center gap-1.5">
+              <svg
+                class="w-3.5 h-3.5"
+                fill="none"
+                viewBox="0 0 24 24"
+                stroke="currentColor"
+                stroke-width="2"
+              >
+                <path
+                  stroke-linecap="round"
+                  stroke-linejoin="round"
+                  d="M11 5H6a2 2 0 00-2 2v11a2 2 0 002 2h11a2 2 0 002-2v-5m-1.414-9.414a2 2 0 112.828 2.828L11.828 15H9v-2.828l8.586-8.586z"
+                />
+              </svg>
+              <span>EDIT</span>
+            </span>
+          {:else if isEditOnlyWithoutImage}
+            <span class="inline-flex items-center gap-1.5">
+              <svg
+                class="w-3.5 h-3.5"
+                fill="none"
+                viewBox="0 0 24 24"
+                stroke="currentColor"
+                stroke-width="2"
+              >
+                <path
+                  stroke-linecap="round"
+                  stroke-linejoin="round"
+                  d="M11 5H6a2 2 0 00-2 2v11a2 2 0 002 2h11a2 2 0 002-2v-5m-1.414-9.414a2 2 0 112.828 2.828L11.828 15H9v-2.828l8.586-8.586z"
+                />
+              </svg>
+              <span>EDIT</span>
+            </span>
+          {:else if isImageModel()}
+            <span class="inline-flex items-center gap-1.5">
+              <svg
+                class="w-3.5 h-3.5"
+                fill="none"
+                viewBox="0 0 24 24"
+                stroke="currentColor"
+                stroke-width="2"
+              >
+                <rect x="3" y="3" width="18" height="18" rx="2" ry="2" />
+                <circle cx="8.5" cy="8.5" r="1.5" />
+                <polyline points="21 15 16 10 5 21" />
+              </svg>
+              <span>GENERATE</span>
+            </span>
+          {:else}
+            SEND
+          {/if}
+        </button>
+      {/if}
     </div>
 
     <!-- Bottom accent line -->

dashboard/src/lib/components/ChatMessages.svelte

@@ -3,16 +3,17 @@
     messages,
     currentResponse,
     isLoading,
+    prefillProgress,
     deleteMessage,
     editAndRegenerate,
     regenerateLastResponse,
     regenerateFromToken,
     setEditingImage,
   } from "$lib/stores/app.svelte";
-  import type { Message } from "$lib/stores/app.svelte";
   import type { MessageAttachment } from "$lib/stores/app.svelte";
   import MarkdownContent from "./MarkdownContent.svelte";
   import TokenHeatmap from "./TokenHeatmap.svelte";
+  import PrefillProgressBar from "./PrefillProgressBar.svelte";
   import ImageLightbox from "./ImageLightbox.svelte";
 
   interface Props {
@@ -25,6 +26,7 @@
   const messageList = $derived(messages());
   const response = $derived(currentResponse());
   const loading = $derived(isLoading());
+  const prefill = $derived(prefillProgress());
 
   // Scroll management - user controls scroll, show button when not at bottom
   const SCROLL_THRESHOLD = 100;
@@ -428,6 +430,9 @@
             {:else}
               <!-- Assistant message styling -->
               <div class="p-3 sm:p-4">
+                {#if loading && isLastAssistantMessage(message.id) && prefill && !message.content}
+                  <PrefillProgressBar progress={prefill} class="mb-3" />
+                {/if}
                 {#if message.thinking && message.thinking.trim().length > 0}
                   <div
                     class="mb-3 rounded border border-exo-yellow/20 bg-exo-black/40"

dashboard/src/lib/components/HuggingFaceResultItem.svelte

@@ -26,7 +26,8 @@
     downloadedOnNodes = [],
   }: HuggingFaceResultItemProps = $props();
 
-  function formatNumber(num: number): string {
+  function formatNumber(num: number | undefined): string {
+    if (num == null) return "0";
     if (num >= 1000000) {
       return `${(num / 1000000).toFixed(1)}M`;
     } else if (num >= 1000) {

dashboard/src/lib/components/MetaInstanceCard.svelte

@@ -0,0 +1,232 @@
+<script lang="ts">
+  import type {
+    MetaInstance,
+    MetaInstanceStatus,
+    NodeInfo,
+  } from "$lib/stores/app.svelte";
+  import {
+    getMetaInstanceStatus,
+    getMetaInstanceBackingNodes,
+    topologyData,
+  } from "$lib/stores/app.svelte";
+
+  interface Props {
+    metaInstance: MetaInstance;
+    onDelete?: (metaInstanceId: string) => void;
+  }
+
+  let { metaInstance, onDelete }: Props = $props();
+
+  const status: MetaInstanceStatus = $derived(
+    getMetaInstanceStatus(metaInstance),
+  );
+  const backingNodeIds: string[] = $derived(
+    getMetaInstanceBackingNodes(metaInstance),
+  );
+
+  const statusConfig = $derived.by(() => {
+    switch (status) {
+      case "active":
+        return {
+          label: "ACTIVE",
+          dotClass: "bg-green-400",
+          borderClass:
+            "border-green-500/30 border-l-green-400",
+          cornerClass: "border-green-500/50",
+          glowClass: "shadow-[0_0_6px_rgba(74,222,128,0.4)]",
+          animate: false,
+        };
+      case "provisioning":
+        return {
+          label: "PROVISIONING",
+          dotClass: "bg-yellow-400",
+          borderClass:
+            "border-exo-yellow/30 border-l-yellow-400",
+          cornerClass: "border-yellow-500/50",
+          glowClass: "shadow-[0_0_6px_rgba(250,204,21,0.4)]",
+          animate: true,
+        };
+      case "error":
+        return {
+          label: "ERROR",
+          dotClass: "bg-red-400",
+          borderClass: "border-red-500/30 border-l-red-400",
+          cornerClass: "border-red-500/50",
+          glowClass: "shadow-[0_0_6px_rgba(248,113,113,0.4)]",
+          animate: false,
+        };
+    }
+  });
+
+  function getNodeName(nodeId: string): string {
+    const topo = topologyData();
+    if (!topo?.nodes) return nodeId.slice(0, 8);
+    const node = topo.nodes[nodeId];
+    return node?.friendly_name || node?.system_info?.model_id || nodeId.slice(0, 8);
+  }
+
+  function formatModelId(modelId: string): string {
+    // Show just the model name part after the org prefix
+    const parts = modelId.split("/");
+    return parts.length > 1 ? parts[parts.length - 1] : modelId;
+  }
+
+  function handleDelete() {
+    if (
+      onDelete &&
+      confirm(
+        `Delete meta-instance for ${formatModelId(metaInstance.modelId)}?`,
+      )
+    ) {
+      onDelete(metaInstance.metaInstanceId);
+    }
+  }
+</script>
+
+<div class="relative group">
+  <!-- Corner accents -->
+  <div
+    class="absolute -top-px -left-px w-2 h-2 border-l border-t {statusConfig.cornerClass}"
+  ></div>
+  <div
+    class="absolute -top-px -right-px w-2 h-2 border-r border-t {statusConfig.cornerClass}"
+  ></div>
+  <div
+    class="absolute -bottom-px -left-px w-2 h-2 border-l border-b {statusConfig.cornerClass}"
+  ></div>
+  <div
+    class="absolute -bottom-px -right-px w-2 h-2 border-r border-b {statusConfig.cornerClass}"
+  ></div>
+
+  <div
+    class="bg-exo-dark-gray/60 border border-l-2 {statusConfig.borderClass} p-3"
+  >
+    <!-- Header: Status + Delete -->
+    <div class="flex justify-between items-start mb-2 pl-2">
+      <div class="flex items-center gap-2">
+        <div
+          class="w-1.5 h-1.5 {statusConfig.dotClass} rounded-full {statusConfig.glowClass} {statusConfig.animate
+            ? 'animate-pulse'
+            : ''}"
+        ></div>
+        <span
+          class="text-xs font-mono tracking-[0.15em] uppercase {status === 'active'
+            ? 'text-green-400'
+            : status === 'error'
+              ? 'text-red-400'
+              : 'text-yellow-400'}"
+        >
+          {statusConfig.label}
+        </span>
+      </div>
+      <button
+        onclick={handleDelete}
+        class="text-xs px-2 py-1 font-mono tracking-wider uppercase border border-red-500/30 text-red-400 hover:bg-red-500/20 hover:text-red-400 hover:border-red-500/50 transition-all duration-200 cursor-pointer"
+      >
+        DELETE
+      </button>
+    </div>
+
+    <!-- Model Info -->
+    <div class="pl-2 space-y-1">
+      <div class="text-exo-yellow text-xs font-mono tracking-wide truncate">
+        {metaInstance.modelId}
+      </div>
+
+      <!-- Sharding + Runtime badges -->
+      <div class="flex items-center gap-2">
+        <span
+          class="inline-flex items-center px-1.5 py-0.5 text-[10px] font-mono tracking-wider uppercase border border-white/10 text-white/50"
+        >
+          {metaInstance.sharding}
+        </span>
+        <span
+          class="inline-flex items-center px-1.5 py-0.5 text-[10px] font-mono tracking-wider uppercase border border-white/10 text-white/50"
+        >
+          {metaInstance.instanceMeta}
+        </span>
+        {#if metaInstance.minNodes > 1}
+          <span
+            class="inline-flex items-center px-1.5 py-0.5 text-[10px] font-mono tracking-wider uppercase border border-white/10 text-white/50"
+          >
+            {metaInstance.minNodes}+ nodes
+          </span>
+        {/if}
+      </div>
+
+      <!-- Node Assignments (when active) -->
+      {#if backingNodeIds.length > 0}
+        <div class="flex items-center gap-1.5 mt-1">
+          <svg
+            class="w-3 h-3 text-green-400/70 flex-shrink-0"
+            viewBox="0 0 24 24"
+            fill="none"
+            stroke="currentColor"
+            stroke-width="2"
+          >
+            <path
+              d="M22 12h-4l-3 9L9 3l-3 9H2"
+              stroke-linecap="round"
+              stroke-linejoin="round"
+            />
+          </svg>
+          <span class="text-white/60 text-xs font-mono truncate">
+            {backingNodeIds.map((id) => getNodeName(id)).join(", ")}
+          </span>
+        </div>
+      {/if}
+
+      <!-- Pinned nodes constraint -->
+      {#if metaInstance.nodeIds && metaInstance.nodeIds.length > 0}
+        <div class="flex items-center gap-1.5">
+          <svg
+            class="w-3 h-3 text-white/40 flex-shrink-0"
+            viewBox="0 0 24 24"
+            fill="none"
+            stroke="currentColor"
+            stroke-width="2"
+          >
+            <rect x="3" y="11" width="18" height="11" rx="2" ry="2" />
+            <path d="M7 11V7a5 5 0 0 1 10 0v4" />
+          </svg>
+          <span class="text-white/40 text-[11px] font-mono">
+            Pinned: {metaInstance.nodeIds
+              .map((id) => getNodeName(id))
+              .join(", ")}
+          </span>
+        </div>
+      {/if}
+
+      <!-- Error details -->
+      {#if metaInstance.placementError}
+        <div
+          class="mt-1.5 p-2 bg-red-500/5 border border-red-500/15 rounded-sm"
+        >
+          <div class="text-red-400 text-[11px] font-mono leading-relaxed">
+            {metaInstance.placementError}
+          </div>
+        </div>
+      {/if}
+
+      <!-- Retry counter -->
+      {#if metaInstance.consecutiveFailures > 0}
+        <div class="flex items-center gap-1.5 mt-1">
+          <svg
+            class="w-3 h-3 text-yellow-500/60 flex-shrink-0"
+            viewBox="0 0 24 24"
+            fill="none"
+            stroke="currentColor"
+            stroke-width="2"
+          >
+            <polyline points="23 4 23 10 17 10" />
+            <path d="M20.49 15a9 9 0 1 1-2.12-9.36L23 10" />
+          </svg>
+          <span class="text-yellow-500/60 text-[11px] font-mono">
+            {metaInstance.consecutiveFailures} consecutive
+            failure{metaInstance.consecutiveFailures !== 1 ? "s" : ""}
+          </span>
+        </div>
+      {/if}
+    </div>
+  </div>
+</div>

dashboard/src/lib/components/PrefillProgressBar.svelte

@@ -0,0 +1,52 @@
+<script lang="ts">
+  import type { PrefillProgress } from "$lib/stores/app.svelte";
+
+  interface Props {
+    progress: PrefillProgress;
+    class?: string;
+  }
+
+  let { progress, class: className = "" }: Props = $props();
+
+  const percentage = $derived(
+    progress.total > 0
+      ? Math.round((progress.processed / progress.total) * 100)
+      : 0,
+  );
+
+  function formatTokenCount(count: number | undefined): string {
+    if (count == null) return "0";
+    if (count >= 1000) {
+      return `${(count / 1000).toFixed(1)}k`;
+    }
+    return count.toString();
+  }
+</script>
+
+<div class="prefill-progress {className}">
+  <div
+    class="flex items-center justify-between text-xs text-exo-light-gray mb-1"
+  >
+    <span>Processing prompt</span>
+    <span class="font-mono">
+      {formatTokenCount(progress.processed)} / {formatTokenCount(
+        progress.total,
+      )} tokens
+    </span>
+  </div>
+  <div class="h-1.5 bg-exo-black/60 rounded-full overflow-hidden">
+    <div
+      class="h-full bg-exo-yellow rounded-full transition-all duration-150 ease-out"
+      style="width: {percentage}%"
+    ></div>
+  </div>
+  <div class="text-right text-xs text-exo-light-gray/70 mt-0.5 font-mono">
+    {percentage}%
+  </div>
+</div>
+
+<style>
+  .prefill-progress {
+    width: 100%;
+  }
+</style>

dashboard/src/lib/components/index.ts

@@ -11,4 +11,5 @@ export { default as FamilySidebar } from "./FamilySidebar.svelte";
 export { default as HuggingFaceResultItem } from "./HuggingFaceResultItem.svelte";
 export { default as ModelFilterPopover } from "./ModelFilterPopover.svelte";
 export { default as ModelPickerGroup } from "./ModelPickerGroup.svelte";
+export { default as MetaInstanceCard } from "./MetaInstanceCard.svelte";
 export { default as ModelPickerModal } from "./ModelPickerModal.svelte";

dashboard/src/lib/stores/app.svelte.ts

@@ -72,8 +72,23 @@ export interface Instance {
     runnerToShard?: Record<string, unknown>;
     nodeToRunner?: Record<string, string>;
   };
+  metaInstanceId?: string | null;
 }
 
+export interface MetaInstance {
+  metaInstanceId: string;
+  modelId: string;
+  sharding: "Pipeline" | "Tensor";
+  instanceMeta: "MlxRing" | "MlxJaccl";
+  minNodes: number;
+  nodeIds: string[] | null;
+  placementError: string | null;
+  consecutiveFailures: number;
+  lastFailureError: string | null;
+}
+
+export type MetaInstanceStatus = "active" | "provisioning" | "error";
+
 // Granular node state types from the new state structure
 interface RawNodeIdentity {
   modelId?: string;
@@ -223,6 +238,7 @@ interface RawStateResponse {
       MlxJacclInstance?: Instance;
     }
   >;
+  metaInstances?: Record<string, MetaInstance>;
   runners?: Record<string, unknown>;
   downloads?: Record<string, unknown[]>;
   // New granular node state fields
@@ -273,6 +289,11 @@ export interface TokenData {
   topLogprobs: TopLogprob[];
 }
 
+export interface PrefillProgress {
+  processed: number;
+  total: number;
+}
+
 export interface Message {
   id: string;
   role: "user" | "assistant" | "system";
@@ -520,10 +541,15 @@ class AppStore {
   ttftMs = $state<number | null>(null); // Time to first token in ms
   tps = $state<number | null>(null); // Tokens per second
   totalTokens = $state<number>(0); // Total tokens in current response
+  prefillProgress = $state<PrefillProgress | null>(null);
+
+  // Abort controller for stopping generation
+  private currentAbortController: AbortController | null = null;
 
   // Topology state
   topologyData = $state<TopologyData | null>(null);
   instances = $state<Record<string, unknown>>({});
+  metaInstances = $state<Record<string, MetaInstance>>({});
   runners = $state<Record<string, unknown>>({});
   downloads = $state<Record<string, unknown[]>>({});
   nodeDisk = $state<
@@ -1259,6 +1285,9 @@ class AppStore {
         this.instances = data.instances;
         this.refreshConversationModelFromInstances();
       }
+      if (data.metaInstances) {
+        this.metaInstances = data.metaInstances;
+      }
       if (data.runners) {
         this.runners = data.runners;
       }
@@ -1284,6 +1313,79 @@ class AppStore {
     }
   }
 
+  async createMetaInstance(
+    modelId: string,
+    sharding: "Pipeline" | "Tensor" = "Pipeline",
+    instanceMeta: "MlxRing" | "MlxJaccl" = "MlxRing",
+    minNodes: number = 1,
+    nodeIds: string[] | null = null,
+  ) {
+    try {
+      const response = await fetch("/meta_instance", {
+        method: "POST",
+        headers: { "Content-Type": "application/json" },
+        body: JSON.stringify({
+          model_id: modelId,
+          sharding,
+          instance_meta: instanceMeta,
+          min_nodes: minNodes,
+          node_ids: nodeIds,
+        }),
+      });
+      if (!response.ok) {
+        console.error("Failed to create meta-instance:", response.status);
+      }
+      await this.fetchState();
+    } catch (error) {
+      console.error("Error creating meta-instance:", error);
+    }
+  }
+
+  async deleteMetaInstance(metaInstanceId: string) {
+    try {
+      const response = await fetch(`/meta_instance/${metaInstanceId}`, {
+        method: "DELETE",
+        headers: { "Content-Type": "application/json" },
+      });
+      if (!response.ok) {
+        console.error("Failed to delete meta-instance:", response.status);
+      }
+      await this.fetchState();
+    } catch (error) {
+      console.error("Error deleting meta-instance:", error);
+    }
+  }
+
+  getMetaInstanceStatus(
+    metaInstance: MetaInstance,
+  ): MetaInstanceStatus {
+    // Check if any running instance is bound to this meta-instance
+    for (const instanceWrapper of Object.values(this.instances)) {
+      if (!instanceWrapper || typeof instanceWrapper !== "object") continue;
+      const keys = Object.keys(instanceWrapper as Record<string, unknown>);
+      if (keys.length !== 1) continue;
+      const inner = (instanceWrapper as Record<string, unknown>)[keys[0]];
+      if (inner && typeof inner === "object" && (inner as Instance).metaInstanceId === metaInstance.metaInstanceId) {
+        return "active";
+      }
+    }
+    if (metaInstance.placementError) return "error";
+    return "provisioning";
+  }
+
+  getMetaInstanceBackingNodes(metaInstance: MetaInstance): string[] {
+    for (const instanceWrapper of Object.values(this.instances)) {
+      if (!instanceWrapper || typeof instanceWrapper !== "object") continue;
+      const keys = Object.keys(instanceWrapper as Record<string, unknown>);
+      if (keys.length !== 1) continue;
+      const inner = (instanceWrapper as Record<string, unknown>)[keys[0]] as Instance;
+      if (inner?.metaInstanceId === metaInstance.metaInstanceId && inner?.shardAssignments?.nodeToRunner) {
+        return Object.keys(inner.shardAssignments.nodeToRunner);
+      }
+    }
+    return [];
+  }
+
   async fetchPlacementPreviews(modelId: string, showLoading = true) {
     if (!modelId) return;
 
@@ -1643,11 +1745,12 @@ class AppStore {
       if (!reader) throw new Error("No response body");
 
       let fullContent = prefixText;
+      let streamedThinking = "";
       const collectedTokens: TokenData[] = [...tokensToKeep];
 
       interface ChatCompletionChunk {
         choices?: Array<{
-          delta?: { content?: string };
+          delta?: { content?: string; reasoning_content?: string };
           logprobs?: {
             content?: Array<{
               token: string;
@@ -1668,6 +1771,7 @@ class AppStore {
         (parsed) => {
           const choice = parsed.choices?.[0];
           const delta = choice?.delta?.content;
+          const thinkingDelta = choice?.delta?.reasoning_content;
 
           // Collect logprobs data
           const logprobsContent = choice?.logprobs?.content;
@@ -1686,7 +1790,11 @@ class AppStore {
             }
           }
 
-          if (delta) {
+          if (thinkingDelta) {
+            streamedThinking += thinkingDelta;
+          }
+
+          if (delta || thinkingDelta) {
             if (firstTokenTime === null) {
               firstTokenTime = performance.now();
               this.ttftMs = firstTokenTime - requestStartTime;
@@ -1700,9 +1808,14 @@ class AppStore {
               this.tps = ((tokenCount - tokensToKeep.length) / elapsed) * 1000;
             }
 
-            fullContent += delta;
-            const { displayContent, thinkingContent } =
+            if (delta) {
+              fullContent += delta;
+            }
+            const { displayContent, thinkingContent: tagThinking } =
               this.stripThinkingTags(fullContent);
+            const combinedThinking = [streamedThinking, tagThinking]
+              .filter(Boolean)
+              .join("\n\n");
 
             if (this.activeConversationId === targetConversationId) {
               this.currentResponse = displayContent;
@@ -1714,7 +1827,7 @@ class AppStore {
               messageId,
               (m) => {
                 m.content = displayContent;
-                m.thinking = thinkingContent || undefined;
+                m.thinking = combinedThinking || undefined;
                 m.tokens = [...collectedTokens];
               },
             );
@@ -1726,11 +1839,14 @@ class AppStore {
 
       // Final update
       if (this.conversationExists(targetConversationId)) {
-        const { displayContent, thinkingContent } =
+        const { displayContent, thinkingContent: tagThinking } =
           this.stripThinkingTags(fullContent);
+        const finalThinking = [streamedThinking, tagThinking]
+          .filter(Boolean)
+          .join("\n\n");
         this.updateConversationMessage(targetConversationId, messageId, (m) => {
           m.content = displayContent;
-          m.thinking = thinkingContent || undefined;
+          m.thinking = finalThinking || undefined;
           m.tokens = [...collectedTokens];
           if (this.ttftMs !== null) m.ttftMs = this.ttftMs;
           if (this.tps !== null) m.tps = this.tps;
@@ -1838,11 +1954,12 @@ class AppStore {
       }
 
       let streamedContent = "";
+      let streamedThinking = "";
       const collectedTokens: TokenData[] = [];
 
       interface ChatCompletionChunk {
         choices?: Array<{
-          delta?: { content?: string };
+          delta?: { content?: string; reasoning_content?: string };
           logprobs?: {
             content?: Array<{
               token: string;
@@ -1863,6 +1980,7 @@ class AppStore {
         (parsed) => {
           const choice = parsed.choices?.[0];
           const delta = choice?.delta?.content;
+          const thinkingDelta = choice?.delta?.reasoning_content;
 
           // Collect logprobs data
           const logprobsContent = choice?.logprobs?.content;
@@ -1881,10 +1999,19 @@ class AppStore {
             }
           }
 
-          if (delta) {
-            streamedContent += delta;
-            const { displayContent, thinkingContent } =
+          if (thinkingDelta) {
+            streamedThinking += thinkingDelta;
+          }
+
+          if (delta || thinkingDelta) {
+            if (delta) {
+              streamedContent += delta;
+            }
+            const { displayContent, thinkingContent: tagThinking } =
               this.stripThinkingTags(streamedContent);
+            const combinedThinking = [streamedThinking, tagThinking]
+              .filter(Boolean)
+              .join("\n\n");
 
             // Only update currentResponse if target conversation is active
             if (this.activeConversationId === targetConversationId) {
@@ -1897,7 +2024,7 @@ class AppStore {
               assistantMessage.id,
               (msg) => {
                 msg.content = displayContent;
-                msg.thinking = thinkingContent || undefined;
+                msg.thinking = combinedThinking || undefined;
                 msg.tokens = [...collectedTokens];
               },
             );
@@ -1909,14 +2036,17 @@ class AppStore {
 
       // Final cleanup of the message (if conversation still exists)
       if (this.conversationExists(targetConversationId)) {
-        const { displayContent, thinkingContent } =
+        const { displayContent, thinkingContent: tagThinking } =
           this.stripThinkingTags(streamedContent);
+        const finalThinking = [streamedThinking, tagThinking]
+          .filter(Boolean)
+          .join("\n\n");
         this.updateConversationMessage(
           targetConversationId,
           assistantMessage.id,
           (msg) => {
             msg.content = displayContent;
-            msg.thinking = thinkingContent || undefined;
+            msg.thinking = finalThinking || undefined;
             msg.tokens = [...collectedTokens];
           },
         );
@@ -2005,6 +2135,7 @@ class AppStore {
     reader: ReadableStreamDefaultReader<Uint8Array>,
     targetConversationId: string,
     onChunk: (parsed: T) => void,
+    onEvent?: Record<string, (data: unknown) => void>,
   ): Promise<void> {
     const decoder = new TextDecoder();
     let buffer = "";
@@ -2025,6 +2156,24 @@ class AppStore {
         const trimmed = line.trim();
         if (!trimmed) continue;
 
+        // Handle SSE comments (": key json") for prefill progress etc.
+        if (trimmed.startsWith(": ") && onEvent) {
+          const comment = trimmed.slice(2);
+          const spaceIdx = comment.indexOf(" ");
+          if (spaceIdx > 0) {
+            const key = comment.slice(0, spaceIdx);
+            if (onEvent[key]) {
+              try {
+                const parsed = JSON.parse(comment.slice(spaceIdx + 1));
+                onEvent[key](parsed);
+              } catch {
+                // Skip malformed JSON in comment
+              }
+            }
+          }
+          continue;
+        }
+
         if (trimmed.startsWith("data: ")) {
           const data = trimmed.slice(6);
           if (data === "[DONE]") continue;
@@ -2256,6 +2405,9 @@ class AppStore {
       let firstTokenTime: number | null = null;
       let tokenCount = 0;
 
+      const abortController = new AbortController();
+      this.currentAbortController = abortController;
+
       const response = await fetch("/v1/chat/completions", {
         method: "POST",
         headers: {
@@ -2272,6 +2424,7 @@ class AppStore {
             enable_thinking: enableThinking,
           }),
         }),
+        signal: abortController.signal,
       });
 
       if (!response.ok) {
@@ -2285,10 +2438,11 @@ class AppStore {
       }
 
       let streamedContent = "";
+      let streamedThinking = "";
 
       interface ChatCompletionChunk {
         choices?: Array<{
-          delta?: { content?: string };
+          delta?: { content?: string; reasoning_content?: string };
           logprobs?: {
             content?: Array<{
               token: string;
@@ -2309,8 +2463,14 @@ class AppStore {
         reader,
         targetConversationId,
         (parsed) => {
+          // Clear prefill progress when first token data arrives
+          if (this.prefillProgress) {
+            this.prefillProgress = null;
+          }
+
           const choice = parsed.choices?.[0];
           const tokenContent = choice?.delta?.content;
+          const thinkingContent = choice?.delta?.reasoning_content;
 
           // Collect logprobs data
           const logprobsContent = choice?.logprobs?.content;
@@ -2329,7 +2489,11 @@ class AppStore {
             }
           }
 
-          if (tokenContent) {
+          if (thinkingContent) {
+            streamedThinking += thinkingContent;
+          }
+
+          if (tokenContent || thinkingContent) {
             // Track first token for TTFT
             if (firstTokenTime === null) {
               firstTokenTime = performance.now();
@@ -2346,11 +2510,16 @@ class AppStore {
               this.tps = (tokenCount / elapsed) * 1000;
             }
 
-            streamedContent += tokenContent;
+            if (tokenContent) {
+              streamedContent += tokenContent;
+            }
 
-            // Strip thinking tags for display and extract thinking content
-            const { displayContent, thinkingContent } =
+            // Use stripThinkingTags as fallback for any <think> tags still in content
+            const { displayContent, thinkingContent: tagThinking } =
               this.stripThinkingTags(streamedContent);
+            const combinedThinking = [streamedThinking, tagThinking]
+              .filter(Boolean)
+              .join("\n\n");
 
             // Only update currentResponse if target conversation is active
             if (this.activeConversationId === targetConversationId) {
@@ -2363,7 +2532,7 @@ class AppStore {
               assistantMessage.id,
               (msg) => {
                 msg.content = displayContent;
-                msg.thinking = thinkingContent || undefined;
+                msg.thinking = combinedThinking || undefined;
                 msg.tokens = [...collectedTokens];
               },
             );
@@ -2371,8 +2540,26 @@ class AppStore {
             this.persistConversation(targetConversationId);
           }
         },
+        {
+          prefill_progress: (data) => {
+            // TaggedModel wraps as {"PrefillProgressChunk": {...}}
+            // model_dump_json() uses snake_case (by_alias defaults to False)
+            const raw = data as Record<string, unknown>;
+            const inner = (raw["PrefillProgressChunk"] ?? raw) as {
+              processed_tokens: number;
+              total_tokens: number;
+            };
+            this.prefillProgress = {
+              processed: inner.processed_tokens,
+              total: inner.total_tokens,
+            };
+          },
+        },
       );
 
+      // Clear prefill progress after stream ends
+      this.prefillProgress = null;
+
       // Calculate final TPS
       if (firstTokenTime !== null && tokenCount > 1) {
         const totalGenerationTime = performance.now() - firstTokenTime;
@@ -2381,14 +2568,17 @@ class AppStore {
 
       // Final cleanup of the message (if conversation still exists)
       if (this.conversationExists(targetConversationId)) {
-        const { displayContent, thinkingContent } =
+        const { displayContent, thinkingContent: tagThinking } =
           this.stripThinkingTags(streamedContent);
+        const finalThinking = [streamedThinking, tagThinking]
+          .filter(Boolean)
+          .join("\n\n");
         this.updateConversationMessage(
           targetConversationId,
           assistantMessage.id,
           (msg) => {
             msg.content = displayContent;
-            msg.thinking = thinkingContent || undefined;
+            msg.thinking = finalThinking || undefined;
             msg.tokens = [...collectedTokens];
             // Store performance metrics on the message
             if (this.ttftMs !== null) {
@@ -2403,20 +2593,31 @@ class AppStore {
         this.persistConversation(targetConversationId);
       }
     } catch (error) {
-      console.error("Error sending message:", error);
-      this.handleStreamingError(
-        error,
-        targetConversationId,
-        assistantMessage.id,
-        "Failed to get response",
-      );
+      if (error instanceof DOMException && error.name === "AbortError") {
+        // User stopped generation — not an error
+      } else {
+        console.error("Error sending message:", error);
+        this.handleStreamingError(
+          error,
+          targetConversationId,
+          assistantMessage.id,
+          "Failed to get response",
+        );
+      }
     } finally {
+      this.currentAbortController = null;
+      this.prefillProgress = null;
       this.isLoading = false;
       this.currentResponse = "";
       this.saveConversationsToStorage();
     }
   }
 
+  stopGeneration(): void {
+    this.currentAbortController?.abort();
+    this.currentAbortController = null;
+  }
+
   /**
    * Generate an image using the image generation API
    */
@@ -3043,8 +3244,10 @@ export const isLoading = () => appStore.isLoading;
 export const ttftMs = () => appStore.ttftMs;
 export const tps = () => appStore.tps;
 export const totalTokens = () => appStore.totalTokens;
+export const prefillProgress = () => appStore.prefillProgress;
 export const topologyData = () => appStore.topologyData;
 export const instances = () => appStore.instances;
+export const metaInstances = () => appStore.metaInstances;
 export const runners = () => appStore.runners;
 export const downloads = () => appStore.downloads;
 export const nodeDisk = () => appStore.nodeDisk;
@@ -3060,6 +3263,7 @@ export const topologyOnlyMode = () => appStore.getTopologyOnlyMode();
 export const chatSidebarVisible = () => appStore.getChatSidebarVisible();
 
 // Actions
+export const stopGeneration = () => appStore.stopGeneration();
 export const startChat = () => appStore.startChat();
 export const sendMessage = (
   content: string,
@@ -3132,6 +3336,21 @@ export const setChatSidebarVisible = (visible: boolean) =>
   appStore.setChatSidebarVisible(visible);
 export const refreshState = () => appStore.fetchState();
 
+// Meta-instance actions
+export const createMetaInstance = (
+  modelId: string,
+  sharding?: "Pipeline" | "Tensor",
+  instanceMeta?: "MlxRing" | "MlxJaccl",
+  minNodes?: number,
+  nodeIds?: string[] | null,
+) => appStore.createMetaInstance(modelId, sharding, instanceMeta, minNodes, nodeIds);
+export const deleteMetaInstance = (metaInstanceId: string) =>
+  appStore.deleteMetaInstance(metaInstanceId);
+export const getMetaInstanceStatus = (metaInstance: MetaInstance) =>
+  appStore.getMetaInstanceStatus(metaInstance);
+export const getMetaInstanceBackingNodes = (metaInstance: MetaInstance) =>
+  appStore.getMetaInstanceBackingNodes(metaInstance);
+
 // Node identities (for OS version mismatch detection)
 export const nodeIdentities = () => appStore.nodeIdentities;
 

dashboard/src/routes/+page.svelte

@@ -47,10 +47,14 @@
     thunderboltBridgeCycles,
     nodeThunderboltBridge,
     nodeIdentities,
+    metaInstances,
+    deleteMetaInstance,
     type DownloadProgress,
     type PlacementPreview,
+    type MetaInstance,
   } from "$lib/stores/app.svelte";
   import HeaderNav from "$lib/components/HeaderNav.svelte";
+  import MetaInstanceCard from "$lib/components/MetaInstanceCard.svelte";
   import { fade, fly } from "svelte/transition";
   import { cubicInOut } from "svelte/easing";
   import { onMount } from "svelte";
@@ -67,6 +71,8 @@
   const loadingPreviews = $derived(isLoadingPreviews());
   const debugEnabled = $derived(debugMode());
   const topologyOnlyEnabled = $derived(topologyOnlyMode());
+  const metaInstanceData = $derived(metaInstances());
+  const metaInstanceCount = $derived(Object.keys(metaInstanceData).length);
   const sidebarVisible = $derived(chatSidebarVisible());
   const tbBridgeCycles = $derived(thunderboltBridgeCycles());
   const tbBridgeData = $derived(nodeThunderboltBridge());
@@ -114,6 +120,74 @@
   });
   let tb5InfoDismissed = $state(false);
 
+  // Detect Mac Studio nodes using RDMA on en2 (the port next to ethernet — RDMA doesn't work there)
+  const macStudioEn2RdmaWarning = $derived.by(() => {
+    const edges = data?.edges;
+    const ids = tbIdentifiers;
+    const rdmaCtl = rdmaCtlData;
+    if (!edges || !ids || !rdmaCtl) return null;
+
+    const affectedConnections: Array<{
+      nodeId: string;
+      nodeName: string;
+      peerNodeId: string;
+      peerNodeName: string;
+      rdmaIface: string;
+    }> = [];
+
+    const isMacStudio = (node: (typeof data.nodes)[string] | undefined) =>
+      node?.system_info?.model_id === "Mac Studio";
+
+    for (const edge of edges) {
+      if (!edge.sourceRdmaIface && !edge.sinkRdmaIface) continue;
+
+      const sourceNode = data?.nodes?.[edge.source];
+      if (
+        isMacStudio(sourceNode) &&
+        edge.sourceRdmaIface === "rdma_en2" &&
+        rdmaCtl[edge.source]?.enabled
+      ) {
+        affectedConnections.push({
+          nodeId: edge.source,
+          nodeName:
+            sourceNode?.friendly_name || edge.source.slice(0, 8) + "...",
+          peerNodeId: edge.target,
+          peerNodeName:
+            data?.nodes?.[edge.target]?.friendly_name ||
+            edge.target.slice(0, 8) + "...",
+          rdmaIface: "en2",
+        });
+      }
+
+      const sinkNode = data?.nodes?.[edge.target];
+      if (
+        isMacStudio(sinkNode) &&
+        edge.sinkRdmaIface === "rdma_en2" &&
+        rdmaCtl[edge.target]?.enabled
+      ) {
+        affectedConnections.push({
+          nodeId: edge.target,
+          nodeName: sinkNode?.friendly_name || edge.target.slice(0, 8) + "...",
+          peerNodeId: edge.source,
+          peerNodeName:
+            sourceNode?.friendly_name || edge.source.slice(0, 8) + "...",
+          rdmaIface: "en2",
+        });
+      }
+    }
+
+    // Deduplicate by nodeId
+    const seen = new Set<string>();
+    const unique = affectedConnections.filter((c) => {
+      if (seen.has(c.nodeId)) return false;
+      seen.add(c.nodeId);
+      return true;
+    });
+
+    return unique.length > 0 ? unique : null;
+  });
+  let macStudioEn2Dismissed = $state(false);
+
   // Helper to get friendly node name from node ID
   function getNodeName(nodeId: string): string {
     const node = data?.nodes?.[nodeId];
@@ -932,13 +1006,6 @@
     };
   }
 
-  // Debug: Log downloads data when it changes
-  $effect(() => {
-    if (downloadsData && Object.keys(downloadsData).length > 0) {
-      console.log("[Download Debug] Current downloads:", downloadsData);
-    }
-  });
-
   // Helper to get download status for an instance
   function getInstanceDownloadStatus(
     instanceId: string,
@@ -1765,7 +1832,7 @@
 </script>
 
 {#snippet clusterWarnings()}
-  {#if tbBridgeCycles.length > 0 || macosVersionMismatch || (tb5WithoutRdma && !tb5InfoDismissed)}
+  {#if tbBridgeCycles.length > 0 || macosVersionMismatch || (tb5WithoutRdma && !tb5InfoDismissed) || (macStudioEn2RdmaWarning && !macStudioEn2Dismissed)}
     <div class="absolute top-4 left-4 flex flex-col gap-2 z-40">
       {#if tbBridgeCycles.length > 0}
         {@const cycle = tbBridgeCycles[0]}
@@ -1930,12 +1997,260 @@
           </button>
         </div>
       {/if}
+
+      {#if macStudioEn2RdmaWarning && !macStudioEn2Dismissed}
+        <div class="group relative" role="alert">
+          <div
+            class="flex items-center gap-2 px-3 py-2 rounded border border-red-500/50 bg-red-500/10 backdrop-blur-sm cursor-help"
+          >
+            <svg
+              class="w-5 h-5 text-red-400 flex-shrink-0"
+              fill="none"
+              viewBox="0 0 24 24"
+              stroke="currentColor"
+              stroke-width="2"
+            >
+              <path
+                stroke-linecap="round"
+                stroke-linejoin="round"
+                d={warningIconPath}
+              />
+            </svg>
+            <span class="text-sm font-mono text-red-200">
+              RDMA INCOMPATIBLE PORT
+            </span>
+            <button
+              type="button"
+              onclick={() => (macStudioEn2Dismissed = true)}
+              class="ml-1 text-red-300/60 hover:text-red-200 transition-colors cursor-pointer"
+              title="Dismiss"
+            >
+              <svg
+                class="w-4 h-4"
+                fill="none"
+                viewBox="0 0 24 24"
+                stroke="currentColor"
+                stroke-width="2"
+              >
+                <path
+                  stroke-linecap="round"
+                  stroke-linejoin="round"
+                  d="M6 18L18 6M6 6l12 12"
+                />
+              </svg>
+            </button>
+          </div>
+
+          <!-- Expanded tooltip on hover -->
+          <div
+            class="absolute top-full left-0 mt-2 w-96 p-4 rounded border border-red-500/30 bg-[#1a1a1a]/95 backdrop-blur-sm opacity-0 invisible group-hover:opacity-100 group-hover:visible transition-all duration-200 z-50 shadow-lg"
+          >
+            <p class="text-xs text-white/80 mb-3">
+              The Thunderbolt 5 port next to the Ethernet port on Mac Studio
+              does
+              <span class="text-red-400 font-semibold">not support RDMA</span>.
+              Move the cable to one of the other three TB5 ports.
+            </p>
+
+            <div class="text-xs text-white/60 mb-3">
+              <span class="text-red-300">Affected:</span>
+              {#each macStudioEn2RdmaWarning as conn}
+                <div class="ml-2 mt-0.5">
+                  <span class="text-white/80">{conn.nodeName}</span>
+                  <span class="text-white/30">&rarr;</span>
+                  <span class="text-white/60">{conn.peerNodeName}</span>
+                  <span class="text-white/30 ml-1">(en2)</span>
+                </div>
+              {/each}
+            </div>
+
+            <!-- Mac Studio back panel illustration -->
+            <div class="bg-black/40 rounded p-3 mb-3">
+              <p
+                class="text-[10px] font-mono text-white/30 uppercase tracking-wider mb-2"
+              >
+                Mac Studio — Rear Panel
+              </p>
+              <svg
+                viewBox="0 0 320 72"
+                class="w-full"
+                xmlns="http://www.w3.org/2000/svg"
+              >
+                <rect
+                  x="1"
+                  y="1"
+                  width="318"
+                  height="70"
+                  rx="6"
+                  ry="6"
+                  fill="none"
+                  stroke="rgba(255,255,255,0.12)"
+                  stroke-width="1"
+                />
+                <!-- TB5 port 1 -->
+                <rect
+                  x="24"
+                  y="22"
+                  width="28"
+                  height="14"
+                  rx="4"
+                  fill="none"
+                  stroke="rgba(255,255,255,0.3)"
+                  stroke-width="1"
+                />
+                <text
+                  x="38"
+                  y="52"
+                  text-anchor="middle"
+                  fill="rgba(255,255,255,0.25)"
+                  style="font-size:7px;font-family:ui-monospace,monospace;"
+                  >TB5</text
+                >
+                <!-- TB5 port 2 -->
+                <rect
+                  x="62"
+                  y="22"
+                  width="28"
+                  height="14"
+                  rx="4"
+                  fill="none"
+                  stroke="rgba(255,255,255,0.3)"
+                  stroke-width="1"
+                />
+                <text
+                  x="76"
+                  y="52"
+                  text-anchor="middle"
+                  fill="rgba(255,255,255,0.25)"
+                  style="font-size:7px;font-family:ui-monospace,monospace;"
+                  >TB5</text
+                >
+                <!-- TB5 port 3 -->
+                <rect
+                  x="100"
+                  y="22"
+                  width="28"
+                  height="14"
+                  rx="4"
+                  fill="none"
+                  stroke="rgba(255,255,255,0.3)"
+                  stroke-width="1"
+                />
+                <text
+                  x="114"
+                  y="52"
+                  text-anchor="middle"
+                  fill="rgba(255,255,255,0.25)"
+                  style="font-size:7px;font-family:ui-monospace,monospace;"
+                  >TB5</text
+                >
+                <!-- TB5 port 4: INCOMPATIBLE (en2) — equally spaced with ports 1-3 -->
+                <rect
+                  x="138"
+                  y="22"
+                  width="28"
+                  height="14"
+                  rx="4"
+                  fill="rgba(239,68,68,0.1)"
+                  stroke="rgba(239,68,68,0.7)"
+                  stroke-width="1.5"
+                />
+                <line
+                  x1="142"
+                  y1="25"
+                  x2="162"
+                  y2="33"
+                  stroke="rgba(239,68,68,0.8)"
+                  stroke-width="1.5"
+                  stroke-linecap="round"
+                />
+                <line
+                  x1="162"
+                  y1="25"
+                  x2="142"
+                  y2="33"
+                  stroke="rgba(239,68,68,0.8)"
+                  stroke-width="1.5"
+                  stroke-linecap="round"
+                />
+                <text
+                  x="152"
+                  y="52"
+                  text-anchor="middle"
+                  fill="rgba(239,68,68,0.6)"
+                  style="font-size:7px;font-family:ui-monospace,monospace;font-weight:600;"
+                  >en2</text
+                >
+                <!-- Ethernet port -->
+                <rect
+                  x="196"
+                  y="19"
+                  width="24"
+                  height="20"
+                  rx="2"
+                  fill="none"
+                  stroke="rgba(255,255,255,0.2)"
+                  stroke-width="1"
+                />
+                <rect
+                  x="200"
+                  y="23"
+                  width="16"
+                  height="12"
+                  rx="1"
+                  fill="none"
+                  stroke="rgba(255,255,255,0.12)"
+                  stroke-width="0.75"
+                />
+                <text
+                  x="208"
+                  y="52"
+                  text-anchor="middle"
+                  fill="rgba(255,255,255,0.25)"
+                  style="font-size:7px;font-family:ui-monospace,monospace;"
+                  >ETH</text
+                >
+                <!-- Green checkmarks on working ports -->
+                <circle
+                  cx="38"
+                  cy="62"
+                  r="3"
+                  fill="none"
+                  stroke="rgba(74,222,128,0.5)"
+                  stroke-width="0.75"
+                />
+                <circle
+                  cx="76"
+                  cy="62"
+                  r="3"
+                  fill="none"
+                  stroke="rgba(74,222,128,0.5)"
+                  stroke-width="0.75"
+                />
+                <circle
+                  cx="114"
+                  cy="62"
+                  r="3"
+                  fill="none"
+                  stroke="rgba(74,222,128,0.5)"
+                  stroke-width="0.75"
+                />
+              </svg>
+            </div>
+
+            <p class="text-xs text-white/50">
+              <span class="text-green-400">Fix:</span> Move the Thunderbolt cable
+              to any of the three leftmost ports (all support RDMA).
+            </p>
+          </div>
+        </div>
+      {/if}
     </div>
   {/if}
 {/snippet}
 
 {#snippet clusterWarningsCompact()}
-  {#if tbBridgeCycles.length > 0 || macosVersionMismatch || (tb5WithoutRdma && !tb5InfoDismissed)}
+  {#if tbBridgeCycles.length > 0 || macosVersionMismatch || (tb5WithoutRdma && !tb5InfoDismissed) || (macStudioEn2RdmaWarning && !macStudioEn2Dismissed)}
     <div class="absolute top-2 left-2 flex flex-col gap-1">
       {#if tbBridgeCycles.length > 0}
         <div
@@ -2003,6 +2318,27 @@
           >
         </div>
       {/if}
+      {#if macStudioEn2RdmaWarning && !macStudioEn2Dismissed}
+        <div
+          class="flex items-center gap-1.5 px-2 py-1 rounded border border-red-500/50 bg-red-500/10 backdrop-blur-sm"
+          title="Mac Studio RDMA incompatible port (en2) — move cable to another TB5 port"
+        >
+          <svg
+            class="w-3.5 h-3.5 text-red-400"
+            fill="none"
+            viewBox="0 0 24 24"
+            stroke="currentColor"
+            stroke-width="2"
+          >
+            <path
+              stroke-linecap="round"
+              stroke-linejoin="round"
+              d={warningIconPath}
+            />
+          </svg>
+          <span class="text-[10px] font-mono text-red-200">BAD RDMA PORT</span>
+        </div>
+      {/if}
     </div>
   {/if}
 {/snippet}
@@ -2726,6 +3062,39 @@
             </div>
           {/if}
 
+          <!-- Meta-Instances Panel -->
+          {#if metaInstanceCount > 0}
+            <div class="p-4 flex-shrink-0 border-t border-exo-yellow/10">
+              <!-- Panel Header -->
+              <div class="flex items-center gap-2 mb-4">
+                <div
+                  class="w-2 h-2 border border-purple-400/60 rotate-45"
+                ></div>
+                <h3
+                  class="text-xs text-purple-400 font-mono tracking-[0.2em] uppercase"
+                >
+                  Meta-Instances
+                </h3>
+                <div
+                  class="flex-1 h-px bg-gradient-to-r from-purple-400/30 to-transparent"
+                ></div>
+                <span class="text-[10px] text-white/40 font-mono"
+                  >{metaInstanceCount}</span
+                >
+              </div>
+
+              <div class="space-y-3">
+                {#each Object.entries(metaInstanceData) as [id, mi]}
+                  <MetaInstanceCard
+                    metaInstance={mi}
+                    onDelete={(metaInstanceId) =>
+                      deleteMetaInstance(metaInstanceId)}
+                  />
+                {/each}
+              </div>
+            </div>
+          {/if}
+
           <!-- Models Panel - Scrollable -->
           <div class="p-4 flex-1 overflow-y-auto">
             <!-- Panel Header -->
@@ -3548,6 +3917,34 @@
                 </div>
               </div>
             {/if}
+
+            <!-- Meta-Instances Section (chat sidebar) -->
+            {#if metaInstanceCount > 0}
+              <div class="p-4 border-t border-exo-yellow/10">
+                <div class="flex items-center gap-2 mb-4">
+                  <div
+                    class="w-2 h-2 border border-purple-400/60 rotate-45"
+                  ></div>
+                  <h3
+                    class="text-xs text-purple-400 font-mono tracking-[0.2em] uppercase"
+                  >
+                    Meta-Instances
+                  </h3>
+                  <div
+                    class="flex-1 h-px bg-gradient-to-r from-purple-400/30 to-transparent"
+                  ></div>
+                </div>
+                <div class="space-y-3">
+                  {#each Object.entries(metaInstanceData) as [id, mi]}
+                    <MetaInstanceCard
+                      metaInstance={mi}
+                      onDelete={(metaInstanceId) =>
+                        deleteMetaInstance(metaInstanceId)}
+                    />
+                  {/each}
+                </div>
+              </div>
+            {/if}
           </aside>
         {/if}
       </div>

flake.nix

@@ -74,7 +74,6 @@
       perSystem =
         { config, self', inputs', pkgs, lib, system, ... }:
         let
-          fenixToolchain = inputs'.fenix.packages.complete;
           # Use pinned nixpkgs for swift-format (swift is broken on x86_64-linux in newer nixpkgs)
           pkgsSwift = import inputs.nixpkgs-swift { inherit system; };
         in

python/parts.nix

@@ -158,6 +158,7 @@
           exo-test-env = testVenv;
         } // {
         exo-bench = mkBenchScript "exo-bench" (inputs.self + /bench/exo_bench.py);
+        exo-eval-tool-calls = mkBenchScript "exo-eval-tool-calls" (inputs.self + /bench/eval_tool_calls.py);
         exo-get-all-models-on-cluster = mkSimplePythonScript "exo-get-all-models-on-cluster" (inputs.self + /tests/get_all_models_on_cluster.py);
       };
 

resources/inference_model_cards/mlx-community--DeepSeek-V3.1-4bit.toml

@@ -6,7 +6,7 @@ tasks = ["TextGeneration"]
 family = "deepseek"
 quantization = "4bit"
 base_model = "DeepSeek V3.1"
-capabilities = ["text", "thinking"]
+capabilities = ["text", "thinking", "thinking_toggle"]
 
 [storage_size]
 in_bytes = 405874409472

resources/inference_model_cards/mlx-community--DeepSeek-V3.1-8bit.toml

@@ -6,7 +6,7 @@ tasks = ["TextGeneration"]
 family = "deepseek"
 quantization = "8bit"
 base_model = "DeepSeek V3.1"
-capabilities = ["text", "thinking"]
+capabilities = ["text", "thinking", "thinking_toggle"]
 
 [storage_size]
 in_bytes = 765577920512

resources/inference_model_cards/mlx-community--GLM-4.5-Air-8bit.toml

@@ -6,7 +6,7 @@ tasks = ["TextGeneration"]
 family = "glm"
 quantization = "8bit"
 base_model = "GLM 4.5 Air"
-capabilities = ["text", "thinking"]
+capabilities = ["text", "thinking", "thinking_toggle"]
 
 [storage_size]
 in_bytes = 122406567936

resources/inference_model_cards/mlx-community--GLM-4.5-Air-bf16.toml

@@ -6,7 +6,7 @@ tasks = ["TextGeneration"]
 family = "glm"
 quantization = "bf16"
 base_model = "GLM 4.5 Air"
-capabilities = ["text", "thinking"]
+capabilities = ["text", "thinking", "thinking_toggle"]
 
 [storage_size]
 in_bytes = 229780750336

resources/inference_model_cards/mlx-community--GLM-4.7-4bit.toml

@@ -6,7 +6,7 @@ tasks = ["TextGeneration"]
 family = "glm"
 quantization = "4bit"
 base_model = "GLM 4.7"
-capabilities = ["text", "thinking"]
+capabilities = ["text", "thinking", "thinking_toggle"]
 
 [storage_size]
 in_bytes = 198556925568

resources/inference_model_cards/mlx-community--GLM-4.7-6bit.toml

@@ -6,7 +6,7 @@ tasks = ["TextGeneration"]
 family = "glm"
 quantization = "6bit"
 base_model = "GLM 4.7"
-capabilities = ["text", "thinking"]
+capabilities = ["text", "thinking", "thinking_toggle"]
 
 [storage_size]
 in_bytes = 286737579648

resources/inference_model_cards/mlx-community--GLM-4.7-8bit-gs32.toml

@@ -6,7 +6,7 @@ tasks = ["TextGeneration"]
 family = "glm"
 quantization = "8bit"
 base_model = "GLM 4.7"
-capabilities = ["text", "thinking"]
+capabilities = ["text", "thinking", "thinking_toggle"]
 
 [storage_size]
 in_bytes = 396963397248

resources/inference_model_cards/mlx-community--GLM-4.7-Flash-4bit.toml

@@ -6,7 +6,7 @@ tasks = ["TextGeneration"]
 family = "glm"
 quantization = "4bit"
 base_model = "GLM 4.7 Flash"
-capabilities = ["text", "thinking"]
+capabilities = ["text", "thinking", "thinking_toggle"]
 
 [storage_size]
 in_bytes = 19327352832

resources/inference_model_cards/mlx-community--GLM-4.7-Flash-5bit.toml

@@ -6,7 +6,7 @@ tasks = ["TextGeneration"]
 family = "glm"
 quantization = "5bit"
 base_model = "GLM 4.7 Flash"
-capabilities = ["text", "thinking"]
+capabilities = ["text", "thinking", "thinking_toggle"]
 
 [storage_size]
 in_bytes = 22548578304

resources/inference_model_cards/mlx-community--GLM-4.7-Flash-6bit.toml

@@ -6,7 +6,7 @@ tasks = ["TextGeneration"]
 family = "glm"
 quantization = "6bit"
 base_model = "GLM 4.7 Flash"
-capabilities = ["text", "thinking"]
+capabilities = ["text", "thinking", "thinking_toggle"]
 
 [storage_size]
 in_bytes = 26843545600

resources/inference_model_cards/mlx-community--GLM-4.7-Flash-8bit.toml

@@ -6,7 +6,7 @@ tasks = ["TextGeneration"]
 family = "glm"
 quantization = "8bit"
 base_model = "GLM 4.7 Flash"
-capabilities = ["text", "thinking"]
+capabilities = ["text", "thinking", "thinking_toggle"]
 
 [storage_size]
 in_bytes = 34359738368

resources/inference_model_cards/mlx-community--Kimi-K2-Thinking.toml

@@ -6,7 +6,7 @@ tasks = ["TextGeneration"]
 family = "kimi"
 quantization = ""
 base_model = "Kimi K2"
-capabilities = ["text", "thinking"]
+capabilities = ["text", "thinking", "thinking_toggle"]
 
 [storage_size]
 in_bytes = 706522120192

resources/inference_model_cards/mlx-community--Kimi-K2.5.toml

@@ -6,7 +6,7 @@ tasks = ["TextGeneration"]
 family = "kimi"
 quantization = ""
 base_model = "Kimi K2.5"
-capabilities = ["text", "thinking"]
+capabilities = ["text", "thinking", "thinking_toggle"]
 
 [storage_size]
 in_bytes = 662498705408

resources/inference_model_cards/mlx-community--MiniMax-M2.1-3bit.toml

@@ -6,7 +6,7 @@ tasks = ["TextGeneration"]
 family = "minimax"
 quantization = "3bit"
 base_model = "MiniMax M2.1"
-capabilities = ["text", "thinking"]
+capabilities = ["text", "thinking", "thinking_toggle"]
 
 [storage_size]
 in_bytes = 100086644736

resources/inference_model_cards/mlx-community--MiniMax-M2.1-8bit.toml

@@ -6,7 +6,7 @@ tasks = ["TextGeneration"]
 family = "minimax"
 quantization = "8bit"
 base_model = "MiniMax M2.1"
-capabilities = ["text", "thinking"]
+capabilities = ["text", "thinking", "thinking_toggle"]
 
 [storage_size]
 in_bytes = 242986745856

resources/inference_model_cards/mlx-community--MiniMax-M2.5-4bit.toml

@@ -0,0 +1,12 @@
+model_id = "mlx-community/MiniMax-M2.5-4bit"
+n_layers = 62
+hidden_size = 3072
+supports_tensor = true
+tasks = ["TextGeneration"]
+family = "minimax"
+quantization = "4bit"
+base_model = "MiniMax M2.5"
+capabilities = ["text", "thinking"]
+
+[storage_size]
+in_bytes = 128666664960

resources/inference_model_cards/mlx-community--MiniMax-M2.5-6bit.toml

@@ -0,0 +1,12 @@
+model_id = "mlx-community/MiniMax-M2.5-6bit"
+n_layers = 62
+hidden_size = 3072
+supports_tensor = true
+tasks = ["TextGeneration"]
+family = "minimax"
+quantization = "6bit"
+base_model = "MiniMax M2.5"
+capabilities = ["text", "thinking"]
+
+[storage_size]
+in_bytes = 185826705408

resources/inference_model_cards/mlx-community--MiniMax-M2.5-8bit.toml

@@ -0,0 +1,12 @@
+model_id = "mlx-community/MiniMax-M2.5-8bit"
+n_layers = 62
+hidden_size = 3072
+supports_tensor = true
+tasks = ["TextGeneration"]
+family = "minimax"
+quantization = "8bit"
+base_model = "MiniMax M2.5"
+capabilities = ["text", "thinking"]
+
+[storage_size]
+in_bytes = 242986745856

resources/inference_model_cards/mlx-community--Qwen3-0.6B-4bit.toml

@@ -6,7 +6,7 @@ tasks = ["TextGeneration"]
 family = "qwen"
 quantization = "4bit"
 base_model = "Qwen3 0.6B"
-capabilities = ["text", "thinking"]
+capabilities = ["text", "thinking", "thinking_toggle"]
 
 [storage_size]
 in_bytes = 342884352

resources/inference_model_cards/mlx-community--Qwen3-0.6B-8bit.toml

@@ -6,7 +6,7 @@ tasks = ["TextGeneration"]
 family = "qwen"
 quantization = "8bit"
 base_model = "Qwen3 0.6B"
-capabilities = ["text", "thinking"]
+capabilities = ["text", "thinking", "thinking_toggle"]
 
 [storage_size]
 in_bytes = 698351616

resources/inference_model_cards/mlx-community--Qwen3-235B-A22B-Instruct-2507-4bit.toml

@@ -6,7 +6,7 @@ tasks = ["TextGeneration"]
 family = "qwen"
 quantization = "4bit"
 base_model = "Qwen3 235B"
-capabilities = ["text", "thinking"]
+capabilities = ["text", "thinking", "thinking_toggle"]
 
 [storage_size]
 in_bytes = 141733920768

resources/inference_model_cards/mlx-community--Qwen3-235B-A22B-Instruct-2507-8bit.toml

@@ -6,7 +6,7 @@ tasks = ["TextGeneration"]
 family = "qwen"
 quantization = "8bit"
 base_model = "Qwen3 235B"
-capabilities = ["text", "thinking"]
+capabilities = ["text", "thinking", "thinking_toggle"]
 
 [storage_size]
 in_bytes = 268435456000

resources/inference_model_cards/mlx-community--Qwen3-30B-A3B-4bit.toml

@@ -6,7 +6,7 @@ tasks = ["TextGeneration"]
 family = "qwen"
 quantization = "4bit"
 base_model = "Qwen3 30B"
-capabilities = ["text", "thinking"]
+capabilities = ["text", "thinking", "thinking_toggle"]
 
 [storage_size]
 in_bytes = 17612931072

resources/inference_model_cards/mlx-community--Qwen3-30B-A3B-8bit.toml

@@ -6,7 +6,7 @@ tasks = ["TextGeneration"]
 family = "qwen"
 quantization = "8bit"
 base_model = "Qwen3 30B"
-capabilities = ["text", "thinking"]
+capabilities = ["text", "thinking", "thinking_toggle"]
 
 [storage_size]
 in_bytes = 33279705088

resources/inference_model_cards/mlx-community--Qwen3-Next-80B-A3B-Thinking-4bit.toml

@@ -6,7 +6,7 @@ tasks = ["TextGeneration"]
 family = "qwen"
 quantization = "4bit"
 base_model = "Qwen3 Next 80B"
-capabilities = ["text", "thinking"]
+capabilities = ["text", "thinking", "thinking_toggle"]
 
 [storage_size]
 in_bytes = 47080074240

resources/inference_model_cards/mlx-community--Qwen3-Next-80B-A3B-Thinking-8bit.toml

@@ -6,7 +6,7 @@ tasks = ["TextGeneration"]
 family = "qwen"
 quantization = "8bit"
 base_model = "Qwen3 Next 80B"
-capabilities = ["text", "thinking"]
+capabilities = ["text", "thinking", "thinking_toggle"]
 
 [storage_size]
 in_bytes = 88814387200

resources/inference_model_cards/mlx-community--Step-3.5-Flash-4bit.toml

@@ -6,7 +6,7 @@ tasks = ["TextGeneration"]
 family = "step"
 quantization = "4bit"
 base_model = "Step 3.5 Flash"
-capabilities = ["text", "thinking"]
+capabilities = ["text", "thinking", "thinking_toggle"]
 
 [storage_size]
 in_bytes = 114572190076

resources/inference_model_cards/mlx-community--Step-3.5-Flash-6bit.toml

@@ -6,7 +6,7 @@ tasks = ["TextGeneration"]
 family = "step"
 quantization = "6bit"
 base_model = "Step 3.5 Flash"
-capabilities = ["text", "thinking"]
+capabilities = ["text", "thinking", "thinking_toggle"]
 
 [storage_size]
 in_bytes = 159039627774

resources/inference_model_cards/mlx-community--Step-3.5-Flash-8Bit.toml

@@ -6,7 +6,7 @@ tasks = ["TextGeneration"]
 family = "step"
 quantization = "8bit"
 base_model = "Step 3.5 Flash"
-capabilities = ["text", "thinking"]
+capabilities = ["text", "thinking", "thinking_toggle"]
 
 [storage_size]
 in_bytes = 209082699847

rust/clippy.toml

@@ -1,2 +0,0 @@
-# we can manually exclude false-positive lint errors for dual packages (if in dependencies)
-#allowed-duplicate-crates = ["hashbrown"]

rust/exo_pyo3_bindings/Cargo.toml

@@ -27,7 +27,7 @@ networking = { workspace = true }
 # interop
 pyo3 = { version = "0.27.2", features = [
     # "abi3-py313", # tells pyo3 (and maturin) to build using the stable ABI with minimum Python version 3.13
-    "nightly", # enables better-supported GIL integration
+    # "nightly", # enables better-supported GIL integration
     "experimental-async", # async support in #[pyfunction] & #[pymethods]
     #"experimental-inspect", # inspection of generated binary => easier to automate type-hint generation
     #"py-clone", # adding Clone-ing of `Py<T>` without GIL (may cause panics - remove if panics happen)
@@ -45,11 +45,10 @@ pyo3-log = "0.13.2"
 # macro dependencies
 extend = { workspace = true }
 delegate = { workspace = true }
-pin-project = { workspace = true }
 
 # async runtime
 tokio = { workspace = true, features = ["full", "tracing"] }
-futures = { workspace = true }
+futures-lite = { workspace = true }
 
 # utility dependencies
 util = { workspace = true }
@@ -60,3 +59,4 @@ env_logger = "0.11"
 
 # Networking
 libp2p = { workspace = true, features = ["full"] }
+pin-project = "1.1.10"

rust/exo_pyo3_bindings/exo_pyo3_bindings.pyi

@@ -19,7 +19,7 @@ class ConnectionUpdate:
         Whether this is a connection or disconnection event
         """
     @property
-    def peer_id(self) -> PeerId:
+    def peer_id(self) -> builtins.str:
         r"""
         Identity of the peer that we have connected to or disconnected from.
         """
@@ -40,92 +40,22 @@ class Keypair:
     Identity keypair of a node.
     """
     @staticmethod
-    def generate_ed25519() -> Keypair:
+    def generate() -> Keypair:
         r"""
         Generate a new Ed25519 keypair.
         """
     @staticmethod
-    def generate_ecdsa() -> Keypair:
+    def from_bytes(bytes: bytes) -> Keypair:
         r"""
-        Generate a new ECDSA keypair.
-        """
-    @staticmethod
-    def generate_secp256k1() -> Keypair:
-        r"""
-        Generate a new Secp256k1 keypair.
-        """
-    @staticmethod
-    def from_protobuf_encoding(bytes: bytes) -> Keypair:
-        r"""
-        Decode a private key from a protobuf structure and parse it as a `Keypair`.
-        """
-    @staticmethod
-    def rsa_from_pkcs8(bytes: bytes) -> Keypair:
-        r"""
-        Decode an keypair from a DER-encoded secret key in PKCS#8 `PrivateKeyInfo`
-        format (i.e. unencrypted) as defined in [RFC5208].
-        
-        [RFC5208]: https://tools.ietf.org/html/rfc5208#section-5
-        """
-    @staticmethod
-    def secp256k1_from_der(bytes: bytes) -> Keypair:
-        r"""
-        Decode a keypair from a DER-encoded Secp256k1 secret key in an `ECPrivateKey`
-        structure as defined in [RFC5915].
-        
-        [RFC5915]: https://tools.ietf.org/html/rfc5915
-        """
-    @staticmethod
-    def ed25519_from_bytes(bytes: bytes) -> Keypair: ...
-    def to_protobuf_encoding(self) -> bytes:
-        r"""
-        Encode a private key as protobuf structure.
-        """
-    def to_peer_id(self) -> PeerId:
-        r"""
-        Convert the `Keypair` into the corresponding `PeerId`.
-        """
-
-@typing.final
-class Multiaddr:
-    r"""
-    Representation of a Multiaddr.
-    """
-    @staticmethod
-    def empty() -> Multiaddr:
-        r"""
-        Create a new, empty multiaddress.
-        """
-    @staticmethod
-    def with_capacity(n: builtins.int) -> Multiaddr:
-        r"""
-        Create a new, empty multiaddress with the given capacity.
-        """
-    @staticmethod
-    def from_bytes(bytes: bytes) -> Multiaddr:
-        r"""
-        Parse a `Multiaddr` value from its byte slice representation.
-        """
-    @staticmethod
-    def from_string(string: builtins.str) -> Multiaddr:
-        r"""
-        Parse a `Multiaddr` value from its string representation.
-        """
-    def len(self) -> builtins.int:
-        r"""
-        Return the length in bytes of this multiaddress.
-        """
-    def is_empty(self) -> builtins.bool:
-        r"""
-        Returns true if the length of this multiaddress is 0.
+        Construct an Ed25519 keypair from secret key bytes
         """
     def to_bytes(self) -> bytes:
         r"""
-        Return a copy of this [`Multiaddr`]'s byte representation.
+        Get the secret key bytes underlying the keypair
         """
-    def to_string(self) -> builtins.str:
+    def to_node_id(self) -> builtins.str:
         r"""
-        Convert a Multiaddr to a string.
+        Convert the `Keypair` into the corresponding `PeerId` string, which we use as our `NodeId`.
         """
 
 @typing.final
@@ -180,37 +110,6 @@ class NoPeersSubscribedToTopicError(builtins.Exception):
     def __repr__(self) -> builtins.str: ...
     def __str__(self) -> builtins.str: ...
 
-@typing.final
-class PeerId:
-    r"""
-    Identifier of a peer of the network.
-    
-    The data is a `CIDv0` compatible multihash of the protobuf encoded public key of the peer
-    as specified in [specs/peer-ids](https://github.com/libp2p/specs/blob/master/peer-ids/peer-ids.md).
-    """
-    @staticmethod
-    def random() -> PeerId:
-        r"""
-        Generates a random peer ID from a cryptographically secure PRNG.
-        
-        This is useful for randomly walking on a DHT, or for testing purposes.
-        """
-    @staticmethod
-    def from_bytes(bytes: bytes) -> PeerId:
-        r"""
-        Parses a `PeerId` from bytes.
-        """
-    def to_bytes(self) -> bytes:
-        r"""
-        Returns a raw bytes representation of this `PeerId`.
-        """
-    def to_base58(self) -> builtins.str:
-        r"""
-        Returns a base-58 encoded string of this `PeerId`.
-        """
-    def __repr__(self) -> builtins.str: ...
-    def __str__(self) -> builtins.str: ...
-
 @typing.final
 class ConnectionUpdateType(enum.Enum):
     r"""

rust/exo_pyo3_bindings/src/allow_threading.rs

@@ -2,7 +2,6 @@
 //!
 
 use pin_project::pin_project;
-use pyo3::marker::Ungil;
 use pyo3::prelude::*;
 use std::{
     future::Future,
@@ -26,8 +25,8 @@ where
 
 impl<F> Future for AllowThreads<F>
 where
-    F: Future + Ungil,
-    F::Output: Ungil,
+    F: Future + Send,
+    F::Output: Send,
 {
     type Output = F::Output;
 

rust/exo_pyo3_bindings/src/ident.rs

@@ -0,0 +1,47 @@
+use crate::ext::ResultExt as _;
+use libp2p::identity::Keypair;
+use pyo3::types::{PyBytes, PyBytesMethods as _};
+use pyo3::{Bound, PyResult, Python, pyclass, pymethods};
+use pyo3_stub_gen::derive::{gen_stub_pyclass, gen_stub_pymethods};
+
+/// Identity keypair of a node.
+#[gen_stub_pyclass]
+#[pyclass(name = "Keypair", frozen)]
+#[repr(transparent)]
+pub struct PyKeypair(pub Keypair);
+
+#[gen_stub_pymethods]
+#[pymethods]
+#[allow(clippy::needless_pass_by_value)]
+impl PyKeypair {
+    /// Generate a new Ed25519 keypair.
+    #[staticmethod]
+    fn generate() -> Self {
+        Self(Keypair::generate_ed25519())
+    }
+
+    /// Construct an Ed25519 keypair from secret key bytes
+    #[staticmethod]
+    fn from_bytes(bytes: Bound<'_, PyBytes>) -> PyResult<Self> {
+        let mut bytes = Vec::from(bytes.as_bytes());
+        Ok(Self(Keypair::ed25519_from_bytes(&mut bytes).pyerr()?))
+    }
+
+    /// Get the secret key bytes underlying the keypair
+    fn to_bytes<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyBytes>> {
+        let bytes = self
+            .0
+            .clone()
+            .try_into_ed25519()
+            .pyerr()?
+            .secret()
+            .as_ref()
+            .to_vec();
+        Ok(PyBytes::new(py, &bytes))
+    }
+
+    /// Convert the `Keypair` into the corresponding `PeerId` string, which we use as our `NodeId`.
+    fn to_node_id(&self) -> String {
+        self.0.public().to_peer_id().to_base58()
+    }
+}

rust/exo_pyo3_bindings/src/lib.rs

@@ -4,26 +4,14 @@
 //!
 //!
 
-// enable Rust-unstable features for convenience
-#![feature(trait_alias)]
-#![feature(tuple_trait)]
-#![feature(unboxed_closures)]
-// #![feature(stmt_expr_attributes)]
-// #![feature(assert_matches)]
-// #![feature(async_fn_in_dyn_trait)]
-// #![feature(async_for_loop)]
-// #![feature(auto_traits)]
-// #![feature(negative_impls)]
-
-extern crate core;
 mod allow_threading;
-pub(crate) mod networking;
-pub(crate) mod pylibp2p;
+mod ident;
+mod networking;
 
+use crate::ident::PyKeypair;
 use crate::networking::networking_submodule;
-use crate::pylibp2p::ident::ident_submodule;
-use crate::pylibp2p::multiaddr::multiaddr_submodule;
 use pyo3::prelude::PyModule;
+use pyo3::types::PyModuleMethods;
 use pyo3::{Bound, PyResult, pyclass, pymodule};
 use pyo3_stub_gen::define_stub_info_gatherer;
 
@@ -32,14 +20,6 @@ pub(crate) mod r#const {
     pub const MPSC_CHANNEL_SIZE: usize = 1024;
 }
 
-/// Namespace for all the type/trait aliases used by this crate.
-pub(crate) mod alias {
-    use std::marker::Tuple;
-
-    pub trait SendFn<Args: Tuple + Send + 'static, Output> =
-        Fn<Args, Output = Output> + Send + 'static;
-}
-
 /// Namespace for crate-wide extension traits/methods
 pub(crate) mod ext {
     use crate::allow_threading::AllowThreads;
@@ -179,8 +159,7 @@ fn main_module(m: &Bound<'_, PyModule>) -> PyResult<()> {
     // 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...
-    ident_submodule(m)?;
-    multiaddr_submodule(m)?;
+    m.add_class::<PyKeypair>()?;
     networking_submodule(m)?;
 
     // top-level constructs

rust/exo_pyo3_bindings/src/networking.rs

@@ -8,8 +8,8 @@
 use crate::r#const::MPSC_CHANNEL_SIZE;
 use crate::ext::{ByteArrayExt as _, FutureExt, PyErrExt as _};
 use crate::ext::{ResultExt as _, TokioMpscReceiverExt as _, TokioMpscSenderExt as _};
+use crate::ident::PyKeypair;
 use crate::pyclass;
-use crate::pylibp2p::ident::{PyKeypair, PyPeerId};
 use libp2p::futures::StreamExt as _;
 use libp2p::gossipsub;
 use libp2p::gossipsub::{IdentTopic, Message, MessageId, PublishError};
@@ -119,7 +119,7 @@ struct PyConnectionUpdate {
 
     /// Identity of the peer that we have connected to or disconnected from.
     #[pyo3(get)]
-    peer_id: PyPeerId,
+    peer_id: String,
 
     /// Remote connection's IPv4 address.
     #[pyo3(get)]
@@ -251,7 +251,7 @@ async fn networking_task(
                         // send connection event to channel (or exit if connection closed)
                         if let Err(e) = connection_update_tx.send(PyConnectionUpdate {
                             update_type: PyConnectionUpdateType::Connected,
-                            peer_id: PyPeerId(peer_id),
+                            peer_id: peer_id.to_base58(),
                             remote_ipv4,
                             remote_tcp_port,
                         }).await {
@@ -272,7 +272,7 @@ async fn networking_task(
                         // send disconnection event to channel (or exit if connection closed)
                         if let Err(e) = connection_update_tx.send(PyConnectionUpdate {
                             update_type: PyConnectionUpdateType::Disconnected,
-                            peer_id: PyPeerId(peer_id),
+                            peer_id: peer_id.to_base58(),
                             remote_ipv4,
                             remote_tcp_port,
                         }).await {

rust/exo_pyo3_bindings/src/pylibp2p/ident.rs

@@ -1,159 +0,0 @@
-use crate::ext::ResultExt as _;
-use libp2p::PeerId;
-use libp2p::identity::Keypair;
-use pyo3::prelude::{PyBytesMethods as _, PyModule, PyModuleMethods as _};
-use pyo3::types::PyBytes;
-use pyo3::{Bound, PyResult, Python, pyclass, pymethods};
-use pyo3_stub_gen::derive::{gen_stub_pyclass, gen_stub_pymethods};
-
-/// Identity keypair of a node.
-#[gen_stub_pyclass]
-#[pyclass(name = "Keypair", frozen)]
-#[repr(transparent)]
-pub struct PyKeypair(pub Keypair);
-
-#[gen_stub_pymethods]
-#[pymethods]
-#[allow(clippy::needless_pass_by_value)]
-impl PyKeypair {
-    /// Generate a new Ed25519 keypair.
-    #[staticmethod]
-    fn generate_ed25519() -> Self {
-        Self(Keypair::generate_ed25519())
-    }
-
-    /// Generate a new ECDSA keypair.
-    #[staticmethod]
-    fn generate_ecdsa() -> Self {
-        Self(Keypair::generate_ecdsa())
-    }
-
-    /// Generate a new Secp256k1 keypair.
-    #[staticmethod]
-    fn generate_secp256k1() -> Self {
-        Self(Keypair::generate_secp256k1())
-    }
-
-    /// Decode a private key from a protobuf structure and parse it as a `Keypair`.
-    #[staticmethod]
-    fn from_protobuf_encoding(bytes: Bound<'_, PyBytes>) -> PyResult<Self> {
-        let bytes = Vec::from(bytes.as_bytes());
-        Ok(Self(Keypair::from_protobuf_encoding(&bytes).pyerr()?))
-    }
-
-    /// Decode an keypair from a DER-encoded secret key in PKCS#8 `PrivateKeyInfo`
-    /// format (i.e. unencrypted) as defined in [RFC5208].
-    ///
-    /// [RFC5208]: https://tools.ietf.org/html/rfc5208#section-5
-    #[staticmethod]
-    fn rsa_from_pkcs8(bytes: Bound<'_, PyBytes>) -> PyResult<Self> {
-        let mut bytes = Vec::from(bytes.as_bytes());
-        Ok(Self(Keypair::rsa_from_pkcs8(&mut bytes).pyerr()?))
-    }
-
-    /// Decode a keypair from a DER-encoded Secp256k1 secret key in an `ECPrivateKey`
-    /// structure as defined in [RFC5915].
-    ///
-    /// [RFC5915]: https://tools.ietf.org/html/rfc5915
-    #[staticmethod]
-    fn secp256k1_from_der(bytes: Bound<'_, PyBytes>) -> PyResult<Self> {
-        let mut bytes = Vec::from(bytes.as_bytes());
-        Ok(Self(Keypair::secp256k1_from_der(&mut bytes).pyerr()?))
-    }
-
-    #[staticmethod]
-    fn ed25519_from_bytes(bytes: Bound<'_, PyBytes>) -> PyResult<Self> {
-        let mut bytes = Vec::from(bytes.as_bytes());
-        Ok(Self(Keypair::ed25519_from_bytes(&mut bytes).pyerr()?))
-    }
-
-    /// Encode a private key as protobuf structure.
-    fn to_protobuf_encoding<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyBytes>> {
-        let bytes = self.0.to_protobuf_encoding().pyerr()?;
-        Ok(PyBytes::new(py, &bytes))
-    }
-
-    /// Convert the `Keypair` into the corresponding `PeerId`.
-    fn to_peer_id(&self) -> PyPeerId {
-        PyPeerId(self.0.public().to_peer_id())
-    }
-
-    // /// Hidden constructor for pickling support. TODO: figure out how to do pickling...
-    // #[gen_stub(skip)]
-    // #[new]
-    // fn py_new(bytes: Bound<'_, PyBytes>) -> PyResult<Self> {
-    //     Self::from_protobuf_encoding(bytes)
-    // }
-    //
-    // #[gen_stub(skip)]
-    // fn __setstate__(&mut self, state: Bound<'_, PyBytes>) -> PyResult<()> {
-    //     *self = Self::from_protobuf_encoding(state)?;
-    //     Ok(())
-    // }
-    //
-    // #[gen_stub(skip)]
-    // fn __getstate__<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyBytes>> {
-    //     self.to_protobuf_encoding(py)
-    // }
-    //
-    // #[gen_stub(skip)]
-    // pub fn __getnewargs__<'py>(&self, py: Python<'py>) -> PyResult<(Bound<'py, PyBytes>,)> {
-    //     Ok((self.to_protobuf_encoding(py)?,))
-    // }
-}
-
-/// Identifier of a peer of the network.
-///
-/// The data is a `CIDv0` compatible multihash of the protobuf encoded public key of the peer
-/// as specified in [specs/peer-ids](https://github.com/libp2p/specs/blob/master/peer-ids/peer-ids.md).
-#[gen_stub_pyclass]
-#[pyclass(name = "PeerId", frozen)]
-#[derive(Debug, Clone)]
-#[repr(transparent)]
-pub struct PyPeerId(pub PeerId);
-
-#[gen_stub_pymethods]
-#[pymethods]
-#[allow(clippy::needless_pass_by_value)]
-impl PyPeerId {
-    /// Generates a random peer ID from a cryptographically secure PRNG.
-    ///
-    /// This is useful for randomly walking on a DHT, or for testing purposes.
-    #[staticmethod]
-    fn random() -> Self {
-        Self(PeerId::random())
-    }
-
-    /// Parses a `PeerId` from bytes.
-    #[staticmethod]
-    fn from_bytes(bytes: Bound<'_, PyBytes>) -> PyResult<Self> {
-        let bytes = Vec::from(bytes.as_bytes());
-        Ok(Self(PeerId::from_bytes(&bytes).pyerr()?))
-    }
-
-    /// Returns a raw bytes representation of this `PeerId`.
-    fn to_bytes<'py>(&self, py: Python<'py>) -> Bound<'py, PyBytes> {
-        let bytes = self.0.to_bytes();
-        PyBytes::new(py, &bytes)
-    }
-
-    /// Returns a base-58 encoded string of this `PeerId`.
-    fn to_base58(&self) -> String {
-        self.0.to_base58()
-    }
-
-    fn __repr__(&self) -> String {
-        format!("PeerId({})", self.to_base58())
-    }
-
-    fn __str__(&self) -> String {
-        self.to_base58()
-    }
-}
-
-pub fn ident_submodule(m: &Bound<'_, PyModule>) -> PyResult<()> {
-    m.add_class::<PyKeypair>()?;
-    m.add_class::<PyPeerId>()?;
-
-    Ok(())
-}

rust/exo_pyo3_bindings/src/pylibp2p/mod.rs

@@ -1,8 +0,0 @@
-//! A module for exposing Rust's libp2p datatypes over Pyo3
-//!
-//! TODO: right now we are coupled to libp2p's identity, but eventually we want to create our own
-//!       independent identity type of some kind or another. This may require handshaking.
-//!
-
-pub mod ident;
-pub mod multiaddr;

rust/exo_pyo3_bindings/src/pylibp2p/multiaddr.rs

@@ -1,81 +0,0 @@
-use crate::ext::ResultExt as _;
-use libp2p::Multiaddr;
-use pyo3::prelude::{PyBytesMethods as _, PyModule, PyModuleMethods as _};
-use pyo3::types::PyBytes;
-use pyo3::{Bound, PyResult, Python, pyclass, pymethods};
-use pyo3_stub_gen::derive::{gen_stub_pyclass, gen_stub_pymethods};
-use std::str::FromStr as _;
-
-/// Representation of a Multiaddr.
-#[gen_stub_pyclass]
-#[pyclass(name = "Multiaddr", frozen)]
-#[derive(Debug, Clone)]
-#[repr(transparent)]
-pub struct PyMultiaddr(pub Multiaddr);
-
-#[gen_stub_pymethods]
-#[pymethods]
-#[allow(clippy::needless_pass_by_value)]
-impl PyMultiaddr {
-    /// Create a new, empty multiaddress.
-    #[staticmethod]
-    fn empty() -> Self {
-        Self(Multiaddr::empty())
-    }
-
-    /// Create a new, empty multiaddress with the given capacity.
-    #[staticmethod]
-    fn with_capacity(n: usize) -> Self {
-        Self(Multiaddr::with_capacity(n))
-    }
-
-    /// Parse a `Multiaddr` value from its byte slice representation.
-    #[staticmethod]
-    fn from_bytes(bytes: Bound<'_, PyBytes>) -> PyResult<Self> {
-        let bytes = Vec::from(bytes.as_bytes());
-        Ok(Self(Multiaddr::try_from(bytes).pyerr()?))
-    }
-
-    /// Parse a `Multiaddr` value from its string representation.
-    #[staticmethod]
-    fn from_string(string: String) -> PyResult<Self> {
-        Ok(Self(Multiaddr::from_str(&string).pyerr()?))
-    }
-
-    /// Return the length in bytes of this multiaddress.
-    fn len(&self) -> usize {
-        self.0.len()
-    }
-
-    /// Returns true if the length of this multiaddress is 0.
-    fn is_empty(&self) -> bool {
-        self.0.is_empty()
-    }
-
-    /// Return a copy of this [`Multiaddr`]'s byte representation.
-    fn to_bytes<'py>(&self, py: Python<'py>) -> Bound<'py, PyBytes> {
-        let bytes = self.0.to_vec();
-        PyBytes::new(py, &bytes)
-    }
-
-    /// Convert a Multiaddr to a string.
-    fn to_string(&self) -> String {
-        self.0.to_string()
-    }
-
-    #[gen_stub(skip)]
-    fn __repr__(&self) -> String {
-        format!("Multiaddr({})", self.0)
-    }
-
-    #[gen_stub(skip)]
-    fn __str__(&self) -> String {
-        self.to_string()
-    }
-}
-
-pub fn multiaddr_submodule(m: &Bound<'_, PyModule>) -> PyResult<()> {
-    m.add_class::<PyMultiaddr>()?;
-
-    Ok(())
-}

rust/networking/Cargo.toml

@@ -22,7 +22,7 @@ delegate = { workspace = true }
 
 # async
 tokio = { workspace = true, features = ["full"] }
-futures = { workspace = true }
+futures-lite = { workspace = true }
 futures-timer = { workspace = true }
 
 # utility dependencies

rust/networking/examples/chatroom.rs

@@ -1,4 +1,4 @@
-use futures::stream::StreamExt as _;
+use futures_lite::StreamExt;
 use libp2p::{gossipsub, identity, swarm::SwarmEvent};
 use networking::{discovery, swarm};
 use tokio::{io, io::AsyncBufReadExt as _, select};
@@ -38,19 +38,19 @@ async fn main() {
                     println!("Publish error: {e:?}");
                 }
             }
-            event = swarm.select_next_some() => match event {
+            event = swarm.next() => match event {
                 // on gossipsub incoming
-                SwarmEvent::Behaviour(swarm::BehaviourEvent::Gossipsub(gossipsub::Event::Message {
+                Some(SwarmEvent::Behaviour(swarm::BehaviourEvent::Gossipsub(gossipsub::Event::Message {
                     propagation_source: peer_id,
                     message_id: id,
                     message,
-                })) => println!(
+                }))) => println!(
                         "\n\nGot message: '{}' with id: {id} from peer: {peer_id}\n\n",
                         String::from_utf8_lossy(&message.data),
                     ),
 
                 // on discovery
-                SwarmEvent::Behaviour(swarm::BehaviourEvent::Discovery(e)) => match e {
+                Some(SwarmEvent::Behaviour(swarm::BehaviourEvent::Discovery(e)) )=> match e {
                     discovery::Event::ConnectionEstablished {
                         peer_id, connection_id, remote_ip, remote_tcp_port
                     } => {
@@ -64,7 +64,7 @@ async fn main() {
                 }
 
                 // ignore outgoing errors: those are normal
-                e@SwarmEvent::OutgoingConnectionError { .. } => { log::debug!("Outgoing connection error: {e:?}"); }
+                e@Some(SwarmEvent::OutgoingConnectionError { .. }) => { log::debug!("Outgoing connection error: {e:?}"); }
 
                 // otherwise log any other event
                 e => { log::info!("Other event {e:?}"); }

rust/networking/examples/chatroom_manual.rs

@@ -1,127 +0,0 @@
-// Copyright 2018 Parity Technologies (UK) Ltd.
-//
-// Permission is hereby granted, free of charge, to any person obtaining a
-// copy of this software and associated documentation files (the "Software"),
-// to deal in the Software without restriction, including without limitation
-// the rights to use, copy, modify, merge, publish, distribute, sublicense,
-// and/or sell copies of the Software, and to permit persons to whom the
-// Software is furnished to do so, subject to the following conditions:
-//
-// The above copyright notice and this permission notice shall be included in
-// all copies or substantial portions of the Software.
-//
-// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS
-// OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
-// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
-// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
-// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
-// FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
-// DEALINGS IN THE SOFTWARE.
-
-use futures::stream::StreamExt;
-use libp2p::{
-    gossipsub, mdns, noise,
-    swarm::{NetworkBehaviour, SwarmEvent},
-    tcp, yamux,
-};
-use std::error::Error;
-use std::time::Duration;
-use tokio::{io, io::AsyncBufReadExt, select};
-use tracing_subscriber::EnvFilter;
-
-// We create a custom network behaviour that combines Gossipsub and Mdns.
-#[derive(NetworkBehaviour)]
-struct MyBehaviour {
-    gossipsub: gossipsub::Behaviour,
-    mdns: mdns::tokio::Behaviour,
-}
-
-#[tokio::main]
-async fn main() -> Result<(), Box<dyn Error>> {
-    let _ = tracing_subscriber::fmt()
-        .with_env_filter(EnvFilter::from_default_env())
-        .try_init();
-
-    let mut swarm = libp2p::SwarmBuilder::with_new_identity()
-        .with_tokio()
-        .with_tcp(
-            tcp::Config::default(),
-            noise::Config::new,
-            yamux::Config::default,
-        )?
-        .with_behaviour(|key| {
-            // Set a custom gossipsub configuration
-            let gossipsub_config = gossipsub::ConfigBuilder::default()
-                .heartbeat_interval(Duration::from_secs(10))
-                .validation_mode(gossipsub::ValidationMode::Strict) // This sets the kind of message validation. The default is Strict (enforce message signing)
-                .build()
-                .map_err(io::Error::other)?; // Temporary hack because `build` does not return a proper `std::error::Error`.
-
-            // build a gossipsub network behaviour
-            let gossipsub = gossipsub::Behaviour::new(
-                gossipsub::MessageAuthenticity::Signed(key.clone()),
-                gossipsub_config,
-            )?;
-
-            let mdns =
-                mdns::tokio::Behaviour::new(mdns::Config::default(), key.public().to_peer_id())?;
-            Ok(MyBehaviour { gossipsub, mdns })
-        })?
-        .build();
-
-    println!("Running swarm with identity {}", swarm.local_peer_id());
-
-    // Create a Gossipsub topic
-    let topic = gossipsub::IdentTopic::new("test-net");
-    // subscribes to our topic
-    swarm.behaviour_mut().gossipsub.subscribe(&topic)?;
-
-    // Read full lines from stdin
-    let mut stdin = io::BufReader::new(io::stdin()).lines();
-
-    // Listen on all interfaces and whatever port the OS assigns
-    swarm.listen_on("/ip4/0.0.0.0/tcp/0".parse()?)?;
-
-    println!("Enter messages via STDIN and they will be sent to connected peers using Gossipsub");
-
-    // Kick it off
-    loop {
-        select! {
-            Ok(Some(line)) = stdin.next_line() => {
-                if let Err(e) = swarm
-                    .behaviour_mut().gossipsub
-                    .publish(topic.clone(), line.as_bytes()) {
-                    println!("Publish error: {e:?}");
-                }
-            }
-            event = swarm.select_next_some() => match event {
-                SwarmEvent::Behaviour(MyBehaviourEvent::Mdns(mdns::Event::Discovered(list))) => {
-                    for (peer_id, multiaddr) in list {
-                        println!("mDNS discovered a new peer: {peer_id} on {multiaddr}");
-                        swarm.behaviour_mut().gossipsub.add_explicit_peer(&peer_id);
-                    }
-                },
-                SwarmEvent::Behaviour(MyBehaviourEvent::Mdns(mdns::Event::Expired(list))) => {
-                    for (peer_id, multiaddr) in list {
-                        println!("mDNS discover peer has expired: {peer_id} on {multiaddr}");
-                        swarm.behaviour_mut().gossipsub.remove_explicit_peer(&peer_id);
-                    }
-                },
-                SwarmEvent::Behaviour(MyBehaviourEvent::Gossipsub(gossipsub::Event::Message {
-                    propagation_source: peer_id,
-                    message_id: id,
-                    message,
-                })) => println!(
-                        "Got message: '{}' with id: {id} from peer: {peer_id}",
-                        String::from_utf8_lossy(&message.data),
-                    ),
-                SwarmEvent::NewListenAddr { address, .. } => {
-                    println!("Local node is listening on {address}");
-                }
-                e => {
-                    println!("Other swarm event: {:?}", e);
-                }
-            }
-        }
-    }
-}

rust/networking/src/discovery.rs

@@ -1,7 +1,7 @@
 use crate::ext::MultiaddrExt;
 use delegate::delegate;
 use either::Either;
-use futures::FutureExt;
+use futures_lite::FutureExt;
 use futures_timer::Delay;
 use libp2p::core::transport::PortUse;
 use libp2p::core::{ConnectedPoint, Endpoint};
@@ -362,7 +362,7 @@ impl NetworkBehaviour for Behaviour {
         }
 
         // retry connecting to all mDNS peers periodically (fails safely if already connected)
-        if self.retry_delay.poll_unpin(cx).is_ready() {
+        if self.retry_delay.poll(cx).is_ready() {
             for (p, mas) in self.mdns_discovered.clone() {
                 for ma in mas {
                     self.dial(p, ma)

rust/networking/src/swarm.rs

@@ -31,7 +31,7 @@ pub fn create_swarm(keypair: identity::Keypair) -> alias::AnyResult<Swarm> {
 mod transport {
     use crate::alias;
     use crate::swarm::{NETWORK_VERSION, OVERRIDE_VERSION_ENV_VAR};
-    use futures::{AsyncRead, AsyncWrite};
+    use futures_lite::{AsyncRead, AsyncWrite};
     use keccak_const::Sha3_256;
     use libp2p::core::muxing;
     use libp2p::core::transport::Boxed;

rust/parts.nix

@@ -1,11 +1,10 @@
 { inputs, ... }:
 {
   perSystem =
-    { config, self', inputs', pkgs, lib, ... }:
+    { inputs', pkgs, lib, ... }:
     let
       # Fenix nightly toolchain with all components
-      fenixPkgs = inputs'.fenix.packages;
-      rustToolchain = fenixPkgs.complete.withComponents [
+      rustToolchain = inputs'.fenix.packages.stable.withComponents [
         "cargo"
         "rustc"
         "clippy"

rust/rust-toolchain.toml

@@ -1,2 +0,0 @@
-[toolchain]
-channel = "nightly"

src/exo/download/coordinator.py

@@ -123,14 +123,17 @@ class DownloadCoordinator:
                 tg.start_soon(self._check_internet_connection)
 
     def _test_internet_connection(self) -> None:
-        try:
-            socket.create_connection(("1.1.1.1", 443), timeout=3).close()
-            self.shard_downloader.set_internet_connection(True)
-        except OSError:
-            self.shard_downloader.set_internet_connection(False)
-        logger.debug(
-            f"Internet connectivity: {self.shard_downloader.internet_connection}"
-        )
+        # Try multiple endpoints since some ISPs/networks block specific IPs
+        for host in ("1.1.1.1", "8.8.8.8", "1.0.0.1"):
+            try:
+                socket.create_connection((host, 443), timeout=3).close()
+                self.shard_downloader.set_internet_connection(True)
+                logger.debug(f"Internet connectivity: True (via {host})")
+                return
+            except OSError:
+                continue
+        self.shard_downloader.set_internet_connection(False)
+        logger.debug("Internet connectivity: False")
 
     async def _check_internet_connection(self) -> None:
         first_connection = True

src/exo/main.py

@@ -45,7 +45,7 @@ class Node:
     @classmethod
     async def create(cls, args: "Args") -> "Self":
         keypair = get_node_id_keypair()
-        node_id = NodeId(keypair.to_peer_id().to_base58())
+        node_id = NodeId(keypair.to_node_id())
         session_id = SessionId(master_node_id=node_id, election_clock=0)
         router = Router.create(keypair)
         await router.register_topic(topics.GLOBAL_EVENTS)

src/exo/master/adapters/chat_completions.py

@@ -19,7 +19,12 @@ from exo.shared.types.api import (
     ToolCall,
     Usage,
 )
-from exo.shared.types.chunks import ErrorChunk, TokenChunk, ToolCallChunk
+from exo.shared.types.chunks import (
+    ErrorChunk,
+    PrefillProgressChunk,
+    TokenChunk,
+    ToolCallChunk,
+)
 from exo.shared.types.common import CommandId
 from exo.shared.types.text_generation import InputMessage, TextGenerationTaskParams
 
@@ -54,7 +59,11 @@ def chat_request_to_text_generation(
             chat_template_messages.append({"role": "system", "content": content})
         else:
             # Skip messages with no meaningful content
-            if msg.content is None and msg.thinking is None and msg.tool_calls is None:
+            if (
+                msg.content is None
+                and msg.reasoning_content is None
+                and msg.tool_calls is None
+            ):
                 continue
 
             if msg.role in ("user", "assistant", "developer"):
@@ -106,6 +115,11 @@ def chunk_to_response(
             ]
         )
 
+    if chunk.is_thinking:
+        delta = ChatCompletionMessage(role="assistant", reasoning_content=chunk.text)
+    else:
+        delta = ChatCompletionMessage(role="assistant", content=chunk.text)
+
     return ChatCompletionResponse(
         id=command_id,
         created=int(time.time()),
@@ -113,7 +127,7 @@ def chunk_to_response(
         choices=[
             StreamingChoiceResponse(
                 index=0,
-                delta=ChatCompletionMessage(role="assistant", content=chunk.text),
+                delta=delta,
                 logprobs=logprobs,
                 finish_reason=chunk.finish_reason,
             )
@@ -123,72 +137,87 @@ def chunk_to_response(
 
 async def generate_chat_stream(
     command_id: CommandId,
-    chunk_stream: AsyncGenerator[ErrorChunk | ToolCallChunk | TokenChunk, None],
+    chunk_stream: AsyncGenerator[
+        PrefillProgressChunk | ErrorChunk | ToolCallChunk | TokenChunk, None
+    ],
 ) -> AsyncGenerator[str, None]:
     """Generate Chat Completions API streaming events from chunks."""
     last_usage: Usage | None = None
 
     async for chunk in chunk_stream:
-        if isinstance(chunk, ErrorChunk):
-            error_response = ErrorResponse(
-                error=ErrorInfo(
-                    message=chunk.error_message or "Internal server error",
-                    type="InternalServerError",
-                    code=500,
-                )
-            )
-            yield f"data: {error_response.model_dump_json()}\n\n"
-            yield "data: [DONE]\n\n"
-            return
+        match chunk:
+            case PrefillProgressChunk():
+                # Use SSE comment so third-party clients ignore it
+                yield f": prefill_progress {chunk.model_dump_json()}\n\n"
 
-        last_usage = chunk.usage or last_usage
-
-        if isinstance(chunk, ToolCallChunk):
-            tool_call_deltas = [
-                ToolCall(
-                    id=tool.id,
-                    index=i,
-                    function=tool,
-                )
-                for i, tool in enumerate(chunk.tool_calls)
-            ]
-            tool_response = ChatCompletionResponse(
-                id=command_id,
-                created=int(time.time()),
-                model=chunk.model,
-                choices=[
-                    StreamingChoiceResponse(
-                        index=0,
-                        delta=ChatCompletionMessage(
-                            role="assistant",
-                            tool_calls=tool_call_deltas,
-                        ),
-                        finish_reason="tool_calls",
+            case ErrorChunk():
+                error_response = ErrorResponse(
+                    error=ErrorInfo(
+                        message=chunk.error_message or "Internal server error",
+                        type="InternalServerError",
+                        code=500,
                     )
-                ],
-                usage=last_usage,
-            )
-            yield f"data: {tool_response.model_dump_json()}\n\n"
-            yield "data: [DONE]\n\n"
-            return
+                )
+                yield f"data: {error_response.model_dump_json()}\n\n"
+                yield "data: [DONE]\n\n"
+                return
 
-        chunk_response = chunk_to_response(chunk, command_id)
-        if chunk.finish_reason is not None:
-            chunk_response = chunk_response.model_copy(update={"usage": last_usage})
-        yield f"data: {chunk_response.model_dump_json()}\n\n"
+            case ToolCallChunk():
+                last_usage = chunk.usage or last_usage
 
-        if chunk.finish_reason is not None:
-            yield "data: [DONE]\n\n"
+                tool_call_deltas = [
+                    ToolCall(
+                        id=tool.id,
+                        index=i,
+                        function=tool,
+                    )
+                    for i, tool in enumerate(chunk.tool_calls)
+                ]
+                tool_response = ChatCompletionResponse(
+                    id=command_id,
+                    created=int(time.time()),
+                    model=chunk.model,
+                    choices=[
+                        StreamingChoiceResponse(
+                            index=0,
+                            delta=ChatCompletionMessage(
+                                role="assistant",
+                                tool_calls=tool_call_deltas,
+                            ),
+                            finish_reason="tool_calls",
+                        )
+                    ],
+                    usage=last_usage,
+                )
+                yield f"data: {tool_response.model_dump_json()}\n\n"
+                yield "data: [DONE]\n\n"
+                return
+
+            case TokenChunk():
+                last_usage = chunk.usage or last_usage
+
+                chunk_response = chunk_to_response(chunk, command_id)
+                if chunk.finish_reason is not None:
+                    chunk_response = chunk_response.model_copy(
+                        update={"usage": last_usage}
+                    )
+                yield f"data: {chunk_response.model_dump_json()}\n\n"
+
+                if chunk.finish_reason is not None:
+                    yield "data: [DONE]\n\n"
 
 
 async def collect_chat_response(
     command_id: CommandId,
-    chunk_stream: AsyncGenerator[ErrorChunk | ToolCallChunk | TokenChunk, None],
+    chunk_stream: AsyncGenerator[
+        ErrorChunk | ToolCallChunk | TokenChunk | PrefillProgressChunk, None
+    ],
 ) -> AsyncGenerator[str]:
     # This is an AsyncGenerator[str] rather than returning a ChatCompletionReponse because
     # FastAPI handles the cancellation better but wouldn't auto-serialize for some reason
     """Collect all token chunks and return a single ChatCompletionResponse."""
     text_parts: list[str] = []
+    thinking_parts: list[str] = []
     tool_calls: list[ToolCall] = []
     logprobs_content: list[LogprobsContentItem] = []
     model: str | None = None
@@ -197,43 +226,52 @@ async def collect_chat_response(
     last_usage: Usage | None = None
 
     async for chunk in chunk_stream:
-        if isinstance(chunk, ErrorChunk):
-            error_message = chunk.error_message or "Internal server error"
-            break
+        match chunk:
+            case PrefillProgressChunk():
+                continue
 
-        if model is None:
-            model = chunk.model
+            case ErrorChunk():
+                error_message = chunk.error_message or "Internal server error"
+                break
 
-        last_usage = chunk.usage or last_usage
-
-        if isinstance(chunk, TokenChunk):
-            text_parts.append(chunk.text)
-            if chunk.logprob is not None:
-                logprobs_content.append(
-                    LogprobsContentItem(
-                        token=chunk.text,
-                        logprob=chunk.logprob,
-                        top_logprobs=chunk.top_logprobs or [],
+            case TokenChunk():
+                if model is None:
+                    model = chunk.model
+                last_usage = chunk.usage or last_usage
+                if chunk.is_thinking:
+                    thinking_parts.append(chunk.text)
+                else:
+                    text_parts.append(chunk.text)
+                if chunk.logprob is not None:
+                    logprobs_content.append(
+                        LogprobsContentItem(
+                            token=chunk.text,
+                            logprob=chunk.logprob,
+                            top_logprobs=chunk.top_logprobs or [],
+                        )
                     )
-                )
+                if chunk.finish_reason is not None:
+                    finish_reason = chunk.finish_reason
 
-        if isinstance(chunk, ToolCallChunk):
-            tool_calls.extend(
-                ToolCall(
-                    id=tool.id,
-                    index=i,
-                    function=tool,
+            case ToolCallChunk():
+                if model is None:
+                    model = chunk.model
+                last_usage = chunk.usage or last_usage
+                tool_calls.extend(
+                    ToolCall(
+                        id=tool.id,
+                        index=i,
+                        function=tool,
+                    )
+                    for i, tool in enumerate(chunk.tool_calls)
                 )
-                for i, tool in enumerate(chunk.tool_calls)
-            )
-
-        if chunk.finish_reason is not None:
-            finish_reason = chunk.finish_reason
+                finish_reason = chunk.finish_reason
 
     if error_message is not None:
         raise ValueError(error_message)
 
     combined_text = "".join(text_parts)
+    combined_thinking = "".join(thinking_parts) if thinking_parts else None
     assert model is not None
 
     yield ChatCompletionResponse(
@@ -246,6 +284,7 @@ async def collect_chat_response(
                 message=ChatCompletionMessage(
                     role="assistant",
                     content=combined_text,
+                    reasoning_content=combined_thinking,
                     tool_calls=tool_calls if tool_calls else None,
                 ),
                 logprobs=Logprobs(content=logprobs_content)

src/exo/master/adapters/claude.py

@@ -1,11 +1,17 @@
 """Claude Messages API adapter for converting requests/responses."""
 
 import json
+import re
 from collections.abc import AsyncGenerator
 from typing import Any
 
 from exo.shared.types.api import FinishReason, Usage
-from exo.shared.types.chunks import ErrorChunk, TokenChunk, ToolCallChunk
+from exo.shared.types.chunks import (
+    ErrorChunk,
+    PrefillProgressChunk,
+    TokenChunk,
+    ToolCallChunk,
+)
 from exo.shared.types.claude_api import (
     ClaudeContentBlock,
     ClaudeContentBlockDeltaEvent,
@@ -23,6 +29,8 @@ from exo.shared.types.claude_api import (
     ClaudeStopReason,
     ClaudeTextBlock,
     ClaudeTextDelta,
+    ClaudeThinkingBlock,
+    ClaudeThinkingDelta,
     ClaudeToolResultBlock,
     ClaudeToolUseBlock,
     ClaudeUsage,
@@ -56,6 +64,22 @@ def _extract_tool_result_text(block: ClaudeToolResultBlock) -> str:
     return "".join(sub_block.text for sub_block in block.content)
 
 
+# Matches "x-anthropic-billing-header: ...;" (with optional trailing newline)
+# or similar telemetry headers that change every request and break KV prefix caching.
+_VOLATILE_HEADER_RE = re.compile(r"^x-anthropic-[^\n]*;\n?", re.MULTILINE)
+
+
+def _strip_volatile_headers(text: str) -> str:
+    """Remove Anthropic billing/telemetry headers from system prompt text.
+
+    Claude Code prepends headers like 'x-anthropic-billing-header: cc_version=...;
+    cc_entrypoint=...; cch=...;' that contain per-request content hashes. These
+    change every request and break KV prefix caching (the prefix diverges at ~20
+    tokens instead of matching thousands of conversation tokens).
+    """
+    return _VOLATILE_HEADER_RE.sub("", text)
+
+
 def claude_request_to_text_generation(
     request: ClaudeMessagesRequest,
 ) -> TextGenerationTaskParams:
@@ -68,6 +92,8 @@ def claude_request_to_text_generation(
             instructions = request.system
         else:
             instructions = "".join(block.text for block in request.system)
+
+        instructions = _strip_volatile_headers(instructions)
         chat_template_messages.append({"role": "system", "content": instructions})
 
     # Convert messages to input
@@ -80,12 +106,15 @@ def claude_request_to_text_generation(
 
         # Process structured content blocks
         text_parts: list[str] = []
+        thinking_parts: list[str] = []
         tool_calls: list[dict[str, Any]] = []
         tool_results: list[ClaudeToolResultBlock] = []
 
         for block in msg.content:
             if isinstance(block, ClaudeTextBlock):
                 text_parts.append(block.text)
+            elif isinstance(block, ClaudeThinkingBlock):
+                thinking_parts.append(block.thinking)
             elif isinstance(block, ClaudeToolUseBlock):
                 tool_calls.append(
                     {
@@ -101,6 +130,7 @@ def claude_request_to_text_generation(
                 tool_results.append(block)
 
         content = "".join(text_parts)
+        reasoning_content = "".join(thinking_parts) if thinking_parts else None
 
         # Build InputMessage from text content
         if msg.role in ("user", "assistant"):
@@ -108,9 +138,14 @@ def claude_request_to_text_generation(
 
         # Build chat_template_messages preserving tool structure
         if tool_calls:
-            chat_template_messages.append(
-                {"role": "assistant", "content": content, "tool_calls": tool_calls}
-            )
+            chat_msg: dict[str, Any] = {
+                "role": "assistant",
+                "content": content,
+                "tool_calls": tool_calls,
+            }
+            if reasoning_content:
+                chat_msg["reasoning_content"] = reasoning_content
+            chat_template_messages.append(chat_msg)
         elif tool_results:
             for tr in tool_results:
                 chat_template_messages.append(
@@ -121,7 +156,10 @@ def claude_request_to_text_generation(
                     }
                 )
         else:
-            chat_template_messages.append({"role": msg.role, "content": content})
+            chat_msg = {"role": msg.role, "content": content}
+            if reasoning_content:
+                chat_msg["reasoning_content"] = reasoning_content
+            chat_template_messages.append(chat_msg)
 
     # Convert Claude tool definitions to OpenAI-style function tools
     tools: list[dict[str, Any]] | None = None
@@ -138,6 +176,10 @@ def claude_request_to_text_generation(
             for tool in request.tools
         ]
 
+    enable_thinking: bool | None = None
+    if request.thinking is not None:
+        enable_thinking = request.thinking.type in ("enabled", "adaptive")
+
     return TextGenerationTaskParams(
         model=request.model,
         input=input_messages
@@ -151,6 +193,7 @@ def claude_request_to_text_generation(
         stop=request.stop_sequences,
         stream=request.stream,
         tools=tools,
+        enable_thinking=enable_thinking,
         chat_template_messages=chat_template_messages
         if chat_template_messages
         else None,
@@ -160,18 +203,24 @@ def claude_request_to_text_generation(
 async def collect_claude_response(
     command_id: CommandId,
     model: str,
-    chunk_stream: AsyncGenerator[ErrorChunk | ToolCallChunk | TokenChunk, None],
+    chunk_stream: AsyncGenerator[
+        ErrorChunk | ToolCallChunk | TokenChunk | PrefillProgressChunk, None
+    ],
 ) -> AsyncGenerator[str]:
     # This is an AsyncGenerator[str] rather than returning a ChatCompletionReponse because
     # FastAPI handles the cancellation better but wouldn't auto-serialize for some reason
     """Collect all token chunks and return a single ClaudeMessagesResponse."""
     text_parts: list[str] = []
+    thinking_parts: list[str] = []
     tool_use_blocks: list[ClaudeToolUseBlock] = []
     stop_reason: ClaudeStopReason | None = None
     last_usage: Usage | None = None
     error_message: str | None = None
 
     async for chunk in chunk_stream:
+        if isinstance(chunk, PrefillProgressChunk):
+            continue
+
         if isinstance(chunk, ErrorChunk):
             error_message = chunk.error_message or "Internal server error"
             break
@@ -190,7 +239,10 @@ async def collect_claude_response(
             stop_reason = "tool_use"
             continue
 
-        text_parts.append(chunk.text)
+        if chunk.is_thinking:
+            thinking_parts.append(chunk.text)
+        else:
+            text_parts.append(chunk.text)
 
         if chunk.finish_reason is not None:
             stop_reason = finish_reason_to_claude_stop_reason(chunk.finish_reason)
@@ -199,9 +251,12 @@ async def collect_claude_response(
         raise ValueError(error_message)
 
     combined_text = "".join(text_parts)
+    combined_thinking = "".join(thinking_parts)
 
     # Build content blocks
     content: list[ClaudeContentBlock] = []
+    if combined_thinking:
+        content.append(ClaudeThinkingBlock(thinking=combined_thinking))
     if combined_text:
         content.append(ClaudeTextBlock(text=combined_text))
     content.extend(tool_use_blocks)
@@ -230,7 +285,9 @@ async def collect_claude_response(
 async def generate_claude_stream(
     command_id: CommandId,
     model: str,
-    chunk_stream: AsyncGenerator[ErrorChunk | ToolCallChunk | TokenChunk, None],
+    chunk_stream: AsyncGenerator[
+        ErrorChunk | ToolCallChunk | TokenChunk | PrefillProgressChunk, None
+    ],
 ) -> AsyncGenerator[str, None]:
     """Generate Claude Messages API streaming events from TokenChunks."""
     # Initial message_start event
@@ -244,18 +301,21 @@ async def generate_claude_stream(
     start_event = ClaudeMessageStartEvent(message=initial_message)
     yield f"event: message_start\ndata: {start_event.model_dump_json()}\n\n"
 
-    # content_block_start for text block at index 0
-    block_start = ClaudeContentBlockStartEvent(
-        index=0, content_block=ClaudeTextBlock(text="")
-    )
-    yield f"event: content_block_start\ndata: {block_start.model_dump_json()}\n\n"
-
     output_tokens = 0
     stop_reason: ClaudeStopReason | None = None
     last_usage: Usage | None = None
-    next_block_index = 1  # text block is 0, tool blocks start at 1
+    next_block_index = 0
+
+    # Track whether we've started thinking/text blocks
+    thinking_block_started = False
+    thinking_block_index = -1
+    text_block_started = False
+    text_block_index = -1
 
     async for chunk in chunk_stream:
+        if isinstance(chunk, PrefillProgressChunk):
+            continue
+
         if isinstance(chunk, ErrorChunk):
             # Close text block and bail
             break
@@ -295,12 +355,45 @@ async def generate_claude_stream(
 
         output_tokens += 1  # Count each chunk as one token
 
-        # content_block_delta
-        delta_event = ClaudeContentBlockDeltaEvent(
-            index=0,
-            delta=ClaudeTextDelta(text=chunk.text),
-        )
-        yield f"event: content_block_delta\ndata: {delta_event.model_dump_json()}\n\n"
+        if chunk.is_thinking:
+            # Start thinking block on first thinking token
+            if not thinking_block_started:
+                thinking_block_started = True
+                thinking_block_index = next_block_index
+                next_block_index += 1
+                block_start = ClaudeContentBlockStartEvent(
+                    index=thinking_block_index,
+                    content_block=ClaudeThinkingBlock(thinking=""),
+                )
+                yield f"event: content_block_start\ndata: {block_start.model_dump_json()}\n\n"
+
+            delta_event = ClaudeContentBlockDeltaEvent(
+                index=thinking_block_index,
+                delta=ClaudeThinkingDelta(thinking=chunk.text),
+            )
+            yield f"event: content_block_delta\ndata: {delta_event.model_dump_json()}\n\n"
+        else:
+            # Close thinking block when transitioning to text
+            if thinking_block_started and text_block_index == -1:
+                block_stop = ClaudeContentBlockStopEvent(index=thinking_block_index)
+                yield f"event: content_block_stop\ndata: {block_stop.model_dump_json()}\n\n"
+
+            # Start text block on first text token
+            if not text_block_started:
+                text_block_started = True
+                text_block_index = next_block_index
+                next_block_index += 1
+                block_start = ClaudeContentBlockStartEvent(
+                    index=text_block_index,
+                    content_block=ClaudeTextBlock(text=""),
+                )
+                yield f"event: content_block_start\ndata: {block_start.model_dump_json()}\n\n"
+
+            delta_event = ClaudeContentBlockDeltaEvent(
+                index=text_block_index,
+                delta=ClaudeTextDelta(text=chunk.text),
+            )
+            yield f"event: content_block_delta\ndata: {delta_event.model_dump_json()}\n\n"
 
         if chunk.finish_reason is not None:
             stop_reason = finish_reason_to_claude_stop_reason(chunk.finish_reason)
@@ -309,9 +402,22 @@ async def generate_claude_stream(
     if last_usage is not None:
         output_tokens = last_usage.completion_tokens
 
-    # content_block_stop for text block
-    block_stop = ClaudeContentBlockStopEvent(index=0)
-    yield f"event: content_block_stop\ndata: {block_stop.model_dump_json()}\n\n"
+    # Close any open blocks
+    if thinking_block_started and text_block_index == -1:
+        block_stop = ClaudeContentBlockStopEvent(index=thinking_block_index)
+        yield f"event: content_block_stop\ndata: {block_stop.model_dump_json()}\n\n"
+
+    if text_block_started:
+        block_stop = ClaudeContentBlockStopEvent(index=text_block_index)
+        yield f"event: content_block_stop\ndata: {block_stop.model_dump_json()}\n\n"
+
+    if not thinking_block_started and not text_block_started:
+        empty_start = ClaudeContentBlockStartEvent(
+            index=0, content_block=ClaudeTextBlock(text="")
+        )
+        yield f"event: content_block_start\ndata: {empty_start.model_dump_json()}\n\n"
+        empty_stop = ClaudeContentBlockStopEvent(index=0)
+        yield f"event: content_block_stop\ndata: {empty_stop.model_dump_json()}\n\n"
 
     # message_delta
     message_delta = ClaudeMessageDeltaEvent(

src/exo/master/adapters/responses.py

@@ -5,7 +5,12 @@ from itertools import count
 from typing import Any
 
 from exo.shared.types.api import Usage
-from exo.shared.types.chunks import ErrorChunk, TokenChunk, ToolCallChunk
+from exo.shared.types.chunks import (
+    ErrorChunk,
+    PrefillProgressChunk,
+    TokenChunk,
+    ToolCallChunk,
+)
 from exo.shared.types.common import CommandId
 from exo.shared.types.openai_responses import (
     FunctionCallInputItem,
@@ -24,8 +29,15 @@ from exo.shared.types.openai_responses import (
     ResponseOutputItemAddedEvent,
     ResponseOutputItemDoneEvent,
     ResponseOutputText,
+    ResponseReasoningItem,
+    ResponseReasoningSummaryPartAddedEvent,
+    ResponseReasoningSummaryPartDoneEvent,
+    ResponseReasoningSummaryText,
+    ResponseReasoningSummaryTextDeltaEvent,
+    ResponseReasoningSummaryTextDoneEvent,
     ResponsesRequest,
     ResponsesResponse,
+    ResponsesStreamEvent,
     ResponseTextDeltaEvent,
     ResponseTextDoneEvent,
     ResponseUsage,
@@ -33,6 +45,11 @@ from exo.shared.types.openai_responses import (
 from exo.shared.types.text_generation import InputMessage, TextGenerationTaskParams
 
 
+def _format_sse(event: ResponsesStreamEvent) -> str:
+    """Format a streaming event as an SSE message."""
+    return f"event: {event.type}\ndata: {event.model_dump_json()}\n\n"
+
+
 def _extract_content(content: str | list[ResponseContentPart]) -> str:
     """Extract plain text from a content field that may be a string or list of parts."""
     if isinstance(content, str):
@@ -121,19 +138,26 @@ def responses_request_to_text_generation(
 async def collect_responses_response(
     command_id: CommandId,
     model: str,
-    chunk_stream: AsyncGenerator[ErrorChunk | ToolCallChunk | TokenChunk, None],
+    chunk_stream: AsyncGenerator[
+        ErrorChunk | ToolCallChunk | TokenChunk | PrefillProgressChunk, None
+    ],
 ) -> AsyncGenerator[str]:
     # This is an AsyncGenerator[str] rather than returning a ChatCompletionReponse because
     # FastAPI handles the cancellation better but wouldn't auto-serialize for some reason
     """Collect all token chunks and return a single ResponsesResponse."""
     response_id = f"resp_{command_id}"
     item_id = f"item_{command_id}"
+    reasoning_id = f"rs_{command_id}"
     accumulated_text = ""
+    thinking_parts: list[str] = []
     function_call_items: list[ResponseFunctionCallItem] = []
     last_usage: Usage | None = None
     error_message: str | None = None
 
     async for chunk in chunk_stream:
+        if isinstance(chunk, PrefillProgressChunk):
+            continue
+
         if isinstance(chunk, ErrorChunk):
             error_message = chunk.error_message or "Internal server error"
             break
@@ -152,6 +176,10 @@ async def collect_responses_response(
                 )
             continue
 
+        if chunk.is_thinking:
+            thinking_parts.append(chunk.text)
+            continue
+
         accumulated_text += chunk.text
 
     if error_message is not None:
@@ -166,13 +194,21 @@ async def collect_responses_response(
             total_tokens=last_usage.total_tokens,
         )
 
-    output: list[ResponseItem] = [
+    output: list[ResponseItem] = []
+    if thinking_parts:
+        output.append(
+            ResponseReasoningItem(
+                id=reasoning_id,
+                summary=[ResponseReasoningSummaryText(text="".join(thinking_parts))],
+            )
+        )
+    output.append(
         ResponseMessageItem(
             id=item_id,
             content=[ResponseOutputText(text=accumulated_text)],
             status="completed",
         )
-    ]
+    )
     output.extend(function_call_items)
 
     yield ResponsesResponse(
@@ -189,11 +225,14 @@ async def collect_responses_response(
 async def generate_responses_stream(
     command_id: CommandId,
     model: str,
-    chunk_stream: AsyncGenerator[ErrorChunk | ToolCallChunk | TokenChunk, None],
+    chunk_stream: AsyncGenerator[
+        ErrorChunk | ToolCallChunk | TokenChunk | PrefillProgressChunk, None
+    ],
 ) -> AsyncGenerator[str, None]:
     """Generate OpenAI Responses API streaming events from TokenChunks."""
     response_id = f"resp_{command_id}"
     item_id = f"item_{command_id}"
+    reasoning_id = f"rs_{command_id}"
     seq = count(1)
 
     # response.created
@@ -207,42 +246,30 @@ async def generate_responses_stream(
     created_event = ResponseCreatedEvent(
         sequence_number=next(seq), response=initial_response
     )
-    yield f"event: response.created\ndata: {created_event.model_dump_json()}\n\n"
+    yield _format_sse(created_event)
 
     # response.in_progress
     in_progress_event = ResponseInProgressEvent(
         sequence_number=next(seq), response=initial_response
     )
-    yield f"event: response.in_progress\ndata: {in_progress_event.model_dump_json()}\n\n"
-
-    # response.output_item.added
-    initial_item = ResponseMessageItem(
-        id=item_id,
-        content=[ResponseOutputText(text="")],
-        status="in_progress",
-    )
-    item_added = ResponseOutputItemAddedEvent(
-        sequence_number=next(seq), output_index=0, item=initial_item
-    )
-    yield f"event: response.output_item.added\ndata: {item_added.model_dump_json()}\n\n"
-
-    # response.content_part.added
-    initial_part = ResponseOutputText(text="")
-    part_added = ResponseContentPartAddedEvent(
-        sequence_number=next(seq),
-        item_id=item_id,
-        output_index=0,
-        content_index=0,
-        part=initial_part,
-    )
-    yield f"event: response.content_part.added\ndata: {part_added.model_dump_json()}\n\n"
+    yield _format_sse(in_progress_event)
 
     accumulated_text = ""
+    accumulated_thinking = ""
     function_call_items: list[ResponseFunctionCallItem] = []
     last_usage: Usage | None = None
-    next_output_index = 1  # message item is at 0
+    next_output_index = 0
+
+    # Track dynamic block creation
+    reasoning_started = False
+    reasoning_output_index = -1
+    message_started = False
+    message_output_index = -1
 
     async for chunk in chunk_stream:
+        if isinstance(chunk, PrefillProgressChunk):
+            continue
+
         if isinstance(chunk, ErrorChunk):
             break
 
@@ -266,7 +293,7 @@ async def generate_responses_stream(
                     output_index=next_output_index,
                     item=fc_item,
                 )
-                yield f"event: response.output_item.added\ndata: {fc_added.model_dump_json()}\n\n"
+                yield _format_sse(fc_added)
 
                 # response.function_call_arguments.delta
                 args_delta = ResponseFunctionCallArgumentsDeltaEvent(
@@ -275,7 +302,7 @@ async def generate_responses_stream(
                     output_index=next_output_index,
                     delta=tool.arguments,
                 )
-                yield f"event: response.function_call_arguments.delta\ndata: {args_delta.model_dump_json()}\n\n"
+                yield _format_sse(args_delta)
 
                 # response.function_call_arguments.done
                 args_done = ResponseFunctionCallArgumentsDoneEvent(
@@ -285,7 +312,7 @@ async def generate_responses_stream(
                     name=tool.name,
                     arguments=tool.arguments,
                 )
-                yield f"event: response.function_call_arguments.done\ndata: {args_done.model_dump_json()}\n\n"
+                yield _format_sse(args_done)
 
                 # response.output_item.done
                 fc_done_item = ResponseFunctionCallItem(
@@ -300,44 +327,205 @@ async def generate_responses_stream(
                     output_index=next_output_index,
                     item=fc_done_item,
                 )
-                yield f"event: response.output_item.done\ndata: {fc_item_done.model_dump_json()}\n\n"
+                yield _format_sse(fc_item_done)
 
                 function_call_items.append(fc_done_item)
                 next_output_index += 1
             continue
 
+        if chunk.is_thinking:
+            # Start reasoning block on first thinking token
+            if not reasoning_started:
+                reasoning_started = True
+                reasoning_output_index = next_output_index
+                next_output_index += 1
+
+                # response.output_item.added for reasoning
+                reasoning_item = ResponseReasoningItem(
+                    id=reasoning_id,
+                    summary=[],
+                    status="in_progress",
+                )
+                rs_added = ResponseOutputItemAddedEvent(
+                    sequence_number=next(seq),
+                    output_index=reasoning_output_index,
+                    item=reasoning_item,
+                )
+                yield _format_sse(rs_added)
+
+                # response.reasoning_summary_part.added
+                part_added = ResponseReasoningSummaryPartAddedEvent(
+                    sequence_number=next(seq),
+                    item_id=reasoning_id,
+                    output_index=reasoning_output_index,
+                    summary_index=0,
+                    part=ResponseReasoningSummaryText(text=""),
+                )
+                yield _format_sse(part_added)
+
+            accumulated_thinking += chunk.text
+
+            # response.reasoning_summary_text.delta
+            rs_delta = ResponseReasoningSummaryTextDeltaEvent(
+                sequence_number=next(seq),
+                item_id=reasoning_id,
+                output_index=reasoning_output_index,
+                summary_index=0,
+                delta=chunk.text,
+            )
+            yield _format_sse(rs_delta)
+            continue
+
+        # Close reasoning block when transitioning to text
+        if reasoning_started and not message_started:
+            # response.reasoning_summary_text.done
+            rs_text_done = ResponseReasoningSummaryTextDoneEvent(
+                sequence_number=next(seq),
+                item_id=reasoning_id,
+                output_index=reasoning_output_index,
+                summary_index=0,
+                text=accumulated_thinking,
+            )
+            yield _format_sse(rs_text_done)
+
+            # response.reasoning_summary_part.done
+            rs_part_done = ResponseReasoningSummaryPartDoneEvent(
+                sequence_number=next(seq),
+                item_id=reasoning_id,
+                output_index=reasoning_output_index,
+                summary_index=0,
+                part=ResponseReasoningSummaryText(text=accumulated_thinking),
+            )
+            yield _format_sse(rs_part_done)
+
+            # response.output_item.done for reasoning
+            rs_item_done = ResponseOutputItemDoneEvent(
+                sequence_number=next(seq),
+                output_index=reasoning_output_index,
+                item=ResponseReasoningItem(
+                    id=reasoning_id,
+                    summary=[ResponseReasoningSummaryText(text=accumulated_thinking)],
+                ),
+            )
+            yield _format_sse(rs_item_done)
+
+        # Start message block on first text token
+        if not message_started:
+            message_started = True
+            message_output_index = next_output_index
+            next_output_index += 1
+
+            initial_item = ResponseMessageItem(
+                id=item_id,
+                content=[ResponseOutputText(text="")],
+                status="in_progress",
+            )
+            item_added = ResponseOutputItemAddedEvent(
+                sequence_number=next(seq),
+                output_index=message_output_index,
+                item=initial_item,
+            )
+            yield _format_sse(item_added)
+
+            initial_part = ResponseOutputText(text="")
+            part_added = ResponseContentPartAddedEvent(
+                sequence_number=next(seq),
+                item_id=item_id,
+                output_index=message_output_index,
+                content_index=0,
+                part=initial_part,
+            )
+            yield _format_sse(part_added)
+
         accumulated_text += chunk.text
 
         # response.output_text.delta
         delta_event = ResponseTextDeltaEvent(
             sequence_number=next(seq),
             item_id=item_id,
-            output_index=0,
+            output_index=message_output_index,
             content_index=0,
             delta=chunk.text,
         )
-        yield f"event: response.output_text.delta\ndata: {delta_event.model_dump_json()}\n\n"
+        yield _format_sse(delta_event)
+
+    # Close reasoning block if it was never followed by text
+    if reasoning_started and not message_started:
+        rs_text_done = ResponseReasoningSummaryTextDoneEvent(
+            sequence_number=next(seq),
+            item_id=reasoning_id,
+            output_index=reasoning_output_index,
+            summary_index=0,
+            text=accumulated_thinking,
+        )
+        yield _format_sse(rs_text_done)
+
+        rs_part_done = ResponseReasoningSummaryPartDoneEvent(
+            sequence_number=next(seq),
+            item_id=reasoning_id,
+            output_index=reasoning_output_index,
+            summary_index=0,
+            part=ResponseReasoningSummaryText(text=accumulated_thinking),
+        )
+        yield _format_sse(rs_part_done)
+
+        rs_item_done = ResponseOutputItemDoneEvent(
+            sequence_number=next(seq),
+            output_index=reasoning_output_index,
+            item=ResponseReasoningItem(
+                id=reasoning_id,
+                summary=[ResponseReasoningSummaryText(text=accumulated_thinking)],
+            ),
+        )
+        yield _format_sse(rs_item_done)
+
+    # If no message block was started, create one now (empty text)
+    if not message_started:
+        message_output_index = next_output_index
+        next_output_index += 1
+
+        initial_item = ResponseMessageItem(
+            id=item_id,
+            content=[ResponseOutputText(text="")],
+            status="in_progress",
+        )
+        item_added = ResponseOutputItemAddedEvent(
+            sequence_number=next(seq),
+            output_index=message_output_index,
+            item=initial_item,
+        )
+        yield _format_sse(item_added)
+
+        initial_part = ResponseOutputText(text="")
+        part_added_evt = ResponseContentPartAddedEvent(
+            sequence_number=next(seq),
+            item_id=item_id,
+            output_index=message_output_index,
+            content_index=0,
+            part=initial_part,
+        )
+        yield _format_sse(part_added_evt)
 
     # response.output_text.done
     text_done = ResponseTextDoneEvent(
         sequence_number=next(seq),
         item_id=item_id,
-        output_index=0,
+        output_index=message_output_index,
         content_index=0,
         text=accumulated_text,
     )
-    yield f"event: response.output_text.done\ndata: {text_done.model_dump_json()}\n\n"
+    yield _format_sse(text_done)
 
     # response.content_part.done
     final_part = ResponseOutputText(text=accumulated_text)
     part_done = ResponseContentPartDoneEvent(
         sequence_number=next(seq),
         item_id=item_id,
-        output_index=0,
+        output_index=message_output_index,
         content_index=0,
         part=final_part,
     )
-    yield f"event: response.content_part.done\ndata: {part_done.model_dump_json()}\n\n"
+    yield _format_sse(part_done)
 
     # response.output_item.done
     final_message_item = ResponseMessageItem(
@@ -346,9 +534,11 @@ async def generate_responses_stream(
         status="completed",
     )
     item_done = ResponseOutputItemDoneEvent(
-        sequence_number=next(seq), output_index=0, item=final_message_item
+        sequence_number=next(seq),
+        output_index=message_output_index,
+        item=final_message_item,
     )
-    yield f"event: response.output_item.done\ndata: {item_done.model_dump_json()}\n\n"
+    yield _format_sse(item_done)
 
     # Create usage from usage data if available
     usage = None
@@ -360,7 +550,15 @@ async def generate_responses_stream(
         )
 
     # response.completed
-    output: list[ResponseItem] = [final_message_item]
+    output: list[ResponseItem] = []
+    if reasoning_started:
+        output.append(
+            ResponseReasoningItem(
+                id=reasoning_id,
+                summary=[ResponseReasoningSummaryText(text=accumulated_thinking)],
+            )
+        )
+    output.append(final_message_item)
     output.extend(function_call_items)
     final_response = ResponsesResponse(
         id=response_id,
@@ -373,4 +571,4 @@ async def generate_responses_stream(
     completed_event = ResponseCompletedEvent(
         sequence_number=next(seq), response=final_response
     )
-    yield f"event: response.completed\ndata: {completed_event.model_dump_json()}\n\n"
+    yield _format_sse(completed_event)

src/exo/master/api.py

@@ -107,6 +107,7 @@ from exo.shared.types.chunks import (
     ErrorChunk,
     ImageChunk,
     InputImageChunk,
+    PrefillProgressChunk,
     TokenChunk,
     ToolCallChunk,
 )
@@ -145,6 +146,7 @@ from exo.shared.types.openai_responses import (
     ResponsesResponse,
 )
 from exo.shared.types.state import State
+from exo.shared.types.worker.downloads import DownloadCompleted
 from exo.shared.types.worker.instances import Instance, InstanceId, InstanceMeta
 from exo.shared.types.worker.shards import Sharding
 from exo.utils.banner import print_startup_banner
@@ -220,7 +222,8 @@ class API:
         )
 
         self._text_generation_queues: dict[
-            CommandId, Sender[TokenChunk | ErrorChunk | ToolCallChunk]
+            CommandId,
+            Sender[TokenChunk | ErrorChunk | ToolCallChunk | PrefillProgressChunk],
         ] = {}
         self._image_generation_queues: dict[
             CommandId, Sender[ImageChunk | ErrorChunk]
@@ -526,19 +529,23 @@ class API:
 
     async def _token_chunk_stream(
         self, command_id: CommandId
-    ) -> AsyncGenerator[ErrorChunk | ToolCallChunk | TokenChunk, None]:
+    ) -> AsyncGenerator[
+        TokenChunk | ErrorChunk | ToolCallChunk | PrefillProgressChunk, None
+    ]:
         """Yield chunks for a given command until completion.
 
         This is the internal low-level stream used by all API adapters.
         """
         try:
             self._text_generation_queues[command_id], recv = channel[
-                ErrorChunk | ToolCallChunk | TokenChunk
+                TokenChunk | ErrorChunk | ToolCallChunk | PrefillProgressChunk
             ]()
 
             with recv as token_chunks:
                 async for chunk in token_chunks:
                     yield chunk
+                    if isinstance(chunk, PrefillProgressChunk):
+                        continue
                     if chunk.finish_reason is not None:
                         break
 
@@ -565,6 +572,9 @@ class API:
         stats: GenerationStats | None = None
 
         async for chunk in self._token_chunk_stream(command_id):
+            if isinstance(chunk, PrefillProgressChunk):
+                continue
+
             if chunk.finish_reason == "error":
                 raise HTTPException(
                     status_code=500,
@@ -1292,8 +1302,18 @@ class API:
 
         return total_available
 
-    async def get_models(self) -> ModelList:
-        """Returns list of available models."""
+    async def get_models(self, status: str | None = Query(default=None)) -> ModelList:
+        """Returns list of available models, optionally filtered by being downloaded."""
+        cards = await get_model_cards()
+
+        if status == "downloaded":
+            downloaded_model_ids: set[str] = set()
+            for node_downloads in self.state.downloads.values():
+                for dl in node_downloads:
+                    if isinstance(dl, DownloadCompleted):
+                        downloaded_model_ids.add(dl.shard_metadata.model_card.model_id)
+            cards = [c for c in cards if c.model_id in downloaded_model_ids]
+
         return ModelList(
             data=[
                 ModelListModel(
@@ -1311,7 +1331,7 @@ class API:
                     base_model=card.base_model,
                     capabilities=card.capabilities,
                 )
-                for card in await get_model_cards()
+                for card in cards
             ]
         )
 
@@ -1434,7 +1454,6 @@ class API:
                                 await queue.send(event.chunk)
                             except BrokenResourceError:
                                 self._text_generation_queues.pop(event.command_id, None)
-
                     if isinstance(event, TracesMerged):
                         self._save_merged_trace(event)
 

src/exo/master/tests/test_claude_tool_use.py

@@ -261,7 +261,7 @@ class TestGenerateClaudeStreamToolUse:
 
         parsed = _parse_sse_events(events)
 
-        # Two tool block starts (at indices 1 and 2)
+        # Two tool block starts (at indices 0 and 1 — no text block when only tools)
         tool_starts = [
             e
             for e in parsed
@@ -270,12 +270,11 @@ class TestGenerateClaudeStreamToolUse:
             == "tool_use"
         ]
         assert len(tool_starts) == 2
-        assert tool_starts[0]["index"] == 1
-        assert tool_starts[1]["index"] == 2
+        assert tool_starts[0]["index"] == 0
+        assert tool_starts[1]["index"] == 1
 
-        # Two tool block stops (at indices 1 and 2), plus text block stop at 0
+        # Two tool block stops (at indices 0 and 1)
         block_stops = [e for e in parsed if e.get("type") == "content_block_stop"]
         stop_indices = [e["index"] for e in block_stops]
         assert 0 in stop_indices
         assert 1 in stop_indices
-        assert 2 in stop_indices

src/exo/master/tests/test_master.py

@@ -42,7 +42,7 @@ from exo.utils.channels import channel
 @pytest.mark.asyncio
 async def test_master():
     keypair = get_node_id_keypair()
-    node_id = NodeId(keypair.to_peer_id().to_base58())
+    node_id = NodeId(keypair.to_node_id())
     session_id = SessionId(master_node_id=node_id, election_clock=0)
 
     ge_sender, global_event_receiver = channel[ForwarderEvent]()
@@ -75,7 +75,7 @@ async def test_master():
     async with anyio.create_task_group() as tg:
         tg.start_soon(master.run)
 
-        sender_node_id = NodeId(f"{keypair.to_peer_id().to_base58()}_sender")
+        sender_node_id = NodeId(f"{keypair.to_node_id()}_sender")
         # inject a NodeGatheredInfo event
         logger.info("inject a NodeGatheredInfo event")
         await local_event_sender.send(

src/exo/routing/connection_message.py

@@ -30,7 +30,7 @@ class ConnectionMessage(CamelCaseModel):
     @classmethod
     def from_update(cls, update: ConnectionUpdate) -> "ConnectionMessage":
         return cls(
-            node_id=NodeId(update.peer_id.to_base58()),
+            node_id=NodeId(update.peer_id),
             connection_type=ConnectionMessageType.from_update_type(update.update_type),
             remote_ipv4=update.remote_ipv4,
             remote_tcp_port=update.remote_tcp_port,

src/exo/routing/router.py

@@ -221,7 +221,7 @@ def get_node_id_keypair(
     Obtain the :class:`PeerId` by from it.
     """
     # TODO(evan): bring back node id persistence once we figure out how to deal with duplicates
-    return Keypair.generate_ed25519()
+    return Keypair.generate()
 
     def lock_path(path: str | bytes | PathLike[str] | PathLike[bytes]) -> Path:
         return Path(str(path) + ".lock")
@@ -235,12 +235,12 @@ def get_node_id_keypair(
                 protobuf_encoded = f.read()
 
                 try:  # if decoded successfully, save & return
-                    return Keypair.from_protobuf_encoding(protobuf_encoded)
+                    return Keypair.from_bytes(protobuf_encoded)
                 except ValueError as e:  # on runtime error, assume corrupt file
                     logger.warning(f"Encountered error when trying to get keypair: {e}")
 
         # if no valid credentials, create new ones and persist
         with open(path, "w+b") as f:
             keypair = Keypair.generate_ed25519()
-            f.write(keypair.to_protobuf_encoding())
+            f.write(keypair.to_bytes())
             return keypair

src/exo/shared/tests/test_node_id_persistence.py

@@ -23,7 +23,7 @@ def _get_keypair_concurrent_subprocess_task(
     sem.release()
     # wait to be told to begin simultaneous read
     ev.wait()
-    queue.put(get_node_id_keypair().to_protobuf_encoding())
+    queue.put(get_node_id_keypair().to_bytes())
 
 
 def _get_keypair_concurrent(num_procs: int) -> bytes:

src/exo/shared/types/api.py

@@ -77,7 +77,7 @@ class ChatCompletionMessage(BaseModel):
     content: (
         str | ChatCompletionMessageText | list[ChatCompletionMessageText] | None
     ) = None
-    thinking: str | None = None  # Added for GPT-OSS harmony format support
+    reasoning_content: str | None = None
     name: str | None = None
     tool_calls: list[ToolCall] | None = None
     tool_call_id: str | None = None

src/exo/shared/types/chunks.py

@@ -27,6 +27,7 @@ class TokenChunk(BaseChunk):
     stats: GenerationStats | None = None
     logprob: float | None = None
     top_logprobs: list[TopLogprobItem] | None = None
+    is_thinking: bool = False
 
 
 class ErrorChunk(BaseChunk):
@@ -76,4 +77,13 @@ class InputImageChunk(BaseChunk):
                 yield name, value
 
 
-GenerationChunk = TokenChunk | ImageChunk | ToolCallChunk | ErrorChunk
+class PrefillProgressChunk(BaseChunk):
+    """Data class for prefill progress events during streaming."""
+
+    processed_tokens: int
+    total_tokens: int
+
+
+GenerationChunk = (
+    TokenChunk | ImageChunk | ToolCallChunk | ErrorChunk | PrefillProgressChunk
+)

src/exo/shared/types/claude_api.py

@@ -47,6 +47,14 @@ class ClaudeImageBlock(BaseModel, frozen=True):
     source: ClaudeImageSource
 
 
+class ClaudeThinkingBlock(BaseModel, frozen=True):
+    """Thinking content block in Claude Messages API."""
+
+    type: Literal["thinking"] = "thinking"
+    thinking: str
+    signature: str | None = None
+
+
 class ClaudeToolUseBlock(BaseModel, frozen=True):
     """Tool use content block in Claude Messages API."""
 
@@ -66,11 +74,17 @@ class ClaudeToolResultBlock(BaseModel, frozen=True):
     cache_control: dict[str, str] | None = None
 
 
-ClaudeContentBlock = ClaudeTextBlock | ClaudeImageBlock | ClaudeToolUseBlock
+ClaudeContentBlock = (
+    ClaudeTextBlock | ClaudeImageBlock | ClaudeThinkingBlock | ClaudeToolUseBlock
+)
 
 # Input content blocks can also include tool_result (sent by user after tool_use)
 ClaudeInputContentBlock = (
-    ClaudeTextBlock | ClaudeImageBlock | ClaudeToolUseBlock | ClaudeToolResultBlock
+    ClaudeTextBlock
+    | ClaudeImageBlock
+    | ClaudeThinkingBlock
+    | ClaudeToolUseBlock
+    | ClaudeToolResultBlock
 )
 
 
@@ -82,6 +96,11 @@ class ClaudeMessage(BaseModel, frozen=True):
     content: str | list[ClaudeInputContentBlock]
 
 
+class ClaudeThinkingConfig(BaseModel, frozen=True):
+    type: Literal["enabled", "disabled", "adaptive"]
+    budget_tokens: int | None = None
+
+
 class ClaudeMessagesRequest(BaseModel):
     """Request body for Claude Messages API."""
 
@@ -96,6 +115,7 @@ class ClaudeMessagesRequest(BaseModel):
     top_k: int | None = None
     tools: list[ClaudeToolDefinition] | None = None
     metadata: dict[str, str] | None = None
+    thinking: ClaudeThinkingConfig | None = None
 
 
 # Response types
@@ -145,7 +165,7 @@ class ClaudeContentBlockStartEvent(BaseModel, frozen=True):
 
     type: Literal["content_block_start"] = "content_block_start"
     index: int
-    content_block: ClaudeTextBlock | ClaudeToolUseBlock
+    content_block: ClaudeTextBlock | ClaudeThinkingBlock | ClaudeToolUseBlock
 
 
 class ClaudeTextDelta(BaseModel, frozen=True):
@@ -155,6 +175,13 @@ class ClaudeTextDelta(BaseModel, frozen=True):
     text: str
 
 
+class ClaudeThinkingDelta(BaseModel, frozen=True):
+    """Delta for thinking content block."""
+
+    type: Literal["thinking_delta"] = "thinking_delta"
+    thinking: str
+
+
 class ClaudeInputJsonDelta(BaseModel, frozen=True):
     """Delta for tool use input JSON content block."""
 
@@ -167,7 +194,7 @@ class ClaudeContentBlockDeltaEvent(BaseModel, frozen=True):
 
     type: Literal["content_block_delta"] = "content_block_delta"
     index: int
-    delta: ClaudeTextDelta | ClaudeInputJsonDelta
+    delta: ClaudeTextDelta | ClaudeThinkingDelta | ClaudeInputJsonDelta
 
 
 class ClaudeContentBlockStopEvent(BaseModel, frozen=True):

src/exo/shared/types/mlx.py

@@ -4,10 +4,13 @@ from collections.abc import Sequence
 
 from mlx_lm.models.cache import (
     ArraysCache,
+    CacheList,
     KVCache,
     QuantizedKVCache,
     RotatingKVCache,
 )
 
 # This list contains one cache entry per transformer layer
-KVCacheType = Sequence[KVCache | RotatingKVCache | QuantizedKVCache | ArraysCache]
+KVCacheType = Sequence[
+    KVCache | RotatingKVCache | QuantizedKVCache | ArraysCache | CacheList
+]

src/exo/shared/types/openai_responses.py

@@ -145,7 +145,23 @@ class ResponseFunctionCallItem(BaseModel, frozen=True):
     status: ResponseStatus = "completed"
 
 
-ResponseItem = ResponseMessageItem | ResponseFunctionCallItem
+class ResponseReasoningSummaryText(BaseModel, frozen=True):
+    """Summary text part in a reasoning output item."""
+
+    type: Literal["summary_text"] = "summary_text"
+    text: str
+
+
+class ResponseReasoningItem(BaseModel, frozen=True):
+    """Reasoning output item in response output array."""
+
+    type: Literal["reasoning"] = "reasoning"
+    id: str
+    summary: list[ResponseReasoningSummaryText] = Field(default_factory=list)
+    status: ResponseStatus = "completed"
+
+
+ResponseItem = ResponseMessageItem | ResponseFunctionCallItem | ResponseReasoningItem
 
 
 class ResponseUsage(BaseModel, frozen=True):
@@ -273,6 +289,58 @@ class ResponseFunctionCallArgumentsDoneEvent(BaseModel, frozen=True):
     arguments: str
 
 
+class ResponseReasoningSummaryPartAddedEvent(BaseModel, frozen=True):
+    """Event sent when a reasoning summary part is added."""
+
+    type: Literal["response.reasoning_summary_part.added"] = (
+        "response.reasoning_summary_part.added"
+    )
+    sequence_number: int
+    item_id: str
+    output_index: int
+    summary_index: int
+    part: ResponseReasoningSummaryText
+
+
+class ResponseReasoningSummaryTextDeltaEvent(BaseModel, frozen=True):
+    """Event sent for reasoning summary text delta during streaming."""
+
+    type: Literal["response.reasoning_summary_text.delta"] = (
+        "response.reasoning_summary_text.delta"
+    )
+    sequence_number: int
+    item_id: str
+    output_index: int
+    summary_index: int
+    delta: str
+
+
+class ResponseReasoningSummaryTextDoneEvent(BaseModel, frozen=True):
+    """Event sent when reasoning summary text is done."""
+
+    type: Literal["response.reasoning_summary_text.done"] = (
+        "response.reasoning_summary_text.done"
+    )
+    sequence_number: int
+    item_id: str
+    output_index: int
+    summary_index: int
+    text: str
+
+
+class ResponseReasoningSummaryPartDoneEvent(BaseModel, frozen=True):
+    """Event sent when a reasoning summary part is done."""
+
+    type: Literal["response.reasoning_summary_part.done"] = (
+        "response.reasoning_summary_part.done"
+    )
+    sequence_number: int
+    item_id: str
+    output_index: int
+    summary_index: int
+    part: ResponseReasoningSummaryText
+
+
 class ResponseCompletedEvent(BaseModel, frozen=True):
     """Event sent when response is completed."""
 
@@ -292,5 +360,9 @@ ResponsesStreamEvent = (
     | ResponseOutputItemDoneEvent
     | ResponseFunctionCallArgumentsDeltaEvent
     | ResponseFunctionCallArgumentsDoneEvent
+    | ResponseReasoningSummaryPartAddedEvent
+    | ResponseReasoningSummaryTextDeltaEvent
+    | ResponseReasoningSummaryTextDoneEvent
+    | ResponseReasoningSummaryPartDoneEvent
     | ResponseCompletedEvent
 )

src/exo/shared/types/worker/runner_response.py

@@ -28,6 +28,7 @@ class GenerationResponse(BaseRunnerResponse):
     finish_reason: FinishReason | None = None
     stats: GenerationStats | None = None
     usage: Usage | None
+    is_thinking: bool = False
 
 
 class ImageGenerationResponse(BaseRunnerResponse):
@@ -67,3 +68,8 @@ class ToolCallResponse(BaseRunnerResponse):
 
 class FinishedResponse(BaseRunnerResponse):
     pass
+
+
+class PrefillProgressResponse(BaseRunnerResponse):
+    processed_tokens: int
+    total_tokens: int

src/exo/utils/channels.py

@@ -1,3 +1,4 @@
+import contextlib
 import multiprocessing as mp
 from dataclasses import dataclass, field
 from math import inf
@@ -132,7 +133,8 @@ class MpSender[T]:
     def close(self) -> None:
         if not self._state.closed.is_set():
             self._state.closed.set()
-        self._state.buffer.put(_MpEndOfStream())
+        with contextlib.suppress(Exception):
+            self._state.buffer.put_nowait(_MpEndOfStream())
         self._state.buffer.close()
 
     # == unique to Mp channels ==
@@ -204,6 +206,8 @@ class MpReceiver[T]:
     def close(self) -> None:
         if not self._state.closed.is_set():
             self._state.closed.set()
+        with contextlib.suppress(Exception):
+            self._state.buffer.put_nowait(_MpEndOfStream())
         self._state.buffer.close()
 
     # == unique to Mp channels ==

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

@@ -5,6 +5,7 @@ import mlx.core as mx
 import psutil
 from mlx_lm.models.cache import (
     ArraysCache,
+    CacheList,
     KVCache,
     QuantizedKVCache,
     RotatingKVCache,
@@ -17,10 +18,22 @@ from exo.worker.engines.mlx import Model
 from exo.worker.engines.mlx.constants import CACHE_GROUP_SIZE, KV_CACHE_BITS
 from exo.worker.runner.bootstrap import logger
 
-# Fraction of device memory above which LRU eviction kicks in
-_DEFAULT_MEMORY_THRESHOLD = 0.9
+
+# Fraction of device memory above which LRU eviction kicks in.
+# Smaller machines need more aggressive eviction.
+def _default_memory_threshold() -> float:
+    total_gb = psutil.virtual_memory().total / (1024**3)
+    if total_gb >= 128:
+        return 0.85
+    if total_gb >= 64:
+        return 0.80
+    if total_gb >= 32:
+        return 0.75
+    return 0.70
+
+
 _MEMORY_THRESHOLD = float(
-    os.environ.get("EXO_MEMORY_THRESHOLD", _DEFAULT_MEMORY_THRESHOLD)
+    os.environ.get("EXO_MEMORY_THRESHOLD", _default_memory_threshold())
 )
 
 
@@ -64,7 +77,7 @@ def has_non_kv_caches(cache: KVCacheType) -> bool:
 
 
 class KVPrefixCache:
-    def __init__(self, group: mx.distributed.Group | None = None):
+    def __init__(self, group: mx.distributed.Group | None):
         self.prompts: list[mx.array] = []  # mx array of tokens (ints)
         self.caches: list[KVCacheType] = []
         self._snapshots: list[list[CacheSnapshot] | None] = []
@@ -156,15 +169,15 @@ class KVPrefixCache:
         best_length = 0
         is_exact = False
 
-        # Find best cache
+        # Find best cache match
         for i, cached_prompt in enumerate(self.prompts):
             length = get_prefix_length(prompt_tokens, cached_prompt)
+            if length >= max_length - 1:
+                best_index, best_length = i, length
+                is_exact = True
+                break
             if length > best_length:
                 best_index, best_length = i, length
-            if length == max_length:
-                is_exact = True
-                best_index, best_length = i, length
-                break
 
         if best_index is None:
             return make_kv_cache(model), prompt_tokens, None
@@ -172,11 +185,12 @@ class KVPrefixCache:
         # For exact match: trim to max_length-1 so remaining has the last token
         # For partial match: trim to best_length, remaining has suffix to prefill
         # This ensures stream_generate always has at least one token to start with
-        target = (max_length - 1) if is_exact else best_length
+        has_ssm = has_non_kv_caches(self.caches[best_index])
+        target = (max_length - 1) if is_exact and not has_ssm else best_length
         restore_pos, restore_snap = self._get_snapshot(best_index, target)
 
         # No usable snapshot — need fresh cache
-        if restore_snap is None and has_non_kv_caches(self.caches[best_index]):
+        if restore_snap is None and has_ssm:
             return make_kv_cache(model), prompt_tokens, None
 
         prompt_cache = deepcopy(self.caches[best_index])
@@ -257,10 +271,21 @@ def encode_prompt(tokenizer: TokenizerWrapper, prompt: str) -> mx.array:
     return mx.array(prompt_tokens)
 
 
+def _entry_length(
+    c: KVCache | RotatingKVCache | QuantizedKVCache | ArraysCache | CacheList,
+) -> int:
+    # Use .offset attribute which KVCache types have (len() not implemented in older QuantizedKVCache).
+    if hasattr(c, "offset"):
+        return c.offset
+    # For CacheList
+    if hasattr(c, "size"):
+        return int(c.size())  # type: ignore
+    return 0
+
+
 def cache_length(cache: KVCacheType) -> int:
     """Get the number of tokens in a KV cache."""
-    # Use .offset attribute which KVCache types have (len() not implemented in older QuantizedKVCache).
-    return max(getattr(c, "offset", 0) for c in cache)
+    return max(_entry_length(c) for c in cache)
 
 
 def get_prefix_length(prompt: mx.array, cached_prompt: mx.array) -> int:

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

@@ -0,0 +1,72 @@
+import json
+import re
+from typing import Any
+
+from mlx_lm.chat_templates import deepseek_v32
+
+from exo.shared.types.api import ToolCallItem
+
+BOS_TOKEN: str = deepseek_v32.bos_token
+EOS_TOKEN: str = deepseek_v32.eos_token
+DSML_TOKEN: str = deepseek_v32.dsml_token
+THINKING_START: str = deepseek_v32.thinking_start_token
+THINKING_END: str = deepseek_v32.thinking_end_token
+USER_TOKEN = "<\uff5cUser\uff5c>"
+ASSISTANT_TOKEN = "<\uff5cAssistant\uff5c>"
+TOOL_CALLS_START = f"<{DSML_TOKEN}function_calls>"
+TOOL_CALLS_END = f"</{DSML_TOKEN}function_calls>"
+encode_messages = deepseek_v32.encode_messages
+
+_INVOKE_PATTERN = re.compile(
+    rf"<{re.escape(DSML_TOKEN)}invoke\s+name=\"([^\"]+)\">"
+    rf"(.*?)"
+    rf"</{re.escape(DSML_TOKEN)}invoke>",
+    re.DOTALL,
+)
+
+_PARAM_PATTERN = re.compile(
+    rf"<{re.escape(DSML_TOKEN)}parameter\s+name=\"([^\"]+)\"\s+string=\"(true|false)\">"
+    rf"(.*?)"
+    rf"</{re.escape(DSML_TOKEN)}parameter>",
+    re.DOTALL,
+)
+
+
+def parse_dsml_output(text: str) -> list[ToolCallItem] | None:
+    """Parse DSML function_calls block from model output text.
+
+    Args:
+        text: The text containing the DSML function_calls block
+              (including the start/end markers).
+
+    Returns:
+        List of ToolCallItem, or None if parsing fails.
+    """
+    tool_calls: list[ToolCallItem] = []
+
+    for invoke_match in _INVOKE_PATTERN.finditer(text):
+        func_name = invoke_match.group(1)
+        invoke_body = invoke_match.group(2)
+
+        args: dict[str, Any] = {}
+        for param_match in _PARAM_PATTERN.finditer(invoke_body):
+            param_name = param_match.group(1)
+            is_string = param_match.group(2) == "true"
+            param_value = param_match.group(3)
+
+            if is_string:
+                args[param_name] = param_value
+            else:
+                try:
+                    args[param_name] = json.loads(param_value)
+                except (json.JSONDecodeError, ValueError):
+                    args[param_name] = param_value
+
+        tool_calls.append(
+            ToolCallItem(
+                name=func_name,
+                arguments=json.dumps(args),
+            )
+        )
+
+    return tool_calls if tool_calls else None

src/exo/worker/engines/mlx/generator/generate.py

@@ -48,7 +48,11 @@ from exo.worker.runner.bootstrap import logger
 
 generation_stream = mx.new_stream(mx.default_device())
 
-_MIN_PREFIX_HIT_TO_UPDATE = 1000
+_MIN_PREFIX_HIT_RATIO_TO_UPDATE = 0.5
+
+
+class PrefillCancelled(BaseException):
+    """Raised when prefill is cancelled via the progress callback."""
 
 
 def prefill(
@@ -57,7 +61,8 @@ def prefill(
     sampler: Callable[[mx.array], mx.array],
     prompt_tokens: mx.array,
     cache: KVCacheType,
-    group: mx.distributed.Group | None = None,
+    group: mx.distributed.Group | None,
+    on_prefill_progress: Callable[[int, int], None] | None,
 ) -> tuple[float, int, list[CacheSnapshot]]:
     """Prefill the KV cache with prompt tokens.
 
@@ -65,7 +70,7 @@ def prefill(
     then trims off the extra generated token.
 
     Returns:
-        tokens_per_sec
+        (tokens_per_sec, num_tokens, snapshots)
     """
     num_tokens = len(prompt_tokens)
     if num_tokens == 0:
@@ -76,6 +81,7 @@ def prefill(
     has_ssm = has_non_kv_caches(cache)
     snapshots: list[CacheSnapshot] = []
 
+    # TODO(evan): kill the callbacks/runner refactor
     def progress_callback(processed: int, total: int) -> None:
         elapsed = time.perf_counter() - start_time
         tok_per_sec = processed / elapsed if elapsed > 0 else 0
@@ -85,6 +91,9 @@ def prefill(
         if has_ssm:
             snapshots.append(snapshot_ssm_states(cache))
 
+        if on_prefill_progress is not None:
+            on_prefill_progress(processed, total)
+
     set_pipeline_prefill(model, is_prefill=True)
 
     mx_barrier(group)
@@ -92,19 +101,23 @@ def prefill(
 
     # Use max_tokens=1 because max_tokens=0 does not work.
     # We just throw away the generated token - we only care about filling the cache
-    for _ in stream_generate(
-        model=model,
-        tokenizer=tokenizer,
-        prompt=prompt_tokens,
-        max_tokens=1,
-        sampler=sampler,
-        prompt_cache=cache,
-        prefill_step_size=8192,
-        kv_group_size=KV_GROUP_SIZE,
-        kv_bits=KV_BITS,
-        prompt_progress_callback=progress_callback,
-    ):
-        break  # Stop after first iteration - cache is now filled
+    try:
+        for _ in stream_generate(
+            model=model,
+            tokenizer=tokenizer,
+            prompt=prompt_tokens,
+            max_tokens=1,
+            sampler=sampler,
+            prompt_cache=cache,
+            prefill_step_size=4096,
+            kv_group_size=KV_GROUP_SIZE,
+            kv_bits=KV_BITS,
+            prompt_progress_callback=progress_callback,
+        ):
+            break  # Stop after first iteration - cache is now filled
+    except PrefillCancelled:
+        set_pipeline_prefill(model, is_prefill=False)
+        raise
 
     set_pipeline_prefill(model, is_prefill=False)
 
@@ -133,7 +146,7 @@ def prefill(
 def warmup_inference(
     model: Model,
     tokenizer: TokenizerWrapper,
-    group: mx.distributed.Group | None = None,
+    group: mx.distributed.Group | None,
 ) -> int:
     content = "Prompt to warm up the inference engine. Repeat this."
 
@@ -255,8 +268,9 @@ def mlx_generate(
     tokenizer: TokenizerWrapper,
     task: TextGenerationTaskParams,
     prompt: str,
-    kv_prefix_cache: KVPrefixCache | None = None,
-    group: mx.distributed.Group | None = None,
+    kv_prefix_cache: KVPrefixCache | None,
+    group: mx.distributed.Group | None,
+    on_prefill_progress: Callable[[int, int], None] | None = None,
 ) -> Generator[GenerationResponse]:
     # Ensure that generation stats only contains peak memory for this generation
     mx.reset_peak_memory()
@@ -311,7 +325,13 @@ def mlx_generate(
 
     # Prefill cache with all tokens except the last one
     prefill_tps, prefill_tokens, ssm_snapshots_list = prefill(
-        model, tokenizer, sampler, prompt_tokens[:-1], caches, group
+        model,
+        tokenizer,
+        sampler,
+        prompt_tokens[:-1],
+        caches,
+        group,
+        on_prefill_progress,
     )
     cache_snapshots: list[CacheSnapshot] | None = ssm_snapshots_list or None
 
@@ -436,9 +456,14 @@ def mlx_generate(
                 full_prompt_tokens = mx.concatenate(
                     [all_prompt_tokens, generated_tokens_array]
                 )
+                hit_ratio = (
+                    prefix_hit_length / len(all_prompt_tokens)
+                    if len(all_prompt_tokens) > 0
+                    else 0.0
+                )
                 if (
                     matched_index is not None
-                    and prefix_hit_length >= _MIN_PREFIX_HIT_TO_UPDATE
+                    and hit_ratio >= _MIN_PREFIX_HIT_RATIO_TO_UPDATE
                 ):
                     kv_prefix_cache.update_kv_cache(
                         matched_index,

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

@@ -1,5 +1,6 @@
 import json
 import os
+import re
 import sys
 import time
 from pathlib import Path
@@ -292,6 +293,8 @@ def get_eos_token_ids_for_model(model_id: ModelId) -> list[int] | None:
     elif "glm" in model_id_lower:
         # For GLM-4.5 and older
         return [151336, 151329, 151338]
+    elif "gpt-oss" in model_id_lower:
+        return [200002, 200012]
     return None
 
 
@@ -405,6 +408,69 @@ def _normalize_tool_calls(msg_dict: dict[str, Any]) -> None:
                 func["arguments"] = json.loads(args)
 
 
+def _collect_nested_property_names(schema: dict[str, Any]) -> set[str]:
+    names: set[str] = set()
+    properties: dict[str, Any] = schema.get("properties", {})  # type: ignore[reportAny]
+    for prop_spec in properties.values():  # pyright: ignore[reportAny]
+        if not isinstance(prop_spec, dict):
+            continue
+        if prop_spec.get("type") == "array":  # type: ignore[reportAny]
+            items: dict[str, Any] | None = prop_spec.get("items")  # type: ignore[reportAny]
+            if isinstance(items, dict) and items.get("type") == "object":  # type: ignore[reportAny]
+                inner_props: dict[str, Any] = items.get("properties", {})  # type: ignore[reportAny]
+                for k in inner_props:  # pyright: ignore[reportUnknownVariableType]
+                    names.add(str(k))  # pyright: ignore[reportUnknownArgumentType]
+                names.update(_collect_nested_property_names(items))  # pyright: ignore[reportUnknownArgumentType]
+    return names
+
+
+def _schemas_lost_in_prompt(prompt: str, tools: list[dict[str, Any]]) -> bool:
+    """Return True if nested property names from any tool schema are absent."""
+    for tool in tools:
+        fn: dict[str, Any] = tool.get("function", {})  # type: ignore
+        params: dict[str, Any] = fn.get("parameters", {})  # type: ignore
+        nested = _collect_nested_property_names(params)
+        if nested and not all(name in prompt for name in nested):
+            return True
+    return False
+
+
+_LOSSY_TEMPLATE_PATTERN = re.compile(
+    r"""inner_type\s*==\s*["']object \| object["']\s*or\s*inner_type\|length\s*>\s*\d+""",
+)
+
+
+def _patch_lossy_chat_template(template: str) -> str | None:
+    """Patch chat templates that collapse nested object schemas to ``any[]``.
+
+    Some templates (e.g., GPT-OSS) have a guard like::
+
+        inner_type == "object | object" or inner_type|length > 50
+
+    The length check silently drops complex array-of-object schemas.
+    We remove the length guard, keeping only the object-union check.
+    Returns the patched template, or *None* if no patch was needed.
+    """
+    patched, n = _LOSSY_TEMPLATE_PATTERN.subn(
+        lambda m: m.group(0).split(" or ")[0],  # keep only the object-union check
+        template,
+    )
+    return patched if n > 0 else None
+
+
+def _needs_dsml_encoding(task_params: TextGenerationTaskParams) -> bool:
+    if "deepseek-v3.2" not in task_params.model.lower():
+        return False
+    # Use DSML encoding when tools are provided or tool results are in the conversation
+    if task_params.tools:
+        return True
+    if task_params.chat_template_messages:
+        return any(
+            msg.get("role") == "tool" for msg in task_params.chat_template_messages
+        )
+    return False
+
+
 def apply_chat_template(
     tokenizer: TokenizerWrapper,
     task_params: TextGenerationTaskParams,
@@ -416,7 +482,6 @@ def apply_chat_template(
 
     When chat_template_messages is available (from Chat Completions API),
     uses those directly to preserve tool_calls, thinking, and other fields.
-    Otherwise builds messages from the task params input/instructions.
     """
     formatted_messages: list[dict[str, Any]] = []
     if task_params.chat_template_messages is not None:
@@ -444,6 +509,19 @@ def apply_chat_template(
         partial_assistant_content = cast(str, formatted_messages[-1].get("content", ""))
         formatted_messages = formatted_messages[:-1]
 
+    if _needs_dsml_encoding(task_params):
+        from exo.worker.engines.mlx.dsml_encoding import encode_messages
+
+        prompt = encode_messages(
+            messages=formatted_messages,
+            thinking_mode="thinking" if task_params.enable_thinking else "chat",
+            tools=task_params.tools,
+        )
+        if partial_assistant_content:
+            prompt += partial_assistant_content
+        logger.info(prompt)
+        return prompt
+
     extra_kwargs: dict[str, Any] = {}
     if task_params.enable_thinking is not None:
         # Qwen3 and GLM use "enable_thinking"; DeepSeek uses "thinking".
@@ -451,14 +529,28 @@ def apply_chat_template(
         extra_kwargs["enable_thinking"] = task_params.enable_thinking
         extra_kwargs["thinking"] = task_params.enable_thinking
 
+    patched_template: str | None = None
+    if task_params.tools:
+        original_template: str | None = getattr(tokenizer, "chat_template", None)
+        if isinstance(original_template, str):
+            patched_template = _patch_lossy_chat_template(original_template)
+            if patched_template is not None:
+                logger.info(
+                    "Patched lossy chat template (removed inner_type length guard)"
+                )
+
     prompt: str = tokenizer.apply_chat_template(
         formatted_messages,
         tokenize=False,
         add_generation_prompt=True,
         tools=task_params.tools,
+        **({"chat_template": patched_template} if patched_template is not None else {}),
         **extra_kwargs,
     )
 
+    if task_params.tools and _schemas_lost_in_prompt(prompt, task_params.tools):
+        logger.warning("Chat template lost nested tool schemas even after patching")
+
     if partial_assistant_content:
         prompt += partial_assistant_content
 

src/exo/worker/main.py

@@ -241,6 +241,11 @@ class Worker:
                     cancelled_task_id=cancelled_task_id, runner_id=runner_id
                 ):
                     await self.runners[runner_id].cancel_task(cancelled_task_id)
+                    await self.event_sender.send(
+                        TaskStatusUpdated(
+                            task_id=task.task_id, task_status=TaskStatus.Complete
+                        )
+                    )
                 case ImageEdits() if task.task_params.total_input_chunks > 0:
                     # Assemble image from chunks and inject into task
                     cmd_id = task.command_id

src/exo/worker/runner/runner.py

@@ -7,10 +7,12 @@ from functools import cache
 from typing import Literal
 
 import mlx.core as mx
+from mlx_lm.models.deepseek_v32 import Model as DeepseekV32Model
 from mlx_lm.models.gpt_oss import Model as GptOssModel
 from mlx_lm.tokenizer_utils import TokenizerWrapper
 from openai_harmony import (  # pyright: ignore[reportMissingTypeStubs]
     HarmonyEncodingName,
+    HarmonyError,  # pyright: ignore[reportUnknownVariableType]
     Role,
     StreamableParser,
     load_harmony_encoding,
@@ -20,7 +22,13 @@ from exo.shared.constants import EXO_MAX_CHUNK_SIZE, EXO_TRACING_ENABLED
 from exo.shared.models.model_cards import ModelId, ModelTask
 from exo.shared.tracing import clear_trace_buffer, get_trace_buffer
 from exo.shared.types.api import ImageGenerationStats
-from exo.shared.types.chunks import ErrorChunk, ImageChunk, TokenChunk, ToolCallChunk
+from exo.shared.types.chunks import (
+    ErrorChunk,
+    ImageChunk,
+    PrefillProgressChunk,
+    TokenChunk,
+    ToolCallChunk,
+)
 from exo.shared.types.common import CommandId
 from exo.shared.types.events import (
     ChunkGenerated,
@@ -80,7 +88,11 @@ from exo.worker.engines.image import (
 )
 from exo.worker.engines.mlx import Model
 from exo.worker.engines.mlx.cache import KVPrefixCache
-from exo.worker.engines.mlx.generator.generate import mlx_generate, warmup_inference
+from exo.worker.engines.mlx.generator.generate import (
+    PrefillCancelled,
+    mlx_generate,
+    warmup_inference,
+)
 from exo.worker.engines.mlx.utils_mlx import (
     apply_chat_template,
     detect_thinking_prompt_suffix,
@@ -297,6 +309,34 @@ def main(
                     assert tokenizer
                     assert check_for_cancel_every
 
+                    # Define callback to send prefill progress events
+                    # and check for cancellation between prefill chunks.
+                    # TODO(evan): kill the callbacks/runner refactor
+                    #  Specifically the part that this is literally duplicated code.
+                    def on_prefill_progress(
+                        processed: int,
+                        total: int,
+                        _task_id: TaskId = task.task_id,
+                        _group: mx.distributed.Group | None = group,
+                    ) -> None:
+                        if device_rank == 0:
+                            event_sender.send(
+                                ChunkGenerated(
+                                    command_id=command_id,
+                                    chunk=PrefillProgressChunk(
+                                        model=shard_metadata.model_card.model_id,
+                                        processed_tokens=processed,
+                                        total_tokens=total,
+                                    ),
+                                )
+                            )
+                        cancelled_tasks.update(cancel_receiver.collect())
+                        want_to_cancel = (_task_id in cancelled_tasks) or (
+                            TaskId("CANCEL_CURRENT_TASK") in cancelled_tasks
+                        )
+                        if mx_any(want_to_cancel, _group):
+                            raise PrefillCancelled()
+
                     try:
                         _check_for_debug_prompts(task_params)
 
@@ -310,19 +350,26 @@ def main(
                             task=task_params,
                             prompt=prompt,
                             kv_prefix_cache=kv_prefix_cache,
+                            on_prefill_progress=on_prefill_progress,
                             group=group,
                         )
 
-                        # For other thinking models (GLM, etc.), check if we need to
-                        # prepend the thinking tag that was consumed by the chat template
-                        if detect_thinking_prompt_suffix(prompt, tokenizer):
+                        if tokenizer.has_thinking:
                             mlx_generator = parse_thinking_models(
-                                mlx_generator, tokenizer
+                                mlx_generator,
+                                tokenizer,
+                                # For other thinking models (GLM, etc.), check if we need to
+                                # prepend the thinking tag that was consumed by the chat template
+                                starts_in_thinking=detect_thinking_prompt_suffix(
+                                    prompt, tokenizer
+                                ),
                             )
 
-                        # GPT-OSS specific parsing to match other model formats.
+                        # Model-specific output parsing for tool calls.
                         if isinstance(inference_model, GptOssModel):
                             mlx_generator = parse_gpt_oss(mlx_generator)
+                        elif isinstance(inference_model, DeepseekV32Model):
+                            mlx_generator = parse_deepseek_v32(mlx_generator)
                         elif tool_parser:
                             mlx_generator = parse_tool_calls(mlx_generator, tool_parser)
 
@@ -374,6 +421,7 @@ def main(
                                                     stats=response.stats,
                                                     logprob=response.logprob,
                                                     top_logprobs=response.top_logprobs,
+                                                    is_thinking=response.is_thinking,
                                                 ),
                                             )
                                         )
@@ -391,6 +439,8 @@ def main(
                                             )
                                         )
 
+                    except PrefillCancelled:
+                        logger.info(f"Prefill cancelled for task {task.task_id}")
                     # can we make this more explicit?
                     except Exception as e:
                         if device_rank == 0:
@@ -588,17 +638,31 @@ def parse_gpt_oss(
 
     for response in responses:
         assert isinstance(response, GenerationResponse)
-        stream.process(response.token)
+        try:
+            stream.process(response.token)
+        except HarmonyError:
+            logger.error("Encountered critical Harmony Error, returning early")
+            return
 
         delta = stream.last_content_delta
         ch = stream.current_channel
         recipient = stream.current_recipient
 
+        # Debug: log every token with state
+        logger.debug(
+            f"parse_gpt_oss token={response.token} text={response.text!r} "
+            f"recipient={recipient!r} ch={ch!r} delta={delta!r} "
+            f"state={stream.state} current_tool={current_tool_name!r}"
+        )
+
         if recipient != current_tool_name:
             if current_tool_name is not None:
                 prefix = "functions."
                 if current_tool_name.startswith(prefix):
                     current_tool_name = current_tool_name[len(prefix) :]
+                logger.info(
+                    f"parse_gpt_oss yielding tool call: name={current_tool_name!r}"
+                )
                 yield ToolCallResponse(
                     tool_calls=[
                         ToolCallItem(
@@ -619,44 +683,208 @@ def parse_gpt_oss(
 
         if ch == "analysis" and not thinking:
             thinking = True
-            yield response.model_copy(update={"text": "<think>"})
 
         if ch != "analysis" and thinking:
             thinking = False
-            yield response.model_copy(update={"text": "</think>"})
 
         if delta:
-            yield response.model_copy(update={"text": delta})
+            yield response.model_copy(update={"text": delta, "is_thinking": thinking})
 
         if response.finish_reason is not None:
-            if thinking:
-                yield response.model_copy(update={"text": "</think>"})
             yield response
 
 
+def parse_deepseek_v32(
+    responses: Generator[GenerationResponse],
+) -> Generator[GenerationResponse | ToolCallResponse]:
+    """Parse DeepSeek V3.2 DSML tool calls from the generation stream.
+
+    Uses accumulated-text matching (not per-token marker checks) because
+    DSML markers like <｜DSML｜function_calls> may span multiple tokens.
+    Also handles <think>...</think> blocks for thinking mode.
+    """
+    from exo.worker.engines.mlx.dsml_encoding import (
+        THINKING_END,
+        THINKING_START,
+        TOOL_CALLS_END,
+        TOOL_CALLS_START,
+        parse_dsml_output,
+    )
+
+    accumulated = ""
+    in_tool_call = False
+    thinking = False
+    # Tokens buffered while we detect the start of a DSML block
+    pending_buffer: list[GenerationResponse] = []
+    # Text accumulated during a tool call block
+    tool_call_text = ""
+
+    for response in responses:
+        assert isinstance(response, GenerationResponse)
+
+        # ── Handle thinking tags ──
+        if not thinking and THINKING_START in response.text:
+            thinking = True
+            # Yield any text before the <think> tag
+            before = response.text[: response.text.index(THINKING_START)]
+            if before:
+                yield response.model_copy(update={"text": before})
+            continue
+
+        if thinking and THINKING_END in response.text:
+            thinking = False
+            # Yield any text after the </think> tag
+            after = response.text[
+                response.text.index(THINKING_END) + len(THINKING_END) :
+            ]
+            if after:
+                yield response.model_copy(update={"text": after, "is_thinking": False})
+            continue
+
+        if thinking:
+            yield response.model_copy(update={"is_thinking": True})
+            continue
+
+        # ── Handle tool call accumulation ──
+        if in_tool_call:
+            tool_call_text += response.text
+            if TOOL_CALLS_END in tool_call_text:
+                # Parse the accumulated DSML block
+                parsed = parse_dsml_output(tool_call_text)
+                if parsed is not None:
+                    logger.info(f"parsed DSML tool calls: {parsed}")
+                    yield ToolCallResponse(
+                        tool_calls=parsed,
+                        usage=response.usage,
+                        stats=response.stats,
+                    )
+                else:
+                    logger.warning(
+                        f"DSML tool call parsing failed for: {tool_call_text}"
+                    )
+                    yield response.model_copy(update={"text": tool_call_text})
+                in_tool_call = False
+                tool_call_text = ""
+                continue
+
+            # EOS reached before end marker — yield buffered text as-is
+            if response.finish_reason is not None:
+                logger.info("DSML tool call parsing interrupted by EOS")
+                yield response.model_copy(update={"text": tool_call_text})
+                in_tool_call = False
+                tool_call_text = ""
+            continue
+
+        # ── Detect start of tool call block ──
+        accumulated += response.text
+
+        if TOOL_CALLS_START in accumulated:
+            # The start marker might be split across pending_buffer + current token
+            start_idx = accumulated.index(TOOL_CALLS_START)
+            # Yield any pending tokens that are purely before the marker
+            pre_text = accumulated[:start_idx]
+            if pre_text:
+                # Flush pending buffer tokens that contributed text before the marker
+                for buf_resp in pending_buffer:
+                    if pre_text:
+                        chunk = buf_resp.text
+                        if len(chunk) <= len(pre_text):
+                            yield buf_resp
+                            pre_text = pre_text[len(chunk) :]
+                        else:
+                            yield buf_resp.model_copy(update={"text": pre_text})
+                            pre_text = ""
+            pending_buffer = []
+            tool_call_text = accumulated[start_idx:]
+            accumulated = ""
+
+            # Check if the end marker is already present (entire tool call in one token)
+            if TOOL_CALLS_END in tool_call_text:
+                parsed = parse_dsml_output(tool_call_text)
+                if parsed is not None:
+                    logger.info(f"parsed DSML tool calls: {parsed}")
+                    yield ToolCallResponse(
+                        tool_calls=parsed,
+                        usage=response.usage,
+                        stats=response.stats,
+                    )
+                else:
+                    logger.warning(
+                        f"DSML tool call parsing failed for: {tool_call_text}"
+                    )
+                    yield response.model_copy(update={"text": tool_call_text})
+                tool_call_text = ""
+            else:
+                in_tool_call = True
+            continue
+
+        # Check if accumulated text might be the start of a DSML marker
+        # Buffer tokens if we see a partial match at the end
+        if _could_be_dsml_prefix(accumulated):
+            pending_buffer.append(response)
+            continue
+
+        # No partial match — flush all pending tokens and the current one
+        for buf_resp in pending_buffer:
+            yield buf_resp
+        pending_buffer = []
+        accumulated = ""
+        yield response
+
+    # Flush any remaining pending buffer at generator end
+    for buf_resp in pending_buffer:
+        yield buf_resp
+
+
+def _could_be_dsml_prefix(text: str) -> bool:
+    """Check if the end of text could be the start of a DSML function_calls marker.
+
+    We look for suffixes of text that are prefixes of the TOOL_CALLS_START pattern.
+    This allows us to buffer tokens until we can determine if a tool call is starting.
+    """
+    from exo.worker.engines.mlx.dsml_encoding import TOOL_CALLS_START
+
+    # Only check the last portion of text that could overlap with the marker
+    max_check = len(TOOL_CALLS_START)
+    tail = text[-max_check:] if len(text) > max_check else text
+
+    # Check if any suffix of tail is a prefix of TOOL_CALLS_START
+    for i in range(len(tail)):
+        suffix = tail[i:]
+        if TOOL_CALLS_START.startswith(suffix):
+            return True
+    return False
+
+
 def parse_thinking_models(
     responses: Generator[GenerationResponse],
     tokenizer: TokenizerWrapper,
+    starts_in_thinking: bool = True,
 ) -> Generator[GenerationResponse]:
+    """Route thinking tokens via is_thinking flag.
+
+    Swallows think tag tokens, sets is_thinking on all others.
+    Always yields tokens with finish_reason to avoid hanging the chunk stream.
     """
-    For models that inject thinking tags in the prompt (like GLM-4.7),
-    prepend the thinking tag to the output stream so the frontend
-    can properly parse thinking content.
-    """
-    first = True
+    in_thinking = starts_in_thinking
     for response in responses:
         if isinstance(response, ToolCallResponse):
             yield response
             continue
-        if first:
-            first = False
-            yield response.model_copy(
-                update={
-                    "text": tokenizer.think_start,
-                    "token": tokenizer.think_start_id,
-                }
-            )
-        yield response
+
+        is_think_tag = (
+            tokenizer.think_end is not None and response.text == tokenizer.think_end
+        ) or (
+            tokenizer.think_start is not None and response.text == tokenizer.think_start
+        )
+
+        if is_think_tag:
+            in_thinking = response.text != tokenizer.think_end
+            # Never swallow finish_reason — the chunk stream needs it to terminate.
+            if response.finish_reason is not None:
+                yield response.model_copy(update={"text": "", "is_thinking": False})
+            continue
+        yield response.model_copy(update={"is_thinking": in_thinking})
 
 
 def _send_image_chunk(

src/exo/worker/runner/runner_supervisor.py

@@ -103,7 +103,7 @@ class RunnerSupervisor:
         self._event_sender.close()
         self._cancel_sender.send(TaskId("CANCEL_CURRENT_TASK"))
         self._cancel_sender.close()
-        self.runner_process.join(1)
+        self.runner_process.join(5)
         if not self.runner_process.is_alive():
             logger.info("Runner process succesfully terminated")
             return

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

@@ -123,7 +123,12 @@ def run_gpt_oss_pipeline_device(
         generated_text = ""
 
         for response in mlx_generate(
-            model=model, tokenizer=tokenizer, task=task, prompt=prompt
+            model=model,
+            tokenizer=tokenizer,
+            task=task,
+            prompt=prompt,
+            kv_prefix_cache=None,
+            group=group,
         ):
             generated_text += response.text
             if response.finish_reason is not None:
@@ -194,6 +199,8 @@ def run_gpt_oss_tensor_parallel_device(
             tokenizer=tokenizer,
             task=task,
             prompt=prompt,
+            kv_prefix_cache=None,
+            group=group,
         ):
             generated_text += response.text
             if response.finish_reason is not None:

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

@@ -88,12 +88,12 @@ class TestKVPrefix:
         return tokenizer
 
     def test_starts_empty(self, mock_tokenizer):
-        cache = KVPrefixCache()
+        cache = KVPrefixCache(None)
         assert len(cache.prompts) == 0
         assert len(cache.caches) == 0
 
     def test_clear_empties_cache(self, mock_tokenizer):
-        cache = KVPrefixCache()
+        cache = KVPrefixCache(None)
         cache.prompts.append(mx.array([1, 2, 3]))
         cache.caches.append([KVCache()])
         cache.clear()
@@ -101,7 +101,7 @@ class TestKVPrefix:
         assert len(cache.caches) == 0
 
     def test_clear_on_empty_cache(self, mock_tokenizer):
-        cache = KVPrefixCache()
+        cache = KVPrefixCache(None)
         cache.clear()
         assert len(cache.prompts) == 0
 
@@ -142,7 +142,9 @@ class TestKVPrefixCacheWithModel:
         tokens = encode_prompt(tokenizer, prompt)
         cache = make_kv_cache(model)
 
-        _, _, snapshots = prefill(model, tokenizer, make_sampler(0.0), tokens, cache)
+        _, _, snapshots = prefill(
+            model, tokenizer, make_sampler(0.0), tokens, cache, group=None
+        )
 
         # Cache should now hold the prompt tokens minus one
         assert cache_length(cache) == len(tokens) - 1
@@ -161,9 +163,11 @@ class TestKVPrefixCacheWithModel:
         tokens = encode_prompt(tokenizer, prompt)
         cache = make_kv_cache(model)
 
-        _, _, snapshots = prefill(model, tokenizer, make_sampler(0.0), tokens, cache)
+        _, _, snapshots = prefill(
+            model, tokenizer, make_sampler(0.0), tokens, cache, group=None
+        )
 
-        kv_prefix_cache = KVPrefixCache()
+        kv_prefix_cache = KVPrefixCache(None)
         kv_prefix_cache.add_kv_cache(tokens, cache, snapshots)
 
         assert len(kv_prefix_cache.prompts) == 1
@@ -176,9 +180,11 @@ class TestKVPrefixCacheWithModel:
         )
         assert matched_index == 0
 
-        # Exact match returns only last token
-        assert len(remaining_tokens) == 1
-        assert mx.array_equal(remaining_tokens, tokens[-1:])
+        # Exact match returns last token(s) — for models with SSM/rotating caches,
+        # snapshot availability constrains how far back we can trim, so remaining
+        # may be 1 or 2 tokens depending on the model.
+        assert len(remaining_tokens) >= 1
+        assert mx.array_equal(remaining_tokens, tokens[-len(remaining_tokens) :])
 
     def test_add_and_get_prefix_match(self, model_and_tokenizer):
         """get_kv_cache with a longer prompt sharing prefix should return partial match."""
@@ -194,10 +200,10 @@ class TestKVPrefixCacheWithModel:
         cache = make_kv_cache(model)
 
         _, _, snapshots = prefill(
-            model, tokenizer, make_sampler(0.0), short_tokens, cache
+            model, tokenizer, make_sampler(0.0), short_tokens, cache, group=None
         )
 
-        kv_prefix_cache = KVPrefixCache()
+        kv_prefix_cache = KVPrefixCache(None)
         kv_prefix_cache.add_kv_cache(short_tokens, cache, snapshots)
 
         # Query with longer prompt that shares the chat template prefix
@@ -238,9 +244,11 @@ class TestKVPrefixCacheWithModel:
         tokens = encode_prompt(tokenizer, prompt)
         cache = make_kv_cache(model)
 
-        _, _, snapshots = prefill(model, tokenizer, make_sampler(0.0), tokens, cache)
+        _, _, snapshots = prefill(
+            model, tokenizer, make_sampler(0.0), tokens, cache, group=None
+        )
 
-        kv_prefix_cache = KVPrefixCache()
+        kv_prefix_cache = KVPrefixCache(None)
         kv_prefix_cache.add_kv_cache(tokens, cache, snapshots)
 
         stored_length = cache_length(kv_prefix_cache.caches[0])
@@ -276,9 +284,11 @@ class TestKVPrefixCacheWithModel:
         tokens = encode_prompt(tokenizer, prompt)
         cache = make_kv_cache(model)
 
-        _, _, snapshots = prefill(model, tokenizer, make_sampler(0.0), tokens, cache)
+        _, _, snapshots = prefill(
+            model, tokenizer, make_sampler(0.0), tokens, cache, group=None
+        )
 
-        kv_prefix_cache = KVPrefixCache()
+        kv_prefix_cache = KVPrefixCache(None)
         kv_prefix_cache.add_kv_cache(tokens, cache, snapshots)
 
         stored_length = cache_length(kv_prefix_cache.caches[0])
@@ -301,7 +311,7 @@ class TestKVPrefixCacheWithModel:
         """mlx_generate should save the cache after generation completes."""
         model, tokenizer = model_and_tokenizer
 
-        kv_prefix_cache = KVPrefixCache()
+        kv_prefix_cache = KVPrefixCache(None)
         task = TextGenerationTaskParams(
             model=DEFAULT_GPT_OSS_MODEL_ID,
             input=[InputMessage(role="user", content="Hello")],
@@ -318,6 +328,7 @@ class TestKVPrefixCacheWithModel:
             task=task,
             prompt=prompt,
             kv_prefix_cache=kv_prefix_cache,
+            group=None,
         ):
             generated_tokens += 1
 
@@ -331,7 +342,7 @@ class TestKVPrefixCacheWithModel:
         """Second mlx_generate call with same prompt should get a prefix hit from stored cache."""
         model, tokenizer = model_and_tokenizer
 
-        kv_prefix_cache = KVPrefixCache()
+        kv_prefix_cache = KVPrefixCache(None)
         task = TextGenerationTaskParams(
             model=DEFAULT_GPT_OSS_MODEL_ID,
             input=[InputMessage(role="user", content="Reuse test")],
@@ -347,6 +358,7 @@ class TestKVPrefixCacheWithModel:
             task=task,
             prompt=prompt,
             kv_prefix_cache=kv_prefix_cache,
+            group=None,
         ):
             pass
 
@@ -368,7 +380,7 @@ class TestKVPrefixCacheWithModel:
         """With a prompt > 1000 tokens, second generation should update the cache entry in-place."""
         model, tokenizer = model_and_tokenizer
 
-        kv_prefix_cache = KVPrefixCache()
+        kv_prefix_cache = KVPrefixCache(None)
 
         # Build a long user message (> 1000 tokens) to exceed _MIN_PREFIX_HIT_TO_UPDATE
         base_text = "The quick brown fox jumps over the lazy dog. "
@@ -395,6 +407,7 @@ class TestKVPrefixCacheWithModel:
             task=task1,
             prompt=prompt1,
             kv_prefix_cache=kv_prefix_cache,
+            group=None,
         ):
             pass
         first_gen_time = time.perf_counter() - t0
@@ -427,6 +440,7 @@ class TestKVPrefixCacheWithModel:
             task=task2,
             prompt=prompt2,
             kv_prefix_cache=kv_prefix_cache,
+            group=None,
         ):
             pass
         second_gen_time = time.perf_counter() - t0
@@ -447,7 +461,7 @@ class TestKVPrefixCacheWithModel:
         """After mlx_generate saves a cache, a second generation must not corrupt the stored copy."""
         model, tokenizer = model_and_tokenizer
 
-        kv_prefix_cache = KVPrefixCache()
+        kv_prefix_cache = KVPrefixCache(None)
         task = TextGenerationTaskParams(
             model=DEFAULT_GPT_OSS_MODEL_ID,
             input=[InputMessage(role="user", content="Immutable test")],
@@ -462,6 +476,7 @@ class TestKVPrefixCacheWithModel:
             task=task,
             prompt=prompt,
             kv_prefix_cache=kv_prefix_cache,
+            group=None,
         ):
             pass
 
@@ -474,6 +489,7 @@ class TestKVPrefixCacheWithModel:
             task=task,
             prompt=prompt,
             kv_prefix_cache=kv_prefix_cache,
+            group=None,
         ):
             pass
 
@@ -484,7 +500,7 @@ class TestKVPrefixCacheWithModel:
         """Under memory pressure, adding a new cache entry evicts the least recently used one."""
         model, tokenizer = model_and_tokenizer
 
-        kv_prefix_cache = KVPrefixCache()
+        kv_prefix_cache = KVPrefixCache(None)
 
         # Add three cache entries with different prompts
         prompts = ["First entry", "Second entry", "Third entry"]
@@ -497,7 +513,7 @@ class TestKVPrefixCacheWithModel:
             prompt = apply_chat_template(tokenizer, task)
             tokens = encode_prompt(tokenizer, prompt)
             cache = make_kv_cache(model)
-            prefill(model, tokenizer, make_sampler(0.0), tokens, cache)
+            prefill(model, tokenizer, make_sampler(0.0), tokens, cache, group=None)
             kv_prefix_cache.add_kv_cache(tokens, cache)
             # Stagger _last_used so LRU order is deterministic
             kv_prefix_cache._last_used[i] = float(i)
@@ -522,7 +538,7 @@ class TestKVPrefixCacheWithModel:
             prompt = apply_chat_template(tokenizer, task)
             tokens = encode_prompt(tokenizer, prompt)
             cache = make_kv_cache(model)
-            prefill(model, tokenizer, make_sampler(0.0), tokens, cache)
+            prefill(model, tokenizer, make_sampler(0.0), tokens, cache, group=None)
             kv_prefix_cache.add_kv_cache(tokens, cache)
 
         # LRU entries should have been evicted (entries 0, 1, 2 in order of _last_used)

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

@@ -0,0 +1,297 @@
+import copy
+import gc
+import importlib
+import json
+import shutil
+import tempfile
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any, cast
+
+import mlx.core as mx
+import mlx.nn as nn
+import pytest
+from mlx.utils import tree_flatten, tree_unflatten
+from mlx_lm.tokenizer_utils import TokenizerWrapper
+
+from exo.shared.types.common import ModelId
+from exo.shared.types.text_generation import InputMessage, TextGenerationTaskParams
+from exo.worker.engines.mlx import Model
+from exo.worker.engines.mlx.cache import KVPrefixCache
+from exo.worker.engines.mlx.generator.generate import mlx_generate
+from exo.worker.engines.mlx.utils_mlx import (
+    apply_chat_template,
+    load_tokenizer_for_model_id,
+)
+
+HF_CACHE = Path.home() / ".cache" / "huggingface" / "hub"
+
+# ── Config reduction ──────────────────────────────────────────────────────── #
+
+_REDUCE = {
+    "num_hidden_layers": 4,
+    "hidden_size": 256,
+    "num_attention_heads": 4,
+    "num_key_value_heads": 4,
+    "intermediate_size": 512,
+    "moe_intermediate_size": 128,
+    "num_experts": 4,
+    "num_experts_per_tok": 2,
+    "n_routed_experts": 4,
+    "num_local_experts": 4,
+    "num_nextn_predict_layers": 0,
+    "first_k_dense_replace": 0,
+    "linear_num_key_heads": 2,
+    "linear_num_value_heads": 2,
+    "num_attention_groups": 4,
+}
+
+
+def _reduce_dict(cfg: dict[str, Any]) -> dict[str, Any]:
+    result = dict(cfg)
+    for key, val in _REDUCE.items():
+        if key in result:
+            result[key] = val
+    return result
+
+
+def _reduce_config(cfg: dict[str, Any]) -> dict[str, Any]:
+    result = _reduce_dict(cfg)
+    n_layers = cast(int, result.get("num_hidden_layers", 4))
+
+    if "text_config" in result and isinstance(result["text_config"], dict):
+        result["text_config"] = _reduce_dict(
+            cast(dict[str, Any], result["text_config"])
+        )
+        tc: dict[str, Any] = result["text_config"]
+        if "num_nextn_predict_layers" in tc:
+            tc["num_nextn_predict_layers"] = 0
+
+    if "layer_types" in result and isinstance(result["layer_types"], list):
+        result["layer_types"] = result["layer_types"][:n_layers]
+
+    if "attention_other_setting" in result and isinstance(
+        result["attention_other_setting"], dict
+    ):
+        aos: dict[str, Any] = dict(
+            cast(dict[str, Any], result["attention_other_setting"])
+        )
+        if "num_attention_heads" in aos:
+            aos["num_attention_heads"] = result.get("num_attention_heads", 4)
+        if "num_attention_groups" in aos:
+            aos["num_attention_groups"] = result.get(
+                "num_attention_groups", cast(int, aos["num_attention_groups"])
+            )
+        result["attention_other_setting"] = aos
+
+    if "moe_layers_enum" in result and isinstance(result["moe_layers_enum"], str):
+        indices = [int(x) for x in result["moe_layers_enum"].split(",") if x.strip()]
+        valid = [i for i in indices if i < n_layers]
+        result["moe_layers_enum"] = ",".join(str(i) for i in valid) if valid else ""
+
+    return result
+
+
+# ── Helpers ───────────────────────────────────────────────────────────────── #
+
+
+def _find_snapshot(hub_name: str) -> Path | None:
+    model_dir = HF_CACHE / f"models--mlx-community--{hub_name}"
+    snaps = model_dir / "snapshots"
+    if not snaps.exists():
+        return None
+    children = sorted(snaps.iterdir())
+    return children[0] if children else None
+
+
+def _copy_tokenizer(src: Path, dst: Path) -> None:
+    for f in src.iterdir():
+        name = f.name
+        if (
+            "tokeniz" in name.lower()
+            or "tiktoken" in name.lower()
+            or name.startswith("vocab")
+            or name.endswith(".jinja")
+            or "tool_declaration" in name
+        ) and f.is_file():
+            shutil.copy2(f, dst / name)
+
+
+def _build_model(module_name: str, cfg: dict[str, Any]) -> Model:
+    mod = importlib.import_module(f"mlx_lm.models.{module_name}")
+    args = mod.ModelArgs.from_dict(cfg)  # pyright: ignore[reportAny]
+    model: nn.Module = mod.Model(args)  # pyright: ignore[reportAny]
+    flat = cast(list[tuple[str, mx.array]], tree_flatten(model.parameters()))
+    random_weights = [
+        (k, mx.random.normal(shape=v.shape, dtype=mx.float16)) for k, v in flat
+    ]
+    model.update(cast(dict[str, Any], tree_unflatten(random_weights)))
+    mx.eval(model.parameters())
+    return cast(Model, model)
+
+
+def _collect_tokens(
+    model: Model,
+    tokenizer: TokenizerWrapper,
+    task: TextGenerationTaskParams,
+    prompt: str,
+    kv_prefix_cache: KVPrefixCache | None,
+) -> list[int]:
+    tokens: list[int] = []
+    for resp in mlx_generate(
+        model=model,
+        tokenizer=tokenizer,
+        task=task,
+        prompt=prompt,
+        kv_prefix_cache=kv_prefix_cache,
+        group=None,
+    ):
+        tokens.append(resp.token)
+        if resp.finish_reason is not None:
+            break
+    return tokens
+
+
+# ── Architecture definitions ──────────────────────────────────────────────── #
+
+
+@dataclass(frozen=True)
+class ArchSpec:
+    name: str
+    hub_name: str
+    module: str
+    tokenizer_hub: str | None = None  # fallback for models without bundled tokenizer
+
+
+ARCHITECTURES: list[ArchSpec] = [
+    ArchSpec("llama", "Llama-3.2-1B-Instruct-4bit", "llama"),
+    ArchSpec("glm_moe_dsa", "GLM-5-MXFP4-Q8", "glm_moe_dsa"),
+    ArchSpec(
+        "glm4_moe", "GLM-4.5-Air-8bit", "glm4_moe", tokenizer_hub="GLM-4.7-8bit-gs32"
+    ),
+    ArchSpec(
+        "glm4_moe_lite",
+        "GLM-4.7-Flash-8bit",
+        "glm4_moe_lite",
+        tokenizer_hub="GLM-4.7-8bit-gs32",
+    ),
+    ArchSpec("glm4_moe_47", "GLM-4.7-8bit-gs32", "glm4_moe"),
+    ArchSpec("qwen3", "Qwen3-4B-Instruct-2507-4bit", "qwen3"),
+    ArchSpec("qwen3_moe", "Qwen3-30B-A3B-4bit", "qwen3_moe"),
+    ArchSpec("qwen3_next", "Qwen3-Next-80B-A3B-Thinking-4bit", "qwen3_next"),
+    ArchSpec("minimax", "MiniMax-M2.1-3bit", "minimax"),
+    ArchSpec("gpt_oss", "gpt-oss-20b-MXFP4-Q8", "gpt_oss"),
+    ArchSpec("step3p5", "Step-3.5-Flash-4bit", "step3p5"),
+    ArchSpec("kimi_k25", "Kimi-K2.5", "kimi_k25"),
+]
+
+
+def _arch_available(spec: ArchSpec) -> bool:
+    snap = _find_snapshot(spec.hub_name)
+    if snap is None:
+        return False
+    if spec.tokenizer_hub is not None:
+        return _find_snapshot(spec.tokenizer_hub) is not None
+    return True
+
+
+def _make_task() -> TextGenerationTaskParams:
+    return TextGenerationTaskParams(
+        model=ModelId("test"),
+        input=[
+            InputMessage(
+                role="user",
+                content="Use the calculator to compute 1847 * 263 + 5921",
+            )
+        ],
+        max_output_tokens=20,
+        temperature=0.0,
+        tools=[
+            {
+                "type": "function",
+                "function": {
+                    "name": "calculate",
+                    "description": "Evaluate a mathematical expression",
+                    "parameters": {
+                        "type": "object",
+                        "properties": {"expression": {"type": "string"}},
+                        "required": ["expression"],
+                    },
+                },
+            }
+        ],
+    )
+
+
+# ── Test class ────────────────────────────────────────────────────────────── #
+
+
+@pytest.mark.slow
+class TestPrefixCacheArchitectures:
+    """Verify prefix cache produces identical output to fresh generation for every architecture."""
+
+    @pytest.fixture(autouse=True)
+    def _cleanup(self):
+        yield
+        mx.clear_cache()
+        gc.collect()
+
+    @pytest.mark.parametrize(
+        "spec",
+        ARCHITECTURES,
+        ids=[a.name for a in ARCHITECTURES],
+    )
+    def test_prefix_cache_exact_hit(self, spec: ArchSpec) -> None:
+        if not _arch_available(spec):
+            pytest.skip(f"Model {spec.hub_name} not cached locally")
+
+        snapshot = _find_snapshot(spec.hub_name)
+        assert snapshot is not None
+
+        tmpdir = Path(tempfile.mkdtemp(prefix=f"exo_test_{spec.name}_"))
+        try:
+            # Build reduced config
+            with open(snapshot / "config.json") as f:
+                cfg = cast(dict[str, Any], json.load(f))
+            reduced = _reduce_config(copy.deepcopy(cfg))
+            (tmpdir / "config.json").write_text(json.dumps(reduced))
+
+            # Copy tokenizer
+            tok_src = snapshot
+            if spec.tokenizer_hub is not None:
+                alt = _find_snapshot(spec.tokenizer_hub)
+                if alt is not None:
+                    tok_src = alt
+            _copy_tokenizer(tok_src, tmpdir)
+
+            # Load tokenizer and model
+            model_id = ModelId(f"mlx-community/{spec.hub_name}")
+            tokenizer = load_tokenizer_for_model_id(model_id, tmpdir)
+            mx.random.seed(0)
+            model = _build_model(spec.module, reduced)
+
+            task = _make_task()
+            prompt = apply_chat_template(tokenizer=tokenizer, task_params=task)
+
+            # Run 1: fresh
+            mx.random.seed(42)
+            fresh = _collect_tokens(model, tokenizer, task, prompt, None)
+            assert len(fresh) > 0, "Fresh generation produced no tokens"
+
+            # Run 2: populate cache
+            kv = KVPrefixCache(None)
+            mx.random.seed(42)
+            populate = _collect_tokens(model, tokenizer, task, prompt, kv)
+
+            # Run 3: exact cache hit
+            mx.random.seed(42)
+            cached = _collect_tokens(model, tokenizer, task, prompt, kv)
+
+            assert fresh == populate, (
+                f"Fresh vs populate mismatch: {fresh[:5]} vs {populate[:5]}"
+            )
+            assert fresh == cached, (
+                f"Fresh vs cached mismatch: {fresh[:5]} vs {cached[:5]}"
+            )
+        finally:
+            shutil.rmtree(tmpdir, ignore_errors=True)

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

@@ -343,8 +343,16 @@ async def test_kimi_tokenizer_specifically():
 @pytest.mark.asyncio
 async def test_glm_tokenizer_specifically():
     """Test GLM tokenizer with its specific EOS tokens."""
+
+    def contains(card: ModelCard, x: str):
+        return x in card.model_id.lower()
+
     glm_model_cards = [
-        card for card in await get_model_cards() if "glm" in card.model_id.lower()
+        card
+        for card in await get_model_cards()
+        if contains(card, "glm")
+        and not contains(card, "-5")
+        and not contains(card, "4.7")
     ]
 
     if not glm_model_cards:

src/exo/worker/tests/unittests/test_runner/test_dsml_e2e.py

@@ -0,0 +1,967 @@
+import json
+from collections.abc import Generator
+from typing import Any
+
+from exo.shared.types.worker.runner_response import (
+    GenerationResponse,
+    ToolCallResponse,
+)
+from exo.worker.engines.mlx.dsml_encoding import (
+    ASSISTANT_TOKEN,
+    BOS_TOKEN,
+    DSML_TOKEN,
+    EOS_TOKEN,
+    THINKING_END,
+    THINKING_START,
+    TOOL_CALLS_END,
+    TOOL_CALLS_START,
+    USER_TOKEN,
+    encode_messages,
+    parse_dsml_output,
+)
+from exo.worker.runner.runner import parse_deepseek_v32
+
+# ── Shared fixtures ──────────────────────────────────────────────
+
+_WEATHER_TOOLS: list[dict[str, Any]] = [
+    {
+        "type": "function",
+        "function": {
+            "name": "get_weather",
+            "description": "Get the current weather in a given city",
+            "parameters": {
+                "type": "object",
+                "properties": {
+                    "city": {"type": "string", "description": "The city name"},
+                    "units": {
+                        "type": "string",
+                        "enum": ["celsius", "fahrenheit"],
+                        "description": "Temperature units",
+                    },
+                },
+                "required": ["city"],
+            },
+        },
+    },
+    {
+        "type": "function",
+        "function": {
+            "name": "get_time",
+            "description": "Get the current time in a timezone",
+            "parameters": {
+                "type": "object",
+                "properties": {
+                    "timezone": {"type": "string"},
+                },
+                "required": ["timezone"],
+            },
+        },
+    },
+]
+
+
+def _simulate_tokens(
+    texts: list[str],
+    finish_on_last: bool = True,
+) -> Generator[GenerationResponse]:
+    """Simulate a model producing tokens from a list of text strings."""
+    for i, text in enumerate(texts):
+        is_last = i == len(texts) - 1
+        yield GenerationResponse(
+            text=text,
+            token=i,
+            finish_reason="stop" if (is_last and finish_on_last) else None,
+            usage=None,
+        )
+
+
+# ── Test: Standard text response (no tool calls) ────────────────
+
+
+class TestE2EStandardResponse:
+    """Model generates a plain text response — no tool calling involved."""
+
+    def test_plain_text_passthrough(self):
+        """Simulate model producing: 'The weather in NYC is 72°F and sunny.'"""
+        # Step 1: Encode the prompt (with tools available)
+        messages: list[dict[str, Any]] = [
+            {"role": "system", "content": "You are a helpful assistant."},
+            {"role": "user", "content": "What's the weather in NYC?"},
+        ]
+        prompt = encode_messages(messages, thinking_mode="chat", tools=_WEATHER_TOOLS)
+
+        # Verify prompt structure
+        assert BOS_TOKEN in prompt
+        assert "## Tools" in prompt
+        assert "get_weather" in prompt
+        assert f"{USER_TOKEN}What's the weather in NYC?{ASSISTANT_TOKEN}" in prompt
+
+        # Step 2: Simulate model response — plain text tokens (no DSML)
+        model_tokens = [
+            "The weather",
+            " in NYC",
+            " is 72",
+            "°F",
+            " and sunny",
+            ".",
+        ]
+        results = list(parse_deepseek_v32(_simulate_tokens(model_tokens)))
+
+        # Step 3: Verify all tokens pass through as GenerationResponse
+        gen_results = [r for r in results if isinstance(r, GenerationResponse)]
+        tool_results = [r for r in results if isinstance(r, ToolCallResponse)]
+
+        assert len(tool_results) == 0
+        assert len(gen_results) == 6
+        full_text = "".join(r.text for r in gen_results)
+        assert full_text == "The weather in NYC is 72°F and sunny."
+        assert gen_results[-1].finish_reason == "stop"
+
+
+# ── Test: Tool call response ─────────────────────────────────────
+
+
+class TestE2EToolCallResponse:
+    """Model generates a DSML tool call — realistic token boundaries."""
+
+    def test_realistic_tool_call_tokens(self):
+        """Simulate model generating a get_weather tool call with realistic token splits.
+
+        Real models split DSML markers across tokens unpredictably.
+        This simulates how DeepSeek V3.2 actually tokenizes DSML output.
+        """
+        # Step 1: Encode prompt
+        messages: list[dict[str, Any]] = [
+            {"role": "system", "content": "You are a helpful assistant."},
+            {"role": "user", "content": "What's the weather in San Francisco?"},
+        ]
+        prompt = encode_messages(messages, thinking_mode="chat", tools=_WEATHER_TOOLS)
+        assert "get_weather" in prompt
+
+        # Step 2: Simulate realistic token-by-token model output
+        # The model first produces some text, then a DSML tool call block
+        model_tokens = [
+            "I'll check the weather for you.",
+            "\n\n",
+            f"<{DSML_TOKEN}",  # marker split across tokens
+            "function_calls>\n",
+            f'<{DSML_TOKEN}invoke name="get_weather">\n',
+            f'<{DSML_TOKEN}parameter name="city" string="true">',
+            "San Francisco",
+            f"</{DSML_TOKEN}parameter>\n",
+            f'<{DSML_TOKEN}parameter name="units" string="false">',
+            '"celsius"',
+            f"</{DSML_TOKEN}parameter>\n",
+            f"</{DSML_TOKEN}invoke>\n",
+            f"</{DSML_TOKEN}function_calls>",
+        ]
+
+        results = list(parse_deepseek_v32(_simulate_tokens(model_tokens)))
+
+        # Step 3: Verify
+        gen_results = [r for r in results if isinstance(r, GenerationResponse)]
+        tool_results = [r for r in results if isinstance(r, ToolCallResponse)]
+
+        # Should have text tokens before tool call + one ToolCallResponse
+        assert len(tool_results) == 1
+        assert len(tool_results[0].tool_calls) == 1
+
+        tc = tool_results[0].tool_calls[0]
+        assert tc.name == "get_weather"
+        args = json.loads(tc.arguments)  # pyright: ignore[reportAny]
+        assert args["city"] == "San Francisco"
+        assert args["units"] == "celsius"
+
+        # The text before the tool call should still be yielded
+        text_before = "".join(r.text for r in gen_results if not r.is_thinking)
+        assert "check the weather" in text_before
+
+    def test_multiple_tool_calls_in_one_block(self):
+        """Model generates two tool calls in a single function_calls block."""
+        messages: list[dict[str, Any]] = [
+            {"role": "system", "content": "You are helpful."},
+            {"role": "user", "content": "Weather in NYC and time in EST?"},
+        ]
+        prompt = encode_messages(messages, thinking_mode="chat", tools=_WEATHER_TOOLS)
+        assert "get_weather" in prompt
+        assert "get_time" in prompt
+
+        # Simulate model output with two invocations
+        model_tokens = [
+            "Let me check both.\n\n",
+            TOOL_CALLS_START,
+            "\n",
+            f'<{DSML_TOKEN}invoke name="get_weather">\n',
+            f'<{DSML_TOKEN}parameter name="city" string="true">NYC</{DSML_TOKEN}parameter>\n',
+            f"</{DSML_TOKEN}invoke>\n",
+            f'<{DSML_TOKEN}invoke name="get_time">\n',
+            f'<{DSML_TOKEN}parameter name="timezone" string="true">EST</{DSML_TOKEN}parameter>\n',
+            f"</{DSML_TOKEN}invoke>\n",
+            TOOL_CALLS_END,
+        ]
+
+        results = list(parse_deepseek_v32(_simulate_tokens(model_tokens)))
+        tool_results = [r for r in results if isinstance(r, ToolCallResponse)]
+
+        assert len(tool_results) == 1
+        assert len(tool_results[0].tool_calls) == 2
+        assert tool_results[0].tool_calls[0].name == "get_weather"
+        assert tool_results[0].tool_calls[1].name == "get_time"
+
+        args0 = json.loads(tool_results[0].tool_calls[0].arguments)  # pyright: ignore[reportAny]
+        args1 = json.loads(tool_results[0].tool_calls[1].arguments)  # pyright: ignore[reportAny]
+        assert args0 == {"city": "NYC"}
+        assert args1 == {"timezone": "EST"}
+
+
+# ── Test: Multi-turn tool use flow ───────────────────────────────
+
+
+class TestE2EMultiTurnToolUse:
+    """Full multi-turn: user asks → model calls tool → tool result → model answers."""
+
+    def test_encode_multi_turn_with_tool_results(self):
+        """Verify the prompt for turn 2 (after tool results) is correctly encoded."""
+        # Turn 1: user asks, model calls tool
+        # Turn 2: tool result provided, model answers
+        messages: list[dict[str, Any]] = [
+            {"role": "system", "content": "You are a weather assistant."},
+            {"role": "user", "content": "What's the weather in NYC?"},
+            {
+                "role": "assistant",
+                "content": "",
+                "tool_calls": [
+                    {
+                        "type": "function",
+                        "function": {
+                            "name": "get_weather",
+                            "arguments": '{"city": "NYC"}',
+                        },
+                    }
+                ],
+            },
+            {"role": "tool", "content": '{"temperature": 72, "condition": "sunny"}'},
+        ]
+
+        prompt = encode_messages(messages, thinking_mode="chat", tools=_WEATHER_TOOLS)
+
+        # Verify multi-turn structure
+        assert BOS_TOKEN in prompt
+        assert "You are a weather assistant." in prompt
+        assert "## Tools" in prompt
+
+        # The assistant's tool call should be encoded as DSML
+        assert TOOL_CALLS_START in prompt
+        assert f'<{DSML_TOKEN}invoke name="get_weather">' in prompt
+        assert EOS_TOKEN in prompt
+
+        # The tool result should be wrapped in function_results
+        assert "<function_results>" in prompt
+        assert "<result>" in prompt
+        assert "72" in prompt
+        assert "</function_results>" in prompt
+
+        # Now simulate model answering after seeing the tool result
+        model_tokens = [
+            "The current",
+            " weather in NYC",
+            " is 72°F",
+            " and sunny.",
+        ]
+        results = list(parse_deepseek_v32(_simulate_tokens(model_tokens)))
+
+        gen_results = [r for r in results if isinstance(r, GenerationResponse)]
+        tool_results = [r for r in results if isinstance(r, ToolCallResponse)]
+
+        assert len(tool_results) == 0
+        full_text = "".join(r.text for r in gen_results)
+        assert full_text == "The current weather in NYC is 72°F and sunny."
+
+    def test_multi_tool_results_encoding(self):
+        """Verify encoding when model called two tools and both return results."""
+        messages: list[dict[str, Any]] = [
+            {"role": "user", "content": "Weather and time?"},
+            {
+                "role": "assistant",
+                "content": "",
+                "tool_calls": [
+                    {
+                        "type": "function",
+                        "function": {
+                            "name": "get_weather",
+                            "arguments": '{"city": "LA"}',
+                        },
+                    },
+                    {
+                        "type": "function",
+                        "function": {
+                            "name": "get_time",
+                            "arguments": '{"timezone": "PST"}',
+                        },
+                    },
+                ],
+            },
+            {"role": "tool", "content": "85F, clear skies"},
+            {"role": "tool", "content": "3:42 PM PST"},
+        ]
+
+        prompt = encode_messages(messages, thinking_mode="chat", tools=_WEATHER_TOOLS)
+
+        # Should have one function_results block with two results
+        assert prompt.count("<function_results>") == 1
+        assert prompt.count("</function_results>") == 1
+        assert "<result>85F, clear skies</result>" in prompt
+        assert "<result>3:42 PM PST</result>" in prompt
+
+
+# ── Test: Thinking + tool call ───────────────────────────────────
+
+
+class TestE2EThinkingAndToolCall:
+    """Model uses thinking mode, reasons, then makes a tool call."""
+
+    def test_thinking_then_tool_call(self):
+        """Model thinks first, then produces a DSML tool call block."""
+        messages: list[dict[str, Any]] = [
+            {"role": "user", "content": "What's the weather?"},
+        ]
+        prompt = encode_messages(
+            messages, tools=_WEATHER_TOOLS, thinking_mode="thinking"
+        )
+        # Thinking mode: prompt should end with <think>
+        assert prompt.endswith(THINKING_START)
+
+        # Simulate: model outputs <think>, thinks, closes thinking, then tool call.
+        # In the full pipeline, parse_thinking_models handles the case where
+        # <think> is in the prompt. Here we test parse_deepseek_v32 directly,
+        # which detects <think>/<think> markers in the stream.
+        model_tokens = [
+            THINKING_START,
+            "The user wants weather",
+            " information. I should use",
+            " the get_weather tool.",
+            THINKING_END,
+            "\n\n",
+            TOOL_CALLS_START,
+            "\n",
+            f'<{DSML_TOKEN}invoke name="get_weather">\n',
+            f'<{DSML_TOKEN}parameter name="city" string="true">',
+            "San Francisco",
+            f"</{DSML_TOKEN}parameter>\n",
+            f"</{DSML_TOKEN}invoke>\n",
+            TOOL_CALLS_END,
+        ]
+
+        results = list(parse_deepseek_v32(_simulate_tokens(model_tokens)))
+
+        gen_results = [r for r in results if isinstance(r, GenerationResponse)]
+        tool_results = [r for r in results if isinstance(r, ToolCallResponse)]
+
+        # Should have thinking tokens + tool call
+        thinking_results = [r for r in gen_results if r.is_thinking]
+
+        assert len(thinking_results) >= 1
+        thinking_text = "".join(r.text for r in thinking_results)
+        assert "get_weather tool" in thinking_text
+
+        assert len(tool_results) == 1
+        assert tool_results[0].tool_calls[0].name == "get_weather"
+        args = json.loads(tool_results[0].tool_calls[0].arguments)  # pyright: ignore[reportAny]
+        assert args["city"] == "San Francisco"
+
+    def test_thinking_prompt_encoding(self):
+        """Verify thinking mode affects prompt encoding correctly."""
+        messages: list[dict[str, Any]] = [
+            {"role": "system", "content": "Be thorough."},
+            {"role": "user", "content": "What's the weather?"},
+        ]
+
+        # With thinking enabled
+        prompt_think = encode_messages(
+            messages, tools=_WEATHER_TOOLS, thinking_mode="thinking"
+        )
+        assert prompt_think.endswith(THINKING_START)
+
+        # With thinking disabled
+        prompt_no_think = encode_messages(
+            messages, tools=_WEATHER_TOOLS, thinking_mode="chat"
+        )
+        assert prompt_no_think.endswith(THINKING_END)
+
+        # Both should have the same tool definitions
+        assert "get_weather" in prompt_think
+        assert "get_weather" in prompt_no_think
+
+
+# ── Test: Round-trip encode → parse ──────────────────────────────
+
+
+class TestE2ERoundTrip:
+    """Verify that DSML we encode can be parsed back correctly."""
+
+    def test_encoded_tool_call_is_parseable(self):
+        """Encode an assistant tool call message, then parse the DSML output."""
+        messages: list[dict[str, Any]] = [
+            {"role": "user", "content": "Weather?"},
+            {
+                "role": "assistant",
+                "content": "",
+                "tool_calls": [
+                    {
+                        "type": "function",
+                        "function": {
+                            "name": "get_weather",
+                            "arguments": '{"city": "Tokyo", "units": "celsius"}',
+                        },
+                    }
+                ],
+            },
+        ]
+
+        prompt = encode_messages(messages, thinking_mode="chat", tools=_WEATHER_TOOLS)
+
+        # Extract the DSML function_calls block from the prompt
+        start = prompt.index(TOOL_CALLS_START)
+        end = prompt.index(TOOL_CALLS_END) + len(TOOL_CALLS_END)
+        dsml_block = prompt[start:end]
+
+        # Parse it back
+        parsed = parse_dsml_output(dsml_block)
+        assert parsed is not None
+        assert len(parsed) == 1
+        assert parsed[0].name == "get_weather"
+        args = json.loads(parsed[0].arguments)  # pyright: ignore[reportAny]
+        assert args["city"] == "Tokyo"
+        assert args["units"] == "celsius"
+
+    def test_encoded_multi_tool_call_round_trips(self):
+        """Encode multiple tool calls, verify they parse back correctly."""
+        messages: list[dict[str, Any]] = [
+            {"role": "user", "content": "Both please"},
+            {
+                "role": "assistant",
+                "content": "",
+                "tool_calls": [
+                    {
+                        "type": "function",
+                        "function": {
+                            "name": "get_weather",
+                            "arguments": '{"city": "Paris"}',
+                        },
+                    },
+                    {
+                        "type": "function",
+                        "function": {
+                            "name": "get_time",
+                            "arguments": '{"timezone": "CET"}',
+                        },
+                    },
+                ],
+            },
+        ]
+
+        prompt = encode_messages(messages, thinking_mode="chat", tools=_WEATHER_TOOLS)
+
+        start = prompt.index(TOOL_CALLS_START)
+        end = prompt.index(TOOL_CALLS_END) + len(TOOL_CALLS_END)
+        dsml_block = prompt[start:end]
+
+        parsed = parse_dsml_output(dsml_block)
+        assert parsed is not None
+        assert len(parsed) == 2
+        assert parsed[0].name == "get_weather"
+        assert parsed[1].name == "get_time"
+        assert json.loads(parsed[0].arguments) == {"city": "Paris"}
+        assert json.loads(parsed[1].arguments) == {"timezone": "CET"}
+
+
+# ── Test: Edge cases with realistic token boundaries ─────────────
+
+
+class TestE2EEdgeCases:
+    """Edge cases that occur in real model inference."""
+
+    def test_dsml_marker_split_at_fullwidth_pipe(self):
+        """The fullwidth pipe character ｜ might be its own token."""
+        # This is a realistic tokenization: the DSML marker is split at the ｜ chars
+        model_tokens = [
+            "Let me help.\n\n",
+            "<\uff5c",  # start of ｜DSML｜
+            "DSML\uff5c",  # rest of DSML token
+            "function_calls>\n",
+            f'<{DSML_TOKEN}invoke name="get_weather">\n',
+            f'<{DSML_TOKEN}parameter name="city" string="true">NYC</{DSML_TOKEN}parameter>\n',
+            f"</{DSML_TOKEN}invoke>\n",
+            TOOL_CALLS_END,
+        ]
+
+        results = list(parse_deepseek_v32(_simulate_tokens(model_tokens)))
+        tool_results = [r for r in results if isinstance(r, ToolCallResponse)]
+
+        assert len(tool_results) == 1
+        assert tool_results[0].tool_calls[0].name == "get_weather"
+
+    def test_tool_call_with_nested_json_object(self):
+        """Model passes a complex JSON object as a non-string parameter."""
+        dsml_block = (
+            f"{TOOL_CALLS_START}\n"
+            f'<{DSML_TOKEN}invoke name="create_event">\n'
+            f'<{DSML_TOKEN}parameter name="title" string="true">Team Standup</{DSML_TOKEN}parameter>\n'
+            f'<{DSML_TOKEN}parameter name="config" string="false">'
+            f'{{"recurring": true, "days": ["mon", "wed", "fri"], "time": "09:00"}}'
+            f"</{DSML_TOKEN}parameter>\n"
+            f"</{DSML_TOKEN}invoke>\n"
+            f"{TOOL_CALLS_END}"
+        )
+
+        # Feed as single token (model might produce it all at once after prefill)
+        results = list(parse_deepseek_v32(_simulate_tokens([dsml_block])))
+        tool_results = [r for r in results if isinstance(r, ToolCallResponse)]
+
+        assert len(tool_results) == 1
+        tc = tool_results[0].tool_calls[0]
+        assert tc.name == "create_event"
+        args = json.loads(tc.arguments)  # pyright: ignore[reportAny]
+        assert args["title"] == "Team Standup"
+        assert args["config"]["recurring"] is True
+        assert args["config"]["days"] == ["mon", "wed", "fri"]
+
+    def test_text_with_angle_brackets_not_mistaken_for_dsml(self):
+        """Angle brackets in normal text should not trigger DSML buffering."""
+        model_tokens = [
+            "The formula is ",
+            "<x, y>",
+            " where x > 0",
+            " and y < 100.",
+        ]
+
+        results = list(parse_deepseek_v32(_simulate_tokens(model_tokens)))
+        gen_results = [r for r in results if isinstance(r, GenerationResponse)]
+        tool_results = [r for r in results if isinstance(r, ToolCallResponse)]
+
+        assert len(tool_results) == 0
+        full_text = "".join(r.text for r in gen_results)
+        assert "formula" in full_text
+        assert "<x, y>" in full_text
+
+    def test_empty_model_response(self):
+        """Model produces only EOS (empty response)."""
+        model_tokens = [""]
+        results = list(parse_deepseek_v32(_simulate_tokens(model_tokens)))
+        gen_results = [r for r in results if isinstance(r, GenerationResponse)]
+        assert len(gen_results) == 1
+        assert gen_results[0].text == ""
+        assert gen_results[0].finish_reason == "stop"
+
+
+# ── Test: Full EPDP spec round-trip ──────────────────────────────
+
+
+class TestE2EFullRoundTrip:
+    """Full round-trip matching the vLLM EPDP spec.
+
+    Simulates the complete multi-turn flow:
+      Turn 1: user asks → think → tool call → tool result → think → answer
+      Turn 2: user asks again → old reasoning stripped → think → answer
+    """
+
+    def test_single_tool_full_flow_with_thinking(self):
+        """Complete flow: user → think → tool call → tool result → think → answer.
+
+        This is the core EPDP flow from the vLLM spec.
+        """
+        # ── Turn 1.1: User asks, encode prompt ──
+        messages: list[dict[str, Any]] = [
+            {"role": "system", "content": "You are a weather assistant."},
+            {"role": "user", "content": "How's the weather in Hangzhou?"},
+        ]
+        prompt_1 = encode_messages(
+            messages, tools=_WEATHER_TOOLS, thinking_mode="thinking"
+        )
+        assert prompt_1.endswith(THINKING_START)
+        assert "## Tools" in prompt_1
+        assert "get_weather" in prompt_1
+
+        # ── Turn 1.1: Model thinks, then calls tool ──
+        model_tokens_1 = [
+            THINKING_START,
+            "The user wants to know the weather in Hangzhou.",
+            " I need to use the get_weather tool.",
+            THINKING_END,
+            "\n\n",
+            TOOL_CALLS_START,
+            "\n",
+            f'<{DSML_TOKEN}invoke name="get_weather">\n',
+            f'<{DSML_TOKEN}parameter name="city" string="true">Hangzhou</{DSML_TOKEN}parameter>\n',
+            f"</{DSML_TOKEN}invoke>\n",
+            TOOL_CALLS_END,
+        ]
+        results_1 = list(parse_deepseek_v32(_simulate_tokens(model_tokens_1)))
+
+        # Verify: thinking tokens + tool call
+        gen_1 = [r for r in results_1 if isinstance(r, GenerationResponse)]
+        tool_1 = [r for r in results_1 if isinstance(r, ToolCallResponse)]
+        thinking_1 = [r for r in gen_1 if r.is_thinking]
+
+        assert len(thinking_1) >= 1
+        assert "get_weather tool" in "".join(r.text for r in thinking_1)
+        assert len(tool_1) == 1
+        assert tool_1[0].tool_calls[0].name == "get_weather"
+        tc_args = json.loads(tool_1[0].tool_calls[0].arguments)  # pyright: ignore[reportAny]
+        assert tc_args == {"city": "Hangzhou"}
+
+        # ── Turn 1.2: Add assistant response + tool result to messages ──
+        messages.append(
+            {
+                "role": "assistant",
+                "content": "",
+                "reasoning_content": "The user wants to know the weather in Hangzhou. I need to use the get_weather tool.",
+                "tool_calls": [
+                    {
+                        "type": "function",
+                        "function": {
+                            "name": "get_weather",
+                            "arguments": '{"city": "Hangzhou"}',
+                        },
+                    }
+                ],
+            }
+        )
+        messages.append(
+            {
+                "role": "tool",
+                "content": '{"temperature": "7~13°C", "condition": "Cloudy"}',
+            }
+        )
+
+        # Encode prompt for turn 1.2
+        prompt_2 = encode_messages(
+            messages, tools=_WEATHER_TOOLS, thinking_mode="thinking"
+        )
+
+        # Verify: prompt has the full conversation structure
+        assert TOOL_CALLS_START in prompt_2  # assistant's encoded tool call
+        assert EOS_TOKEN in prompt_2  # assistant turn ends with EOS
+        assert "<function_results>" in prompt_2
+        assert "<result>" in prompt_2
+        assert "Cloudy" in prompt_2
+        assert "</function_results>" in prompt_2
+        # After tool results with thinking enabled → <think> appended
+        assert prompt_2.endswith(THINKING_START)
+        # The assistant's reasoning_content should appear (it's after last_user_idx)
+        assert "get_weather tool" in prompt_2
+
+        # ── Turn 1.2: Model thinks about results, then answers ──
+        model_tokens_2 = [
+            THINKING_START,
+            "The weather in Hangzhou is Cloudy, 7~13°C.",
+            " I'll tell the user.",
+            THINKING_END,
+            "The weather in Hangzhou is currently cloudy with temperatures between 7°C and 13°C.",
+        ]
+        results_2 = list(parse_deepseek_v32(_simulate_tokens(model_tokens_2)))
+
+        gen_2 = [r for r in results_2 if isinstance(r, GenerationResponse)]
+        tool_2 = [r for r in results_2 if isinstance(r, ToolCallResponse)]
+        thinking_2 = [r for r in gen_2 if r.is_thinking]
+        non_thinking_2 = [r for r in gen_2 if not r.is_thinking]
+
+        assert len(tool_2) == 0  # No more tool calls
+        assert len(thinking_2) >= 1
+        assert "Cloudy" in "".join(r.text for r in thinking_2)
+        assert len(non_thinking_2) >= 1
+        final_text = "".join(r.text for r in non_thinking_2)
+        assert "7°C" in final_text
+        assert "13°C" in final_text
+
+    def test_multi_tool_full_flow(self):
+        """Flow with two tools: user → think → 2 tool calls → 2 results → think → answer."""
+        # ── Initial prompt ──
+        messages: list[dict[str, Any]] = [
+            {"role": "system", "content": "You help with weather and time."},
+            {"role": "user", "content": "Weather in NYC and time in EST?"},
+        ]
+        prompt_1 = encode_messages(
+            messages, tools=_WEATHER_TOOLS, thinking_mode="thinking"
+        )
+        assert prompt_1.endswith(THINKING_START)
+
+        # ── Model thinks, calls both tools ──
+        model_tokens_1 = [
+            THINKING_START,
+            "Two requests: weather and time. I'll call both.",
+            THINKING_END,
+            "\n\n",
+            TOOL_CALLS_START,
+            "\n",
+            f'<{DSML_TOKEN}invoke name="get_weather">\n',
+            f'<{DSML_TOKEN}parameter name="city" string="true">NYC</{DSML_TOKEN}parameter>\n',
+            f"</{DSML_TOKEN}invoke>\n",
+            f'<{DSML_TOKEN}invoke name="get_time">\n',
+            f'<{DSML_TOKEN}parameter name="timezone" string="true">EST</{DSML_TOKEN}parameter>\n',
+            f"</{DSML_TOKEN}invoke>\n",
+            TOOL_CALLS_END,
+        ]
+        results_1 = list(parse_deepseek_v32(_simulate_tokens(model_tokens_1)))
+        tool_1 = [r for r in results_1 if isinstance(r, ToolCallResponse)]
+
+        assert len(tool_1) == 1
+        assert len(tool_1[0].tool_calls) == 2
+        assert tool_1[0].tool_calls[0].name == "get_weather"
+        assert tool_1[0].tool_calls[1].name == "get_time"
+
+        # ── Add assistant + both tool results ──
+        messages.append(
+            {
+                "role": "assistant",
+                "content": "",
+                "reasoning_content": "Two requests: weather and time. I'll call both.",
+                "tool_calls": [
+                    {
+                        "type": "function",
+                        "function": {
+                            "name": "get_weather",
+                            "arguments": '{"city": "NYC"}',
+                        },
+                    },
+                    {
+                        "type": "function",
+                        "function": {
+                            "name": "get_time",
+                            "arguments": '{"timezone": "EST"}',
+                        },
+                    },
+                ],
+            }
+        )
+        messages.append({"role": "tool", "content": "72°F, sunny"})
+        messages.append({"role": "tool", "content": "2:30 PM EST"})
+
+        prompt_2 = encode_messages(
+            messages, tools=_WEATHER_TOOLS, thinking_mode="thinking"
+        )
+
+        # Verify multi-tool result encoding
+        # Count is 2: 1 in _TOOLS_SYSTEM_TEMPLATE example + 1 in conversation
+        assert prompt_2.count("<function_results>") == 2
+        assert prompt_2.count("</function_results>") == 2
+        assert "<result>72°F, sunny</result>" in prompt_2
+        assert "<result>2:30 PM EST</result>" in prompt_2
+        assert prompt_2.endswith(THINKING_START)
+
+        # ── Model thinks about results, answers ──
+        model_tokens_2 = [
+            THINKING_START,
+            "Got both results. Weather is 72°F sunny, time is 2:30 PM.",
+            THINKING_END,
+            "In NYC it's currently 72°F and sunny. The time in EST is 2:30 PM.",
+        ]
+        results_2 = list(parse_deepseek_v32(_simulate_tokens(model_tokens_2)))
+
+        tool_2 = [r for r in results_2 if isinstance(r, ToolCallResponse)]
+        gen_2 = [r for r in results_2 if isinstance(r, GenerationResponse)]
+        non_thinking_2 = [r for r in gen_2 if not r.is_thinking]
+
+        assert len(tool_2) == 0
+        final_text = "".join(r.text for r in non_thinking_2)
+        assert "72°F" in final_text
+        assert "2:30 PM" in final_text
+
+    def test_two_user_turns_reasoning_stripped(self):
+        """Turn 2: old reasoning_content is stripped from history.
+
+        Per the vLLM spec, clear_reasoning_content is called between user turns
+        to save bandwidth. Our _drop_old_thinking handles this.
+        """
+        # Full turn 1 conversation (already completed)
+        messages: list[dict[str, Any]] = [
+            {"role": "system", "content": "You are helpful."},
+            {"role": "user", "content": "Weather in Hangzhou?"},
+            {
+                "role": "assistant",
+                "content": "",
+                "reasoning_content": "I need to call get_weather for Hangzhou.",
+                "tool_calls": [
+                    {
+                        "type": "function",
+                        "function": {
+                            "name": "get_weather",
+                            "arguments": '{"city": "Hangzhou"}',
+                        },
+                    }
+                ],
+            },
+            {"role": "tool", "content": "Cloudy 7~13°C"},
+            {
+                "role": "assistant",
+                "content": "The weather in Hangzhou is cloudy, 7-13°C.",
+                "reasoning_content": "The tool returned cloudy weather. I'll summarize.",
+            },
+            # Turn 2: user asks again
+            {"role": "user", "content": "What about Beijing?"},
+        ]
+
+        prompt = encode_messages(
+            messages, tools=_WEATHER_TOOLS, thinking_mode="thinking"
+        )
+
+        # Old reasoning_content from turn 1 assistants should be STRIPPED
+        # (they're before the last user message at index 5)
+        assert "I need to call get_weather" not in prompt
+        assert "tool returned cloudy" not in prompt
+
+        # But the assistant's content and tool calls should still be there
+        assert "cloudy, 7-13°C" in prompt
+        assert TOOL_CALLS_START in prompt
+
+        # Prompt ends with <think> for the new turn
+        assert prompt.endswith(THINKING_START)
+
+        # ── Turn 2: Model thinks, calls tool for Beijing ──
+        model_tokens = [
+            THINKING_START,
+            "Now the user wants Beijing weather.",
+            THINKING_END,
+            "\n\n",
+            TOOL_CALLS_START,
+            "\n",
+            f'<{DSML_TOKEN}invoke name="get_weather">\n',
+            f'<{DSML_TOKEN}parameter name="city" string="true">Beijing</{DSML_TOKEN}parameter>\n',
+            f"</{DSML_TOKEN}invoke>\n",
+            TOOL_CALLS_END,
+        ]
+        results = list(parse_deepseek_v32(_simulate_tokens(model_tokens)))
+        tool_results = [r for r in results if isinstance(r, ToolCallResponse)]
+
+        assert len(tool_results) == 1
+        assert tool_results[0].tool_calls[0].name == "get_weather"
+        args = json.loads(tool_results[0].tool_calls[0].arguments)  # pyright: ignore[reportAny]
+        assert args == {"city": "Beijing"}
+
+    def test_chained_tool_calls_loop(self):
+        """Model calls tool, gets result, calls another tool, gets result, answers.
+
+        This simulates the inner while loop from the vLLM spec where the model
+        may need multiple sub-turns of tool calling before it has enough info.
+        """
+        # ── Sub-turn 1: user asks, model calls get_time ──
+        messages: list[dict[str, Any]] = [
+            {"role": "system", "content": "You are helpful."},
+            {"role": "user", "content": "What's the weather in Hangzhou tomorrow?"},
+        ]
+
+        prompt_1 = encode_messages(
+            messages, tools=_WEATHER_TOOLS, thinking_mode="thinking"
+        )
+        assert prompt_1.endswith(THINKING_START)
+
+        # Model first calls get_time to figure out the date
+        model_tokens_1 = [
+            THINKING_START,
+            "I need the current date first to calculate tomorrow.",
+            THINKING_END,
+            "\n\n",
+            TOOL_CALLS_START,
+            "\n",
+            f'<{DSML_TOKEN}invoke name="get_time">\n',
+            f'<{DSML_TOKEN}parameter name="timezone" string="true">Asia/Shanghai</{DSML_TOKEN}parameter>\n',
+            f"</{DSML_TOKEN}invoke>\n",
+            TOOL_CALLS_END,
+        ]
+        results_1 = list(parse_deepseek_v32(_simulate_tokens(model_tokens_1)))
+        tool_1 = [r for r in results_1 if isinstance(r, ToolCallResponse)]
+        assert len(tool_1) == 1
+        assert tool_1[0].tool_calls[0].name == "get_time"
+
+        # ── Sub-turn 2: add tool result, model calls get_weather ──
+        messages.append(
+            {
+                "role": "assistant",
+                "content": "",
+                "reasoning_content": "I need the current date first to calculate tomorrow.",
+                "tool_calls": [
+                    {
+                        "type": "function",
+                        "function": {
+                            "name": "get_time",
+                            "arguments": '{"timezone": "Asia/Shanghai"}',
+                        },
+                    }
+                ],
+            }
+        )
+        messages.append({"role": "tool", "content": "2025-12-01 14:30 CST"})
+
+        prompt_2 = encode_messages(
+            messages, tools=_WEATHER_TOOLS, thinking_mode="thinking"
+        )
+        assert "<result>2025-12-01 14:30 CST</result>" in prompt_2
+        assert prompt_2.endswith(THINKING_START)
+
+        # Model now knows the date, calls get_weather
+        model_tokens_2 = [
+            THINKING_START,
+            "Today is 2025-12-01, so tomorrow is 2025-12-02.",
+            " Now I can check weather for Hangzhou.",
+            THINKING_END,
+            "\n\n",
+            TOOL_CALLS_START,
+            "\n",
+            f'<{DSML_TOKEN}invoke name="get_weather">\n',
+            f'<{DSML_TOKEN}parameter name="city" string="true">Hangzhou</{DSML_TOKEN}parameter>\n',
+            f"</{DSML_TOKEN}invoke>\n",
+            TOOL_CALLS_END,
+        ]
+        results_2 = list(parse_deepseek_v32(_simulate_tokens(model_tokens_2)))
+        tool_2 = [r for r in results_2 if isinstance(r, ToolCallResponse)]
+        assert len(tool_2) == 1
+        assert tool_2[0].tool_calls[0].name == "get_weather"
+
+        # ── Sub-turn 3: add weather result, model answers ──
+        messages.append(
+            {
+                "role": "assistant",
+                "content": "",
+                "reasoning_content": "Today is 2025-12-01, so tomorrow is 2025-12-02. Now I can check weather for Hangzhou.",
+                "tool_calls": [
+                    {
+                        "type": "function",
+                        "function": {
+                            "name": "get_weather",
+                            "arguments": '{"city": "Hangzhou"}',
+                        },
+                    }
+                ],
+            }
+        )
+        messages.append({"role": "tool", "content": "Sunny, 5~12°C"})
+
+        prompt_3 = encode_messages(
+            messages, tools=_WEATHER_TOOLS, thinking_mode="thinking"
+        )
+        # Should have both function_results blocks (one per tool round)
+        # Count is 3: 1 in _TOOLS_SYSTEM_TEMPLATE example + 2 in conversation
+        assert prompt_3.count("<function_results>") == 3
+        assert prompt_3.count("</function_results>") == 3
+        assert "<result>2025-12-01 14:30 CST</result>" in prompt_3
+        assert "<result>Sunny, 5~12°C</result>" in prompt_3
+        assert prompt_3.endswith(THINKING_START)
+
+        # Model finally answers
+        model_tokens_3 = [
+            THINKING_START,
+            "I have the weather for tomorrow in Hangzhou.",
+            THINKING_END,
+            "Tomorrow in Hangzhou will be sunny with temperatures between 5°C and 12°C.",
+        ]
+        results_3 = list(parse_deepseek_v32(_simulate_tokens(model_tokens_3)))
+
+        tool_3 = [r for r in results_3 if isinstance(r, ToolCallResponse)]
+        gen_3 = [r for r in results_3 if isinstance(r, GenerationResponse)]
+        non_thinking_3 = [r for r in gen_3 if not r.is_thinking]
+
+        assert len(tool_3) == 0  # No more tool calls — loop ends
+        final_text = "".join(r.text for r in non_thinking_3)
+        assert "sunny" in final_text.lower()
+        assert "5°C" in final_text
+        assert "12°C" in final_text

src/exo/worker/tests/unittests/test_runner/test_event_ordering.py

@@ -148,6 +148,7 @@ class MockTokenizer:
     tool_call_start = None
     tool_call_end = None
     has_tool_calling = False
+    has_thinking = False
 
 
 class MockGroup:

src/exo/worker/tests/unittests/test_runner/test_parse_gpt_oss.py

@@ -0,0 +1,173 @@
+from collections.abc import Generator
+
+from exo.shared.types.worker.runner_response import (
+    GenerationResponse,
+    ToolCallResponse,
+)
+from exo.worker.runner.runner import parse_gpt_oss
+
+# Token IDs from mlx-community/gpt-oss-20b-MXFP4-Q8 tokenizer.
+# These are stable since they come from the model's vocabulary.
+_CHANNEL = 200005  # <|channel|>
+_START = 200006  # <|start|>
+_MESSAGE = 200008  # <|message|>
+_CALL = 200012  # <|call|>
+_END = 200007  # <|end|>
+_ASSISTANT = 173781  # "assistant"
+
+# fmt: off
+# " to=functions.get_current_weather<|channel|>commentary json<|message|>{\"location\": \"Tokyo\"}<|call|>"
+FORMAT_A_TOKENS: list[tuple[int, str]] = [
+    (316,    " to"),
+    (28,     "="),
+    (44580,  "functions"),
+    (775,    ".get"),
+    (23981,  "_current"),
+    (170154, "_weather"),
+    (_CHANNEL, "<|channel|>"),
+    (12606,  "comment"),
+    (815,    "ary"),
+    (5701,   " json"),
+    (_MESSAGE, "<|message|>"),
+    (10848,  '{"'),
+    (7693,   "location"),
+    (1243,   '":'),
+    (392,    ' "'),
+    (173844, "Tokyo"),
+    (18583,  '"}'),
+    (_CALL,  "<|call|>"),
+]
+
+# "<|channel|>commentary to=functions.get_current_weather json<|message|>{\"location\": \"Tokyo\"}<|call|>"
+FORMAT_B_TOKENS: list[tuple[int, str]] = [
+    (_CHANNEL, "<|channel|>"),
+    (12606,  "comment"),
+    (815,    "ary"),
+    (316,    " to"),
+    (28,     "="),
+    (44580,  "functions"),
+    (775,    ".get"),
+    (23981,  "_current"),
+    (170154, "_weather"),
+    (5701,   " json"),
+    (_MESSAGE, "<|message|>"),
+    (10848,  '{"'),
+    (7693,   "location"),
+    (1243,   '":'),
+    (392,    ' "'),
+    (173844, "Tokyo"),
+    (18583,  '"}'),
+    (_CALL,  "<|call|>"),
+]
+
+# "<|channel|>analysis<|message|>Let me think...<|end|><|start|>assistant<|channel|>commentary to=functions.X ..."
+# Full analysis-then-tool-call as the model actually generates it.
+THINKING_THEN_TOOL_TOKENS: list[tuple[int, str]] = [
+    (_CHANNEL, "<|channel|>"),
+    (35644,  "analysis"),
+    (_MESSAGE, "<|message|>"),
+    (12845,  "Let"),
+    (668,    " me"),
+    (2411,   " think"),
+    (1078,   " about"),
+    (495,    " this"),
+    (13,     "."),
+    (_END,   "<|end|>"),
+    # Model generates a new message header for the tool call:
+    (_START, "<|start|>"),
+    (_ASSISTANT, "assistant"),
+    *FORMAT_B_TOKENS,
+]
+# fmt: on
+
+
+def _make_gen_responses(
+    tokens: list[tuple[int, str]],
+) -> list[GenerationResponse]:
+    """Build GenerationResponse list from (token_id, text) pairs."""
+    responses: list[GenerationResponse] = []
+    for i, (tid, text) in enumerate(tokens):
+        is_last = i == len(tokens) - 1
+        responses.append(
+            GenerationResponse(
+                text=text,
+                token=tid,
+                finish_reason="stop" if is_last else None,
+                usage=None,
+            )
+        )
+    return responses
+
+
+def _collect(
+    tokens: list[tuple[int, str]],
+) -> list[GenerationResponse | ToolCallResponse]:
+    """Feed tokens through parse_gpt_oss and collect all yielded responses."""
+
+    def _gen() -> Generator[GenerationResponse, None, None]:
+        yield from _make_gen_responses(tokens)
+
+    return list(parse_gpt_oss(_gen()))
+
+
+def _get_tool_call(
+    results: list[GenerationResponse | ToolCallResponse],
+) -> ToolCallResponse:
+    """Extract the single ToolCallResponse from results."""
+    tool_calls = [r for r in results if isinstance(r, ToolCallResponse)]
+    assert len(tool_calls) == 1, f"Expected 1 ToolCallResponse, got {len(tool_calls)}"
+    return tool_calls[0]
+
+
+class TestParseGptOssRecipientPlacement:
+    """Both Harmony recipient placements must produce identical tool calls."""
+
+    def test_format_a_yields_tool_call(self):
+        results = _collect(FORMAT_A_TOKENS)
+        tc = _get_tool_call(results)
+        assert tc.tool_calls[0].name == "get_current_weather"
+        assert '"location"' in tc.tool_calls[0].arguments
+        assert "Tokyo" in tc.tool_calls[0].arguments
+
+    def test_format_b_yields_tool_call(self):
+        results = _collect(FORMAT_B_TOKENS)
+        tc = _get_tool_call(results)
+        assert tc.tool_calls[0].name == "get_current_weather"
+        assert '"location"' in tc.tool_calls[0].arguments
+        assert "Tokyo" in tc.tool_calls[0].arguments
+
+    def test_both_formats_produce_identical_tool_calls(self):
+        tc_a = _get_tool_call(_collect(FORMAT_A_TOKENS))
+        tc_b = _get_tool_call(_collect(FORMAT_B_TOKENS))
+        assert tc_a.tool_calls[0].name == tc_b.tool_calls[0].name
+        assert tc_a.tool_calls[0].arguments == tc_b.tool_calls[0].arguments
+
+
+class TestParseGptOssThinkingThenToolCall:
+    """Analysis (thinking) followed by a tool call must yield both."""
+
+    def test_thinking_then_tool_call(self):
+        results = _collect(THINKING_THEN_TOOL_TOKENS)
+
+        # Thinking tokens should have is_thinking=True and no <think> tags
+        thinking_responses = [
+            r for r in results if isinstance(r, GenerationResponse) and r.is_thinking
+        ]
+        thinking_text = "".join(r.text for r in thinking_responses)
+        assert "Let me think about this." in thinking_text
+        assert "<think>" not in thinking_text
+        assert "</think>" not in thinking_text
+
+        # Non-thinking tokens should have is_thinking=False
+        non_thinking = [
+            r
+            for r in results
+            if isinstance(r, GenerationResponse) and not r.is_thinking
+        ]
+        non_thinking_text = "".join(r.text for r in non_thinking)
+        assert "<think>" not in non_thinking_text
+
+        # And the tool call
+        tc = _get_tool_call(results)
+        assert tc.tool_calls[0].name == "get_current_weather"
+        assert "Tokyo" in tc.tool_calls[0].arguments

tests/eval_tool_calls.sh

@@ -0,0 +1,55 @@
+#!/usr/bin/env bash
+
+[ $# -lt 1 ] && {
+  echo "Usage: $0 host1 [host2 ...]"
+  exit 1
+}
+
+[ -z "$(git status --porcelain)" ] || {
+  echo "Uncommitted changes"
+  exit 1
+}
+
+commit=$(git rev-parse HEAD)
+git fetch -q origin
+git branch -r --contains "$commit" | grep -qE '^\s*origin/' || {
+  echo "Not pushed to origin"
+  exit 1
+}
+hosts=("$@")
+cleanup() {
+  for host in "${hosts[@]}"; do
+    ssh -T -o BatchMode=yes "$host@$host" "pkill -f bin/exo" &
+  done
+  sleep 1
+  jobs -pr | xargs -r kill 2>/dev/null || true
+}
+trap 'cleanup' EXIT INT TERM
+
+for host; do
+  ssh -T -o BatchMode=yes -o ServerAliveInterval=30 "$host@$host" \
+    "EXO_LIBP2P_NAMESPACE=$commit /nix/var/nix/profiles/default/bin/nix build github:exo-explore/exo/$commit" &
+done
+wait
+for host; do
+  ssh -T -o BatchMode=yes -o ServerAliveInterval=30 "$host@$host" \
+    "EXO_LIBP2P_NAMESPACE=$commit /nix/var/nix/profiles/default/bin/nix run github:exo-explore/exo/$commit" &>/dev/null &
+done
+
+for host; do
+  echo "Waiting for $host..." 1>&2
+  until curl -sf "http://$host:52415/models" &>/dev/null; do sleep 1; done
+done
+
+echo "Waiting 30s for cluster setup" 1>&2
+sleep 30
+echo "EXO loaded" 1>&2
+eval_runner="${hosts[0]}"
+mkdir -p "./bench/$commit"
+nix run .#exo-get-all-models-on-cluster -- "$eval_runner" | while IFS= read -r model; do
+  echo "running eval for $model" 1>&2
+  ssh -Tn -o BatchMode=yes -o ServerAliveInterval=30 "$eval_runner@$eval_runner" \
+    "/nix/var/nix/profiles/default/bin/nix run github:exo-explore/exo/$commit#exo-eval-tool-calls -- --model $model --stdout" \
+    >>"./bench/$commit/${model//\//--}-eval.json"
+  echo
+done

tmp/config_examples/claude_code.sh

@@ -0,0 +1,8 @@
+#!/bin/bash
+# Run Claude Code against a local exo cluster! (Here, GPT OSS 120B)
+ANTHROPIC_BASE_URL="http://localhost:52415/" \
+  ANTHROPIC_AUTH_TOKEN="dummy" \
+  ANTHROPIC_MODEL="mlx-community/gpt-oss-120b-MXFP4-Q8" \
+  ANTHROPIC_SMALL_FAST_MODEL="mlx-community/gpt-oss-120b-MXFP4-Q8" \
+  CLAUDE_CODE_DISABLE_NONESSENTIAL_TRAFFIC=1 \
+  claude