implement engine interface for mlx and mflux

Add more model cards (#1970 )
chore(app): hardcode bug report presigned-URL endpoint (#1971 )
2026-09-11 21:08:48 -04:00 · 2026-04-25 01:52:29 +01:00 · 2026-04-23 15:28:40 +01:00 · 2026-04-23 14:08:16 +00:00 · 2026-04-23 11:39:36 +01:00 · 2026-04-23 01:50:39 +00:00
57 changed files with 1975 additions and 1353 deletions
@@ -32,7 +32,6 @@ jobs:
      SPARKLE_ED25519_PRIVATE: ${{ secrets.SPARKLE_ED25519_PRIVATE }}
      SPARKLE_S3_BUCKET: ${{ secrets.SPARKLE_S3_BUCKET }}
      SPARKLE_S3_PREFIX: ${{ secrets.SPARKLE_S3_PREFIX }}
-      EXO_BUG_REPORT_PRESIGNED_URL_ENDPOINT: ${{ secrets.EXO_BUG_REPORT_PRESIGNED_URL_ENDPOINT }}
      AWS_REGION: ${{ secrets.AWS_REGION }}
      EXO_BUILD_NUMBER: ${{ github.run_number }}
      EXO_LIBP2P_NAMESPACE: ${{ github.ref_name }}
@@ -347,7 +346,6 @@ jobs:
            EXO_BUILD_COMMIT="$GITHUB_SHA" \
            SPARKLE_FEED_URL="$SPARKLE_FEED_URL" \
            SPARKLE_ED25519_PUBLIC="$SPARKLE_ED25519_PUBLIC" \
-            EXO_BUG_REPORT_PRESIGNED_URL_ENDPOINT="$EXO_BUG_REPORT_PRESIGNED_URL_ENDPOINT" \
            CODE_SIGNING_IDENTITY="$SIGNING_IDENTITY" \
            CODE_SIGN_INJECT_BASE_ENTITLEMENTS=YES
          mkdir -p ../../output
@@ -1767,12 +1767,12 @@ def clip(
        array: The clipped array.
    """

-def compile(
-    fun: Callable,
+def compile[F: Callable[..., object]](
+    fun: F,
    inputs: object | None = ...,
    outputs: object | None = ...,
    shapeless: bool = ...,
-) -> Callable:
+) -> F:
    """
    Returns a compiled function which produces the same output as ``fun``.

@@ -2915,8 +2915,8 @@ def gather_mm(
    a: array,
    b: array,
    /,
-    lhs_indices: array,
-    rhs_indices: array,
+    lhs_indices: array | None = ...,
+    rhs_indices: array | None = ...,
    *,
    sorted_indices: bool = ...,
    stream: Stream | Device | None = ...,
@@ -4707,6 +4707,7 @@ def softmax(
    /,
    axis: int | Sequence[int] | None = ...,
    *,
+    precise: bool = ...,
    stream: Stream | Device | None = ...,
 ) -> array:
    """
@@ -57,6 +57,10 @@ class Module(dict):
    def __init__(self) -> None:
        """Should be called by the subclasses of ``Module``."""

+    def __getitem__(self, key: str) -> mx.array | Module: ...
+    def get(
+        self, key: str, default: mx.array | Module | None = ...
+    ) -> mx.array | Module | None: ...
    @property
    def training(self):  # -> bool:
        """Boolean indicating if the model is in training mode."""
@@ -3,7 +3,7 @@ This type stub file was generated by pyright.
 """

 from dataclasses import dataclass
-from typing import Optional
+from typing import Any, Optional

 import mlx.core as mx

@@ -37,10 +37,10 @@ def quantized_scaled_dot_product_attention(
    bits: int = ...,
 ) -> mx.array: ...
 def scaled_dot_product_attention(
-    queries,
-    keys,
-    values,
-    cache,
+    queries: mx.array,
+    keys: mx.array,
+    values: mx.array,
+    cache: Optional[Any],
    scale: float,
    mask: Optional[mx.array],
    sinks: Optional[mx.array] = ...,
@@ -0,0 +1,103 @@
+"""Type stubs for mlx_lm.models.gpt_oss"""
+
+from dataclasses import dataclass
+from typing import Any, List, Optional
+
+import mlx.core as mx
+import mlx.nn as nn
+
+from .base import BaseModelArgs
+from .cache import KVCache
+from .switch_layers import SwitchGLU
+
+@dataclass
+class ModelArgs(BaseModelArgs):
+    model_type: str
+    hidden_size: int
+    intermediate_size: int
+    num_hidden_layers: int
+    num_attention_heads: int
+    num_key_value_heads: int
+    num_local_experts: int
+    num_experts_per_tok: int
+    vocab_size: int
+    rms_norm_eps: float
+    sliding_window: int
+    layer_types: Optional[List[str]]
+
+def mlx_topk(a: mx.array, k: int, axis: int = -1) -> tuple[mx.array, mx.array]: ...
+
+class AttentionBlock(nn.Module):
+    head_dim: int
+    num_attention_heads: int
+    num_key_value_heads: int
+    num_key_value_groups: int
+    sinks: mx.array
+    q_proj: nn.Linear
+    k_proj: nn.Linear
+    v_proj: nn.Linear
+    o_proj: nn.Linear
+    sm_scale: float
+    rope: nn.Module
+
+    def __init__(self, config: ModelArgs) -> None: ...
+    def __call__(
+        self,
+        x: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[Any] = None,
+    ) -> mx.array: ...
+
+class TransformerBlock(nn.Module):
+    self_attn: AttentionBlock
+    mlp: MLPBlock
+
+    def __init__(self, config: ModelArgs) -> None: ...
+    def __call__(
+        self,
+        x: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[Any] = None,
+    ) -> mx.array: ...
+
+class MLPBlock(nn.Module):
+    hidden_size: int
+    num_local_experts: int
+    num_experts_per_tok: int
+    experts: SwitchGLU
+    router: nn.Linear
+    sharding_group: Optional[mx.distributed.Group]
+
+    def __init__(self, config: ModelArgs) -> None: ...
+    def __call__(self, x: mx.array) -> mx.array: ...
+
+class GptOssMoeModel(nn.Module):
+    embed_tokens: nn.Embedding
+    norm: nn.RMSNorm
+    layer_types: List[str]
+    layers: list[TransformerBlock]
+    window_size: int
+    swa_idx: int
+    ga_idx: int
+
+    def __init__(self, args: ModelArgs) -> None: ...
+    def __call__(
+        self,
+        inputs: mx.array,
+        cache: Optional[Any] = None,
+    ) -> mx.array: ...
+
+class Model(nn.Module):
+    model_type: str
+    model: GptOssMoeModel
+    lm_head: nn.Linear
+
+    def __init__(self, args: ModelArgs) -> None: ...
+    def __call__(
+        self,
+        inputs: mx.array,
+        cache: Optional[Any] = None,
+    ) -> mx.array: ...
+    @property
+    def layers(self) -> list[nn.Module]: ...
+    def make_cache(self) -> list[KVCache]: ...
@@ -0,0 +1,94 @@
+"""Type stubs for mlx_lm.models.minimax"""
+
+from dataclasses import dataclass
+from typing import Any, Optional
+
+import mlx.core as mx
+import mlx.nn as nn
+
+from .base import BaseModelArgs
+from .switch_layers import SwitchGLU
+
+@dataclass
+class ModelArgs(BaseModelArgs):
+    model_type: str
+    hidden_size: int
+    intermediate_size: int
+    num_hidden_layers: int
+    num_attention_heads: int
+    num_key_value_heads: int
+    num_local_experts: int
+    num_experts_per_tok: int
+    max_position_embeddings: int
+
+class MiniMaxAttention(nn.Module):
+    num_heads: int
+    num_attention_heads: int
+    num_key_value_heads: int
+    head_dim: int
+    scale: float
+    q_proj: nn.Linear
+    k_proj: nn.Linear
+    v_proj: nn.Linear
+    o_proj: nn.Linear
+    q_norm: nn.Module
+    k_norm: nn.Module
+    rope: nn.Module
+
+    def __init__(self, args: ModelArgs) -> None: ...
+    def __call__(
+        self,
+        x: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[Any] = None,
+    ) -> mx.array: ...
+
+class MiniMaxSparseMoeBlock(nn.Module):
+    num_experts_per_tok: int
+    gate: nn.Linear
+    switch_mlp: SwitchGLU
+    e_score_correction_bias: mx.array
+    sharding_group: Optional[mx.distributed.Group]
+
+    def __init__(self, args: ModelArgs) -> None: ...
+    def __call__(self, x: mx.array) -> mx.array: ...
+
+class MiniMaxDecoderLayer(nn.Module):
+    self_attn: MiniMaxAttention
+    block_sparse_moe: MiniMaxSparseMoeBlock
+    input_layernorm: nn.RMSNorm
+    post_attention_layernorm: nn.RMSNorm
+
+    def __init__(self, args: ModelArgs) -> None: ...
+    def __call__(
+        self,
+        x: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[Any] = None,
+    ) -> mx.array: ...
+
+class MiniMaxModel(nn.Module):
+    embed_tokens: nn.Embedding
+    layers: list[MiniMaxDecoderLayer]
+    norm: nn.RMSNorm
+
+    def __init__(self, args: ModelArgs) -> None: ...
+    def __call__(
+        self,
+        inputs: mx.array,
+        cache: Optional[Any] = None,
+    ) -> mx.array: ...
+
+class Model(nn.Module):
+    model_type: str
+    model: MiniMaxModel
+    lm_head: nn.Linear
+
+    def __init__(self, args: ModelArgs) -> None: ...
+    def __call__(
+        self,
+        inputs: mx.array,
+        cache: Optional[Any] = None,
+    ) -> mx.array: ...
+    @property
+    def layers(self) -> list[MiniMaxDecoderLayer]: ...
@@ -92,6 +92,15 @@ class NemotronHAttention(nn.Module):
        cache: Optional[KVCache] = None,
    ) -> mx.array: ...

+class MoEGate(nn.Module):
+    config: ModelArgs
+    top_k: int
+    norm_topk_prob: bool
+    weight: mx.array
+
+    def __init__(self, config: ModelArgs) -> None: ...
+    def __call__(self, x: mx.array) -> tuple[mx.array, mx.array]: ...
+
 class NemotronHMLP(nn.Module):
    up_proj: nn.Linear
    down_proj: nn.Linear
@@ -102,9 +111,14 @@ class NemotronHMLP(nn.Module):
    def __call__(self, x: mx.array) -> mx.array: ...

 class NemotronHMoE(nn.Module):
+    config: ModelArgs
    num_experts_per_tok: int
+    moe_latent_size: Optional[int]
    switch_mlp: SwitchMLP
+    gate: MoEGate
    shared_experts: NemotronHMLP
+    fc1_latent_proj: nn.Linear
+    fc2_latent_proj: nn.Linear

    def __init__(self, config: ModelArgs) -> None: ...
    def __call__(self, x: mx.array) -> mx.array: ...
@@ -71,6 +71,7 @@ class Qwen3NextAttention(nn.Module):
 class Qwen3NextSparseMoeBlock(nn.Module):
    norm_topk_prob: bool
    num_experts: int
+    num_experts_per_tok: int
    top_k: int
    gate: nn.Linear
    switch_mlp: SwitchGLU
@@ -584,9 +584,18 @@ struct ContentView: View {

            case .prompting:
                VStack(alignment: .leading, spacing: 6) {
-                    Text("What's the issue? (optional)")
+                    VStack(alignment: .leading, spacing: 2) {
+                        Text("Tell us what went wrong (optional)")
+                            .font(.caption2)
+                            .foregroundColor(.secondary)
+                        Text(
+                            "A quick description of what you were doing and what happened helps us track down the bug for you."
+                        )
                        .font(.caption2)
                        .foregroundColor(.secondary)
+                        .opacity(0.8)
+                        .fixedSize(horizontal: false, vertical: true)
+                    }
                    TextEditor(text: $bugReportUserDescription)
                        .font(.caption2)
                        .frame(height: 60)
@@ -9,7 +9,7 @@
 	<key>EXOBuildCommit</key>
 	<string>$(EXO_BUILD_COMMIT)</string>
 	<key>EXOBugReportPresignedUrlEndpoint</key>
-	<string>$(EXO_BUG_REPORT_PRESIGNED_URL_ENDPOINT)</string>
+	<string>https://reports.exolabs.net/presigned-urls</string>
 	<key>NSLocalNetworkUsageDescription</key>
 	<string>EXO needs local network access to discover and connect to other devices in your cluster for distributed AI inference.</string>
 	<key>NSBonjourServices</key>
@@ -3,11 +3,13 @@ from __future__ import annotations

 import argparse
 import contextlib
+import io
 import json
 import os
 import sys
 import time
 import tomllib
+from concurrent.futures import ThreadPoolExecutor, as_completed
 from dataclasses import dataclass, field
 from pathlib import Path
 from typing import Any, Literal
@@ -209,7 +211,7 @@ def _openai_build_request(
        "model": model,
        "messages": messages,
        "tools": tools,
-        "max_tokens": 16384,
+        "max_tokens": 4096,
        "temperature": 0.0,
    }
    return "/v1/chat/completions", body
@@ -276,7 +278,7 @@ def _openai_build_followup(
        "model": model,
        "messages": followup_messages,
        "tools": tools,
-        "max_tokens": 16384,
+        "max_tokens": 4096,
        "temperature": 0.0,
    }
    return "/v1/chat/completions", body
@@ -379,7 +381,7 @@ def _claude_build_request(
        "model": model,
        "messages": claude_messages,
        "tools": claude_tools,
-        "max_tokens": 16384,
+        "max_tokens": 4096,
        "temperature": 0.0,
    }
    if system_content is not None:
@@ -489,7 +491,7 @@ def _claude_build_followup(
        "model": model,
        "messages": claude_messages,
        "tools": claude_tools,
-        "max_tokens": 16384,
+        "max_tokens": 4096,
        "temperature": 0.0,
    }
    if system_content is not None:
@@ -913,6 +915,12 @@ Examples:
        default=1,
        help="Repeat each scenario N times (default: 1)",
    )
+    parser.add_argument(
+        "--concurrency",
+        type=int,
+        default=1,
+        help="Run up to N scenarios in parallel against the same instance (default: 1)",
+    )
    parser.add_argument(
        "--scenarios",
        nargs="*",
@@ -935,6 +943,13 @@ Examples:
    )
    args = parser.parse_args()

+    if args.concurrency < 1:
+        print(
+            f"--concurrency must be >= 1 (got {args.concurrency})",
+            file=sys.stderr,
+        )
+        sys.exit(2)
+
    all_scenarios = load_scenarios(SCENARIOS_PATH)
    if args.scenarios:
        scenarios = [s for s in all_scenarios if s.name in args.scenarios]
@@ -1010,42 +1025,72 @@ Examples:
    cluster_snapshot = capture_cluster_snapshot(exo)
    all_results: list[ScenarioResult] = []

+    tasks: list[tuple[int, Scenario, ApiName]] = [
+        (run_idx, scenario, api_name)
+        for run_idx in range(args.repeat)
+        for scenario in scenarios
+        for api_name in api_names
+    ]
+
+    def _run_one(
+        http_client: httpx.Client,
+        task: tuple[int, Scenario, ApiName],
+    ) -> tuple[tuple[int, Scenario, ApiName], list[ScenarioResult], str]:
+        run_idx, scenario, api_name = task
+        buf = io.StringIO()
+        run_tag = f"[run {run_idx + 1}/{args.repeat}]" if args.repeat > 1 else ""
+        print(
+            f"\n  {run_tag}[{api_name:>9}] {scenario.name}: {scenario.description}",
+            file=buf,
+        )
+        scenario_results = run_scenario(
+            http_client,
+            args.host,
+            args.port,
+            full_model_id,
+            scenario,
+            api_name,
+            args.timeout,
+            args.verbose,
+        )
+        for r in scenario_results:
+            status = "PASS" if r.passed else "FAIL"
+            print(
+                f"    [{r.phase:>10}] {status}  ({r.latency_ms:.0f}ms)",
+                file=buf,
+            )
+            for check_name, check_ok in r.checks.items():
+                mark = "+" if check_ok else "-"
+                print(f"      {mark} {check_name}", file=buf)
+            if r.error:
+                print(f"      ! {r.error}", file=buf)
+        return task, scenario_results, buf.getvalue()
+
    try:
        with httpx.Client() as http_client:
-            for run_idx in range(args.repeat):
-                if args.repeat > 1:
-                    print(f"\n--- Run {run_idx + 1}/{args.repeat} ---", file=log)
-
-                for scenario in scenarios:
-                    for api_name in api_names:
-                        print(
-                            f"\n  [{api_name:>9}] {scenario.name}: {scenario.description}",
-                            file=log,
-                        )
-
-                        scenario_results = run_scenario(
-                            http_client,
-                            args.host,
-                            args.port,
-                            full_model_id,
-                            scenario,
-                            api_name,
-                            args.timeout,
-                            args.verbose,
-                        )
+            if args.concurrency == 1:
+                current_run = -1
+                for task in tasks:
+                    run_idx = task[0]
+                    if args.repeat > 1 and run_idx != current_run:
+                        print(f"\n--- Run {run_idx + 1}/{args.repeat} ---", file=log)
+                        current_run = run_idx
+                    _, scenario_results, buffered = _run_one(http_client, task)
+                    all_results.extend(scenario_results)
+                    log.write(buffered)
+                    log.flush()
+            else:
+                print(
+                    f"Running {len(tasks)} tasks with concurrency={args.concurrency}",
+                    file=log,
+                )
+                with ThreadPoolExecutor(max_workers=args.concurrency) as pool:
+                    futures = [pool.submit(_run_one, http_client, t) for t in tasks]
+                    for fut in as_completed(futures):
+                        _, scenario_results, buffered = fut.result()
                        all_results.extend(scenario_results)
-
-                        for r in scenario_results:
-                            status = "PASS" if r.passed else "FAIL"
-                            print(
-                                f"    [{r.phase:>10}] {status}  ({r.latency_ms:.0f}ms)",
-                                file=log,
-                            )
-                            for check_name, check_ok in r.checks.items():
-                                mark = "+" if check_ok else "-"
-                                print(f"      {mark} {check_name}", file=log)
-                            if r.error:
-                                print(f"      ! {r.error}", file=log)
+                        log.write(buffered)
+                        log.flush()
    finally:
        try:
            exo.request_json("DELETE", f"/instance/{instance_id}")
@@ -564,7 +564,7 @@ def add_common_instance_args(ap: argparse.ArgumentParser) -> None:
    ap.add_argument(
        "--settle-timeout",
        type=float,
-        default=0,
+        default=60.0,
        help="Max seconds to wait for the cluster to produce valid placements (0 = try once).",
    )
    ap.add_argument(
@@ -88,10 +88,12 @@
  let codexModel = $state("");
  let codexMcpPath = $state("/Users/username");
  let openClawModel = $state("");
+  let piModel = $state("");
  $effect(() => {
    const def = modelsBySize.length > 0 ? modelsBySize[0] : "your-model-id";
    codexModel = def;
    openClawModel = def;
+    piModel = def;
  });

  const claudeShellCommand = $derived(
@@ -218,6 +220,55 @@
    ),
  );

+  const piModelsJson = $derived.by(() => {
+    const models: Record<string, unknown>[] = [];
+    for (const modelId of runningModels) {
+      const caps = modelCapabilities[modelId] || [];
+      const ctxLen = modelContextLengths[modelId] || 0;
+      const entry: Record<string, unknown> = { id: modelId };
+      if (caps.includes("vision")) {
+        entry.input = ["text", "image"];
+      }
+      // Mark thinking-capable models so pi surfaces its thinking-level selector
+      // for them. exo capability strings: "thinking" (model emits reasoning
+      // content) and "thinking_toggle" (user can turn it on/off).
+      if (caps.includes("thinking") || caps.includes("thinking_toggle")) {
+        entry.reasoning = true;
+      }
+      if (ctxLen > 0) {
+        entry.contextWindow = ctxLen;
+      }
+      models.push(entry);
+    }
+    if (models.length === 0) {
+      models.push({ id: "your-model-id" });
+    }
+    return JSON.stringify(
+      {
+        providers: {
+          exo: {
+            baseUrl: `${apiUrl}/v1`,
+            api: "openai-completions",
+            apiKey: "exo",
+            compat: {
+              supportsDeveloperRole: false,
+              // exo's OpenAI surface takes a boolean `enable_thinking` toggle,
+              // not graded effort levels, so disable pi's `reasoning_effort`
+              // parameter and use the matching top-level-boolean format.
+              supportsReasoningEffort: false,
+              thinkingFormat: "qwen",
+            },
+            models,
+          },
+        },
+      },
+      null,
+      2,
+    );
+  });
+
+  const piShellCommand = $derived(`pi --provider exo --model ${piModel}`);
+
  const ollamaCommand = $derived(
    `OLLAMA_HOST=${apiUrl}/ollama ollama run ${modelsBySize.length > 0 ? modelsBySize[0] : "your-model-id"}`,
  );
@@ -277,6 +328,7 @@
    "OpenCode",
    "Codex",
    "OpenClaw",
+    "Pi",
    "Open WebUI",
    "n8n",
    "Firefox",
@@ -515,6 +567,33 @@
          config={`openclaw doctor --fix${(modelCapabilities[openClawModel] || []).includes("vision") ? `\nopenclaw models set-image exo/${openClawModel}` : ""}\nopenclaw gateway &\nopenclaw dashboard`}
          language="bash"
        />
+      {:else if activeTab === "Pi"}
+        {#if runningModels.length > 1}
+          <div class="text-xs">
+            <span
+              class="text-exo-light-gray/50 text-[10px] uppercase tracking-wider block mb-1"
+              >Model</span
+            >
+            <select bind:value={piModel} class={selectClass}>
+              {#each runningModels as model}
+                <option value={model}>{model.split("/").pop()}</option>
+              {/each}
+            </select>
+          </div>
+        {/if}
+        <IntegrationCard
+          title="Models Config"
+          subtitle="~/.pi/agent/models.json"
+          description="Register exo as a custom provider in pi. Create or edit this file, then run pi and pick an exo model via /model. Install pi with: npm install -g @mariozechner/pi-coding-agent"
+          config={piModelsJson}
+        />
+        <IntegrationCard
+          title="Shell Command"
+          subtitle="Run in terminal"
+          description="Launch pi directly with the exo provider and model selected."
+          config={piShellCommand}
+          language="bash"
+        />
      {:else if activeTab === "Open WebUI"}
        <IntegrationCard
          title="1. Start Open WebUI"
@@ -0,0 +1,6 @@
+{
+  "name": "exo",
+  "lockfileVersion": 3,
+  "requires": true,
+  "packages": {}
+}
@@ -3,7 +3,7 @@ name = "exo"
 version = "0.3.70"
 description = "Exo"
 readme = "README.md"
-requires-python = ">=3.13"
+requires-python = "==3.13.*"
 dependencies = [
  "aiofiles>=24.1.0",
  "aiohttp>=3.12.14",
@@ -17,8 +17,8 @@ dependencies = [
  "loguru>=0.7.3",
  "exo-pyo3-bindings",                       # rust bindings
  "anyio==4.11.0",
-  "mlx==0.31.1; sys_platform == 'darwin'",
-  "mlx-lm",
+  "mlx==0.31.2; sys_platform == 'darwin'",
+  "mlx-lm; sys_platform=='darwin'",
  "tiktoken>=0.12.0",                        # required for kimi k2 tokenizer
  "hypercorn>=0.18.0",
  "openai-harmony>=0.0.8",
@@ -49,15 +49,21 @@ dev = [
 [project.optional-dependencies]
 build = ["nanobind"]
 cpu = [
+  "mlx==0.31.1; sys_platform == 'linux'",
  "mlx-cpu==0.31.1; sys_platform == 'linux'",
+  "mlx-lm; sys_platform == 'linux'",
  "torch>=2.10.0; sys_platform == 'linux'",
 ]
 cuda12 = [
+  "mlx==0.31.1; sys_platform == 'linux'",
  "mlx-cuda-12==0.31.1; sys_platform == 'linux'",
+  "mlx-lm; sys_platform == 'linux'",
  "torch>=2.10.0; sys_platform == 'linux'",
 ]
 cuda13 = [
+  "mlx==0.31.1; sys_platform == 'linux'",
  "mlx-cuda-13==0.31.1; sys_platform == 'linux'",
+  "mlx-lm; sys_platform == 'linux'",
  "torch>=2.10.0; sys_platform == 'linux'",
 ]

@@ -75,7 +81,7 @@ mlx-lm = { git = "https://github.com/rltakashige/mlx-lm", branch = "leo/fix-arra
 torch = [
  { index = "pytorch-cu130", marker = "sys_platform == 'linux' and extra == 'cuda13' and extra != 'cpu' and extra != 'cuda12'" },
  { index = "pytorch-cu120", marker = "sys_platform == 'linux' and extra == 'cuda12' and extra != 'cpu' and extra != 'cuda13'" },
-  { index = "pytorch-cpu", marker = "sys_platform == 'linux' and extra == 'cpu' and extra != 'cuda12' and extra != 'cuda13'" },
+  { index = "pytorch-cpu", marker = "(extra != 'cuda12' and extra != 'cuda13' and sys_platform == 'linux') or sys_platform == 'darwin'" },
 ]
 vllm = { git = "https://github.com/hmellor/vllm.git", branch = "transformers-v5" }

@@ -149,6 +155,10 @@ prerelease = "allow"
 environments = ["sys_platform == 'darwin'", "sys_platform == 'linux'"]
 conflicts = [[{ extra = "cuda12" }, { extra = "cuda13" }, { extra = "cpu" }]]
 constraint-dependencies = ["transformers>=5.0.0,<5.4.0"]
+override-dependencies = [
+  "mlx==0.31.1; sys_platform=='linux'",
+  "mlx; sys_platform=='darwin'",
+]

 [tool.uv.extra-build-dependencies]
 miniaudio = ["setuptools", "cffi", "pycparser"]
@@ -5,14 +5,13 @@ let
    workspaceRoot = ../.;
  };

-  mkPythonSet = { pkgs, lib, self' }:
+  mkPythonSet = { pkgs, lib, self', members }:
    let
      inherit (pkgs.stdenv.hostPlatform) isLinux isDarwin isx86_64;
      inherit (pkgs.config) cudaSupport;
      inherit (pkgs) cudaPackages;
      cuda13Support = cudaSupport && cudaPackages.cudaMajorVersion == "13";
      libmlx_source = if cuda13Support then "mlx-cuda-13" else if cudaSupport then "mlx-cuda-12" else "mlx-cpu";
-      uv_extra = if cuda13Support then "cuda13" else if cudaSupport then "cuda12" else "cpu";
      python = pkgs.python313;
      cudaLibs = with cudaPackages; [
        cuda_cudart
@@ -51,7 +50,7 @@ let
          '';
        };
      };
-      buildSystemsOverlay = final: prev: { } //
+      buildSystemsOverlay = final: prev:
        lib.optionalAttrs isDarwin
          {
            mlx = prev.mlx.overrideAttrs (old:
@@ -81,7 +80,7 @@ let
                nativeBuildInputs = (old.nativeBuildInputs or [ ]) ++ [ pkgs.cmake self'.packages.metal-toolchain ];
                # TODO: non-sdk_26 support
                buildInputs = (old.buildInputs or [ ])
-                  ++ [ gguf-tools pkgs.fmt pkgs.nlohmann_json pkgs.apple-sdk_26 ];
+                ++ [ gguf-tools pkgs.fmt pkgs.nlohmann_json pkgs.apple-sdk_26 ];
                patches = [
                  (pkgs.replaceVars ../nix/darwin-build-fixes.patch {
                    sdkVersion = pkgs.apple-sdk_26.version;
@@ -113,42 +112,42 @@ let
                MACOSX_DEPLOYMENT_TARGET = pkgs.apple-sdk_26.version;
              });
          } // lib.optionalAttrs isLinux {
-        mlx = prev.mlx.overrideAttrs (old: {
-          buildInputs = old.buildInputs ++ lib.optionals cudaSupport cudaLibs;
-          autoPatchelfIgnoreMissingDeps = lib.optionals cudaSupport [ "libcuda.so.1" ];
-          postInstall = (old.postInstall or "") + ''
-            cp -r "${final.${libmlx_source}}/${final.python.sitePackages}/mlx" "$out/${final.python.sitePackages}/mlx/"
-          '';
-        });
-      } // lib.optionalAttrs cudaSupport {
-        "${libmlx_source}" = prev."${libmlx_source}".overrideAttrs (old: {
-          buildInputs = old.buildInputs ++ cudaLibs;
-          autoPatchelfIgnoreMissingDeps = [ "libcuda.so.1" ];
-        });
-        nvidia-cufile = prev.nvidia-cufile.overrideAttrs (old: {
-          buildInputs = old.buildInputs ++ [ pkgs.rdma-core ];
-          autoPatchelfIgnoreMissingDeps = [ "libcuda.so.1" ];
-        });
-        nvidia-cusolver = prev.nvidia-cusolver.overrideAttrs (old: {
-          buildInputs = old.buildInputs ++ cudaLibs;
-          autoPatchelfIgnoreMissingDeps = [ "libcuda.so.1" ];
-        });
-        nvidia-nvshmem-cu13 = prev.nvidia-nvshmem-cu13.overrideAttrs (old: {
-          buildInputs = old.buildInputs ++ [ pkgs.rdma-core pkgs.pmix pkgs.libfabric pkgs.ucx pkgs.openmpi ];
-          autoPatchelfIgnoreMissingDeps = [ "libcuda.so.1" ];
-        });
-        nvidia-cusparse = prev.nvidia-cusparse.overrideAttrs (old: {
-          buildInputs = old.buildInputs ++ [ cudaLibs ];
-          autoPatchelfIgnoreMissingDeps = [ "libcuda.so.1" ];
-        });
-        torch = prev.torch.overrideAttrs (old: {
-          buildInputs = old.buildInputs ++ cudaLibs;
-          autoPatchelfIgnoreMissingDeps = [ "libcuda.so.1" ];
-        });
-      };
+          mlx = prev.mlx.overrideAttrs (old: {
+            buildInputs = old.buildInputs ++ lib.optionals cudaSupport cudaLibs;
+            autoPatchelfIgnoreMissingDeps = lib.optionals cudaSupport [ "libcuda.so.1" ];
+            postInstall = ''
+              cp -r "${final.${libmlx_source}}/${final.python.sitePackages}/mlx" "$out/${final.python.sitePackages}/mlx/"
+            '';
+          });
+        } // lib.optionalAttrs cudaSupport {
+          "${libmlx_source}" = prev."${libmlx_source}".overrideAttrs (old: {
+            buildInputs = old.buildInputs ++ cudaLibs;
+            autoPatchelfIgnoreMissingDeps = [ "libcuda.so.1" ];
+          });
+          nvidia-cufile = prev.nvidia-cufile.overrideAttrs (old: {
+            buildInputs = old.buildInputs ++ [ pkgs.rdma-core ];
+            autoPatchelfIgnoreMissingDeps = [ "libcuda.so.1" ];
+          });
+          nvidia-cusolver = prev.nvidia-cusolver.overrideAttrs (old: {
+            buildInputs = old.buildInputs ++ cudaLibs;
+            autoPatchelfIgnoreMissingDeps = [ "libcuda.so.1" ];
+          });
+          nvidia-nvshmem-cu13 = prev.nvidia-nvshmem-cu13.overrideAttrs (old: {
+            buildInputs = old.buildInputs ++ [ pkgs.rdma-core pkgs.pmix pkgs.libfabric pkgs.ucx pkgs.openmpi ];
+            autoPatchelfIgnoreMissingDeps = [ "libcuda.so.1" ];
+          });
+          nvidia-cusparse = prev.nvidia-cusparse.overrideAttrs (old: {
+            buildInputs = old.buildInputs ++ [ cudaLibs ];
+            autoPatchelfIgnoreMissingDeps = [ "libcuda.so.1" ];
+          });
+          torch = prev.torch.overrideAttrs (old: {
+            buildInputs = old.buildInputs ++ cudaLibs;
+            autoPatchelfIgnoreMissingDeps = [ "libcuda.so.1" ];
+          });
+        };
      pyprojectOverlay = workspace.mkPyprojectOverlay {
        sourcePreference = "wheel";
-        dependencies = { exo = [ uv_extra ]; exo-bench = [ ]; };
+        dependencies = members;
      };
      editableOverlay = workspace.mkEditablePyprojectOverlay {
        # Use environment variable pointing to editable root directory
@@ -165,8 +164,8 @@ let
          buildSystemsOverlay
        ]
      );
-
-      mkApp = cmd: name: members: pkgs.writeShellApplication {
+      venv = name: (pythonSet.mkVirtualEnv "${name}-env" members).overrideAttrs (_: { venvSkip = [ "lib/python${python.pythonVersion}/site-packages/mlx/share/cmake/*" ]; });
+      mkApp = cmd: name: pkgs.writeShellApplication {
        inherit name;
        runtimeEnv = {
          EXO_DASHBOARD_DIR = self'.packages.dashboard;
@@ -174,17 +173,17 @@ let
        };
        runtimeInputs = [
          # mlx and mlx-cuda ship clashing cmake files - we dont need them at runtime anyway
-          ((pythonSet.mkVirtualEnv "${name}-env" members).overrideAttrs (_: { venvSkip = [ "lib/python${python.pythonVersion}/site-packages/mlx/share/cmake/*" ]; }))
+          (venv name)
        ]
        ++ lib.optionals isDarwin [ pkgs.macmon ];
        text = "exec " + lib.optionalString cudaSupport "${lib.getExe pkgs.nix-gl-host} " + cmd;
      };
    in
    {
-      inherit pythonSet;
+      inherit venv;
      editablePythonSet = pythonSet.overrideScope editableOverlay;
-      mkPythonScript = members: name: path: mkApp ''python ${path} "$@"'' name members;
-      mkExo = name: members: mkApp ''exo "$@"'' name members;
+      mkPythonScript = path: mkApp ''python ${path} "$@"'';
+      mkExo = mkApp ''exo "$@"'';
    };
 in
 {
@@ -192,16 +191,21 @@ in
    { self', pkgs, unfreePkgs, lib, ... }:
    let
      inherit (pkgs.stdenv.hostPlatform) isLinux;
-      inherit (mkPythonSet { inherit self' pkgs lib; }) pythonSet editablePythonSet mkPythonScript mkExo;
-
-      exoVenv = pythonSet.mkVirtualEnv "exo-env" { exo = lib.optionals isLinux [ "cpu" ]; };
+      inherit (mkPythonSet { inherit self' pkgs lib; members = { exo = [ "cpu" ]; }; }) editablePythonSet mkExo;

      # Virtual environment with dev dependencies for testing
-      testVenv = pythonSet.mkVirtualEnv "exo-test-env" {
-        exo = [ "dev" ] ++ lib.optionals isLinux [ "cpu" ]; # Include pytest, pytest-asyncio, pytest-env
+      testVenv = (mkPythonSet {
+        inherit self' pkgs lib; members = {
+        exo = [ "dev" "cpu" ]; # Include pytest, pytest-asyncio, pytest-env
      };
+      }).venv "exo-test";

-      mkBenchScript = mkPythonScript { exo-bench = [ ]; };
+      mkBenchScript = (mkPythonSet {
+        inherit self' pkgs lib; members = {
+        exo = [ "cpu" ];
+        exo-bench = [ ]; # Include pytest, pytest-asyncio, pytest-env
+      };
+      }).mkPythonScript;

      mkSimplePythonScript = name: path: pkgs.writeShellApplication {
        inherit name;
@@ -212,9 +216,7 @@ in
    in
    {
      packages = {
-        exo = mkExo "exo" { exo = lib.optionals isLinux [ "cpu" ]; };
-        # for devShell
-        exo-venv = exoVenv;
+        exo = mkExo "exo";
        editableVenv = editablePythonSet.mkVirtualEnv "exo-dev-env" { exo = [ "dev" ]; };
        # for running tests in ci
        exo-test-env = testVenv;
@@ -224,8 +226,8 @@ in
        # used by ./tests/run_exo_on.sh
        exo-get-all-models-on-cluster = mkSimplePythonScript "exo-get-all-models-on-cluster" (inputs.self + /tests/get_all_models_on_cluster.py);
      } // lib.optionalAttrs isLinux {
-        exo-cuda-12 = (mkPythonSet { inherit self' lib; inherit (unfreePkgs.pkgsCuda.cudaPackages_12) pkgs; }).mkExo "exo-cuda-12" { exo = [ "cuda12" ]; };
-        exo-cuda-13 = (mkPythonSet { inherit self' lib; inherit (unfreePkgs.pkgsCuda.cudaPackages_13) pkgs; }).mkExo "exo-cuda-13" { exo = [ "cuda13" ]; };
+        exo-cuda-12 = (mkPythonSet { inherit self' lib; inherit (unfreePkgs.pkgsCuda.cudaPackages_12) pkgs; members = { exo = [ "cuda12" ]; }; }).mkExo "exo-cuda-12";
+        exo-cuda-13 = (mkPythonSet { inherit self' lib; inherit (unfreePkgs.pkgsCuda.cudaPackages_13) pkgs; members = { exo = [ "cuda13" ]; }; }).mkExo "exo-cuda-13";
      };

      checks = {
@@ -235,19 +237,11 @@ in
          touch $out
        '';

-        typecheck = pkgs.runCommand "typecheck"
-          {
-            nativeBuildInputs = [
-              testVenv
-              pkgs.basedpyright
-            ];
-          }
-          ''
-            cd ${inputs.self}
-            export HOME=$TMPDIR
-            basedpyright --pythonpath ${testVenv}/bin/python --project ${inputs.self}/pyproject.toml
-            touch $out
-          '';
+        typecheck = pkgs.runCommand "typecheck" { nativeBuildInputs = [ testVenv ]; } ''
+          cd ${inputs.self}
+          basedpyright
+          touch $out
+        '';
      };
    };
 }
@@ -0,0 +1,21 @@
+model_id = "mlx-community/GLM-5.1-DQ4plus-q8"
+n_layers = 78
+hidden_size = 6144
+num_key_value_heads = 64
+supports_tensor = true
+tasks = ["TextGeneration"]
+family = "glm"
+quantization = "8bit"
+base_model = "GLM-5.1"
+capabilities = ["text", "thinking"]
+
+context_length = 202752
+
+[storage_size]
+in_bytes = 465173655552
+
+# Source: https://huggingface.co/zai-org/GLM-5.1
+# Source: https://docs.z.ai/api-reference/llm/chat-completion
+[sampling_defaults]
+temperature = 1.0
+top_p = 0.95
@@ -0,0 +1,21 @@
+model_id = "mlx-community/GLM-5.1-MXFP4-Q8"
+n_layers = 78
+hidden_size = 6144
+num_key_value_heads = 64
+supports_tensor = true
+tasks = ["TextGeneration"]
+family = "glm"
+quantization = "MXFP4-Q8"
+base_model = "GLM-5.1"
+capabilities = ["text", "thinking"]
+
+context_length = 202752
+
+[storage_size]
+in_bytes = 405480321024
+
+# Source: https://huggingface.co/zai-org/GLM-5.1
+# Source: https://docs.z.ai/api-reference/llm/chat-completion
+[sampling_defaults]
+temperature = 1.0
+top_p = 0.95
@@ -0,0 +1,21 @@
+model_id = "mlx-community/GLM-5.1"
+n_layers = 78
+hidden_size = 6144
+num_key_value_heads = 64
+supports_tensor = true
+tasks = ["TextGeneration"]
+family = "glm"
+quantization = "bf16"
+base_model = "GLM-5.1"
+capabilities = ["text", "thinking"]
+
+context_length = 202752
+
+[storage_size]
+in_bytes = 1487822475264
+
+# Source: https://huggingface.co/zai-org/GLM-5.1
+# Source: https://docs.z.ai/api-reference/llm/chat-completion
+[sampling_defaults]
+temperature = 1.0
+top_p = 0.95
@@ -0,0 +1,33 @@
+model_id = "mlx-community/Kimi-K2.6-mlx-DQ3_K_M-q8"
+n_layers = 61
+hidden_size = 7168
+num_key_value_heads = 64
+supports_tensor = true
+tasks = ["TextGeneration"]
+family = "kimi"
+quantization = "3bit"
+base_model = "Kimi K2.6"
+capabilities = ["text", "thinking", "thinking_toggle", "vision"]
+
+context_length = 262144
+
+[storage_size]
+in_bytes = 470628683776
+
+[vision]
+image_token_id = 163605
+model_type = "kimi_vl"
+weights_repo = "exolabs/Kimi-K2.6-vision"
+processor_repo = "moonshotai/Kimi-K2.6"
+
+# Source: https://huggingface.co/moonshotai/Kimi-K2.6
+[sampling_defaults]
+temperature = 1.0
+top_p = 0.95
+min_p = 0.01
+
+# Source: https://huggingface.co/moonshotai/Kimi-K2.6
+[sampling_defaults.non_thinking]
+temperature = 0.6
+top_p = 0.95
+min_p = 0.01
@@ -0,0 +1,35 @@
+model_id = "mlx-community/Qwen3.6-27B-4bit"
+n_layers = 64
+hidden_size = 5120
+num_key_value_heads = 4
+supports_tensor = true
+tasks = ["TextGeneration"]
+family = "qwen"
+quantization = "4bit"
+base_model = "Qwen3.6 27B"
+capabilities = ["text", "thinking", "thinking_toggle", "vision"]
+
+context_length = 262144
+
+[storage_size]
+in_bytes = 16054262240
+
+# Source: https://huggingface.co/Qwen/Qwen3.6-27B#best-practices
+# Source: https://unsloth.ai/docs/models/qwen3.5
+[sampling_defaults]
+temperature = 1.0
+top_p = 0.95
+top_k = 20
+min_p = 0.0
+repetition_penalty = 1.0
+presence_penalty = 1.5
+
+# Source: https://huggingface.co/Qwen/Qwen3.6-27B#best-practices
+# Source: https://unsloth.ai/docs/models/qwen3.5
+[sampling_defaults.non_thinking]
+temperature = 0.7
+top_p = 0.8
+top_k = 20
+min_p = 0.0
+repetition_penalty = 1.0
+presence_penalty = 1.5
@@ -0,0 +1,35 @@
+model_id = "mlx-community/Qwen3.6-27B-8bit"
+n_layers = 64
+hidden_size = 5120
+num_key_value_heads = 4
+supports_tensor = true
+tasks = ["TextGeneration"]
+family = "qwen"
+quantization = "8bit"
+base_model = "Qwen3.6 27B"
+capabilities = ["text", "thinking", "thinking_toggle", "vision"]
+
+context_length = 262144
+
+[storage_size]
+in_bytes = 29500938720
+
+# Source: https://huggingface.co/Qwen/Qwen3.6-27B#best-practices
+# Source: https://unsloth.ai/docs/models/qwen3.5
+[sampling_defaults]
+temperature = 1.0
+top_p = 0.95
+top_k = 20
+min_p = 0.0
+repetition_penalty = 1.0
+presence_penalty = 1.5
+
+# Source: https://huggingface.co/Qwen/Qwen3.6-27B#best-practices
+# Source: https://unsloth.ai/docs/models/qwen3.5
+[sampling_defaults.non_thinking]
+temperature = 0.7
+top_p = 0.8
+top_k = 20
+min_p = 0.0
+repetition_penalty = 1.0
+presence_penalty = 1.5
@@ -0,0 +1,35 @@
+model_id = "mlx-community/Qwen3.6-27B-bf16"
+n_layers = 64
+hidden_size = 5120
+num_key_value_heads = 4
+supports_tensor = true
+tasks = ["TextGeneration"]
+family = "qwen"
+quantization = "bf16"
+base_model = "Qwen3.6 27B"
+capabilities = ["text", "thinking", "thinking_toggle", "vision"]
+
+context_length = 262144
+
+[storage_size]
+in_bytes = 54713457120
+
+# Source: https://huggingface.co/Qwen/Qwen3.6-27B#best-practices
+# Source: https://unsloth.ai/docs/models/qwen3.5
+[sampling_defaults]
+temperature = 1.0
+top_p = 0.95
+top_k = 20
+min_p = 0.0
+repetition_penalty = 1.0
+presence_penalty = 1.5
+
+# Source: https://huggingface.co/Qwen/Qwen3.6-27B#best-practices
+# Source: https://unsloth.ai/docs/models/qwen3.5
+[sampling_defaults.non_thinking]
+temperature = 0.7
+top_p = 0.8
+top_k = 20
+min_p = 0.0
+repetition_penalty = 1.0
+presence_penalty = 1.5
@@ -0,0 +1,33 @@
+model_id = "moonshotai/Kimi-K2.6"
+n_layers = 61
+hidden_size = 7168
+num_key_value_heads = 64
+supports_tensor = true
+tasks = ["TextGeneration"]
+family = "kimi"
+quantization = ""
+base_model = "Kimi K2.6"
+capabilities = ["text", "thinking", "thinking_toggle", "vision"]
+
+context_length = 262144
+
+[storage_size]
+in_bytes = 595148192736
+
+[vision]
+image_token_id = 163605
+model_type = "kimi_vl"
+weights_repo = "exolabs/Kimi-K2.6-vision"
+processor_repo = "moonshotai/Kimi-K2.6"
+
+# Source: https://huggingface.co/moonshotai/Kimi-K2.6
+[sampling_defaults]
+temperature = 1.0
+top_p = 0.95
+min_p = 0.01
+
+# Source: https://huggingface.co/moonshotai/Kimi-K2.6
+[sampling_defaults.non_thinking]
+temperature = 0.6
+top_p = 0.95
+min_p = 0.01
@@ -113,6 +113,23 @@ def _extract_content(content: str | list[ResponseContentPart]) -> str:
    )


+def _append_tool_call(
+    chat_template_messages: list[dict[str, Any]], tool_call: dict[str, Any]
+) -> None:
+    if chat_template_messages:
+        prev = chat_template_messages[-1]
+        if prev.get("role") == "assistant" and isinstance(prev.get("content"), str):
+            existing: list[dict[str, Any]] | None = prev.get("tool_calls")
+            if existing is None:
+                prev["tool_calls"] = [tool_call]
+            else:
+                existing.append(tool_call)
+            return
+    chat_template_messages.append(
+        {"role": "assistant", "content": "", "tool_calls": [tool_call]}
+    )
+
+
 async def responses_request_to_text_generation(
    request: ResponsesRequest,
 ) -> TextGenerationTaskParams:
@@ -182,59 +199,44 @@ async def responses_request_to_text_generation(
                    | McpCallInputItem()
                    | CustomToolCallInputItem()
                ):
-                    chat_template_messages.append(
+                    _append_tool_call(
+                        chat_template_messages,
                        {
-                            "role": "assistant",
-                            "content": "",
-                            "tool_calls": [
-                                {
-                                    "id": item.call_id,
-                                    "type": "function",
-                                    "function": {
-                                        "name": item.name,
-                                        "arguments": item.arguments,
-                                    },
-                                }
-                            ],
-                        }
+                            "id": item.call_id,
+                            "type": "function",
+                            "function": {
+                                "name": item.name,
+                                "arguments": item.arguments,
+                            },
+                        },
                    )
                case (
                    LocalShellCallInputItem()
                    | ShellCallInputItem()
                    | ComputerCallInputItem()
                ):
-                    chat_template_messages.append(
+                    _append_tool_call(
+                        chat_template_messages,
                        {
-                            "role": "assistant",
-                            "content": "",
-                            "tool_calls": [
-                                {
-                                    "id": item.call_id,
-                                    "type": "function",
-                                    "function": {
-                                        "name": item.type,
-                                        "arguments": json.dumps(item.action),
-                                    },
-                                }
-                            ],
-                        }
+                            "id": item.call_id,
+                            "type": "function",
+                            "function": {
+                                "name": item.type,
+                                "arguments": json.dumps(item.action),
+                            },
+                        },
                    )
                case ApplyPatchCallInputItem():
-                    chat_template_messages.append(
+                    _append_tool_call(
+                        chat_template_messages,
                        {
-                            "role": "assistant",
-                            "content": "",
-                            "tool_calls": [
-                                {
-                                    "id": item.call_id,
-                                    "type": "function",
-                                    "function": {
-                                        "name": "apply_patch",
-                                        "arguments": json.dumps({"patch": item.patch}),
-                                    },
-                                }
-                            ],
-                        }
+                            "id": item.call_id,
+                            "type": "function",
+                            "function": {
+                                "name": "apply_patch",
+                                "arguments": json.dumps({"patch": item.patch}),
+                            },
+                        },
                    )
                case (
                    WebSearchCallInputItem()
@@ -254,21 +256,16 @@ async def responses_request_to_text_generation(
                        args = {"prompt": item.prompt}
                    else:
                        args = {"query": item.query}
-                    chat_template_messages.append(
+                    _append_tool_call(
+                        chat_template_messages,
                        {
-                            "role": "assistant",
-                            "content": "",
-                            "tool_calls": [
-                                {
-                                    "id": item.call_id,
-                                    "type": "function",
-                                    "function": {
-                                        "name": item.type,
-                                        "arguments": json.dumps(args),
-                                    },
-                                }
-                            ],
-                        }
+                            "id": item.call_id,
+                            "type": "function",
+                            "function": {
+                                "name": item.type,
+                                "arguments": json.dumps(args),
+                            },
+                        },
                    )
                case (
                    FunctionCallOutputInputItem()
@@ -320,21 +317,16 @@ async def responses_request_to_text_generation(
                            }
                        )
                case McpApprovalRequestInputItem():
-                    chat_template_messages.append(
+                    _append_tool_call(
+                        chat_template_messages,
                        {
-                            "role": "assistant",
-                            "content": "",
-                            "tool_calls": [
-                                {
-                                    "id": item.call_id,
-                                    "type": "function",
-                                    "function": {
-                                        "name": item.name,
-                                        "arguments": item.arguments,
-                                    },
-                                }
-                            ],
-                        }
+                            "id": item.call_id,
+                            "type": "function",
+                            "function": {
+                                "name": item.name,
+                                "arguments": item.arguments,
+                            },
+                        },
                    )
                case McpApprovalResponseInputItem():
                    chat_template_messages.append(
@@ -119,7 +119,7 @@ from exo.api.types.openai_responses import (
 )
 from exo.master.image_store import ImageStore
 from exo.master.placement import place_instance as get_instance_placements
-from exo.routing.state_manager import StateManager
+from exo.shared.apply import apply
 from exo.shared.constants import (
    DASHBOARD_DIR,
    EXO_CACHE_HOME,
@@ -185,7 +185,10 @@ from exo.shared.types.tasks import (
 from exo.shared.types.tasks import (
    TextGeneration as TextGenerationTask,
 )
-from exo.shared.types.text_generation import Base64Image, TextGenerationTaskParams
+from exo.shared.types.text_generation import (
+    Base64ImageHash,
+    TextGenerationTaskParams,
+)
 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
@@ -223,9 +226,8 @@ class API:
        download_command_sender: Sender[ForwarderDownloadCommand],
        # This lets us pause the API if an election is running
        election_receiver: Receiver[ElectionMessage],
-        state_manager: StateManager[State],
    ) -> None:
-        self.state_manager = state_manager
+        self.state = State()
        self._event_log = DiskEventLog(_API_EVENT_LOG_DIR)
        self._system_id = SystemId()
        self.command_sender = command_sender
@@ -235,6 +237,7 @@ class API:
        self.node_id: NodeId = node_id
        self.last_completed_election: int = 0
        self.port = port
+        self._sent_image_hashes: set[str] = set()

        self.paused: bool = False
        self.paused_ev: anyio.Event = anyio.Event()
@@ -272,16 +275,11 @@ class API:
        self._image_store = ImageStore(EXO_IMAGE_CACHE_DIR)
        self._tg: TaskGroup = TaskGroup()

-    def reset(
-        self,
-        result_clock: int,
-        event_receiver: Receiver[IndexedEvent],
-        state_manager: StateManager[State],
-    ):
+    def reset(self, result_clock: int, event_receiver: Receiver[IndexedEvent]):
        logger.info("Resetting API State")
        self._event_log.close()
        self._event_log = DiskEventLog(_API_EVENT_LOG_DIR)
-        self.state_manager = state_manager
+        self.state = State()
        self._system_id = SystemId()
        self._text_generation_queues = {}
        self._image_generation_queues = {}
@@ -289,6 +287,7 @@ class API:
        self.event_receiver.close()
        self.event_receiver = event_receiver
        self._tg.start_soon(self._apply_state)
+        self._sent_image_hashes = set()

    def unpause(self, result_clock: int):
        logger.info("Unpausing API")
@@ -378,12 +377,11 @@ class API:
        self.app.get("/onboarding")(self.get_onboarding)
        self.app.post("/onboarding")(self.complete_onboarding)

-    def get_state(self, path: str = "") -> Any:  # pyright: ignore[reportAny]
-        state = self.state_manager.get()
+    def get_state(self, path: str = ""):
        if path == "":
-            return state
+            return self.state
        try:
-            x = state.model_dump(by_alias=True)
+            x = self.state.model_dump(by_alias=True)
            for attr in path.split("/"):
                if attr != "":
                    if isinstance(x, dict):
@@ -445,7 +443,6 @@ class API:
        min_nodes: int = 1,
    ) -> Instance:
        model_card = await ModelCard.load(model_id)
-        state = self.state_manager.get()

        try:
            placements = get_instance_placements(
@@ -455,16 +452,16 @@ class API:
                    instance_meta=instance_meta,
                    min_nodes=min_nodes,
                ),
-                node_memory=state.node_memory,
-                node_network=state.node_network,
-                topology=state.topology,
-                current_instances=state.instances,
-                download_status=state.downloads,
+                node_memory=self.state.node_memory,
+                node_network=self.state.node_network,
+                topology=self.state.topology,
+                current_instances=self.state.instances,
+                download_status=self.state.downloads,
            )
        except ValueError as exc:
            raise HTTPException(status_code=400, detail=str(exc)) from exc

-        current_ids = set(state.instances.keys())
+        current_ids = set(self.state.instances.keys())
        new_ids = [
            instance_id for instance_id in placements if instance_id not in current_ids
        ]
@@ -484,9 +481,8 @@ class API:
        seen: set[tuple[ModelId, Sharding, InstanceMeta, int]] = set()
        previews: list[PlacementPreview] = []
        required_nodes = set(node_ids) if node_ids else None
-        state = self.state_manager.get()

-        if len(list(state.topology.list_nodes())) == 0:
+        if len(list(self.state.topology.list_nodes())) == 0:
            return PlacementPreviewResponse(previews=[])

        try:
@@ -501,7 +497,9 @@ class API:
                instance_combinations.extend(
                    [
                        (sharding, instance_meta, i)
-                        for i in range(1, len(list(state.topology.list_nodes())) + 1)
+                        for i in range(
+                            1, len(list(self.state.topology.list_nodes())) + 1
+                        )
                    ]
                )
        # TODO: PDD
@@ -516,12 +514,12 @@ class API:
                        instance_meta=instance_meta,
                        min_nodes=min_nodes,
                    ),
-                    node_memory=state.node_memory,
-                    node_network=state.node_network,
-                    topology=state.topology,
-                    current_instances=state.instances,
+                    node_memory=self.state.node_memory,
+                    node_network=self.state.node_network,
+                    topology=self.state.topology,
+                    current_instances=self.state.instances,
                    required_nodes=required_nodes,
-                    download_status=state.downloads,
+                    download_status=self.state.downloads,
                )
            except ValueError as exc:
                if (model_card.model_id, sharding, instance_meta, 0) not in seen:
@@ -537,7 +535,7 @@ class API:
                seen.add((model_card.model_id, sharding, instance_meta, 0))
                continue

-            current_ids = set(state.instances.keys())
+            current_ids = set(self.state.instances.keys())
            new_instances = [
                instance
                for instance_id, instance in placements.items()
@@ -599,14 +597,12 @@ class API:
        return PlacementPreviewResponse(previews=previews)

    def get_instance(self, instance_id: InstanceId) -> Instance:
-        state = self.state_manager.get()
-        if instance_id not in state.instances:
+        if instance_id not in self.state.instances:
            raise HTTPException(status_code=404, detail="Instance not found")
-        return state.instances[instance_id]
+        return self.state.instances[instance_id]

    async def delete_instance(self, instance_id: InstanceId) -> DeleteInstanceResponse:
-        state = self.state_manager.get()
-        if instance_id not in state.instances:
+        if instance_id not in self.state.instances:
            raise HTTPException(status_code=404, detail="Instance not found")

        command = DeleteInstance(
@@ -675,9 +671,7 @@ class API:
    async def _collect_text_generation_with_stats(
        self, command_id: CommandId
    ) -> BenchChatCompletionResponse:
-        sampler = PowerSampler(
-            get_node_system=lambda: self.state_manager.get().node_system
-        )
+        sampler = PowerSampler(get_node_system=lambda: self.state.node_system)
        text_parts: list[str] = []
        tool_calls: list[ToolCall] = []
        model: ModelId | None = None
@@ -748,8 +742,6 @@ class API:
            "TODO: we should send a notification to the user to download the model"
        )

-    _sent_image_hashes: set[str] = set()
-
    async def _send_text_generation_with_images(
        self, task_params: TextGenerationTaskParams
    ) -> TextGeneration:
@@ -761,23 +753,19 @@ class API:
            return command

        hashes = [hashlib.sha256(img.encode("ascii")).hexdigest() for img in images]
+        all_hashes = {idx: Base64ImageHash(h) for idx, h in enumerate(hashes)}
+        task_params = task_params.model_copy(
+            update={"images": [], "image_hashes": all_hashes}
+        )
+        command = TextGeneration(task_params=task_params)

-        cached_hashes: dict[int, str] = {}
        new_images: list[tuple[int, str]] = []
        for idx, (img, h) in enumerate(zip(images, hashes, strict=True)):
-            if h in self._sent_image_hashes:
-                cached_hashes[idx] = h
-            else:
+            if h not in self._sent_image_hashes:
                self._sent_image_hashes.add(h)
                new_images.append((idx, img))

-        wrapped_hashes = {idx: Base64Image(h) for idx, h in cached_hashes.items()}
-
        if not new_images:
-            task_params = task_params.model_copy(
-                update={"images": [], "image_hashes": wrapped_hashes}
-            )
-            command = TextGeneration(task_params=task_params)
            await self._send(command)
            return command

@@ -786,16 +774,6 @@ class API:
            for i in range(0, len(img_data), EXO_MAX_CHUNK_SIZE):
                all_chunks.append((img_idx, img_data[i : i + EXO_MAX_CHUNK_SIZE]))

-        task_params = task_params.model_copy(
-            update={
-                "images": [],
-                "image_hashes": wrapped_hashes,
-                "total_input_chunks": len(all_chunks),
-                "image_count": len(new_images),
-            }
-        )
-        command = TextGeneration(task_params=task_params)
-
        for global_idx, (img_idx, chunk_data) in enumerate(all_chunks):
            await self._send(
                SendInputChunk(
@@ -875,10 +853,9 @@ class API:

        Raises HTTPException 404 if no instance is found for the model.
        """
-        state = self.state_manager.get()
        if not any(
            instance.shard_assignments.model_id == model_id
-            for instance in state.instances.values()
+            for instance in self.state.instances.values()
        ):
            await self._trigger_notify_user_to_download_model(model_id)
            raise HTTPException(
@@ -894,10 +871,9 @@ class API:
        """
        model_card = await ModelCard.load(model)
        resolved_model = model_card.model_id
-        state = self.state_manager.get()
        if not any(
            instance.shard_assignments.model_id == resolved_model
-            for instance in state.instances.values()
+            for instance in self.state.instances.values()
        ):
            await self._trigger_notify_user_to_download_model(resolved_model)
            raise HTTPException(
@@ -1207,9 +1183,7 @@ class API:
        num_images: int,
        response_format: str,
    ) -> BenchImageGenerationResponse:
-        sampler = PowerSampler(
-            get_node_system=lambda: self.state_manager.get().node_system
-        )
+        sampler = PowerSampler(get_node_system=lambda: self.state.node_system)
        images: list[ImageData] = []
        stats: ImageGenerationStats | None = None
        async with anyio.create_task_group() as tg:
@@ -1579,9 +1553,8 @@ class API:
        def none_if_empty(value: str) -> str | None:
            return value or None

-        state = self.state_manager.get()
        downloaded_model_ids: set[str] = set()
-        for node_downloads in state.downloads.values():
+        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)
@@ -1635,8 +1608,7 @@ class API:
        """Returns list of running models (active instances)."""
        models: list[OllamaPsModel] = []
        seen: set[str] = set()
-        state = self.state_manager.get()
-        for instance in state.instances.values():
+        for instance in self.state.instances.values():
            model_id = str(instance.shard_assignments.model_id)
            if model_id in seen:
                continue
@@ -1658,8 +1630,7 @@ class API:
        """Calculate total available memory across all nodes in bytes."""
        total_available = Memory()

-        state = self.state_manager.get()
-        for memory in state.node_memory.values():
+        for memory in self.state.node_memory.values():
            total_available += memory.ram_available

        return total_available
@@ -1667,11 +1638,10 @@ class API:
    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()
-        state = self.state_manager.get()

        if status == "downloaded":
            downloaded_model_ids: set[str] = set()
-            for node_downloads in state.downloads.values():
+            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)
@@ -1827,6 +1797,7 @@ class API:
        with self.event_receiver as events:
            async for i_event in events:
                self._event_log.append(i_event.event)
+                self.state = apply(self.state, i_event)
                event = i_event.event

                if isinstance(event, ChunkGenerated):
@@ -1853,8 +1824,7 @@ class API:

    def _close_streams_for_instance(self, instance_id: InstanceId) -> None:
        """Close any active generation streams for commands running on the given instance."""
-        state = self.state_manager.get()
-        for task in state.tasks.values():
+        for task in self.state.tasks.values():
            if task.instance_id != instance_id:
                continue
            if not isinstance(
@@ -88,7 +88,9 @@ class DownloadCoordinator:

        try:
            if progress.status == "complete":
-                found = await to_thread.run_sync(resolve_existing_model, model_id)
+                found = await to_thread.run_sync(
+                    resolve_existing_model, model_id, callback_shard.model_card
+                )
                if found is not None:
                    completed = self._completed_from_path(
                        callback_shard, found, progress.total
@@ -193,7 +195,9 @@ class DownloadCoordinator:
                return

        # Check all model directories for pre-existing complete models
-        found_path = await to_thread.run_sync(resolve_existing_model, model_id)
+        found_path = await to_thread.run_sync(
+            resolve_existing_model, model_id, shard.model_card
+        )
        if found_path is not None:
            logger.info(f"DownloadCoordinator: Model {model_id} found at {found_path}")
            completed = self._completed_from_path(
@@ -220,7 +224,9 @@ class DownloadCoordinator:
        )

        if initial_progress.status == "complete":
-            found = await to_thread.run_sync(resolve_existing_model, model_id)
+            found = await to_thread.run_sync(
+                resolve_existing_model, model_id, shard.model_card
+            )
            if found is not None:
                completed = self._completed_from_path(
                    shard, found, initial_progress.total
@@ -351,7 +357,9 @@ class DownloadCoordinator:

                    if progress.status == "complete":
                        found = await to_thread.run_sync(
-                            resolve_existing_model, model_id
+                            resolve_existing_model,
+                            model_id,
+                            progress.shard.model_card,
                        )
                        if found is not None:
                            status: DownloadProgress = self._completed_from_path(
@@ -380,7 +388,9 @@ class DownloadCoordinator:
                        # (is_model_directory_complete) which validates that all
                        # safetensors weight files are present.
                        found = await to_thread.run_sync(
-                            resolve_existing_model, model_id
+                            resolve_existing_model,
+                            model_id,
+                            progress.shard.model_card,
                        )
                        if found is not None:
                            status = self._completed_from_path(
@@ -421,7 +431,9 @@ class DownloadCoordinator:
                            (DownloadCompleted, DownloadOngoing, DownloadFailed),
                        ):
                            continue
-                        found = await to_thread.run_sync(resolve_existing_model, mid)
+                        found = await to_thread.run_sync(
+                            resolve_existing_model, mid, card
+                        )
                        if found is not None and is_read_only_model_dir(found):
                            path_shard = PipelineShardMetadata(
                                model_card=card,
@@ -35,7 +35,7 @@ from exo.shared.constants import (
    EXO_MODELS_DIRS,
    EXO_MODELS_READ_ONLY_DIRS,
 )
-from exo.shared.models.model_cards import ModelTask
+from exo.shared.models.model_cards import ModelCard, ModelTask
 from exo.shared.types.common import ModelId
 from exo.shared.types.memory import Memory
 from exo.shared.types.worker.downloads import (
@@ -118,7 +118,9 @@ class InsufficientDiskSpaceError(Exception):
    """Raised when no writable model directory has enough free space."""


-def resolve_existing_model(model_id: ModelId) -> Path | None:
+def resolve_existing_model(
+    model_id: ModelId, card: ModelCard | None = None
+) -> Path | None:
    """Search all model directories for a complete, pre-existing model.

    Checks read-only directories first, then writable directories.
@@ -128,7 +130,7 @@ def resolve_existing_model(model_id: ModelId) -> Path | None:
    normalized = model_id.normalize()
    for search_dir in (*EXO_MODELS_READ_ONLY_DIRS, *EXO_MODELS_DIRS):
        candidate = search_dir / normalized
-        if candidate.is_dir() and is_model_directory_complete(candidate):
+        if candidate.is_dir() and is_model_directory_complete(candidate, card):
            return candidate
    return None

@@ -165,6 +167,29 @@ def select_download_dir(required_bytes: int) -> Path:
    )


+async def select_download_dir_for_shard(
+    model_id: ModelId,
+    filtered_file_list: list[FileListEntry],
+    total_size: int,
+) -> Path:
+    for candidate_dir in EXO_MODELS_DIRS:
+        if not candidate_dir.exists():
+            continue
+        sub = candidate_dir / model_id.normalize()
+        if not await aios.path.isdir(sub):
+            continue
+        existing_bytes = 0
+        for file_entry in filtered_file_list:
+            existing_bytes += await get_downloaded_size(sub / file_entry.path)
+        remaining = max(total_size - existing_bytes, 0)
+        try:
+            if shutil.disk_usage(candidate_dir).free >= remaining:
+                return candidate_dir
+        except OSError:
+            continue
+    return select_download_dir(total_size)
+
+
 async def resolve_model_dir(model_id: ModelId) -> Path:
    """Return the directory for a model's files, creating it if needed.

@@ -279,10 +304,26 @@ def _scan_model_directory(
    return list(entries_by_path.values())


-def is_model_directory_complete(model_dir: Path) -> bool:
-    """Check if a model directory contains all required weight files."""
+def is_model_directory_complete(model_dir: Path, card: ModelCard | None = None) -> bool:
+    """Check if a model directory contains all required weight files.
+    Also checks for sibling weights repo.
+    """
    file_list = _scan_model_directory(model_dir, recursive=True)
-    return file_list is not None and all(f.size is not None for f in file_list)
+    if file_list is None or not all(f.size is not None for f in file_list):
+        return False
+    if (
+        card is not None
+        and card.vision is not None
+        and card.vision.weights_repo != str(card.model_id)
+    ):
+        vision_id = ModelId(card.vision.weights_repo)
+        normalized = vision_id.normalize()
+        for search_dir in (*EXO_MODELS_READ_ONLY_DIRS, *EXO_MODELS_DIRS):
+            candidate = search_dir / normalized
+            if candidate.is_dir() and is_model_directory_complete(candidate):
+                return True
+        return False
+    return True


 async def _build_file_list_from_local_directory(
@@ -834,7 +875,9 @@ async def download_shard(
            else EXO_DEFAULT_MODELS_DIR / model_id.normalize()
        )
    else:
-        models_dir = select_download_dir(total_size)
+        models_dir = await select_download_dir_for_shard(
+            model_id, filtered_file_list, total_size
+        )
        target_dir = models_dir / model_id.normalize()
        await aios.makedirs(target_dir, exist_ok=True)
    file_progress: dict[str, RepoFileDownloadProgress] = {}
@@ -117,40 +117,39 @@ class ResumableShardDownloader(ShardDownloader):
    ) -> Path:
        allow_patterns = ["config.json"] if config_only else None

+        has_vision_sibling = (
+            not config_only
+            and not self.offline
+            and shard.model_card.vision is not None
+            and shard.model_card.vision.weights_repo != str(shard.model_card.model_id)
+        )
+
+        async def main_progress(
+            cb_shard: ShardMetadata, progress: RepoDownloadProgress
+        ) -> None:
+            if has_vision_sibling and progress.status == "complete":
+                return
+            await self.on_progress_wrapper(cb_shard, progress)
+
        target_dir, _ = await download_shard(
            shard,
-            self.on_progress_wrapper,
+            main_progress,
            max_parallel_downloads=self.max_parallel_downloads,
            allow_patterns=allow_patterns,
            skip_internet=self.offline,
        )

-        if (
-            not config_only
-            and not self.offline
-            and shard.model_card.vision
-            and shard.model_card.vision.weights_repo != str(shard.model_card.model_id)
-        ):
-            vision_repo = shard.model_card.vision.weights_repo
-            vision_card = ModelCard(
-                model_id=ModelId(vision_repo),
-                storage_size=Memory.from_bytes(0),
-                n_layers=1,
-                hidden_size=1,
-                supports_tensor=False,
-                tasks=[ModelTask.TextGeneration],
-            )
-            vision_shard = PipelineShardMetadata(
-                model_card=vision_card,
-                device_rank=0,
-                world_size=1,
-                start_layer=0,
-                end_layer=1,
-                n_layers=1,
-            )
+        if has_vision_sibling:
+            vision_shard = self._build_vision_shard(shard)
+
+            async def vision_progress(
+                _cb_shard: ShardMetadata, progress: RepoDownloadProgress
+            ) -> None:
+                await self.on_progress_wrapper(shard, progress)
+
            await download_shard(
                vision_shard,
-                self.on_progress_wrapper,
+                vision_progress,
                max_parallel_downloads=self.max_parallel_downloads,
                allow_patterns=["*.safetensors", "config.json"],
                skip_internet=self.offline,
@@ -158,6 +157,87 @@ class ResumableShardDownloader(ShardDownloader):

        return target_dir

+    async def _status_for_shard(
+        self, shard: ShardMetadata
+    ) -> tuple[Path, RepoDownloadProgress]:
+        async def _noop(
+            _cb_shard: ShardMetadata, _progress: RepoDownloadProgress
+        ) -> None:
+            return
+
+        path, main_progress = await download_shard(
+            shard,
+            _noop,
+            skip_download=True,
+            skip_internet=self.offline,
+        )
+
+        has_vision_sibling = (
+            shard.model_card.vision is not None
+            and shard.model_card.vision.weights_repo != str(shard.model_card.model_id)
+        )
+        if not has_vision_sibling:
+            return path, main_progress
+
+        vision_shard = self._build_vision_shard(shard)
+        _, vision_progress = await download_shard(
+            vision_shard,
+            _noop,
+            skip_download=True,
+            skip_internet=self.offline,
+        )
+        combined = self._combine_progress(shard, main_progress, vision_progress)
+        return path, combined
+
+    @staticmethod
+    def _build_vision_shard(shard: ShardMetadata) -> PipelineShardMetadata:
+        assert shard.model_card.vision is not None
+        vision_card = ModelCard(
+            model_id=ModelId(shard.model_card.vision.weights_repo),
+            storage_size=Memory.from_bytes(0),
+            n_layers=1,
+            hidden_size=1,
+            supports_tensor=False,
+            tasks=[ModelTask.TextGeneration],
+        )
+        return PipelineShardMetadata(
+            model_card=vision_card,
+            device_rank=0,
+            world_size=1,
+            start_layer=0,
+            end_layer=1,
+            n_layers=1,
+        )
+
+    @staticmethod
+    def _combine_progress(
+        shard: ShardMetadata,
+        main: RepoDownloadProgress,
+        vision: RepoDownloadProgress,
+    ) -> RepoDownloadProgress:
+        status_rank = {"not_started": 0, "in_progress": 1, "complete": 2}
+        combined_status = min(
+            (main.status, vision.status), key=lambda s: status_rank[s]
+        )
+        file_progress = dict(main.file_progress)
+        for file_path, fp in vision.file_progress.items():
+            file_progress[f"{vision.repo_id}/{file_path}"] = fp
+        return RepoDownloadProgress(
+            repo_id=main.repo_id,
+            repo_revision=main.repo_revision,
+            shard=shard,
+            completed_files=main.completed_files + vision.completed_files,
+            total_files=main.total_files + vision.total_files,
+            downloaded=main.downloaded + vision.downloaded,
+            downloaded_this_session=main.downloaded_this_session
+            + vision.downloaded_this_session,
+            total=main.total + vision.total,
+            overall_speed=main.overall_speed + vision.overall_speed,
+            overall_eta=max(main.overall_eta, vision.overall_eta),
+            status=combined_status,
+            file_progress=file_progress,
+        )
+
    async def get_shard_download_status(
        self,
    ) -> AsyncIterator[tuple[Path, RepoDownloadProgress]]:
@@ -166,12 +246,7 @@ class ResumableShardDownloader(ShardDownloader):
        ) -> tuple[Path, RepoDownloadProgress]:
            """Helper coroutine that builds the shard for a model and gets its download status."""
            shard = await build_full_shard(model_id)
-            return await download_shard(
-                shard,
-                self.on_progress_wrapper,
-                skip_download=True,
-                skip_internet=self.offline,
-            )
+            return await self._status_for_shard(shard)

        semaphore = asyncio.Semaphore(self.max_parallel_downloads)

@@ -195,10 +270,5 @@ class ResumableShardDownloader(ShardDownloader):
    async def get_shard_download_status_for_shard(
        self, shard: ShardMetadata
    ) -> RepoDownloadProgress:
-        _, progress = await download_shard(
-            shard,
-            self.on_progress_wrapper,
-            skip_download=True,
-            skip_internet=self.offline,
-        )
+        _, progress = await self._status_for_shard(shard)
        return progress
@@ -17,7 +17,6 @@ from exo.download.impl_shard_downloader import exo_shard_downloader
 from exo.master.main import Master
 from exo.routing.event_router import EventRouter
 from exo.routing.router import Router, get_node_id_keypair
-from exo.routing.state_manager import state_manager_from_routers
 from exo.shared.constants import EXO_LOG
 from exo.shared.election import Election, ElectionResult
 from exo.shared.logging import logger_cleanup, logger_setup
@@ -65,7 +64,6 @@ class Node:
            command_sender=router.sender(topics.COMMANDS),
            external_outbound=router.sender(topics.LOCAL_EVENTS),
            external_inbound=router.receiver(topics.GLOBAL_EVENTS),
-            state_receiver=router.receiver(topics.STATE_SNAPSHOTS),
        )

        logger.info(f"Starting node {node_id}")
@@ -90,7 +88,6 @@ class Node:
                command_sender=router.sender(topics.COMMANDS),
                download_command_sender=router.sender(topics.DOWNLOAD_COMMANDS),
                election_receiver=router.receiver(topics.ELECTION_MESSAGES),
-                state_manager=state_manager_from_routers(router, event_router),
            )
        else:
            api = None
@@ -103,7 +100,6 @@ class Node:
                command_sender=router.sender(topics.COMMANDS),
                download_command_sender=router.sender(topics.DOWNLOAD_COMMANDS),
                api_port=args.api_port,
-                state_manager=state_manager_from_routers(router, event_router),
            )
        else:
            worker = None
@@ -117,7 +113,6 @@ class Node:
            local_event_receiver=router.receiver(topics.LOCAL_EVENTS),
            command_receiver=router.receiver(topics.COMMANDS),
            download_command_sender=router.sender(topics.DOWNLOAD_COMMANDS),
-            state_manager=state_manager_from_routers(router, event_router),
        )

        er_send, er_recv = channel[ElectionResult]()
@@ -196,7 +191,6 @@ class Node:
                        self.router.sender(topics.COMMANDS),
                        self.router.receiver(topics.GLOBAL_EVENTS),
                        self.router.sender(topics.LOCAL_EVENTS),
-                        self.router.receiver(topics.STATE_SNAPSHOTS),
                    )

                if (
@@ -219,9 +213,6 @@ class Node:
                        download_command_sender=self.router.sender(
                            topics.DOWNLOAD_COMMANDS
                        ),
-                        state_manager=state_manager_from_routers(
-                            self.router, self.event_router
-                        ),
                    )
                    self._tg.start_soon(self.master.run)
                elif (
@@ -262,19 +253,10 @@ class Node:
                                topics.DOWNLOAD_COMMANDS
                            ),
                            api_port=self._api_port,
-                            state_manager=state_manager_from_routers(
-                                self.router, self.event_router
-                            ),
                        )
                        self._tg.start_soon(self.worker.run)
                    if self.api:
-                        self.api.reset(
-                            result.won_clock,
-                            self.event_router.receiver(),
-                            state_manager=state_manager_from_routers(
-                                self.router, self.event_router
-                            ),
-                        )
+                        self.api.reset(result.won_clock, self.event_router.receiver())
                    self._tg.start_soon(self.event_router.run)
                else:
                    if self.api:
@@ -10,7 +10,7 @@ from exo.master.placement import (
    get_transition_events,
    place_instance,
 )
-from exo.routing.state_manager import StateManager
+from exo.shared.apply import apply
 from exo.shared.constants import EXO_EVENT_LOG_DIR, EXO_TRACING_ENABLED
 from exo.shared.types.commands import (
    AddCustomModelCard,
@@ -80,11 +80,10 @@ class Master:
        local_event_receiver: Receiver[LocalForwarderEvent],
        global_event_sender: Sender[GlobalForwarderEvent],
        download_command_sender: Sender[ForwarderDownloadCommand],
-        state_manager: StateManager[State],
    ):
        self.node_id = node_id
        self.session_id = session_id
-        self.state_manager = state_manager
+        self.state = State()
        self._tg: TaskGroup = TaskGroup()
        self.command_task_mapping: dict[CommandId, TaskId] = {}
        self.command_receiver = command_receiver
@@ -119,7 +118,6 @@ class Master:
    async def _command_processor(self) -> None:
        with self.command_receiver as commands:
            async for forwarder_command in commands:
-                state = self.state_manager.get()
                try:
                    logger.info(f"Executing command: {forwarder_command.command}")

@@ -130,14 +128,14 @@ class Master:
                        case TestCommand():
                            pass
                        case TextGeneration():
-                            for instance in state.instances.values():
+                            for instance in self.state.instances.values():
                                if (
                                    instance.shard_assignments.model_id
                                    == command.task_params.model
                                ):
                                    task_count = sum(
                                        1
-                                        for task in state.tasks.values()
+                                        for task in self.state.tasks.values()
                                        if task.instance_id == instance.instance_id
                                    )
                                    instance_task_counts[instance.instance_id] = (
@@ -172,14 +170,14 @@ class Master:

                            self.command_task_mapping[command.command_id] = task_id
                        case ImageGeneration():
-                            for instance in state.instances.values():
+                            for instance in self.state.instances.values():
                                if (
                                    instance.shard_assignments.model_id
                                    == command.task_params.model
                                ):
                                    task_count = sum(
                                        1
-                                        for task in state.tasks.values()
+                                        for task in self.state.tasks.values()
                                        if task.instance_id == instance.instance_id
                                    )
                                    instance_task_counts[instance.instance_id] = (
@@ -216,7 +214,7 @@ class Master:
                            self.command_task_mapping[command.command_id] = task_id

                            if EXO_TRACING_ENABLED:
-                                selected_instance = state.instances.get(
+                                selected_instance = self.state.instances.get(
                                    selected_instance_id
                                )
                                if selected_instance:
@@ -226,14 +224,14 @@ class Master:
                                    )
                                    self._expected_ranks[task_id] = ranks
                        case ImageEdits():
-                            for instance in state.instances.values():
+                            for instance in self.state.instances.values():
                                if (
                                    instance.shard_assignments.model_id
                                    == command.task_params.model
                                ):
                                    task_count = sum(
                                        1
-                                        for task in state.tasks.values()
+                                        for task in self.state.tasks.values()
                                        if task.instance_id == instance.instance_id
                                    )
                                    instance_task_counts[instance.instance_id] = (
@@ -270,7 +268,7 @@ class Master:
                            self.command_task_mapping[command.command_id] = task_id

                            if EXO_TRACING_ENABLED:
-                                selected_instance = state.instances.get(
+                                selected_instance = self.state.instances.get(
                                    selected_instance_id
                                )
                                if selected_instance:
@@ -280,12 +278,12 @@ class Master:
                                    )
                                    self._expected_ranks[task_id] = ranks
                        case DeleteInstance():
-                            placement = delete_instance(command, state.instances)
+                            placement = delete_instance(command, self.state.instances)
                            transition_events = get_transition_events(
-                                state.instances, placement, state.tasks
+                                self.state.instances, placement, self.state.tasks
                            )
                            for cmd in cancel_unnecessary_downloads(
-                                placement, state.downloads
+                                placement, self.state.downloads
                            ):
                                await self.download_command_sender.send(
                                    ForwarderDownloadCommand(
@@ -296,24 +294,24 @@ class Master:
                        case PlaceInstance():
                            placement = place_instance(
                                command,
-                                state.topology,
-                                state.instances,
-                                state.node_memory,
-                                state.node_network,
-                                download_status=state.downloads,
+                                self.state.topology,
+                                self.state.instances,
+                                self.state.node_memory,
+                                self.state.node_network,
+                                download_status=self.state.downloads,
                            )
                            transition_events = get_transition_events(
-                                state.instances, placement, state.tasks
+                                self.state.instances, placement, self.state.tasks
                            )
                            generated_events.extend(transition_events)
                        case CreateInstance():
                            placement = add_instance_to_placements(
                                command,
-                                state.topology,
-                                state.instances,
+                                self.state.topology,
+                                self.state.instances,
                            )
                            transition_events = get_transition_events(
-                                state.instances, placement, state.tasks
+                                self.state.instances, placement, self.state.tasks
                            )
                            generated_events.extend(transition_events)
                        case SendInputChunk(chunk=chunk):
@@ -376,10 +374,9 @@ class Master:
    # These plan loops are the cracks showing in our event sourcing architecture - more things could be commands
    async def _plan(self) -> None:
        while True:
-            state = self.state_manager.get()
            # kill broken instances
-            connected_node_ids = set(state.topology.list_nodes())
-            for instance_id, instance in state.instances.items():
+            connected_node_ids = set(self.state.topology.list_nodes())
+            for instance_id, instance in self.state.instances.items():
                for node_id in instance.shard_assignments.node_to_runner:
                    if node_id not in connected_node_ids:
                        await self.event_sender.send(
@@ -388,7 +385,7 @@ class Master:
                        break

            # time out dead nodes
-            for node_id, time in state.last_seen.items():
+            for node_id, time in self.state.last_seen.items():
                now = datetime.now(tz=timezone.utc)
                if now - time > timedelta(seconds=30):
                    logger.info(f"Manually removing node {node_id} due to inactivity")
@@ -413,7 +410,6 @@ class Master:
                        continue

                    logger.debug(f"Master indexing event: {str(event)[:100]}")
-                    indexed = IndexedEvent(event=event, idx=len(self._event_log))

                    event = event.model_copy(
                        update={"_master_time_stamp": datetime.now(tz=timezone.utc)}
@@ -423,6 +419,9 @@ class Master:
                            update={"when": str(datetime.now(tz=timezone.utc))}
                        )

+                    indexed = IndexedEvent(event=event, idx=len(self._event_log))
+                    self.state = apply(self.state, indexed)
+
                    self._event_log.append(event)
                    await self._send_event(indexed)

@@ -15,7 +15,6 @@ from exo.shared.types.events import (
    IndexedEvent,
    LocalForwarderEvent,
 )
-from exo.shared.types.state import ForwarderState
 from exo.utils.channels import Receiver, Sender, channel
 from exo.utils.event_buffer import OrderedBuffer
 from exo.utils.task_group import TaskGroup
@@ -27,7 +26,6 @@ class EventRouter:
    command_sender: Sender[ForwarderCommand]
    external_inbound: Receiver[GlobalForwarderEvent]
    external_outbound: Sender[LocalForwarderEvent]
-    state_receiver: Receiver[ForwarderState]
    _system_id: SystemId = field(init=False, default_factory=SystemId)
    internal_outbound: list[Sender[IndexedEvent]] = field(
        init=False, default_factory=list
@@ -1,67 +0,0 @@
-from collections.abc import Callable, Iterable
-from dataclasses import dataclass
-
-import exo.routing.topics as topics
-from exo.shared.apply import apply
-from exo.shared.types.common import SessionId
-from exo.shared.types.events import (
-    IndexedEvent,
-)
-from exo.shared.types.state import BaseState, ForwarderState, State
-from exo.utils import fmap, fold
-from exo.utils.channels import Receiver
-
-from .event_router import EventRouter
-from .router import Router
-
-
-@dataclass
-class StateManager[T: BaseState]:
-    _apply: Callable[[T, IndexedEvent], T]
-    _event_recv: Receiver[IndexedEvent]
-    _state_recv: Receiver[T]
-    _state: T
-
-    def get(self) -> T:
-        return self._state
-
-    async def run(self):
-        async with self._event_recv, self._state_recv:
-            async for event in self._event_recv:
-                # apply new states eagerly
-                def order_state(current: T, other: T) -> T:
-                    return (
-                        current
-                        if other.last_event_idx() < current.last_event_idx()
-                        else other
-                    )
-
-                self._state = fold(self._state, order_state, self._state_recv.collect())
-                # catch up / ignore stale
-                if event.idx <= self._state.last_event_idx():
-                    continue
-                # apply state
-                self._state = self._apply(self._state, event)
-
-
-@dataclass
-class _Hack:
-    recv: Receiver[ForwarderState]
-    session_id: SessionId
-
-    def collect(self) -> Iterable[State]:
-        return fmap(self._matches, self.recv.collect())
-
-    def _matches(self, s: ForwarderState) -> State | None:
-        return s.state if s.session_id == self.session_id else None
-
-
-def state_manager_from_routers(
-    router: Router, event_router: EventRouter
-) -> StateManager[State]:
-    return StateManager(
-        apply,
-        event_router.receiver(),
-        _Hack(router.receiver(topics.STATE_SNAPSHOTS), event_router.session_id),  # type: ignore
-        State(),
-    )
@@ -8,7 +8,6 @@ from exo.shared.types.events import (
    GlobalForwarderEvent,
    LocalForwarderEvent,
 )
-from exo.shared.types.state import ForwarderState
 from exo.utils.pydantic_ext import FrozenModel


@@ -50,4 +49,3 @@ CONNECTION_MESSAGES = TypedTopic(
 DOWNLOAD_COMMANDS = TypedTopic(
    "download_commands", PublishPolicy.Always, ForwarderDownloadCommand
 )
-STATE_SNAPSHOTS = TypedTopic("state_snapshots", PublishPolicy.Always, ForwarderState)
@@ -1,11 +1,12 @@
-from abc import ABC, abstractmethod
 from collections.abc import Mapping, Sequence
 from datetime import datetime
+from typing import Any, cast

-from pydantic import Field, field_serializer, field_validator
+from pydantic import ConfigDict, Field, field_serializer, field_validator
+from pydantic.alias_generators import to_camel

 from exo.shared.topology import Topology, TopologySnapshot
-from exo.shared.types.common import NodeId, SessionId
+from exo.shared.types.common import NodeId
 from exo.shared.types.profiling import (
    DiskUsage,
    MemoryUsage,
@@ -23,12 +24,7 @@ from exo.shared.types.worker.runners import RunnerId, RunnerStatus
 from exo.utils.pydantic_ext import FrozenModel


-class BaseState(ABC):
-    @abstractmethod
-    def last_event_idx(self) -> int: ...
-
-
-class State(BaseState, FrozenModel, arbitrary_types_allowed=True):
+class State(FrozenModel):
    """Global system state.

    The :class:`Topology` instance is encoded/decoded via an immutable
@@ -36,6 +32,14 @@ class State(BaseState, FrozenModel, arbitrary_types_allowed=True):
    standard JSON serialisation.
    """

+    model_config = ConfigDict(
+        alias_generator=to_camel,
+        validate_by_name=True,
+        extra="forbid",
+        # I want to reenable this ASAP, but it's causing an issue with TaskStatus
+        strict=True,
+        arbitrary_types_allowed=True,
+    )
    instances: Mapping[InstanceId, Instance] = {}
    runners: Mapping[RunnerId, RunnerStatus] = {}
    downloads: Mapping[NodeId, Sequence[DownloadProgress]] = {}
@@ -57,9 +61,6 @@ class State(BaseState, FrozenModel, arbitrary_types_allowed=True):
    # Detected cycles where all nodes have Thunderbolt bridge enabled (>2 nodes)
    thunderbolt_bridge_cycles: Sequence[Sequence[NodeId]] = []

-    def last_event_idx(self) -> int:
-        return self.last_event_applied_idx
-
    @field_serializer("topology", mode="plain")
    def _encode_topology(self, value: Topology) -> TopologySnapshot:
        return value.to_snapshot()
@@ -77,12 +78,7 @@ class State(BaseState, FrozenModel, arbitrary_types_allowed=True):
            return value

        if isinstance(value, Mapping):  # likely a snapshot-dict coming from JSON
-            snapshot = TopologySnapshot.model_validate(value)
+            snapshot = TopologySnapshot(**cast(dict[str, Any], value))  # type: ignore[arg-type]
            return Topology.from_snapshot(snapshot)

        raise TypeError("Invalid representation for Topology field in State")
-
-
-class ForwarderState(FrozenModel):
-    state: State
-    session_id: SessionId
@@ -114,8 +114,6 @@ class TextGenerationTaskParams(BaseModel, frozen=True):
    frequency_penalty: float | None = None
    images: list[Base64Image] = Field(default_factory=list)
    image_hashes: dict[int, Base64ImageHash] = Field(default_factory=dict)
-    total_input_chunks: int = 0
-    image_count: int = 0

    def with_card_sampling_defaults(self) -> "TextGenerationTaskParams":
        from exo.shared.models.model_cards import get_card
@@ -1,5 +1,4 @@
-from collections.abc import Callable, Iterable
-from typing import Any, Type, TypeGuard
+from typing import Any, Type

 from .phantom import PhantomData

@@ -15,17 +14,3 @@ def todo[T](
    _phantom: PhantomData[T] = None,
 ) -> T:
    raise NotImplementedError(msg)
-
-
-def fold[T, U](acc: T, fn: Callable[[T, U], T], iterator: Iterable[U]) -> T:
-    for it in iterator:
-        acc = fn(acc, it)
-    return acc
-
-
-def _filter_none[U](item: U | None) -> TypeGuard[U]:
-    return item is not None
-
-
-def fmap[T, U](fn: Callable[[T], U | None], iterator: Iterable[T]) -> Iterable[U]:
-    return filter(_filter_none, map(fn, iterator))
@@ -12,7 +12,7 @@ from mlx.nn.layers.distributed import (
    sum_gradients,
 )
 from mlx_lm.models.base import (
-    scaled_dot_product_attention,  # pyright: ignore[reportUnknownVariableType]
+    scaled_dot_product_attention,
 )
 from mlx_lm.models.cache import ArraysCache, KVCache
 from mlx_lm.models.deepseek_v3 import DeepseekV3MLP
@@ -306,24 +306,20 @@ def pipeline_auto_parallel(
    )

    if isinstance(inner_model_instance, GptOssMoeModel):
-        inner_model_instance.layer_types = inner_model_instance.layer_types[  # type: ignore
+        inner_model_instance.layer_types = inner_model_instance.layer_types[
            start_layer:end_layer
        ]
        # We can assume the model has at least one layer thanks to placement.
        # If a layer type doesn't exist, we can set it to 0.
        inner_model_instance.swa_idx = (
            0
-            if "sliding_attention" not in inner_model_instance.layer_types  # type: ignore
-            else inner_model_instance.layer_types.index(  # type: ignore
-                "sliding_attention"
-            )
+            if "sliding_attention" not in inner_model_instance.layer_types
+            else inner_model_instance.layer_types.index("sliding_attention")
        )
        inner_model_instance.ga_idx = (
            0
-            if "full_attention" not in inner_model_instance.layer_types  # type: ignore
-            else inner_model_instance.layer_types.index(  # type: ignore
-                "full_attention"
-            )
+            if "full_attention" not in inner_model_instance.layer_types
+            else inner_model_instance.layer_types.index("full_attention")
        )

    if isinstance(inner_model_instance, Step35InnerModel):
@@ -883,7 +879,7 @@ class WrappedMiniMaxAttention(CustomMlxLayer):
            keys,
            values,
            cache=cache,
-            scale=self._original_layer.scale,  # type: ignore
+            scale=self._original_layer.scale,
            mask=mask,
        )

@@ -922,8 +918,8 @@ class MiniMaxShardingStrategy(TensorParallelShardingStrategy):
            self.all_to_sharded_linear_in_place(
                layer.block_sparse_moe.switch_mlp.up_proj
            )
-            layer.block_sparse_moe = ShardedMoE(layer.block_sparse_moe)  # pyright: ignore[reportAttributeAccessIssue, reportArgumentType]
-            layer.block_sparse_moe.sharding_group = self.group  # pyright: ignore[reportAttributeAccessIssue]
+            layer.block_sparse_moe = ShardedMoE(layer.block_sparse_moe)  # type: ignore
+            layer.block_sparse_moe.sharding_group = self.group
            mx.eval(layer)

            yield ModelLoadingResponse(layers_loaded=i, total=total)
@@ -1172,7 +1168,7 @@ class GptOssShardingStrategy(TensorParallelShardingStrategy):
            self.all_to_sharded_linear_in_place(layer.mlp.experts.up_proj)

            layer.mlp = ShardedMoE(layer.mlp)  # type: ignore
-            layer.mlp.sharding_group = self.group  # pyright: ignore[reportAttributeAccessIssue]
+            layer.mlp.sharding_group = self.group
            mx.eval(layer)

            yield ModelLoadingResponse(layers_loaded=i, total=total)
@@ -1374,9 +1370,11 @@ class Gemma4ShardingStrategy(TensorParallelShardingStrategy):

            attn = layer.self_attn
            attn.q_proj = self.all_to_sharded_linear(attn.q_proj)
-            attn.k_proj = self.all_to_sharded_linear(attn.k_proj)
-            if not attn.use_k_eq_v:
-                attn.v_proj = self.all_to_sharded_linear(attn.v_proj)
+            has_kv: bool = cast(bool, attn.has_kv)
+            if has_kv:
+                attn.k_proj = self.all_to_sharded_linear(attn.k_proj)
+                if not attn.use_k_eq_v:
+                    attn.v_proj = self.all_to_sharded_linear(attn.v_proj)
            attn.o_proj = self.sharded_to_all_linear(attn.o_proj)
            attn.n_heads //= self.N
            attn.n_kv_heads //= self.N
@@ -249,7 +249,12 @@ class KVPrefixCache:
        # 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
        has_ssm = has_non_kv_caches(self.caches[best_index])
-        target = (max_length - 1) if is_exact and not has_ssm else best_length
+        cached_length = cache_length(self.caches[best_index])
+        if has_ssm:
+            target = best_length
+        else:
+            desired = (max_length - 1) if is_exact else best_length
+            target = min(cached_length, desired)
        restore_pos, restore_snap = self._get_snapshot(best_index, target)

        # No usable snapshot — need fresh cache
@@ -257,7 +262,6 @@ class KVPrefixCache:
            return make_kv_cache(model), prompt_tokens, None, False

        prompt_cache = deepcopy(self.caches[best_index])
-        cached_length = cache_length(self.caches[best_index])
        tokens_to_trim = cached_length - restore_pos
        if tokens_to_trim > 0:
            trim_cache(prompt_cache, tokens_to_trim, restore_snap)
@@ -17,6 +17,13 @@ TOOL_CALLS_START = f"<{DSML_TOKEN}function_calls>"
 TOOL_CALLS_END = f"</{DSML_TOKEN}function_calls>"
 _ORPHAN_THINK_END = ASSISTANT_TOKEN + THINKING_END
 _FIXED_THINK_BLOCK = ASSISTANT_TOKEN + THINKING_START + "\n" + THINKING_END
+_FUNCTION_RESULTS_CLOSE = "</function_results>"
+_ORPHAN_TOOL_RESULT_SUFFIX = _FUNCTION_RESULTS_CLOSE + "\n\n" + THINKING_END
+_EMPTY_THINK_BLOCKS = (
+    THINKING_START + "\n\n" + THINKING_END,
+    THINKING_START + "\n" + THINKING_END,
+    THINKING_START + THINKING_END,
+)


 def encode_messages(
@@ -35,7 +42,11 @@ def encode_messages(
        add_default_bos_token=add_default_bos_token,
        tools=tools,
    )
-    return prompt.replace(_ORPHAN_THINK_END, _FIXED_THINK_BLOCK)
+    prompt = prompt.replace(_ORPHAN_TOOL_RESULT_SUFFIX, _FUNCTION_RESULTS_CLOSE)
+    prompt = prompt.replace(_ORPHAN_THINK_END, _FIXED_THINK_BLOCK)
+    for empty in _EMPTY_THINK_BLOCKS:
+        prompt = prompt.replace(empty, "")
+    return prompt


 _INVOKE_PATTERN = re.compile(
@@ -208,9 +208,20 @@ def load_mlx_items(
    if vision_config is not None:
        from exo.worker.engines.mlx.vision import VisionProcessor

-        vision_processor: VisionProcessor | None = VisionProcessor(
-            vision_config, bound_instance.bound_shard.model_card.model_id
-        )
+        vision_start_time = time.perf_counter()
+        try:
+            vision_processor: VisionProcessor | None = VisionProcessor(
+                vision_config, bound_instance.bound_shard.model_card.model_id
+            )
+            vision_processor.load()
+            logger.info(
+                f"Time taken to load vision weights: {(time.perf_counter() - vision_start_time):.2f}s"
+            )
+        except Exception as e:
+            logger.opt(exception=e).error(
+                "Failed to load vision weights — disabling vision for this runner"
+            )
+            vision_processor = None
    else:
        vision_processor = None

@@ -536,7 +547,6 @@ def render_chat_template(
        )
        if partial_assistant_content:
            prompt += partial_assistant_content
-        logger.info(prompt)
        return prompt

    for msg in formatted_messages:
@@ -36,6 +36,19 @@ from exo.worker.runner.bootstrap import logger

 _video_processor_patched = False

+_MLX_VLM_MODEL_TYPE_ALIASES: dict[str, str] = {
+    "kimi_k25": "kimi_vl",
+    "kimi_k26": "kimi_vl",
+}
+
+
+def _torch_tensor_to_mx(
+    tensor: Any,  # pyright: ignore[reportAny]
+) -> mx.array:
+    if str(tensor.dtype) == "torch.bfloat16":  # type: ignore
+        return mx.array(tensor.float().numpy(), dtype=mx.bfloat16)  # type: ignore
+    return mx.array(tensor.numpy())  # type: ignore
+

 def _filter_config(cls: type, d: dict[str, Any]) -> dict[str, Any]:
    valid = set(inspect.signature(cls.__init__).parameters.keys()) - {"self"}
@@ -85,6 +98,8 @@ def _instantiate_projector(
    params = {n: p for n, p in init_sig.parameters.items() if n != "self"}
    kwargs: dict[str, Any] = {}

+    if "config" in params:
+        kwargs["config"] = model_config
    if "embedding_dim" in params:
        kwargs["embedding_dim"] = vision_config.hidden_size  # pyright: ignore[reportAny]
    if "text_hidden_size" in params:
@@ -205,7 +220,9 @@ class VisionEncoder:
        return {}

    def _import_mlx_vlm(self, *submodules: str) -> Any:  # type: ignore
-        mt = self._config.model_type
+        mt = _MLX_VLM_MODEL_TYPE_ALIASES.get(
+            self._config.model_type, self._config.model_type
+        )
        results: list[Any] = []
        for sub in submodules:
            name = f"mlx_vlm.models.{mt}.{sub}"
@@ -238,7 +255,7 @@ class VisionEncoder:
    def _load_image_processor_from_module(self, repo: str) -> "ImageProcessor | None":
        # mlx_vlm.utils.load_image_processor only works for models that set
        # `Model.ImageProcessor = <cls>`, but Gemma4 just uses
-        # `Gemma4ImageProcessor` from the package `__init__.py`
+        # `Gemma4ImageProcessor` from the package `__init__.py`.
        try:
            pkg: Any = importlib.import_module(
                f"mlx_vlm.models.{self._config.model_type}"
@@ -319,10 +336,16 @@ class VisionEncoder:
        else:
            self._load_weights_from_model_repo()

-        repo = processor_repo or str(self._model_path)
-        image_proc = load_image_processor(
-            repo
-        ) or self._load_image_processor_from_module(repo)
+        if processor_repo:
+            repo = str(build_model_path(ModelId(processor_repo)))
+        else:
+            repo = str(self._model_path)
+        try:
+            image_proc = load_image_processor(repo)
+        except ValueError:
+            image_proc = None
+        if image_proc is None:
+            image_proc = self._load_image_processor_from_module(repo)
        if image_proc is not None:
            self._processor = image_proc
        else:
@@ -339,39 +362,42 @@ class VisionEncoder:
        if not safetensors_files:
            raise FileNotFoundError(f"No safetensors files found in {self._model_path}")

-        weights: dict[str, mx.array] = {}
-        for sf_path in safetensors_files:
-            with safe_open(str(sf_path), framework="pt") as f:
-                keys = f.keys()
-                for key in keys:
-                    tensor = f.get_tensor(key)  # type: ignore
-                    np_tensor = tensor.float().numpy()  # type: ignore
-                    weights[key] = mx.array(np_tensor, dtype=mx.bfloat16)  # type: ignore
-
        vision_weights: dict[str, mx.array] = {}
        projector_weights: dict[str, mx.array] = {}
-        for key, val in weights.items():
-            if key.startswith("vision_tower."):
-                short_key = key[len("vision_tower.") :]
-                if short_key.startswith("encoder."):
-                    short_key = short_key[len("encoder.") :]
-                m = re.match(r"^(blocks\.\d+)\.(wqkv|wo)\.(weight|bias)$", short_key)
-                if m:
-                    short_key = f"{m.group(1)}.attn.{m.group(2)}.{m.group(3)}"
-                if short_key == "patch_embed.proj.weight" and val.ndim == 4:
-                    val = val.transpose(0, 2, 3, 1)
-                vision_weights[short_key] = val
-            elif key.startswith(("mm_projector.", "multi_modal_projector.")):
-                if key.startswith("multi_modal_projector."):
-                    short_key = key[len("multi_modal_projector.") :]
-                    if short_key.startswith("mm_projector."):
-                        short_key = short_key[len("mm_projector.") :]
-                else:
-                    short_key = key[len("mm_projector.") :]
-                short_key = short_key.replace("proj.0.", "linear_1.").replace(
-                    "proj.2.", "linear_2."
-                )
-                projector_weights[short_key] = val
+
+        for sf_path in safetensors_files:
+            with safe_open(str(sf_path), framework="pt") as f:
+                keys = cast(list[str], list(f.keys()))  # type: ignore
+                for key in keys:
+                    if key.startswith("vision_tower."):
+                        short_key = key[len("vision_tower.") :]
+                        if short_key.startswith("encoder."):
+                            short_key = short_key[len("encoder.") :]
+                        m = re.match(
+                            r"^(blocks\.\d+)\.(wqkv|wo)\.(weight|bias)$", short_key
+                        )
+                        if m:
+                            short_key = f"{m.group(1)}.attn.{m.group(2)}.{m.group(3)}"
+                        tensor = f.get_tensor(key)  # type: ignore
+                        val = mx.array(tensor.float().numpy(), dtype=mx.bfloat16)  # type: ignore
+                        if short_key == "patch_embed.proj.weight" and val.ndim == 4:
+                            val = val.transpose(0, 2, 3, 1)
+                        vision_weights[short_key] = val
+                    elif key.startswith(("mm_projector.", "multi_modal_projector.")):
+                        if key.startswith("multi_modal_projector."):
+                            short_key = key[len("multi_modal_projector.") :]
+                            if short_key.startswith("mm_projector."):
+                                short_key = short_key[len("mm_projector.") :]
+                        else:
+                            short_key = key[len("mm_projector.") :]
+                        short_key = short_key.replace("proj.0.", "linear_1.").replace(
+                            "proj.2.", "linear_2."
+                        )
+                        tensor = f.get_tensor(key)  # type: ignore
+                        projector_weights[short_key] = mx.array(
+                            tensor.float().numpy(),  # type: ignore
+                            dtype=mx.bfloat16,
+                        )

        assert self._vision_tower is not None
        self._vision_tower.load_weights(list(vision_weights.items()))
@@ -407,18 +433,26 @@ class VisionEncoder:
        needs_sanitize = False

        for sf_path in safetensors_files:
-            file_weights: dict[str, mx.array] = mx.load(str(sf_path))  # type: ignore
-            for key, val in file_weights.items():
-                for prefix in vision_prefixes:
-                    if key.startswith(prefix):
-                        vision_weights[key[len(prefix) :]] = val
-                        if prefix == "model.visual.":
-                            needs_sanitize = True
-                        break
-                else:
+            with safe_open(str(sf_path), framework="pt") as f:
+                keys = cast(list[str], list(f.keys()))  # type: ignore
+                for key in keys:
+                    matched = False
+                    for prefix in vision_prefixes:
+                        if key.startswith(prefix):
+                            vision_weights[key[len(prefix) :]] = _torch_tensor_to_mx(
+                                f.get_tensor(key)
+                            )
+                            if prefix == "model.visual.":
+                                needs_sanitize = True
+                            matched = True
+                            break
+                    if matched:
+                        continue
                    for prefix in projector_prefixes:
                        if key.startswith(prefix):
-                            projector_weights[key[len(prefix) :]] = val
+                            projector_weights[key[len(prefix) :]] = _torch_tensor_to_mx(
+                                f.get_tensor(key)
+                            )
                            break

        if not vision_weights:
@@ -463,7 +497,12 @@ class VisionEncoder:
        grid_thw: mx.array | None
        n_tokens_per_image: list[int]

-        if self._config.processor_repo:
+        is_kimi_vl_processor = any(
+            "mlx_vlm.models.kimi_vl" in cls.__module__
+            for cls in type(self._processor).__mro__
+        )
+
+        if self._config.processor_repo and not is_kimi_vl_processor:
            processed = self._processor.preprocess(
                [{"type": "image", "image": img} for img in pil_images],
                return_tensors="np",
@@ -481,6 +520,24 @@ class VisionEncoder:
                int(mx.prod(grid_thw[i]).item()) // merge_length
                for i in range(grid_thw.shape[0])
            ]
+        elif is_kimi_vl_processor:
+            proc: Any = self._processor
+            raw_processed = proc.preprocess(pil_images, return_tensors="np")  # type: ignore
+            stacked_pixels = mx.array(raw_processed["pixel_values"])  # type: ignore
+            if stacked_pixels.ndim == 3:
+                stacked_pixels = stacked_pixels[None]
+            per_image_pixels = [
+                stacked_pixels[i : i + 1] for i in range(stacked_pixels.shape[0])
+            ]
+            grid_raw = raw_processed.get("image_grid_hws")  # type: ignore
+            if grid_raw is None:
+                grid_raw = raw_processed["grid_thws"]  # type: ignore
+            grid_thw = mx.array(grid_raw)  # type: ignore
+            merge_length = int(np.prod(self._merge_kernel_size or [2, 2]))
+            n_tokens_per_image = [
+                int(mx.prod(grid_thw[i]).item()) // merge_length
+                for i in range(grid_thw.shape[0])
+            ]
        else:
            batch, tokens_override = _run_processor(self._processor, pil_images)
            # `Gemma4ImageProcessor` returns pixel_values as a plain ndarray
@@ -8,7 +8,7 @@ from loguru import logger

 from exo.api.types import ImageEditsTaskParams
 from exo.download.download_utils import is_read_only_model_dir, resolve_existing_model
-from exo.routing.state_manager import StateManager
+from exo.shared.apply import apply
 from exo.shared.constants import EXO_MAX_INSTANCE_RETRIES
 from exo.shared.models.model_cards import ModelId, add_to_card_cache, delete_custom_card
 from exo.shared.types.chunks import InputImageChunk
@@ -71,7 +71,6 @@ class Worker:
        # but I think it's the correct way to be thinking about commands
        command_sender: Sender[ForwarderCommand],
        download_command_sender: Sender[ForwarderDownloadCommand],
-        state_manager: StateManager[State],
        api_port: int,
    ):
        self.node_id: NodeId = node_id
@@ -81,7 +80,7 @@ class Worker:
        self.download_command_sender = download_command_sender
        self.api_port = api_port

-        self.state_manager = state_manager
+        self.state: State = State()
        self.runners: dict[RunnerId, RunnerSupervisor] = {}
        self._tg: TaskGroup = TaskGroup()

@@ -135,6 +134,8 @@ class Worker:
    async def _event_applier(self):
        with self.event_receiver as events:
            async for event in events:
+                # 2. for each event, apply it to the state
+                self.state = apply(self.state, event=event)
                event = event.event

                if isinstance(event, InstanceDeleted):
@@ -151,6 +152,26 @@ class Worker:
                        event.chunk
                    )

+                    if (
+                        len(self.input_chunk_buffer[cmd_id])
+                        == self.input_chunk_counts[cmd_id]
+                    ):
+                        per_image: defaultdict[int, list[InputImageChunk]] = (
+                            defaultdict(list)
+                        )
+                        for chunk in self.input_chunk_buffer[cmd_id].values():
+                            per_image[chunk.image_index].append(chunk)
+                        for chunks_for_image in per_image.values():
+                            sorted_chunks = sorted(
+                                chunks_for_image, key=lambda c: c.chunk_index
+                            )
+                            img = Base64Image("".join(c.data for c in sorted_chunks))
+                            self.image_cache[
+                                Base64ImageHash(
+                                    hashlib.sha256(img.encode("ascii")).hexdigest()
+                                )
+                            ] = img
+
                if isinstance(event, CustomModelCardAdded):
                    await event.model_card.save_to_custom_dir()
                    add_to_card_cache(event.model_card)
@@ -160,16 +181,16 @@ class Worker:

    async def plan_step(self):
        while True:
-            state = self.state_manager.get()
            await anyio.sleep(0.1)
            task: Task | None = plan(
                self.node_id,
                self.runners,
-                state.downloads,
-                state.instances,
-                state.runners,
-                state.tasks,
+                self.state.downloads,
+                self.state.instances,
+                self.state.runners,
+                self.state.tasks,
                self.input_chunk_buffer,
+                self.image_cache,
                self._instance_backoff,
                self._download_backoff,
            )
@@ -209,7 +230,7 @@ class Worker:
                    self._download_backoff.record_attempt(model_id)

                    found_path = await to_thread.run_sync(
-                        resolve_existing_model, model_id
+                        resolve_existing_model, model_id, shard.model_card
                    )
                    if found_path is not None:
                        logger.info(f"Model {model_id} found at {found_path}")
@@ -305,44 +326,13 @@ class Worker:
                        del self.input_chunk_buffer[cmd_id]
                    if cmd_id in self.input_chunk_counts:
                        del self.input_chunk_counts[cmd_id]
-                    await self._start_runner_task(state, modified_task)
+                    await self._start_runner_task(modified_task)

-                case TextGeneration() if (
-                    task.task_params.image_hashes
-                    or task.task_params.total_input_chunks > 0
-                ):
+                case TextGeneration() if task.task_params.image_hashes:
                    cmd_id = task.command_id
-                    by_index: dict[int, Base64Image] = {}
-
-                    for idx, h in task.task_params.image_hashes.items():
-                        assert h in self.image_cache
-                        by_index[idx] = self.image_cache[h]
-
-                    if task.task_params.total_input_chunks > 0:
-                        chunk_buffer = self.input_chunk_buffer.get(cmd_id, {})
-                        per_image: defaultdict[int, list[InputImageChunk]] = (
-                            defaultdict(list)
-                        )
-                        for chunk in chunk_buffer.values():
-                            per_image[chunk.image_index].append(chunk)
-                        for img_idx in sorted(per_image):
-                            sorted_chunks = sorted(
-                                per_image[img_idx], key=lambda c: c.chunk_index
-                            )
-                            img = Base64Image("".join(c.data for c in sorted_chunks))
-                            self.image_cache[
-                                Base64ImageHash(
-                                    hashlib.sha256(img.encode("ascii")).hexdigest()
-                                )
-                            ] = img
-                            by_index[img_idx] = img
-                        logger.info(
-                            f"Assembled {len(per_image)} VLM image(s) "
-                            f"from {len(chunk_buffer)} chunks"
-                        )
-
                    resolved_images = [
-                        Base64Image(by_index[i]) for i in sorted(by_index)
+                        self.image_cache[h]
+                        for _, h in sorted(task.task_params.image_hashes.items())
                    ]
                    modified_task = task.model_copy(
                        update={
@@ -355,22 +345,22 @@ class Worker:
                        del self.input_chunk_buffer[cmd_id]
                    if cmd_id in self.input_chunk_counts:
                        del self.input_chunk_counts[cmd_id]
-                    await self._start_runner_task(state, modified_task)
+                    await self._start_runner_task(modified_task)
                case LoadModel(instance_id=instance_id):
-                    if (instance := state.instances.get(instance_id)) is not None:
+                    if (instance := self.state.instances.get(instance_id)) is not None:
                        model_id = instance.shard_assignments.model_id
                        self._download_backoff.reset(model_id)

-                    await self._start_runner_task(state, task)
+                    await self._start_runner_task(task)
                case task:
-                    await self._start_runner_task(state, task)
+                    await self._start_runner_task(task)

    async def shutdown(self):
        self._tg.cancel_tasks()
        await self._stopped.wait()

-    async def _start_runner_task(self, state: State, task: Task):
-        if (instance := state.instances.get(task.instance_id)) is not None:
+    async def _start_runner_task(self, task: Task):
+        if (instance := self.state.instances.get(task.instance_id)) is not None:
            await self.runners[
                instance.shard_assignments.node_to_runner[self.node_id]
            ].start_task(task)
@@ -387,13 +377,14 @@ class Worker:

    async def _poll_connection_updates(self):
        while True:
-            state = self.state_manager.get()
-            edges = set(conn.edge for conn in state.topology.out_edges(self.node_id))
+            edges = set(
+                conn.edge for conn in self.state.topology.out_edges(self.node_id)
+            )
            conns: defaultdict[NodeId, set[str]] = defaultdict(set)
            async for ip, nid in check_reachable(
-                state.topology,
+                self.state.topology,
                self.node_id,
-                state.node_network,
+                self.state.node_network,
                api_port=self.api_port,
            ):
                if ip in conns[nid]:
@@ -414,7 +405,7 @@ class Worker:
                        )
                    )

-            for conn in state.topology.out_edges(self.node_id):
+            for conn in self.state.topology.out_edges(self.node_id):
                if not isinstance(conn.edge, SocketConnection):
                    continue
                # ignore mDNS discovered connections
@@ -19,6 +19,7 @@ from exo.shared.types.tasks import (
    TaskStatus,
    TextGeneration,
 )
+from exo.shared.types.text_generation import Base64Image, Base64ImageHash
 from exo.shared.types.worker.downloads import (
    DownloadCompleted,
    DownloadFailed,
@@ -52,6 +53,7 @@ def plan(
    all_runners: Mapping[RunnerId, RunnerStatus],  # all global
    tasks: Mapping[TaskId, Task],
    input_chunk_buffer: Mapping[CommandId, Mapping[int, InputImageChunk]],
+    image_cache: Mapping[Base64ImageHash, Base64Image],
    instance_backoff: KeyedBackoff[InstanceId],
    download_backoff: KeyedBackoff[ModelId],
 ) -> Task | None:
@@ -66,7 +68,7 @@ def plan(
        or _init_distributed_backend(runners, all_runners)
        or _load_model(runners, all_runners, global_download_status)
        or _ready_to_warmup(runners, all_runners)
-        or _pending_tasks(runners, tasks, all_runners, input_chunk_buffer)
+        or _pending_tasks(runners, tasks, all_runners, input_chunk_buffer, image_cache)
    )


@@ -300,6 +302,7 @@ def _pending_tasks(
    tasks: Mapping[TaskId, Task],
    all_runners: Mapping[RunnerId, RunnerStatus],
    input_chunk_buffer: Mapping[CommandId, Mapping[int, InputImageChunk]],
+    image_cache: Mapping[Base64ImageHash, Base64Image],
 ) -> Task | None:
    for task in tasks.values():
        # for now, just forward chat completions
@@ -309,16 +312,20 @@ def _pending_tasks(
        if task.task_status not in (TaskStatus.Pending, TaskStatus.Running):
            continue

-        # For tasks with images, verify all input chunks have been received
-        expected_image_chunks = 0
-        if isinstance(task, (ImageEdits, TextGeneration)):
-            expected_image_chunks = task.task_params.total_input_chunks
-        if expected_image_chunks > 0:
-            cmd_id = task.command_id
-            received = len(input_chunk_buffer.get(cmd_id, {}))
-            if received < expected_image_chunks:
+        if isinstance(task, ImageEdits) and task.task_params.total_input_chunks > 0:
+            received = len(input_chunk_buffer.get(task.command_id, {}))
+            if received < task.task_params.total_input_chunks:
                continue  # Wait for all chunks to arrive

+        if (
+            isinstance(task, TextGeneration)
+            and task.task_params.image_hashes
+            and not all(
+                h in image_cache for h in task.task_params.image_hashes.values()
+            )
+        ):
+            continue  # Wait for all images to be assembled into the cache
+
        for runner in runners.values():
            if task.instance_id != runner.bound_instance.instance.instance_id:
                continue
@@ -79,6 +79,13 @@ def apply_all_parsers(
        issubclass(model_type, DeepseekV32Model)
        and "deepseek" in model_id.normalize().lower()
    ):
+        if tokenizer.has_thinking:
+            generator = parse_thinking_models(
+                generator,
+                tokenizer.think_start,
+                tokenizer.think_end,
+                starts_in_thinking=detect_thinking_prompt_suffix(prompt, tokenizer),
+            )
        generator = parse_deepseek_v32(generator)
    else:
        if tokenizer.has_thinking:
@@ -210,11 +217,10 @@ def parse_deepseek_v32(

    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.
+    Thinking tag handling is delegated to parse_thinking_models, which
+    wraps this parser in apply_all_parsers.
    """
    from exo.worker.engines.mlx.dsml_encoding import (
-        THINKING_END,
-        THINKING_START,
        TOOL_CALLS_END,
        TOOL_CALLS_START,
        parse_dsml_output,
@@ -222,7 +228,6 @@ def parse_deepseek_v32(

    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
@@ -264,29 +269,6 @@ def parse_deepseek_v32(
                yield response
            break

-        # ── 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
@@ -96,7 +96,12 @@ def run_gpt_oss_pipeline_device(
            n_layers=24,
        )

-        model, tokenizer = shard_and_load(shard_meta, group, on_layer_loaded=None)
+        gen = shard_and_load(shard_meta, group)
+        try:
+            while True:
+                next(gen)
+        except StopIteration as stop:
+            model, tokenizer = stop.value
        model = cast(Model, model)

        # Generate a prompt of exact token length
@@ -172,7 +177,12 @@ def run_gpt_oss_tensor_parallel_device(
            n_layers=24,
        )

-        model, tokenizer = shard_and_load(shard_meta, group, on_layer_loaded=None)
+        gen = shard_and_load(shard_meta, group)
+        try:
+            while True:
+                next(gen)
+        except StopIteration as stop:
+            model, tokenizer = stop.value
        model = cast(Model, model)

        base_text = "The quick brown fox jumps over the lazy dog. "
@@ -343,7 +343,7 @@ class TestKVPrefixCacheWithModel:
            )

    def test_mlx_generate_populates_cache(self, model_and_tokenizer):
-        """mlx_generate should save the cache after generation completes."""
+        """mlx_generate should save the post-prefill cache (before the decode loop)."""
        model, tokenizer = model_and_tokenizer

        kv_prefix_cache = KVPrefixCache(None)
@@ -356,7 +356,6 @@ class TestKVPrefixCacheWithModel:
        prompt_tokens = encode_prompt(tokenizer, prompt)

        # Consume the entire generator so the cache-saving code after yield runs
-        generated_tokens = 0
        for _response in mlx_generate(
            model=model,
            tokenizer=tokenizer,
@@ -365,13 +364,14 @@ class TestKVPrefixCacheWithModel:
            kv_prefix_cache=kv_prefix_cache,
            group=None,
        ):
-            generated_tokens += 1
+            pass

        assert len(kv_prefix_cache.prompts) == 1
        assert len(kv_prefix_cache.caches) == 1
-        # Cache should contain prompt + generated tokens
-        expected_length = len(prompt_tokens) + generated_tokens
-        assert cache_length(kv_prefix_cache.caches[0]) == expected_length
+        # add_kv_cache is called before the decode loop and stores a deepcopy of
+        # the cache as it is just after prefill + trim(2). Generation tokens are
+        # never written into the stored entry.
+        assert cache_length(kv_prefix_cache.caches[0]) == len(prompt_tokens) - 2

    def test_mlx_generate_second_call_gets_prefix_hit(self, model_and_tokenizer):
        """Second mlx_generate call with same prompt should get a prefix hit from stored cache."""
@@ -174,7 +174,12 @@ def _run_pipeline_device(
            n_layers=TOTAL_LAYERS,
        )

-        model, tokenizer = shard_and_load(shard_meta, group, on_layer_loaded=None)
+        gen = shard_and_load(shard_meta, group)
+        try:
+            while True:
+                next(gen)
+        except StopIteration as stop:
+            model, tokenizer = stop.value
        model = cast(Any, model)

        prompt, task = _build_prompt(tokenizer, prompt_tokens)
@@ -14,6 +14,8 @@ import pytest
 from mlx.utils import tree_flatten, tree_unflatten
 from mlx_lm.tokenizer_utils import TokenizerWrapper

+from exo.download.download_utils import resolve_existing_model
+from exo.shared.constants import EXO_MODELS_DIRS, EXO_MODELS_READ_ONLY_DIRS
 from exo.shared.types.common import ModelId
 from exo.shared.types.mlx import Model
 from exo.shared.types.text_generation import (
@@ -28,8 +30,6 @@ from exo.worker.engines.mlx.utils_mlx import (
    load_tokenizer_for_model_id,
 )

-HF_CACHE = Path.home() / ".cache" / "huggingface" / "hub"
-
 # ── Config reduction ──────────────────────────────────────────────────────── #

 _REDUCE = {
@@ -100,12 +100,21 @@ def _reduce_config(cfg: dict[str, Any]) -> dict[str, Any]:


 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
+    """Locate a model directory under exo's models dirs.
+
+    Uses resolve_existing_model for fully-downloaded models; falls back to any
+    existing directory (even partial) so that tokenizer-only copies still work.
+    """
+    model_id = ModelId(f"mlx-community/{hub_name}")
+    found = resolve_existing_model(model_id)
+    if found is not None:
+        return found
+    normalized = model_id.normalize()
+    for search_dir in (*EXO_MODELS_READ_ONLY_DIRS, *EXO_MODELS_DIRS):
+        candidate = search_dir / normalized
+        if candidate.is_dir():
+            return candidate
+    return None


 def _copy_tokenizer(src: Path, dst: Path) -> None:
@@ -192,13 +201,31 @@ ARCHITECTURES: list[ArchSpec] = [
 ]


+def _has_chat_template(model_dir: Path) -> bool:
+    """Check if a model dir has a usable chat template (inline or separate)."""
+    if (model_dir / "chat_template.jinja").exists():
+        return True
+    cfg = model_dir / "tokenizer_config.json"
+    if not cfg.exists():
+        return False
+    try:
+        data = cast(dict[str, Any], json.loads(cfg.read_text()))
+    except (OSError, json.JSONDecodeError):
+        return False
+    return bool(data.get("chat_template"))
+
+
 def _arch_available(spec: ArchSpec) -> bool:
    snap = _find_snapshot(spec.hub_name)
    if snap is None or not (snap / "config.json").exists():
        return False
+    tokenizer_snap = snap
    if spec.tokenizer_hub is not None:
-        return _find_snapshot(spec.tokenizer_hub) is not None
-    return True
+        alt = _find_snapshot(spec.tokenizer_hub)
+        if alt is None:
+            return False
+        tokenizer_snap = alt
+    return _has_chat_template(tokenizer_snap)


 def _make_task() -> TextGenerationTaskParams:
@@ -0,0 +1,395 @@
+# type: ignore
+"""uv run pytest -v -m "" src/exo/worker/tests/unittests/test_mlx/test_tp_bit_exact.py"""
+
+import importlib
+import json
+import multiprocessing as mp
+import os
+import sys
+import tempfile
+import traceback
+
+import numpy as np
+import pytest
+
+MODEL_CONFIGS = {
+    "llama": dict(
+        module="mlx_lm.models.llama",
+        args=dict(
+            model_type="llama",
+            hidden_size=512,
+            intermediate_size=1024,
+            num_hidden_layers=2,
+            num_attention_heads=16,
+            num_key_value_heads=4,
+            rms_norm_eps=1e-6,
+            vocab_size=512,
+            max_position_embeddings=128,
+            head_dim=32,
+            rope_theta=10000.0,
+        ),
+    ),
+    "qwen3_5_moe": dict(
+        module="mlx_lm.models.qwen3_5_moe",
+        args=dict(
+            model_type="qwen3_5_moe",
+            text_config=dict(
+                model_type="qwen3_5_moe",
+                vocab_size=512,
+                hidden_size=512,
+                intermediate_size=1024,
+                num_hidden_layers=4,
+                num_attention_heads=16,
+                num_key_value_heads=4,
+                head_dim=32,
+                max_position_embeddings=128,
+                rms_norm_eps=1e-6,
+                tie_word_embeddings=False,
+                attention_bias=False,
+                full_attention_interval=2,
+                linear_num_value_heads=32,
+                linear_num_key_heads=16,
+                linear_key_head_dim=32,
+                linear_value_head_dim=32,
+                linear_conv_kernel_dim=4,
+                num_experts=16,
+                num_experts_per_tok=2,
+                decoder_sparse_step=1,
+                shared_expert_intermediate_size=256,
+                moe_intermediate_size=256,
+                norm_topk_prob=True,
+                rope_parameters={
+                    "type": "default",
+                    "rope_theta": 10000.0,
+                    "partial_rotary_factor": 0.25,
+                    "mrope_section": [11, 11, 10],
+                },
+            ),
+        ),
+    ),
+    "qwen3_next": dict(
+        module="mlx_lm.models.qwen3_next",
+        args=dict(
+            model_type="qwen3_next",
+            hidden_size=512,
+            intermediate_size=1024,
+            num_hidden_layers=4,
+            num_attention_heads=16,
+            num_key_value_heads=4,
+            head_dim=32,
+            max_position_embeddings=128,
+            rms_norm_eps=1e-6,
+            vocab_size=512,
+            attention_bias=False,
+            full_attention_interval=2,
+            linear_num_value_heads=32,
+            linear_num_key_heads=16,
+            linear_key_head_dim=32,
+            linear_value_head_dim=32,
+            linear_conv_kernel_dim=4,
+            num_experts=16,
+            num_experts_per_tok=2,
+            decoder_sparse_step=1,
+            shared_expert_intermediate_size=256,
+            moe_intermediate_size=256,
+            norm_topk_prob=True,
+            mlp_only_layers=[],
+            rope_theta=10000.0,
+            partial_rotary_factor=0.25,
+        ),
+    ),
+    "deepseek_v3": dict(
+        module="mlx_lm.models.deepseek_v3",
+        args=dict(
+            model_type="deepseek_v3",
+            hidden_size=512,
+            intermediate_size=1024,
+            num_hidden_layers=2,
+            num_attention_heads=16,
+            num_key_value_heads=16,
+            vocab_size=512,
+            max_position_embeddings=128,
+            rms_norm_eps=1e-6,
+            n_routed_experts=8,
+            n_shared_experts=1,
+            num_experts_per_tok=2,
+            moe_intermediate_size=256,
+            moe_layer_freq=1,
+            first_k_dense_replace=0,
+            n_group=1,
+            topk_group=1,
+            routed_scaling_factor=1.0,
+            q_lora_rank=None,
+            kv_lora_rank=16,
+            qk_nope_head_dim=16,
+            qk_rope_head_dim=16,
+            v_head_dim=32,
+            rope_theta=10000.0,
+            rope_scaling={},
+            attention_bias=False,
+            norm_topk_prob=True,
+            scoring_func="sigmoid",
+            topk_method="noaux_tc",
+        ),
+    ),
+    "deepseek_v3_q4": dict(
+        module="mlx_lm.models.deepseek_v3",
+        quantize=dict(group_size=32, bits=4, mode="affine"),
+        args=dict(
+            model_type="deepseek_v3",
+            hidden_size=512,
+            intermediate_size=1024,
+            num_hidden_layers=2,
+            num_attention_heads=16,
+            num_key_value_heads=16,
+            vocab_size=512,
+            max_position_embeddings=128,
+            rms_norm_eps=1e-6,
+            n_routed_experts=8,
+            n_shared_experts=1,
+            num_experts_per_tok=2,
+            moe_intermediate_size=256,
+            moe_layer_freq=1,
+            first_k_dense_replace=0,
+            n_group=1,
+            topk_group=1,
+            routed_scaling_factor=1.0,
+            q_lora_rank=None,
+            kv_lora_rank=64,
+            qk_nope_head_dim=32,
+            qk_rope_head_dim=32,
+            v_head_dim=32,
+            rope_theta=10000.0,
+            rope_scaling={},
+            attention_bias=False,
+            norm_topk_prob=True,
+            scoring_func="sigmoid",
+            topk_method="noaux_tc",
+        ),
+    ),
+    "glm4_moe_lite": dict(
+        module="mlx_lm.models.glm4_moe_lite",
+        args=dict(
+            model_type="glm4_moe_lite",
+            hidden_size=512,
+            intermediate_size=1024,
+            num_hidden_layers=2,
+            num_attention_heads=16,
+            num_key_value_heads=16,
+            vocab_size=512,
+            max_position_embeddings=128,
+            rms_norm_eps=1e-6,
+            n_routed_experts=8,
+            n_shared_experts=1,
+            num_experts_per_tok=2,
+            moe_intermediate_size=256,
+            first_k_dense_replace=1,
+            n_group=1,
+            topk_group=1,
+            routed_scaling_factor=1.0,
+            rope_theta=10000.0,
+            attention_bias=False,
+            q_lora_rank=None,
+            kv_lora_rank=16,
+            qk_rope_head_dim=16,
+            qk_nope_head_dim=16,
+            v_head_dim=32,
+        ),
+    ),
+    "minimax": dict(
+        module="mlx_lm.models.minimax",
+        args=dict(
+            model_type="minimax",
+            hidden_size=512,
+            intermediate_size=1024,
+            num_attention_heads=16,
+            num_key_value_heads=4,
+            max_position_embeddings=128,
+            num_experts_per_tok=2,
+            num_local_experts=8,
+            shared_intermediate_size=256,
+            num_hidden_layers=2,
+            rms_norm_eps=1e-6,
+            rope_theta=10000.0,
+            rotary_dim=32,
+            vocab_size=512,
+        ),
+    ),
+    "gpt_oss": dict(
+        module="mlx_lm.models.gpt_oss",
+        args=dict(
+            model_type="gpt_oss",
+            hidden_size=512,
+            intermediate_size=256,
+            num_hidden_layers=2,
+            num_attention_heads=16,
+            num_key_value_heads=4,
+            vocab_size=512,
+            head_dim=32,
+            rms_norm_eps=1e-6,
+            num_local_experts=8,
+            num_experts_per_tok=2,
+            layer_types=["sliding_attention", "full_attention"],
+            sliding_window=64,
+            rope_theta=10000.0,
+        ),
+    ),
+    "gemma4": dict(
+        module="mlx_lm.models.gemma4",
+        args=dict(
+            model_type="gemma4",
+            vocab_size=512,
+            text_config=dict(
+                vocab_size=512,
+                hidden_size=512,
+                intermediate_size=1024,
+                num_hidden_layers=4,
+                num_attention_heads=16,
+                num_key_value_heads=4,
+                head_dim=32,
+                global_head_dim=32,
+                num_kv_shared_layers=0,
+                vocab_size_per_layer_input=512,
+                hidden_size_per_layer_input=512,
+                rms_norm_eps=1e-6,
+                max_position_embeddings=128,
+                sliding_window=64,
+                sliding_window_pattern=2,
+                layer_types=[
+                    "sliding_attention",
+                    "full_attention",
+                    "sliding_attention",
+                    "full_attention",
+                ],
+                enable_moe_block=True,
+                num_experts=8,
+                top_k_experts=2,
+                moe_intermediate_size=256,
+            ),
+        ),
+    ),
+}
+
+_PROMPT = [[1, 23, 45, 67, 89, 12, 34, 56]]
+
+
+def _build(name):
+    import mlx.core as mx
+    import mlx.nn as nn
+    from mlx.utils import tree_map_with_path
+
+    import exo.worker.engines.mlx.auto_parallel  # noqa: F401
+
+    cfg = MODEL_CONFIGS[name]
+    module = importlib.import_module(cfg["module"])
+    model_cls = module.Model
+    model_args_cls = module.ModelArgs
+
+    mx.random.seed(0)
+    args = model_args_cls(**cfg["args"])
+    m = model_cls(args)
+
+    def _to_bf16(_p, v):
+        if hasattr(v, "dtype") and v.dtype in (mx.float16, mx.float32, mx.bfloat16):
+            return v.astype(mx.bfloat16)
+        return v
+
+    m.update(tree_map_with_path(_to_bf16, m.parameters()))
+    if "quantize" in cfg:
+        nn.quantize(m, **cfg["quantize"])
+    mx.eval(m.parameters())
+    return mx, m
+
+
+def _run(name, out_path, shard):
+    import mlx.core as mx
+
+    if shard:
+        g = mx.distributed.init(backend="ring", strict=True)
+    mx_, m = _build(name)
+    if shard:
+        from exo.worker.engines.mlx.auto_parallel import tensor_auto_parallel
+
+        m = tensor_auto_parallel(m, g, on_layer_loaded=None)
+        mx_.eval(m.parameters())
+    inputs = mx_.array(_PROMPT, dtype=mx_.int32)
+    logits = m(inputs)
+    mx_.eval(logits)
+    np.savez(out_path, logits=np.asarray(logits.astype(mx_.float32)))
+
+
+def _ref_worker(name, out_path, q):
+    try:
+        _run(name, out_path, shard=False)
+        q.put(True)
+    except BaseException as e:
+        q.put(f"{e}\n{traceback.format_exc()}")
+
+
+def _tp_worker(name, rank, hf, out_path, q):
+    os.environ["MLX_HOSTFILE"] = hf
+    os.environ["MLX_RANK"] = str(rank)
+    try:
+        path = out_path if rank == 0 else out_path + f".r{rank}"
+        _run(name, path, shard=True)
+        q.put((rank, True, None))
+    except BaseException as e:
+        q.put((rank, False, f"{e}\n{traceback.format_exc()}"))
+
+
+def _run_compare(name, world_size, port_base):
+    d = tempfile.mkdtemp()
+    ref_path = f"{d}/ref.npz"
+    tp_path = f"{d}/tp.npz"
+    ctx = mp.get_context("spawn")
+    q = ctx.Queue()
+
+    p = ctx.Process(target=_ref_worker, args=(name, ref_path, q))
+    p.start()
+    p.join(300)
+    r = q.get(timeout=10)
+    if r is not True:
+        pytest.fail(f"[{name}] ref FAIL: {str(r)[:500]}")
+
+    hosts = [f"127.0.0.1:{port_base + i}" for i in range(world_size)]
+    with tempfile.NamedTemporaryFile(mode="w", suffix=".json", delete=False) as f:
+        json.dump(hosts, f)
+        hf = f.name
+    ps = [
+        ctx.Process(target=_tp_worker, args=(name, rank, hf, tp_path, q))
+        for rank in range(world_size)
+    ]
+    for pp in ps:
+        pp.start()
+    results = [q.get(timeout=300) for _ in range(world_size)]
+    for pp in ps:
+        pp.join(60)
+    for rank, ok, payload in results:
+        if not ok:
+            pytest.fail(f"[{name}] rank {rank} FAIL: {payload[:500]}")
+
+    ref = np.load(ref_path)["logits"]
+    tp = np.load(tp_path)["logits"]
+    diff = np.abs(ref - tp)
+    max_diff = float(diff.max())
+    mean_diff = float(diff.mean())
+    assert max_diff == 0.0, (
+        f"[{name} TP={world_size}] not bit-exact: max={max_diff} mean={mean_diff}"
+    )
+
+
+pytestmark = [
+    pytest.mark.slow,
+    pytest.mark.skipif(
+        sys.platform != "darwin", reason="MLX distributed requires Metal"
+    ),
+]
+
+
+@pytest.mark.skip("TP=2 is currently very different to TP=1. This test will not pass")
+@pytest.mark.parametrize("world_size", [2, 4])
+@pytest.mark.parametrize("name", list(MODEL_CONFIGS))
+def test_tp_bit_exact(name, world_size):
+    name_idx = list(MODEL_CONFIGS).index(name)
+    port = 32000 + name_idx * 20 + world_size
+    _run_compare(name, world_size, port)
@@ -54,6 +54,7 @@ def test_plan_requests_download_when_waiting_and_shard_not_downloaded():
        all_runners=all_runners,
        tasks={},
        input_chunk_buffer={},
+        image_cache={},
        instance_backoff=KeyedBackoff(),
        download_backoff=KeyedBackoff(),
    )
@@ -109,6 +110,7 @@ def test_plan_loads_model_when_all_shards_downloaded_and_waiting():
        all_runners=all_runners,
        tasks={},
        input_chunk_buffer={},
+        image_cache={},
        instance_backoff=KeyedBackoff(),
        download_backoff=KeyedBackoff(),
    )
@@ -154,6 +156,7 @@ def test_plan_does_not_request_download_when_shard_already_downloaded():
        all_runners=all_runners,
        tasks={},
        input_chunk_buffer={},
+        image_cache={},
        instance_backoff=KeyedBackoff(),
        download_backoff=KeyedBackoff(),
    )
@@ -204,6 +207,7 @@ def test_plan_does_not_load_model_until_all_shards_downloaded_globally():
        all_runners=all_runners,
        tasks={},
        input_chunk_buffer={},
+        image_cache={},
        instance_backoff=KeyedBackoff(),
        download_backoff=KeyedBackoff(),
    )
@@ -227,6 +231,7 @@ def test_plan_does_not_load_model_until_all_shards_downloaded_globally():
        all_runners=all_runners,
        tasks={},
        input_chunk_buffer={},
+        image_cache={},
        instance_backoff=KeyedBackoff(),
        download_backoff=KeyedBackoff(),
    )
@@ -54,6 +54,7 @@ def test_plan_kills_runner_when_instance_missing():
        all_runners=all_runners,
        tasks={},
        input_chunk_buffer={},
+        image_cache={},
        instance_backoff=KeyedBackoff(),
        download_backoff=KeyedBackoff(),
    )
@@ -96,6 +97,7 @@ def test_plan_kills_runner_when_sibling_failed():
        all_runners=all_runners,
        tasks={},
        input_chunk_buffer={},
+        image_cache={},
        instance_backoff=KeyedBackoff(),
        download_backoff=KeyedBackoff(),
    )
@@ -130,6 +132,7 @@ def test_plan_creates_runner_when_missing_for_node():
        all_runners=all_runners,
        tasks={},
        input_chunk_buffer={},
+        image_cache={},
        instance_backoff=KeyedBackoff(),
        download_backoff=KeyedBackoff(),
    )
@@ -171,6 +174,7 @@ def test_plan_does_not_create_runner_when_supervisor_already_present():
        all_runners=all_runners,
        tasks={},
        input_chunk_buffer={},
+        image_cache={},
        instance_backoff=KeyedBackoff(),
        download_backoff=KeyedBackoff(),
    )
@@ -203,6 +207,7 @@ def test_plan_does_not_create_runner_for_unassigned_node():
        all_runners=all_runners,
        tasks={},
        input_chunk_buffer={},
+        image_cache={},
        instance_backoff=KeyedBackoff(),
        download_backoff=KeyedBackoff(),
    )
@@ -78,6 +78,7 @@ def test_plan_forwards_pending_chat_completion_when_runner_ready():
        all_runners=all_runners,
        tasks={TASK_1_ID: task},
        input_chunk_buffer={},
+        image_cache={},
        instance_backoff=KeyedBackoff(),
        download_backoff=KeyedBackoff(),
    )
@@ -131,6 +132,7 @@ def test_plan_does_not_forward_chat_completion_if_any_runner_not_ready():
        all_runners=all_runners,
        tasks={TASK_1_ID: task},
        input_chunk_buffer={},
+        image_cache={},
        instance_backoff=KeyedBackoff(),
        download_backoff=KeyedBackoff(),
    )
@@ -181,6 +183,7 @@ def test_plan_does_not_forward_tasks_for_other_instances():
        all_runners=all_runners,
        tasks={foreign_task.task_id: foreign_task},
        input_chunk_buffer={},
+        image_cache={},
        instance_backoff=KeyedBackoff(),
        download_backoff=KeyedBackoff(),
    )
@@ -249,6 +252,7 @@ def test_plan_ignores_non_pending_or_non_chat_tasks():
        all_runners=all_runners,
        tasks={TASK_1_ID: completed_task, other_task_id: other_task},
        input_chunk_buffer={},
+        image_cache={},
        instance_backoff=KeyedBackoff(),
        download_backoff=KeyedBackoff(),
    )
@@ -291,6 +295,7 @@ def test_plan_returns_none_when_nothing_to_do():
        all_runners=all_runners,
        tasks={},
        input_chunk_buffer={},
+        image_cache={},
        instance_backoff=KeyedBackoff(),
        download_backoff=KeyedBackoff(),
    )
@@ -63,6 +63,7 @@ def test_plan_starts_warmup_for_accepting_rank_when_all_loaded_or_warming():
        all_runners=all_runners,
        tasks={},
        input_chunk_buffer={},
+        image_cache={},
        instance_backoff=KeyedBackoff(),
        download_backoff=KeyedBackoff(),
    )
@@ -107,6 +108,7 @@ def test_plan_starts_warmup_for_rank_zero_after_others_warming():
        all_runners=all_runners,
        tasks={},
        input_chunk_buffer={},
+        image_cache={},
        instance_backoff=KeyedBackoff(),
        download_backoff=KeyedBackoff(),
    )
@@ -150,6 +152,7 @@ def test_plan_does_not_start_warmup_for_non_zero_rank_until_all_loaded_or_warmin
        all_runners=all_runners,
        tasks={},
        input_chunk_buffer={},
+        image_cache={},
        instance_backoff=KeyedBackoff(),
        download_backoff=KeyedBackoff(),
    )
@@ -197,6 +200,7 @@ def test_plan_does_not_start_warmup_for_rank_zero_until_others_warming():
        all_runners=all_runners,
        tasks={},
        input_chunk_buffer={},
+        image_cache={},
        instance_backoff=KeyedBackoff(),
        download_backoff=KeyedBackoff(),
    )
@@ -216,6 +220,7 @@ def test_plan_does_not_start_warmup_for_rank_zero_until_others_warming():
        all_runners=all_runners,
        tasks={},
        input_chunk_buffer={},
+        image_cache={},
        instance_backoff=KeyedBackoff(),
        download_backoff=KeyedBackoff(),
    )
@@ -262,6 +267,7 @@ def test_plan_starts_warmup_for_connecting_rank_after_others_warming():
        all_runners=all_runners,
        tasks={},
        input_chunk_buffer={},
+        image_cache={},
        instance_backoff=KeyedBackoff(),
        download_backoff=KeyedBackoff(),
    )
@@ -307,6 +313,7 @@ def test_plan_does_not_start_warmup_for_accepting_rank_until_all_loaded_or_warmi
        all_runners=all_runners,
        tasks={},
        input_chunk_buffer={},
+        image_cache={},
        instance_backoff=KeyedBackoff(),
        download_backoff=KeyedBackoff(),
    )
@@ -351,6 +358,7 @@ def test_plan_does_not_start_warmup_for_connecting_rank_until_others_warming():
        all_runners=all_runners,
        tasks={},
        input_chunk_buffer={},
+        image_cache={},
        instance_backoff=KeyedBackoff(),
        download_backoff=KeyedBackoff(),
    )
@@ -20,7 +20,25 @@ from exo.worker.engines.mlx.dsml_encoding import (
    encode_messages,
    parse_dsml_output,
 )
-from exo.worker.runner.llm_inference.model_output_parsers import parse_deepseek_v32
+from exo.worker.runner.llm_inference.model_output_parsers import (
+    parse_deepseek_v32,
+    parse_thinking_models,
+)
+
+
+def _parse_deepseek_with_thinking(
+    source: Generator[GenerationResponse | None],
+    starts_in_thinking: bool = False,
+) -> Generator[GenerationResponse | ToolCallResponse | None]:
+    return parse_deepseek_v32(
+        parse_thinking_models(
+            source,
+            think_start=THINKING_START,
+            think_end=THINKING_END,
+            starts_in_thinking=starts_in_thinking,
+        )
+    )
+

 # ── Shared fixtures ──────────────────────────────────────────────

@@ -333,9 +351,7 @@ class TestE2EThinkingAndToolCall:
        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.
+        # Use the full production chain (parse_thinking_models → parse_deepseek_v32).
        model_tokens = [
            THINKING_START,
            "The user wants weather",
@@ -353,7 +369,7 @@ class TestE2EThinkingAndToolCall:
            TOOL_CALLS_END,
        ]

-        results = list(parse_deepseek_v32(_simulate_tokens(model_tokens)))
+        results = list(_parse_deepseek_with_thinking(_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)]
@@ -387,7 +403,7 @@ class TestE2EThinkingAndToolCall:
        prompt_no_think = encode_messages(
            messages, tools=_WEATHER_TOOLS, thinking_mode="chat"
        )
-        assert prompt_no_think.endswith(THINKING_END)
+        assert not prompt_no_think.endswith(THINKING_START)

        # Both should have the same tool definitions
        assert "get_weather" in prompt_think
@@ -597,7 +613,9 @@ class TestE2EFullRoundTrip:
            f"</{DSML_TOKEN}invoke>\n",
            TOOL_CALLS_END,
        ]
-        results_1 = list(parse_deepseek_v32(_simulate_tokens(model_tokens_1)))
+        results_1 = list(
+            _parse_deepseek_with_thinking(_simulate_tokens(model_tokens_1))
+        )

        # Verify: thinking tokens + tool call
        gen_1 = [r for r in results_1 if isinstance(r, GenerationResponse)]
@@ -660,7 +678,9 @@ class TestE2EFullRoundTrip:
            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)))
+        results_2 = list(
+            _parse_deepseek_with_thinking(_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)]
@@ -380,6 +380,110 @@ class TestGenericToolCallsFinishReason:
 # ── Double parser chain (parse_thinking_models → parse_deepseek_v32) ──


+class TestDeepSeekV32StartsInThinking:
+    """Regression tests for deepseek v3.2 where the chat template appends
+    <think> to the prompt so the model starts already inside a thinking block.
+    """
+
+    def test_reasoning_tagged_when_starts_in_thinking(self):
+        tokens = [
+            _make_response("let me", 0),
+            _make_response(" think", 1),
+            _make_response(THINKING_END, 2),
+            _make_response("\n", 3),
+            _make_response("42", 4, finish_reason="stop"),
+        ]
+        thinking = parse_thinking_models(
+            _queue_source(tokens),
+            think_start=THINKING_START,
+            think_end=THINKING_END,
+            starts_in_thinking=True,
+        )
+        results = _step_until_finish(parse_deepseek_v32(thinking))
+        gens = [
+            r
+            for r in results
+            if isinstance(r, GenerationResponse) and r.finish_reason is None
+        ]
+        texts = [(r.text, r.is_thinking) for r in gens]
+        assert texts == [("let me", True), (" think", True), ("\n", False)]
+        final = [
+            r
+            for r in results
+            if isinstance(r, GenerationResponse) and r.finish_reason is not None
+        ]
+        assert len(final) == 1
+        assert final[0].text == "42"
+        assert final[0].is_thinking is False
+
+    def test_starts_in_thinking_then_tool_call(self):
+        tokens = [
+            _make_response("need weather", 0),
+            _make_response(THINKING_END, 1),
+            _make_response("\n\n", 2),
+            _make_response(TOOL_CALLS_START, 3),
+            _make_response("\n", 4),
+            _make_response(f'<{DSML_TOKEN}invoke name="get_weather">\n', 5),
+            _make_response(
+                f'<{DSML_TOKEN}parameter name="city" string="true">NYC</{DSML_TOKEN}parameter>\n',
+                6,
+            ),
+            _make_response(f"</{DSML_TOKEN}invoke>\n", 7),
+            _make_response(TOOL_CALLS_END, 8, finish_reason="stop"),
+        ]
+        thinking = parse_thinking_models(
+            _queue_source(tokens),
+            think_start=THINKING_START,
+            think_end=THINKING_END,
+            starts_in_thinking=True,
+        )
+        results = _step_until_finish(parse_deepseek_v32(thinking))
+        reasoning_gens = [
+            r
+            for r in results
+            if isinstance(r, GenerationResponse)
+            and r.finish_reason is None
+            and r.is_thinking
+        ]
+        assert [r.text for r in reasoning_gens] == ["need weather"]
+        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_reasoning_tokens_counted_starts_in_thinking(self):
+        usage = Usage(
+            prompt_tokens=10,
+            completion_tokens=5,
+            total_tokens=15,
+            prompt_tokens_details=PromptTokensDetails(cached_tokens=0),
+            completion_tokens_details=CompletionTokensDetails(reasoning_tokens=0),
+        )
+        tokens = [
+            _make_response("reasoning", 0),
+            _make_response(" more", 1),
+            _make_response(THINKING_END, 2),
+            _make_response("\n", 3),
+            GenerationResponse(text="42", token=4, finish_reason="stop", usage=usage),
+        ]
+        thinking = parse_thinking_models(
+            _queue_source(tokens),
+            think_start=THINKING_START,
+            think_end=THINKING_END,
+            starts_in_thinking=True,
+        )
+        results = _step_until_finish(
+            count_reasoning_tokens(parse_deepseek_v32(thinking))
+        )
+        final = [
+            r
+            for r in results
+            if isinstance(r, GenerationResponse) and r.finish_reason is not None
+        ]
+        assert len(final) == 1
+        assert final[0].usage is not None
+        assert final[0].usage.completion_tokens_details.reasoning_tokens == 2
+
+
 class TestBatchGeneratorSingleNext:
    def test_finish_reason_with_buffered_tokens_drain_loop(self):
        from exo.worker.runner.llm_inference.batch_generator import GeneratorQueue
Author	SHA1	Message	Date
Evan	4216ca541a	implement engine interface for mlx and mflux	2026-04-25 01:52:29 +01:00
rltakashige	fd707de30b	Add more model cards (#1970 )	2026-04-23 15:28:40 +01:00
Alex CheemaandClaude Opus 4.7	45248c5c85	chore(app): hardcode bug report presigned-URL endpoint (#1971 ) ## Motivation The bug-report presigned-URL endpoint (`https://reports.exolabs.net/presigned-urls`) was injected at build time from the `EXO_BUG_REPORT_PRESIGNED_URL_ENDPOINT` GitHub Actions secret into `Info.plist`, then read at runtime by `BugReportService`. It isn't actually a secret — the POST body is just `{"keys":[...]}` with no credential (see `app/EXO/EXO/Services/BugReportService.swift:136-142`), abuse prevention lives server-side on the lambda, and the URL is already visible in every publicly-distributed DMG's `Info.plist`. Treating it as a repo secret added plumbing with no security benefit and broke local dev builds — hitting Send Bug Report on an uncustomised `just build-app` raised "Bug report endpoint is invalid". ## Changes - `app/EXO/EXO/Info.plist`: replace `$(EXO_BUG_REPORT_PRESIGNED_URL_ENDPOINT)` with the literal URL. - `.github/workflows/build-app.yml`: drop the `EXO_BUG_REPORT_PRESIGNED_URL_ENDPOINT` job-level env var and the xcodebuild build-setting passthrough. No other workflow changes. Swift code is unchanged — `BugReportService` still reads from `Info.plist`, which leaves an escape hatch if anyone ever needs to override via `xcodebuild EXOBugReportPresignedUrlEndpoint=...` without recompiling. Follow-up: the `EXO_BUG_REPORT_PRESIGNED_URL_ENDPOINT` repo secret can now be deleted in the GitHub Actions settings UI. ## Why It Works `Info.plist` variable substitution turns `$(FOO)` into whatever build setting `FOO` resolves to. CI was setting `FOO` via xcodebuild; local dev wasn't, so the key resolved to an empty string, which `BugReportService.fetchPresignedUploadUrls` rejects via the `!trimmedEndpointString.isEmpty` guard at `BugReportService.swift:131`. Hardcoding the literal string removes the substitution entirely, so every build — local or CI — gets the right value. ## Test Plan ### Manual Testing <!-- Hardware: MacBook Pro (macOS app build via Xcode) --> - `just build-app` with no extra env vars (reproduces the failure path on `main`). - `/usr/libexec/PlistBuddy -c "Print :EXOBugReportPresignedUrlEndpoint" app/EXO/build/Build/Products/Debug/EXO.app/Contents/Info.plist` → returns `https://reports.exolabs.net/presigned-urls` (was empty before this change). - `open app/EXO/build/Build/Products/Debug/EXO.app` → menubar → Debug Info → Send Bug Report → type a description → Send → upload succeeds and the Create GitHub Issue button appears (was failing with "Bug report endpoint is invalid" before). - Cross-check on the Slack side that the uploaded `report.json` lands under `reports/YYYY/MM/DD/<ts>/` as before. ### Automated Testing <!-- Describe changes to automated tests, or how existing tests cover this change --> - No new tests. This is a single-string change to `Info.plist` plus a workflow cleanup. `nix flake check` in CI verifies formatting/lint for the rest of the tree. 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>	2026-04-23 14:08:16 +00:00
rltakashige	290e3fd927	Keep image cache fresh (temporary fix) (#1961 ) ## Motivation When a new node joins, it might not have the cache. Caveat: This is potentially fallible if a new node joins and updates real topology, but the API topology hasn't caught up with this fact and the user queues up a new text generation. In practice, there is only a split second where this is the case, and this is only for users of the dashboard interface. We should fix this properly after the release.	2026-04-23 11:39:36 +01:00
rltakashige	3894cf134e	Fix Gemma 4 E2B TP + DeepSeek V32 thinking parsing (#1967 )	2026-04-23 01:50:39 +00:00
Alex CheemaandClaude Opus 4.7	8993ccaf09	feat(app): add friendly context message to bug report prompt (#1959 ) ## Motivation When a user clicks Send Bug Report in the macOS app, we already give them the option to add more context via an optional text field. But the current prompt is just a terse label — `"What's the issue? (optional)"` — which doesn't tell the user why bothering to fill it in matters. A friendly one-line explanation increases the chance they'll describe what went wrong, which is the single most useful signal when we triage the resulting diagnostic bundle. ## Changes - `app/EXO/EXO/ContentView.swift`: In the `.prompting` phase of `sendBugReportButton`, replace the single label with a two-line hierarchy: - Primary: `Tell us what went wrong (optional)` - Helper: `A quick description of what you were doing and what happened helps us track down the bug for you.` - The helper uses `.caption2` + `.secondary` + `.opacity(0.8)` + `.fixedSize(horizontal: false, vertical: true)` so it stays visually subordinate and wraps cleanly inside the 340pt popover. No changes to `BugReportService`, the `user_description` payload, or any other flow. ## Why It Works The optional description is already plumbed end-to-end (text editor → `bugReportUserDescription` state → `BugReportService.sendReport(..., userDescription:)` → `report.json`'s `user_description` field → GitHub issue pre-fill). The only gap was user-facing motivation, so this is purely a copy/layout tweak inside the existing `.prompting` case — no new state, bindings, or service changes. ## Test Plan ### Manual Testing <!-- Hardware: MacBook Pro (macOS app build via Xcode) --> - Build the macOS app in Xcode (`app/EXO/EXO.xcodeproj`) and launch it. - Open the menubar popover → expand Debug Info → click Send Bug Report. - Verify the new primary label and helper sentence both appear above the text editor and wrap cleanly within the popover width. - Leave the field empty → click Send → upload should succeed (no `user_description` in payload, same as before). - Fill in a description → click Send → upload succeeds and the success card with Create GitHub Issue appears; clicking it opens GitHub with the description pre-filled. - Click Cancel from the prompting state → returns to idle. ### Automated Testing <!-- Describe changes to automated tests, or how existing tests cover this change --> - No new automated tests. This is a SwiftUI copy/layout change; existing `EXOTests` are smoke-level and don't cover `ContentView` view bodies, and UI snapshot tests aren't worth adding for a two-line copy tweak. - `nix fmt` reports 0 files changed after the edit; `nix flake check` in CI will verify formatting/lint for the rest of the tree. 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>	2026-04-22 17:39:36 +00:00
Nadeem Hilal Wani	4939fbe995	feat(dashboard): add Pi integration tab (#1925 ) ## Summary Adds a new Pi tab to the Integrations page (`/#/integrations`) alongside the existing Claude Code, OpenCode, Codex, OpenClaw, Open WebUI, n8n, and Firefox tabs. [pi](https://pi.dev) (`@mariozechner/pi-coding-agent`) is a terminal coding agent that supports custom OpenAI-compatible providers via `~/.pi/agent/models.json`. This tab gives users a copy-pasteable config to wire pi up to their exo cluster. ## What's in the tab - Model selector (shown when multiple models are running) — picks the default model for the generated shell command. - Models Config card — generates `~/.pi/agent/models.json` registering `exo` as a custom provider: - `baseUrl` → `<apiUrl>/v1` - `api` → `openai-completions` - `apiKey` → `"exo"` (placeholder; exo ignores it) - `compat.supportsDeveloperRole: false` and `compat.supportsReasoningEffort: false`, per pi docs recommendation for local OpenAI-compatible servers - Auto-populates every running model with `id`, `contextWindow` (from `/v1/models`), and `input: ["text", "image"]` for vision-capable models - Shell Command card — `pi --provider exo --model <model>` for quick launch. The tab gracefully falls back to `your-model-id` when no models are running, matching the behavior of the other tabs. ## Usage 1. `npm install -g @mariozechner/pi-coding-agent` 2. Paste the generated config into `~/.pi/agent/models.json` 3. Run `pi` and pick an exo model via `/model` — or run the shell command directly ## Changes - `dashboard/src/routes/integrations/+page.svelte` — adds `"Pi"` to the `tabs` tuple, `piModel` state, `piModelsJson` + `piShellCommand` derivations, and the tab content block. Single-file, scoped change — no backend or type changes. ## Testing - `cd dashboard && npm run build` — ✅ builds cleanly - `svelte-check` on the edited file — no new errors - Manually verified the tab renders, the model selector updates the generated JSON, and the config reflects `/v1/models` capabilities (vision → `input: ["text","image"]`, `context_length` → `contextWindow`). ## Screenshots <img width="1545" height="1236" alt="pi-tab" src="https://github.com/user-attachments/assets/38aa179f-4ed9-4a1e-9783-d3baa7738263" />	2026-04-22 17:29:47 +00:00
rltakashige	73782ecc65	Fix event mutation causing indexed vs event mismatch (#1964 ) Fixes small issue with #1957	2026-04-22 16:12:24 +00:00
rltakashige	f6e418ed23	Cleanup on #1952 (#1960 )	2026-04-22 17:05:49 +01:00
rltakashigeandEvan	7a312a177b	Misc fixes: upstream JACCL all_sum, API, etc. + Add Kimi K2.6 (#1952 ) ## Motivation This fixes a bunch of observed model quality issues introduced upstream in JACCL, as well as API issues and prefix cache calculation. ## Test Plan ### Manual Testing Tested a bunch ### Automated Testing Added a test, automated eval tool calls on Kimi K2.6, Minimax M2.7, GPT OSS and Qwen3.6 models. --------- Co-authored-by: Evan <evanev7@gmail.com>	2026-04-22 15:43:27 +00:00
Evan Quiney	0a549f8846	remove layer loading callback (#1890 ) first part of modularising the backend is simplifying some of the control flow. more tbd.	2026-04-22 14:03:31 +01:00