disconnect on mdns expiry

Closes #1586

## Summary
- track in-flight tasks in `RunnerSupervisor` (not only unacknowledged
pending tasks)
- when `_check_runner()` detects a crashed runner, emit
`ChunkGenerated(ErrorChunk)` for each in-flight command task
(`TextGeneration`, `ImageGeneration`, `ImageEdits`)
- keep existing `RunnerStatusUpdated(RunnerFailed)` emission so
planner/state still transition correctly
- add a unit test for supervisor crash path to ensure an error chunk is
emitted before failed runner status

## Why
`#1586` reports streams that can hang forever when runners crash during
warmup/loading. This keeps failure signaling at the runner-supervisor
layer, matching maintainer guidance in the issue thread.

## Validation
- attempted: `uv run pytest
src/exo/worker/tests/unittests/test_runner/test_runner_supervisor.py`
- blocked locally by environment disk exhaustion while uv tried to
materialize heavy CUDA wheels (`No space left on device` during
`nvidia-cudnn-cu13` extraction)

I kept the change scoped and added a targeted unit test for the failure
path.

---------

Co-authored-by: Evan <evanev7@gmail.com>

.mlx_typings/mlx_lm/models/cache.pyi

@@ -164,8 +164,9 @@ class KVCache(_BaseCache):
     def to_quantized(
         self, group_size: int = ..., bits: int = ...
     ) -> QuantizedKVCache: ...
-    def make_mask(self, *args, **kwargs):  # -> array | Literal['causal'] | None:
-        ...
+    def make_mask(
+        self, *args: Any, **kwargs: Any
+    ) -> mx.array | Literal["causal"] | None: ...
 
 class RotatingKVCache(_BaseCache):
     step = ...
@@ -218,8 +219,7 @@ class ArraysCache(_BaseCache):
         In-place extend this cache with the other cache.
         """
 
-    def make_mask(self, N: int):  # -> array | None:
-        ...
+    def make_mask(self, N: int) -> mx.array | None: ...
 
 class MambaCache(ArraysCache):
     def __init__(self, left_padding: Optional[List[int]] = ...) -> None: ...

.mlx_typings/mlx_lm/models/qwen3_5.pyi

@@ -0,0 +1,153 @@
+from dataclasses import dataclass
+from typing import Any, Optional
+
+import mlx.core as mx
+import mlx.nn as nn
+
+from .cache import ArraysCache, KVCache
+from .qwen3_next import (
+    Qwen3NextAttention as Attention,
+    Qwen3NextMLP as MLP,
+    Qwen3NextRMSNormGated as RMSNormGated,
+    Qwen3NextSparseMoeBlock,
+)
+
+SparseMoeBlock = Qwen3NextSparseMoeBlock
+from .switch_layers import SwitchGLU
+
+@dataclass
+class TextModelArgs:
+    model_type: str
+    hidden_size: int
+    intermediate_size: int
+    num_hidden_layers: int
+    num_attention_heads: int
+    rms_norm_eps: float
+    vocab_size: int
+    num_key_value_heads: int
+    max_position_embeddings: int
+    linear_num_value_heads: int
+    linear_num_key_heads: int
+    linear_key_head_dim: int
+    linear_value_head_dim: int
+    linear_conv_kernel_dim: int
+    tie_word_embeddings: bool
+    attention_bias: bool
+    head_dim: Optional[int]
+    full_attention_interval: int
+    num_experts: int
+    num_experts_per_tok: int
+    decoder_sparse_step: int
+    shared_expert_intermediate_size: int
+    moe_intermediate_size: int
+    norm_topk_prob: bool
+    rope_parameters: Optional[dict[str, Any]]
+    partial_rotary_factor: float
+    rope_theta: float
+    rope_scaling: Optional[dict[str, Any]]
+
+    @classmethod
+    def from_dict(cls, params: dict[str, Any]) -> TextModelArgs: ...
+    def __post_init__(self) -> None: ...
+
+class GatedDeltaNet(nn.Module):
+    hidden_size: int
+    num_v_heads: int
+    num_k_heads: int
+    head_k_dim: int
+    head_v_dim: int
+    key_dim: int
+    value_dim: int
+    conv_kernel_size: int
+    conv_dim: int
+    conv1d: nn.Conv1d
+    in_proj_qkv: nn.Linear
+    in_proj_z: nn.Linear
+    in_proj_b: nn.Linear
+    in_proj_a: nn.Linear
+    dt_bias: mx.array
+    A_log: mx.array
+    norm: RMSNormGated
+    out_proj: nn.Linear
+
+    def __init__(self, config: TextModelArgs) -> None: ...
+    def __call__(
+        self,
+        inputs: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[Any] = None,
+    ) -> mx.array: ...
+
+class DecoderLayer(nn.Module):
+    is_linear: bool
+    linear_attn: GatedDeltaNet
+    self_attn: Attention
+    input_layernorm: nn.RMSNorm
+    post_attention_layernorm: nn.RMSNorm
+    mlp: MLP | SparseMoeBlock
+
+    def __init__(self, args: TextModelArgs, layer_idx: int) -> None: ...
+    def __call__(
+        self,
+        x: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[Any] = None,
+    ) -> mx.array: ...
+
+class Qwen3_5TextModel(nn.Module):
+    embed_tokens: nn.Embedding
+    layers: list[DecoderLayer]
+    norm: nn.RMSNorm
+    ssm_idx: int
+    fa_idx: int
+
+    def __init__(self, args: TextModelArgs) -> None: ...
+    def __call__(
+        self,
+        inputs: mx.array,
+        cache: Optional[Any] = None,
+        input_embeddings: Optional[mx.array] = None,
+    ) -> mx.array: ...
+
+class TextModel(nn.Module):
+    args: TextModelArgs
+    model_type: str
+    model: Qwen3_5TextModel
+    lm_head: nn.Linear
+
+    def __init__(self, args: TextModelArgs) -> None: ...
+    def __call__(
+        self,
+        inputs: mx.array,
+        cache: Optional[Any] = None,
+        input_embeddings: Optional[mx.array] = None,
+    ) -> mx.array: ...
+    @property
+    def layers(self) -> list[DecoderLayer]: ...
+    def make_cache(self) -> list[ArraysCache | KVCache]: ...
+    def sanitize(self, weights: dict[str, Any]) -> dict[str, Any]: ...
+
+@dataclass
+class ModelArgs:
+    model_type: str
+    text_config: dict[str, Any]
+
+    @classmethod
+    def from_dict(cls, params: dict[str, Any]) -> ModelArgs: ...
+
+class Model(nn.Module):
+    args: ModelArgs
+    model_type: str
+    language_model: TextModel
+
+    def __init__(self, args: ModelArgs) -> None: ...
+    def __call__(
+        self,
+        inputs: mx.array,
+        cache: Optional[Any] = None,
+        input_embeddings: Optional[mx.array] = None,
+    ) -> mx.array: ...
+    def sanitize(self, weights: dict[str, Any]) -> dict[str, Any]: ...
+    @property
+    def layers(self) -> list[DecoderLayer]: ...
+    def make_cache(self) -> list[ArraysCache | KVCache]: ...

.mlx_typings/mlx_lm/models/qwen3_5_moe.pyi

@@ -0,0 +1,19 @@
+from dataclasses import dataclass
+from typing import Any, Optional
+
+import mlx.core as mx
+import mlx.nn as nn
+
+from .cache import ArraysCache, KVCache
+from .qwen3_5 import DecoderLayer, Model as Qwen3_5Model, TextModel
+
+@dataclass
+class ModelArgs:
+    model_type: str
+    text_config: dict[str, Any]
+
+    @classmethod
+    def from_dict(cls, params: dict[str, Any]) -> ModelArgs: ...
+
+class Model(Qwen3_5Model):
+    def sanitize(self, weights: dict[str, Any]) -> dict[str, Any]: ...

.mlx_typings/mlx_lm/models/qwen3_next.pyi

@@ -7,6 +7,15 @@ import mlx.nn as nn
 
 from .switch_layers import SwitchGLU
 
+class Qwen3NextRMSNormGated(nn.Module):
+    eps: float
+    weight: mx.array
+
+    def __init__(self, hidden_size: int, eps: float = ...) -> None: ...
+    def __call__(
+        self, hidden_states: mx.array, gate: mx.array | None = None
+    ) -> mx.array: ...
+
 class Qwen3NextMLP(nn.Module):
     gate_proj: nn.Linear
     down_proj: nn.Linear

Cargo.lock

@@ -922,7 +922,6 @@ dependencies = [
  "pyo3-log",
  "pyo3-stub-gen",
  "tokio",
- "util",
 ]
 
 [[package]]
@@ -2793,7 +2792,6 @@ dependencies = [
  "pin-project",
  "tokio",
  "tracing-subscriber",
- "util",
 ]
 
 [[package]]
@@ -4623,10 +4621,6 @@ version = "0.2.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "06abde3611657adf66d383f00b093d7faecc7fa57071cce2578660c9f1010821"
 
-[[package]]
-name = "util"
-version = "0.0.1"
-
 [[package]]
 name = "uuid"
 version = "1.19.0"

Cargo.toml

@@ -1,6 +1,6 @@
 [workspace]
 resolver = "3"
-members = ["rust/networking", "rust/exo_pyo3_bindings", "rust/util"]
+members = ["rust/networking", "rust/exo_pyo3_bindings"]
 
 [workspace.package]
 version = "0.0.1"
@@ -20,7 +20,6 @@ opt-level = 3
 [workspace.dependencies]
 ## Crate members as common dependencies
 networking = { path = "rust/networking" }
-util = { path = "rust/util" }
 
 # Macro dependecies
 extend = "1.2"

pyproject.toml

@@ -19,7 +19,7 @@ dependencies = [
   "anyio==4.11.0",
   "mlx; sys_platform == 'darwin'",
   "mlx[cpu]==0.30.6; sys_platform == 'linux'",
-  "mlx-lm==0.30.7",
+  "mlx-lm",
   "tiktoken>=0.12.0",                          # required for kimi k2 tokenizer
   "hypercorn>=0.18.0",
   "openai-harmony>=0.0.8",
@@ -62,7 +62,7 @@ members = ["rust/exo_pyo3_bindings", "bench"]
 [tool.uv.sources]
 exo_pyo3_bindings = { workspace = true }
 mlx = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git", branch = "address-rdma-gpu-locks", marker = "sys_platform == 'darwin'" }
-#mlx-lm = { git = "https://github.com/davidmcc73/mlx-lm", branch = "stable" }
+mlx-lm = { git = "https://github.com/ml-explore/mlx-lm", rev = "834fac934c4e04de9b3d723e2b9287a2c60cfd4a" }
 # Uncomment to use local mlx/mlx-lm development versions:
 # mlx = { path = "/Users/Shared/mlx", editable=true }
 # mlx-lm = { path = "/Users/Shared/mlx-lm", editable=true }

resources/inference_model_cards/mlx-community--Qwen3.5-122B-A10B-4bit.toml

@@ -0,0 +1,12 @@
+model_id = "mlx-community/Qwen3.5-122B-A10B-4bit"
+n_layers = 48
+hidden_size = 3072
+supports_tensor = true
+tasks = ["TextGeneration"]
+family = "qwen"
+quantization = "4bit"
+base_model = "Qwen3.5 122B A10B"
+capabilities = ["text", "thinking", "thinking_toggle"]
+
+[storage_size]
+in_bytes = 69593314272

resources/inference_model_cards/mlx-community--Qwen3.5-122B-A10B-6bit.toml

@@ -0,0 +1,12 @@
+model_id = "mlx-community/Qwen3.5-122B-A10B-6bit"
+n_layers = 48
+hidden_size = 3072
+supports_tensor = true
+tasks = ["TextGeneration"]
+family = "qwen"
+quantization = "6bit"
+base_model = "Qwen3.5 122B A10B"
+capabilities = ["text", "thinking", "thinking_toggle"]
+
+[storage_size]
+in_bytes = 100120675296

resources/inference_model_cards/mlx-community--Qwen3.5-122B-A10B-8bit.toml

@@ -0,0 +1,12 @@
+model_id = "mlx-community/Qwen3.5-122B-A10B-8bit"
+n_layers = 48
+hidden_size = 3072
+supports_tensor = true
+tasks = ["TextGeneration"]
+family = "qwen"
+quantization = "8bit"
+base_model = "Qwen3.5 122B A10B"
+capabilities = ["text", "thinking", "thinking_toggle"]
+
+[storage_size]
+in_bytes = 130648036320

resources/inference_model_cards/mlx-community--Qwen3.5-122B-A10B-bf16.toml

@@ -0,0 +1,12 @@
+model_id = "mlx-community/Qwen3.5-122B-A10B-bf16"
+n_layers = 48
+hidden_size = 3072
+supports_tensor = true
+tasks = ["TextGeneration"]
+family = "qwen"
+quantization = "bf16"
+base_model = "Qwen3.5 122B A10B"
+capabilities = ["text", "thinking", "thinking_toggle"]
+
+[storage_size]
+in_bytes = 245125640160

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

@@ -0,0 +1,12 @@
+model_id = "mlx-community/Qwen3.5-27B-4bit"
+n_layers = 64
+hidden_size = 5120
+supports_tensor = true
+tasks = ["TextGeneration"]
+family = "qwen"
+quantization = "4bit"
+base_model = "Qwen3.5 27B"
+capabilities = ["text", "thinking"]
+
+[storage_size]
+in_bytes = 16054266848

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

@@ -0,0 +1,12 @@
+model_id = "mlx-community/Qwen3.5-27B-8bit"
+n_layers = 64
+hidden_size = 5120
+supports_tensor = true
+tasks = ["TextGeneration"]
+family = "qwen"
+quantization = "8bit"
+base_model = "Qwen3.5 27B"
+capabilities = ["text", "thinking"]
+
+[storage_size]
+in_bytes = 29500943328

resources/inference_model_cards/mlx-community--Qwen3.5-2B-MLX-8bit.toml

@@ -0,0 +1,12 @@
+model_id = "mlx-community/Qwen3.5-2B-MLX-8bit"
+n_layers = 24
+hidden_size = 2048
+supports_tensor = true
+tasks = ["TextGeneration"]
+family = "qwen"
+quantization = "8bit"
+base_model = "Qwen3.5 2B"
+capabilities = ["text", "thinking"]
+
+[storage_size]
+in_bytes = 2662787264

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

@@ -0,0 +1,12 @@
+model_id = "mlx-community/Qwen3.5-35B-A3B-4bit"
+n_layers = 40
+hidden_size = 2048
+supports_tensor = true
+tasks = ["TextGeneration"]
+family = "qwen"
+quantization = "4bit"
+base_model = "Qwen3.5 35B A3B"
+capabilities = ["text", "thinking", "thinking_toggle"]
+
+[storage_size]
+in_bytes = 20391405152

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

@@ -0,0 +1,12 @@
+model_id = "mlx-community/Qwen3.5-35B-A3B-8bit"
+n_layers = 40
+hidden_size = 2048
+supports_tensor = true
+tasks = ["TextGeneration"]
+family = "qwen"
+quantization = "8bit"
+base_model = "Qwen3.5 35B A3B"
+capabilities = ["text", "thinking", "thinking_toggle"]
+
+[storage_size]
+in_bytes = 37721130592

resources/inference_model_cards/mlx-community--Qwen3.5-397B-A17B-4bit.toml

@@ -0,0 +1,12 @@
+model_id = "mlx-community/Qwen3.5-397B-A17B-4bit"
+n_layers = 60
+hidden_size = 4096
+supports_tensor = true
+tasks = ["TextGeneration"]
+family = "qwen"
+quantization = "4bit"
+base_model = "Qwen3.5 397B A17B"
+capabilities = ["text", "thinking", "thinking_toggle"]
+
+[storage_size]
+in_bytes = 223860768352

resources/inference_model_cards/mlx-community--Qwen3.5-397B-A17B-6bit.toml

@@ -0,0 +1,12 @@
+model_id = "mlx-community/Qwen3.5-397B-A17B-6bit"
+n_layers = 60
+hidden_size = 4096
+supports_tensor = true
+tasks = ["TextGeneration"]
+family = "qwen"
+quantization = "6bit"
+base_model = "Qwen3.5 397B A17B"
+capabilities = ["text", "thinking", "thinking_toggle"]
+
+[storage_size]
+in_bytes = 322946674272

resources/inference_model_cards/mlx-community--Qwen3.5-397B-A17B-8bit.toml

@@ -0,0 +1,12 @@
+model_id = "mlx-community/Qwen3.5-397B-A17B-8bit"
+n_layers = 60
+hidden_size = 4096
+supports_tensor = true
+tasks = ["TextGeneration"]
+family = "qwen"
+quantization = "8bit"
+base_model = "Qwen3.5 397B A17B"
+capabilities = ["text", "thinking", "thinking_toggle"]
+
+[storage_size]
+in_bytes = 422032580192

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

@@ -0,0 +1,12 @@
+model_id = "mlx-community/Qwen3.5-9B-4bit"
+n_layers = 32
+hidden_size = 4096
+supports_tensor = true
+tasks = ["TextGeneration"]
+family = "qwen"
+quantization = "4bit"
+base_model = "Qwen3.5 9B"
+capabilities = ["text", "thinking"]
+
+[storage_size]
+in_bytes = 5950062560

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

@@ -0,0 +1,12 @@
+model_id = "mlx-community/Qwen3.5-9B-8bit"
+n_layers = 32
+hidden_size = 4096
+supports_tensor = true
+tasks = ["TextGeneration"]
+family = "qwen"
+quantization = "8bit"
+base_model = "Qwen3.5 9B"
+capabilities = ["text", "thinking"]
+
+[storage_size]
+in_bytes = 10426433504

rust/exo_pyo3_bindings/Cargo.toml

@@ -54,9 +54,6 @@ delegate = { workspace = true }
 tokio = { workspace = true, features = ["full", "tracing"] }
 futures-lite = { workspace = true }
 
-# utility dependencies
-util = { workspace = true }
-
 # Tracing
 log = { workspace = true }
 env_logger = "0.11"

rust/exo_pyo3_bindings/src/networking.rs

@@ -11,7 +11,7 @@ use crate::networking::exception::{
 use crate::pyclass;
 use futures_lite::{Stream, StreamExt as _};
 use libp2p::gossipsub::PublishError;
-use networking::swarm::{FromSwarm, ToSwarm, create_swarm};
+use networking::{FromSwarm, ToSwarm, create_swarm};
 use pyo3::exceptions::PyRuntimeError;
 use pyo3::prelude::{PyModule, PyModuleMethods as _};
 use pyo3::types::PyBytes;

rust/exo_pyo3_bindings/tests/test_python.py

@@ -1,16 +1,20 @@
 import asyncio
 
 import pytest
-from exo_pyo3_bindings import Keypair, NetworkingHandle, NoPeersSubscribedToTopicError
+from exo_pyo3_bindings import (
+    Keypair,
+    NetworkingHandle,
+    NoPeersSubscribedToTopicError,
+    PyFromSwarm,
+)
 
 
 @pytest.mark.asyncio
 async def test_sleep_on_multiple_items() -> None:
     print("PYTHON: starting handle")
-    h = NetworkingHandle(Keypair.generate_ed25519())
+    h = NetworkingHandle(Keypair.generate())
 
-    ct = asyncio.create_task(_await_cons(h))
-    mt = asyncio.create_task(_await_msg(h))
+    rt = asyncio.create_task(_await_recv(h))
 
     # sleep for 4 ticks
     for i in range(4):
@@ -22,13 +26,11 @@ async def test_sleep_on_multiple_items() -> None:
             print("caught it", e)
 
 
-async def _await_cons(h: NetworkingHandle):
+async def _await_recv(h: NetworkingHandle):
     while True:
-        c = await h.connection_update_recv()
-        print(f"PYTHON: connection update: {c}")
-
-
-async def _await_msg(h: NetworkingHandle):
-    while True:
-        m = await h.gossipsub_recv()
-        print(f"PYTHON: message: {m}")
+        event = await h.recv()
+        match event:
+            case PyFromSwarm.Connection() as c:
+                print(f"PYTHON: connection update: {c}")
+            case PyFromSwarm.Message() as m:
+                print(f"PYTHON: message: {m}")

rust/networking/Cargo.toml

@@ -27,7 +27,6 @@ futures-timer = { workspace = true }
 tokio = { workspace = true, features = ["full"] }
 
 # utility dependencies
-util = { workspace = true }
 tracing-subscriber = { version = "0.3.19", features = [
   "default",
   "env-filter",

rust/networking/examples/chatroom.rs

@@ -1,7 +1,6 @@
-use futures_lite::StreamExt;
+use futures_lite::StreamExt as _;
 use libp2p::identity;
-use networking::swarm;
-use networking::swarm::{FromSwarm, ToSwarm};
+use networking::{FromSwarm, ToSwarm};
 use tokio::sync::{mpsc, oneshot};
 use tokio::{io, io::AsyncBufReadExt as _};
 use tracing_subscriber::EnvFilter;
@@ -9,23 +8,23 @@ use tracing_subscriber::filter::LevelFilter;
 
 #[tokio::main]
 async fn main() {
-    let _ = tracing_subscriber::fmt()
+    tracing_subscriber::fmt()
         .with_env_filter(EnvFilter::from_default_env().add_directive(LevelFilter::INFO.into()))
-        .try_init();
+        .init();
 
     let (to_swarm, from_client) = mpsc::channel(20);
 
     // Configure swarm
-    let mut swarm = swarm::create_swarm(identity::Keypair::generate_ed25519(), from_client)
+    let mut swarm = networking::create_swarm(identity::Keypair::generate_ed25519(), from_client)
         .expect("Swarm creation failed")
         .into_stream();
 
     // Create a Gossipsub topic & subscribe
-    let (tx, rx) = oneshot::channel();
-    _ = to_swarm
+    let (sub_tx, sub_rx) = oneshot::channel();
+    to_swarm
         .send(ToSwarm::Subscribe {
-            topic: "test-net".to_string(),
-            result_sender: tx,
+            topic: "test-net".to_owned(),
+            result_sender: sub_tx,
         })
         .await
         .expect("should send");
@@ -35,15 +34,16 @@ async fn main() {
     println!("Enter messages via STDIN and they will be sent to connected peers using Gossipsub");
 
     tokio::task::spawn(async move {
-        rx.await
+        sub_rx
+            .await
             .expect("tx not dropped")
             .expect("subscribe shouldn't fail");
         loop {
             if let Ok(Some(line)) = stdin.next_line().await {
                 let (tx, rx) = oneshot::channel();
                 if let Err(e) = to_swarm
-                    .send(swarm::ToSwarm::Publish {
-                        topic: "test-net".to_string(),
+                    .send(ToSwarm::Publish {
+                        topic: "test-net".to_owned(),
                         data: line.as_bytes().to_vec(),
                         result_sender: tx,
                     })
@@ -51,7 +51,7 @@ async fn main() {
                 {
                     println!("Send error: {e:?}");
                     return;
-                };
+                }
                 match rx.await {
                     Ok(Err(e)) => println!("Publish error: {e:?}"),
                     Err(e) => println!("Publish error: {e:?}"),

rust/networking/src/discovery.rs

@@ -1,382 +0,0 @@
-use crate::ext::MultiaddrExt;
-use delegate::delegate;
-use either::Either;
-use futures_lite::FutureExt;
-use futures_timer::Delay;
-use libp2p::core::transport::PortUse;
-use libp2p::core::{ConnectedPoint, Endpoint};
-use libp2p::swarm::behaviour::ConnectionEstablished;
-use libp2p::swarm::dial_opts::DialOpts;
-use libp2p::swarm::{
-    CloseConnection, ConnectionClosed, ConnectionDenied, ConnectionHandler,
-    ConnectionHandlerSelect, ConnectionId, FromSwarm, NetworkBehaviour, THandler, THandlerInEvent,
-    THandlerOutEvent, ToSwarm, dummy,
-};
-use libp2p::{Multiaddr, PeerId, identity, mdns};
-use std::collections::{BTreeSet, HashMap};
-use std::convert::Infallible;
-use std::io;
-use std::net::IpAddr;
-use std::task::{Context, Poll};
-use std::time::Duration;
-use util::wakerdeque::WakerDeque;
-
-const RETRY_CONNECT_INTERVAL: Duration = Duration::from_secs(5);
-
-mod managed {
-    use libp2p::swarm::NetworkBehaviour;
-    use libp2p::{identity, mdns, ping};
-    use std::io;
-    use std::time::Duration;
-
-    const MDNS_RECORD_TTL: Duration = Duration::from_secs(2_500);
-    const MDNS_QUERY_INTERVAL: Duration = Duration::from_secs(1_500);
-    const PING_TIMEOUT: Duration = Duration::from_millis(2_500);
-    const PING_INTERVAL: Duration = Duration::from_millis(2_500);
-
-    #[derive(NetworkBehaviour)]
-    pub struct Behaviour {
-        mdns: mdns::tokio::Behaviour,
-        ping: ping::Behaviour,
-    }
-
-    impl Behaviour {
-        pub fn new(keypair: &identity::Keypair) -> io::Result<Self> {
-            Ok(Self {
-                mdns: mdns_behaviour(keypair)?,
-                ping: ping_behaviour(),
-            })
-        }
-    }
-
-    fn mdns_behaviour(keypair: &identity::Keypair) -> io::Result<mdns::tokio::Behaviour> {
-        use mdns::{Config, tokio};
-
-        // mDNS config => enable IPv6
-        let mdns_config = Config {
-            ttl: MDNS_RECORD_TTL,
-            query_interval: MDNS_QUERY_INTERVAL,
-
-            // enable_ipv6: true, // TODO: for some reason, TCP+mDNS don't work well with ipv6?? figure out how to make work
-            ..Default::default()
-        };
-
-        let mdns_behaviour = tokio::Behaviour::new(mdns_config, keypair.public().to_peer_id());
-        Ok(mdns_behaviour?)
-    }
-
-    fn ping_behaviour() -> ping::Behaviour {
-        ping::Behaviour::new(
-            ping::Config::new()
-                .with_timeout(PING_TIMEOUT)
-                .with_interval(PING_INTERVAL),
-        )
-    }
-}
-
-/// Events for when a listening connection is truly established and truly closed.
-#[derive(Debug, Clone)]
-pub enum Event {
-    ConnectionEstablished {
-        peer_id: PeerId,
-        connection_id: ConnectionId,
-        remote_ip: IpAddr,
-        remote_tcp_port: u16,
-    },
-    ConnectionClosed {
-        peer_id: PeerId,
-        connection_id: ConnectionId,
-        remote_ip: IpAddr,
-        remote_tcp_port: u16,
-    },
-}
-
-/// Discovery behavior that wraps mDNS to produce truly discovered durable peer-connections.
-///
-/// The behaviour operates as such:
-///  1) All true (listening) connections/disconnections are tracked, emitting corresponding events
-///     to the swarm.
-///  1) mDNS discovered/expired peers are tracked; discovered but not connected peers are dialed
-///     immediately, and expired but connected peers are disconnected from immediately.
-///  2) Every fixed interval: discovered but not connected peers are dialed, and expired but
-///     connected peers are disconnected from.
-pub struct Behaviour {
-    // state-tracking for managed behaviors & mDNS-discovered peers
-    managed: managed::Behaviour,
-    mdns_discovered: HashMap<PeerId, BTreeSet<Multiaddr>>,
-
-    retry_delay: Delay, // retry interval
-
-    // pending events to emmit => waker-backed Deque to control polling
-    pending_events: WakerDeque<ToSwarm<Event, Infallible>>,
-}
-
-impl Behaviour {
-    pub fn new(keypair: &identity::Keypair) -> io::Result<Self> {
-        Ok(Self {
-            managed: managed::Behaviour::new(keypair)?,
-            mdns_discovered: HashMap::new(),
-            retry_delay: Delay::new(RETRY_CONNECT_INTERVAL),
-            pending_events: WakerDeque::new(),
-        })
-    }
-
-    fn dial(&mut self, peer_id: PeerId, addr: Multiaddr) {
-        self.pending_events.push_back(ToSwarm::Dial {
-            opts: DialOpts::peer_id(peer_id).addresses(vec![addr]).build(),
-        })
-    }
-
-    fn close_connection(&mut self, peer_id: PeerId, connection: ConnectionId) {
-        // push front to make this IMMEDIATE
-        self.pending_events.push_front(ToSwarm::CloseConnection {
-            peer_id,
-            connection: CloseConnection::One(connection),
-        })
-    }
-
-    fn handle_mdns_discovered(&mut self, peers: Vec<(PeerId, Multiaddr)>) {
-        for (p, ma) in peers {
-            self.dial(p, ma.clone()); // always connect
-
-            // get peer's multi-addresses or insert if missing
-            let Some(mas) = self.mdns_discovered.get_mut(&p) else {
-                self.mdns_discovered.insert(p, BTreeSet::from([ma]));
-                continue;
-            };
-
-            // multiaddress should never already be present - else something has gone wrong
-            let is_new_addr = mas.insert(ma);
-            assert!(is_new_addr, "cannot discover a discovered peer");
-        }
-    }
-
-    fn handle_mdns_expired(&mut self, peers: Vec<(PeerId, Multiaddr)>) {
-        for (p, ma) in peers {
-            // at this point, we *must* have the peer
-            let mas = self
-                .mdns_discovered
-                .get_mut(&p)
-                .expect("nonexistent peer cannot expire");
-
-            // at this point, we *must* have the multiaddress
-            let was_present = mas.remove(&ma);
-            assert!(was_present, "nonexistent multiaddress cannot expire");
-
-            // if empty, remove the peer-id entirely
-            if mas.is_empty() {
-                self.mdns_discovered.remove(&p);
-            }
-        }
-    }
-
-    fn on_connection_established(
-        &mut self,
-        peer_id: PeerId,
-        connection_id: ConnectionId,
-        remote_ip: IpAddr,
-        remote_tcp_port: u16,
-    ) {
-        // send out connected event
-        self.pending_events
-            .push_back(ToSwarm::GenerateEvent(Event::ConnectionEstablished {
-                peer_id,
-                connection_id,
-                remote_ip,
-                remote_tcp_port,
-            }));
-    }
-
-    fn on_connection_closed(
-        &mut self,
-        peer_id: PeerId,
-        connection_id: ConnectionId,
-        remote_ip: IpAddr,
-        remote_tcp_port: u16,
-    ) {
-        // send out disconnected event
-        self.pending_events
-            .push_back(ToSwarm::GenerateEvent(Event::ConnectionClosed {
-                peer_id,
-                connection_id,
-                remote_ip,
-                remote_tcp_port,
-            }));
-    }
-}
-
-impl NetworkBehaviour for Behaviour {
-    type ConnectionHandler =
-        ConnectionHandlerSelect<dummy::ConnectionHandler, THandler<managed::Behaviour>>;
-    type ToSwarm = Event;
-
-    // simply delegate to underlying mDNS behaviour
-
-    delegate! {
-        to self.managed {
-            fn handle_pending_inbound_connection(&mut self, connection_id: ConnectionId, local_addr: &Multiaddr, remote_addr: &Multiaddr) -> Result<(), ConnectionDenied>;
-            fn handle_pending_outbound_connection(&mut self, connection_id: ConnectionId, maybe_peer: Option<PeerId>, addresses: &[Multiaddr], effective_role: Endpoint) -> Result<Vec<Multiaddr>, ConnectionDenied>;
-        }
-    }
-
-    fn handle_established_inbound_connection(
-        &mut self,
-        connection_id: ConnectionId,
-        peer: PeerId,
-        local_addr: &Multiaddr,
-        remote_addr: &Multiaddr,
-    ) -> Result<THandler<Self>, ConnectionDenied> {
-        Ok(ConnectionHandler::select(
-            dummy::ConnectionHandler,
-            self.managed.handle_established_inbound_connection(
-                connection_id,
-                peer,
-                local_addr,
-                remote_addr,
-            )?,
-        ))
-    }
-
-    #[allow(clippy::needless_question_mark)]
-    fn handle_established_outbound_connection(
-        &mut self,
-        connection_id: ConnectionId,
-        peer: PeerId,
-        addr: &Multiaddr,
-        role_override: Endpoint,
-        port_use: PortUse,
-    ) -> Result<THandler<Self>, ConnectionDenied> {
-        Ok(ConnectionHandler::select(
-            dummy::ConnectionHandler,
-            self.managed.handle_established_outbound_connection(
-                connection_id,
-                peer,
-                addr,
-                role_override,
-                port_use,
-            )?,
-        ))
-    }
-
-    fn on_connection_handler_event(
-        &mut self,
-        peer_id: PeerId,
-        connection_id: ConnectionId,
-        event: THandlerOutEvent<Self>,
-    ) {
-        match event {
-            Either::Left(ev) => libp2p::core::util::unreachable(ev),
-            Either::Right(ev) => {
-                self.managed
-                    .on_connection_handler_event(peer_id, connection_id, ev)
-            }
-        }
-    }
-
-    // hook into these methods to drive behavior
-
-    fn on_swarm_event(&mut self, event: FromSwarm) {
-        self.managed.on_swarm_event(event); // let mDNS handle swarm events
-
-        // handle swarm events to update internal state:
-        match event {
-            FromSwarm::ConnectionEstablished(ConnectionEstablished {
-                peer_id,
-                connection_id,
-                endpoint,
-                ..
-            }) => {
-                let remote_address = match endpoint {
-                    ConnectedPoint::Dialer { address, .. } => address,
-                    ConnectedPoint::Listener { send_back_addr, .. } => send_back_addr,
-                };
-
-                if let Some((ip, port)) = remote_address.try_to_tcp_addr() {
-                    // handle connection established event which is filtered correctly
-                    self.on_connection_established(peer_id, connection_id, ip, port)
-                }
-            }
-            FromSwarm::ConnectionClosed(ConnectionClosed {
-                peer_id,
-                connection_id,
-                endpoint,
-                ..
-            }) => {
-                let remote_address = match endpoint {
-                    ConnectedPoint::Dialer { address, .. } => address,
-                    ConnectedPoint::Listener { send_back_addr, .. } => send_back_addr,
-                };
-
-                if let Some((ip, port)) = remote_address.try_to_tcp_addr() {
-                    // handle connection closed event which is filtered correctly
-                    self.on_connection_closed(peer_id, connection_id, ip, port)
-                }
-            }
-
-            // since we are running TCP/IP transport layer, we are assuming that
-            // no address changes can occur, hence encountering one is a fatal error
-            FromSwarm::AddressChange(a) => {
-                unreachable!("unhandlable: address change encountered: {:?}", a)
-            }
-            _ => {}
-        }
-    }
-
-    fn poll(&mut self, cx: &mut Context) -> Poll<ToSwarm<Self::ToSwarm, THandlerInEvent<Self>>> {
-        // delegate to managed behaviors for any behaviors they need to perform
-        match self.managed.poll(cx) {
-            Poll::Ready(ToSwarm::GenerateEvent(e)) => {
-                match e {
-                    // handle discovered and expired events from mDNS
-                    managed::BehaviourEvent::Mdns(e) => match e.clone() {
-                        mdns::Event::Discovered(peers) => {
-                            self.handle_mdns_discovered(peers);
-                        }
-                        mdns::Event::Expired(peers) => {
-                            self.handle_mdns_expired(peers);
-                        }
-                    },
-
-                    // handle ping events => if error then disconnect
-                    managed::BehaviourEvent::Ping(e) => {
-                        if let Err(_) = e.result {
-                            self.close_connection(e.peer, e.connection.clone())
-                        }
-                    }
-                }
-
-                // since we just consumed an event, we should immediately wake just in case
-                // there are more events to come where that came from
-                cx.waker().wake_by_ref();
-            }
-
-            // forward any other mDNS event to the swarm or its connection handler(s)
-            Poll::Ready(e) => {
-                return Poll::Ready(
-                    e.map_out(|_| unreachable!("events returning to swarm already handled"))
-                        .map_in(Either::Right),
-                );
-            }
-
-            Poll::Pending => {}
-        }
-
-        // retry connecting to all mDNS peers periodically (fails safely if already connected)
-        if self.retry_delay.poll(cx).is_ready() {
-            for (p, mas) in self.mdns_discovered.clone() {
-                for ma in mas {
-                    self.dial(p, ma)
-                }
-            }
-            self.retry_delay.reset(RETRY_CONNECT_INTERVAL) // reset timeout
-        }
-
-        // send out any pending events from our own service
-        if let Some(e) = self.pending_events.pop_front(cx) {
-            return Poll::Ready(e.map_in(Either::Left));
-        }
-
-        // wait for pending events
-        Poll::Pending
-    }
-}

rust/networking/src/lib.rs

@@ -3,8 +3,6 @@
 //! this is here as a placeholder documentation
 //!
 //!
-pub mod discovery;
-pub mod swarm;
 
 /// Namespace for all the type/trait aliases used by this crate.
 pub(crate) mod alias {
@@ -13,32 +11,302 @@ pub(crate) mod alias {
     pub type AnyError = Box<dyn Error + Send + Sync + 'static>;
     pub type AnyResult<T> = Result<T, AnyError>;
 }
+use std::collections::{HashMap, HashSet};
+use std::pin::Pin;
 
-/// Namespace for crate-wide extension traits/methods
-pub(crate) mod ext {
-    use extend::ext;
-    use libp2p::Multiaddr;
-    use libp2p::multiaddr::Protocol;
-    use std::net::IpAddr;
+pub use behaviour::{Behaviour, BehaviourEvent};
+use futures_lite::{Stream, StreamExt};
+use libp2p::swarm::dial_opts::DialOpts;
+use libp2p::{Multiaddr, mdns};
+use libp2p::{PeerId, SwarmBuilder, gossipsub, identity, swarm::SwarmEvent};
+use tokio::sync::{mpsc, oneshot};
+use transport::tcp_transport;
 
-    #[ext(pub, name = MultiaddrExt)]
-    impl Multiaddr {
-        /// If the multiaddress corresponds to a TCP address, extracts it
-        fn try_to_tcp_addr(&self) -> Option<(IpAddr, u16)> {
-            let mut ps = self.into_iter();
-            let ip = if let Some(p) = ps.next() {
-                match p {
-                    Protocol::Ip4(ip) => IpAddr::V4(ip),
-                    Protocol::Ip6(ip) => IpAddr::V6(ip),
-                    _ => return None,
+/// The current version of the network: this prevents devices running different versions of the
+/// software from interacting with each other.
+///
+/// TODO: right now this is a hardcoded constant; figure out what the versioning semantics should
+///       even be, and how to inject the right version into this config/initialization. E.g. should
+///       this be passed in as a parameter? What about rapidly changing versions in debug builds?
+///       this is all VERY very hard to figure out and needs to be mulled over as a team.
+pub const NETWORK_VERSION: &[u8] = b"v0.0.1";
+pub const OVERRIDE_VERSION_ENV_VAR: &str = "EXO_LIBP2P_NAMESPACE";
+
+// Uses oneshot senders to emulate function calling apis while avoiding requiring unique ownership
+// of the Swarm.
+pub enum ToSwarm {
+    Unsubscribe {
+        topic: String,
+        result_sender: oneshot::Sender<bool>,
+    },
+    Subscribe {
+        topic: String,
+        result_sender: oneshot::Sender<Result<bool, gossipsub::SubscriptionError>>,
+    },
+    Publish {
+        topic: String,
+        data: Vec<u8>,
+        result_sender: oneshot::Sender<Result<gossipsub::MessageId, gossipsub::PublishError>>,
+    },
+}
+pub enum FromSwarm {
+    Message {
+        from: PeerId,
+        topic: String,
+        data: Vec<u8>,
+    },
+    Discovered {
+        peer_id: PeerId,
+    },
+    Expired {
+        peer_id: PeerId,
+    },
+}
+
+pub struct Swarm {
+    swarm: libp2p::Swarm<Behaviour>,
+    from_client: mpsc::Receiver<ToSwarm>,
+}
+
+impl Swarm {
+    #[inline]
+    pub fn into_stream(self) -> Pin<Box<dyn Stream<Item = FromSwarm> + Send>> {
+        let Self {
+            mut swarm,
+            mut from_client,
+        } = self;
+        Box::pin(async_stream::stream! {
+            loop {
+                tokio::select! {
+                    msg = from_client.recv() => {
+                        let Some(msg) = msg else { break };
+                        on_message(&mut swarm, msg);
+                    }
+                    event = swarm.next() => {
+                        let Some(event) = event else { break };
+                        for item in on_swarm_event(&mut swarm, event) {
+                            yield item;
+                        }
+                    }
                 }
-            } else {
-                return None;
-            };
-            let Some(Protocol::Tcp(port)) = ps.next() else {
-                return None;
-            };
-            Some((ip, port))
+            }
+        })
+    }
+}
+
+fn on_swarm_event(
+    swarm: &mut libp2p::Swarm<Behaviour>,
+    event: SwarmEvent<BehaviourEvent>,
+) -> Vec<FromSwarm> {
+    match event {
+        SwarmEvent::Behaviour(BehaviourEvent::Gossipsub(gossipsub::Event::Message {
+            message:
+                gossipsub::Message {
+                    source: Some(peer_id),
+                    topic,
+                    data,
+                    ..
+                },
+            ..
+        })) => vec![FromSwarm::Message {
+            from: peer_id,
+            topic: topic.into_string(),
+            data,
+        }],
+        SwarmEvent::Behaviour(BehaviourEvent::Discovery(mdns::Event::Discovered(peer_ids))) => {
+            let mut base = HashMap::<PeerId, Vec<Multiaddr>>::default();
+            for (pid, addr) in peer_ids {
+                base.entry(pid).or_default().push(addr)
+            }
+            let ret = base
+                .keys()
+                .map(|&peer_id| FromSwarm::Discovered { peer_id })
+                .collect::<Vec<_>>();
+            for (pid, addrs) in base {
+                _ = swarm.dial(DialOpts::peer_id(pid).addresses(addrs).build());
+            }
+            ret
+        }
+        SwarmEvent::Behaviour(BehaviourEvent::Discovery(mdns::Event::Expired(peers))) => peers
+            .into_iter()
+            .map(|(peer_id, _)| peer_id)
+            .collect::<HashSet<PeerId>>()
+            .into_iter()
+            .map(|peer_id| FromSwarm::Expired { peer_id })
+            .collect(),
+        _ => vec![],
+    }
+}
+
+fn on_message(swarm: &mut libp2p::Swarm<Behaviour>, message: ToSwarm) {
+    match message {
+        ToSwarm::Subscribe {
+            topic,
+            result_sender,
+        } => {
+            let result = swarm
+                .behaviour_mut()
+                .gossipsub
+                .subscribe(&gossipsub::IdentTopic::new(topic));
+            _ = result_sender.send(result);
+        }
+        ToSwarm::Unsubscribe {
+            topic,
+            result_sender,
+        } => {
+            let result = swarm
+                .behaviour_mut()
+                .gossipsub
+                .unsubscribe(&gossipsub::IdentTopic::new(topic));
+            _ = result_sender.send(result);
+        }
+        ToSwarm::Publish {
+            topic,
+            data,
+            result_sender,
+        } => {
+            let result = swarm
+                .behaviour_mut()
+                .gossipsub
+                .publish(gossipsub::IdentTopic::new(topic), data);
+            _ = result_sender.send(result);
         }
     }
 }
+
+/// Create and configure a swarm which listens to all ports on OS
+pub fn create_swarm(
+    keypair: identity::Keypair,
+    from_client: mpsc::Receiver<ToSwarm>,
+) -> alias::AnyResult<Swarm> {
+    let mut swarm = SwarmBuilder::with_existing_identity(keypair)
+        .with_tokio()
+        .with_other_transport(tcp_transport)?
+        .with_behaviour(Behaviour::new)?
+        .build();
+
+    // Listen on all interfaces and whatever port the OS assigns
+    swarm.listen_on("/ip4/0.0.0.0/tcp/0".parse()?)?;
+    Ok(Swarm { swarm, from_client })
+}
+
+mod transport {
+    use crate::alias;
+    use crate::{NETWORK_VERSION, OVERRIDE_VERSION_ENV_VAR};
+    use futures_lite::{AsyncRead, AsyncWrite};
+    use keccak_const::Sha3_256;
+    use libp2p::core::muxing;
+    use libp2p::core::transport::Boxed;
+    use libp2p::pnet::{PnetError, PnetOutput};
+    use libp2p::{PeerId, Transport, identity, noise, pnet, yamux};
+    use std::{env, sync::LazyLock};
+
+    /// Key used for networking's private network; parametrized on the [`NETWORK_VERSION`].
+    /// See [`pnet_upgrade`] for more.
+    static PNET_PRESHARED_KEY: LazyLock<[u8; 32]> = LazyLock::new(|| {
+        let builder = Sha3_256::new().update(b"exo_discovery_network");
+
+        if let Ok(var) = env::var(OVERRIDE_VERSION_ENV_VAR) {
+            let bytes = var.into_bytes();
+            builder.update(&bytes)
+        } else {
+            builder.update(NETWORK_VERSION)
+        }
+        .finalize()
+    });
+
+    /// Make the Swarm run on a private network, as to not clash with public libp2p nodes and
+    /// also different-versioned instances of this same network.
+    /// This is implemented as an additional "upgrade" ontop of existing [`libp2p::Transport`] layers.
+    async fn pnet_upgrade<TSocket>(
+        socket: TSocket,
+        _: impl Sized,
+    ) -> Result<PnetOutput<TSocket>, PnetError>
+    where
+        TSocket: AsyncRead + AsyncWrite + Send + Unpin + 'static,
+    {
+        use pnet::{PnetConfig, PreSharedKey};
+        PnetConfig::new(PreSharedKey::new(*PNET_PRESHARED_KEY))
+            .handshake(socket)
+            .await
+    }
+
+    /// TCP/IP transport layer configuration.
+    pub fn tcp_transport(
+        keypair: &identity::Keypair,
+    ) -> alias::AnyResult<Boxed<(PeerId, muxing::StreamMuxerBox)>> {
+        use libp2p::{
+            core::upgrade::Version,
+            tcp::{Config, tokio},
+        };
+
+        // `TCP_NODELAY` enabled => avoid latency
+        let tcp_config = Config::default().nodelay(true);
+
+        // V1 + lazy flushing => 0-RTT negotiation
+        let upgrade_version = Version::V1Lazy;
+
+        // Noise is faster than TLS + we don't care much for security
+        let noise_config = noise::Config::new(keypair)?;
+
+        // Use default Yamux config for multiplexing
+        let yamux_config = yamux::Config::default();
+
+        // Create new Tokio-driven TCP/IP transport layer
+        let base_transport = tokio::Transport::new(tcp_config)
+            .and_then(pnet_upgrade)
+            .upgrade(upgrade_version)
+            .authenticate(noise_config)
+            .multiplex(yamux_config);
+
+        // Return boxed transport (to flatten complex type)
+        Ok(base_transport.boxed())
+    }
+}
+
+mod behaviour {
+    use crate::alias;
+    use libp2p::swarm::NetworkBehaviour;
+    use libp2p::{gossipsub, identity, mdns};
+
+    /// Behavior of the Swarm which composes all desired behaviors:
+    /// Right now its just [`mdns::tokio::Behaviour`] and [`gossipsub::Behaviour`].
+    #[derive(NetworkBehaviour)]
+    pub struct Behaviour {
+        pub discovery: mdns::tokio::Behaviour,
+        pub gossipsub: gossipsub::Behaviour,
+    }
+
+    impl Behaviour {
+        pub fn new(keypair: &identity::Keypair) -> alias::AnyResult<Self> {
+            Ok(Self {
+                discovery: mdns_behaviour(keypair)?,
+                gossipsub: gossipsub_behaviour(keypair),
+            })
+        }
+    }
+
+    fn mdns_behaviour(keypair: &identity::Keypair) -> alias::AnyResult<mdns::tokio::Behaviour> {
+        Ok(mdns::tokio::Behaviour::new(
+            mdns::Config::default(),
+            keypair.public().to_peer_id(),
+        )?)
+    }
+
+    fn gossipsub_behaviour(keypair: &identity::Keypair) -> gossipsub::Behaviour {
+        use gossipsub::{ConfigBuilder, MessageAuthenticity, ValidationMode};
+
+        // build a gossipsub network behaviour
+        //  => signed message authenticity + strict validation mode means the message-ID is
+        //     automatically provided by gossipsub w/out needing to provide custom message-ID function
+        gossipsub::Behaviour::new(
+            MessageAuthenticity::Signed(keypair.clone()),
+            ConfigBuilder::default()
+                .max_transmit_size(1024 * 1024)
+                .validation_mode(ValidationMode::Strict)
+                .build()
+                .expect("the configuration should always be valid"),
+        )
+        .expect("creating gossipsub behavior should always work")
+    }
+}

rust/networking/src/swarm.rs

@@ -1,273 +0,0 @@
-use std::pin::Pin;
-
-use crate::swarm::transport::tcp_transport;
-use crate::{alias, discovery};
-pub use behaviour::{Behaviour, BehaviourEvent};
-use futures_lite::{Stream, StreamExt};
-use libp2p::{PeerId, SwarmBuilder, gossipsub, identity, swarm::SwarmEvent};
-use tokio::sync::{mpsc, oneshot};
-
-/// The current version of the network: this prevents devices running different versions of the
-/// software from interacting with each other.
-///
-/// TODO: right now this is a hardcoded constant; figure out what the versioning semantics should
-///       even be, and how to inject the right version into this config/initialization. E.g. should
-///       this be passed in as a parameter? What about rapidly changing versions in debug builds?
-///       this is all VERY very hard to figure out and needs to be mulled over as a team.
-pub const NETWORK_VERSION: &[u8] = b"v0.0.1";
-pub const OVERRIDE_VERSION_ENV_VAR: &str = "EXO_LIBP2P_NAMESPACE";
-
-// Uses oneshot senders to emulate function calling apis while avoiding requiring unique ownership
-// of the Swarm.
-pub enum ToSwarm {
-    Unsubscribe {
-        topic: String,
-        result_sender: oneshot::Sender<bool>,
-    },
-    Subscribe {
-        topic: String,
-        result_sender: oneshot::Sender<Result<bool, gossipsub::SubscriptionError>>,
-    },
-    Publish {
-        topic: String,
-        data: Vec<u8>,
-        result_sender: oneshot::Sender<Result<gossipsub::MessageId, gossipsub::PublishError>>,
-    },
-}
-pub enum FromSwarm {
-    Message {
-        from: PeerId,
-        topic: String,
-        data: Vec<u8>,
-    },
-    Discovered {
-        peer_id: PeerId,
-    },
-    Expired {
-        peer_id: PeerId,
-    },
-}
-
-pub struct Swarm {
-    swarm: libp2p::Swarm<Behaviour>,
-    from_client: mpsc::Receiver<ToSwarm>,
-}
-
-impl Swarm {
-    pub fn into_stream(self) -> Pin<Box<dyn Stream<Item = FromSwarm> + Send>> {
-        let Swarm {
-            mut swarm,
-            mut from_client,
-        } = self;
-        let stream = async_stream::stream! {
-            loop {
-                tokio::select! {
-                    msg = from_client.recv() => {
-                        let Some(msg) = msg else { break };
-                        on_message(&mut swarm, msg);
-                    }
-                    event = swarm.next() => {
-                        let Some(event) = event else { break };
-                        if let Some(item) = filter_swarm_event(event) {
-                            yield item;
-                        }
-                    }
-                }
-            }
-        };
-        Box::pin(stream)
-    }
-}
-
-fn on_message(swarm: &mut libp2p::Swarm<Behaviour>, message: ToSwarm) {
-    match message {
-        ToSwarm::Subscribe {
-            topic,
-            result_sender,
-        } => {
-            let result = swarm
-                .behaviour_mut()
-                .gossipsub
-                .subscribe(&gossipsub::IdentTopic::new(topic));
-            _ = result_sender.send(result);
-        }
-        ToSwarm::Unsubscribe {
-            topic,
-            result_sender,
-        } => {
-            let result = swarm
-                .behaviour_mut()
-                .gossipsub
-                .unsubscribe(&gossipsub::IdentTopic::new(topic));
-            _ = result_sender.send(result);
-        }
-        ToSwarm::Publish {
-            topic,
-            data,
-            result_sender,
-        } => {
-            let result = swarm
-                .behaviour_mut()
-                .gossipsub
-                .publish(gossipsub::IdentTopic::new(topic), data);
-            _ = result_sender.send(result);
-        }
-    }
-}
-
-fn filter_swarm_event(event: SwarmEvent<BehaviourEvent>) -> Option<FromSwarm> {
-    match event {
-        SwarmEvent::Behaviour(BehaviourEvent::Gossipsub(gossipsub::Event::Message {
-            message:
-                gossipsub::Message {
-                    source: Some(peer_id),
-                    topic,
-                    data,
-                    ..
-                },
-            ..
-        })) => Some(FromSwarm::Message {
-            from: peer_id,
-            topic: topic.into_string(),
-            data,
-        }),
-        SwarmEvent::Behaviour(BehaviourEvent::Discovery(
-            discovery::Event::ConnectionEstablished { peer_id, .. },
-        )) => Some(FromSwarm::Discovered { peer_id }),
-        SwarmEvent::Behaviour(BehaviourEvent::Discovery(discovery::Event::ConnectionClosed {
-            peer_id,
-            ..
-        })) => Some(FromSwarm::Expired { peer_id }),
-        _ => None,
-    }
-}
-
-/// Create and configure a swarm which listens to all ports on OS
-pub fn create_swarm(
-    keypair: identity::Keypair,
-    from_client: mpsc::Receiver<ToSwarm>,
-) -> alias::AnyResult<Swarm> {
-    let mut swarm = SwarmBuilder::with_existing_identity(keypair)
-        .with_tokio()
-        .with_other_transport(tcp_transport)?
-        .with_behaviour(Behaviour::new)?
-        .build();
-
-    // Listen on all interfaces and whatever port the OS assigns
-    swarm.listen_on("/ip4/0.0.0.0/tcp/0".parse()?)?;
-    Ok(Swarm { swarm, from_client })
-}
-
-mod transport {
-    use crate::alias;
-    use crate::swarm::{NETWORK_VERSION, OVERRIDE_VERSION_ENV_VAR};
-    use futures_lite::{AsyncRead, AsyncWrite};
-    use keccak_const::Sha3_256;
-    use libp2p::core::muxing;
-    use libp2p::core::transport::Boxed;
-    use libp2p::pnet::{PnetError, PnetOutput};
-    use libp2p::{PeerId, Transport, identity, noise, pnet, yamux};
-    use std::{env, sync::LazyLock};
-
-    /// Key used for networking's private network; parametrized on the [`NETWORK_VERSION`].
-    /// See [`pnet_upgrade`] for more.
-    static PNET_PRESHARED_KEY: LazyLock<[u8; 32]> = LazyLock::new(|| {
-        let builder = Sha3_256::new().update(b"exo_discovery_network");
-
-        if let Ok(var) = env::var(OVERRIDE_VERSION_ENV_VAR) {
-            let bytes = var.into_bytes();
-            builder.update(&bytes)
-        } else {
-            builder.update(NETWORK_VERSION)
-        }
-        .finalize()
-    });
-
-    /// Make the Swarm run on a private network, as to not clash with public libp2p nodes and
-    /// also different-versioned instances of this same network.
-    /// This is implemented as an additional "upgrade" ontop of existing [`libp2p::Transport`] layers.
-    async fn pnet_upgrade<TSocket>(
-        socket: TSocket,
-        _: impl Sized,
-    ) -> Result<PnetOutput<TSocket>, PnetError>
-    where
-        TSocket: AsyncRead + AsyncWrite + Send + Unpin + 'static,
-    {
-        use pnet::{PnetConfig, PreSharedKey};
-        PnetConfig::new(PreSharedKey::new(*PNET_PRESHARED_KEY))
-            .handshake(socket)
-            .await
-    }
-
-    /// TCP/IP transport layer configuration.
-    pub fn tcp_transport(
-        keypair: &identity::Keypair,
-    ) -> alias::AnyResult<Boxed<(PeerId, muxing::StreamMuxerBox)>> {
-        use libp2p::{
-            core::upgrade::Version,
-            tcp::{Config, tokio},
-        };
-
-        // `TCP_NODELAY` enabled => avoid latency
-        let tcp_config = Config::default().nodelay(true);
-
-        // V1 + lazy flushing => 0-RTT negotiation
-        let upgrade_version = Version::V1Lazy;
-
-        // Noise is faster than TLS + we don't care much for security
-        let noise_config = noise::Config::new(keypair)?;
-
-        // Use default Yamux config for multiplexing
-        let yamux_config = yamux::Config::default();
-
-        // Create new Tokio-driven TCP/IP transport layer
-        let base_transport = tokio::Transport::new(tcp_config)
-            .and_then(pnet_upgrade)
-            .upgrade(upgrade_version)
-            .authenticate(noise_config)
-            .multiplex(yamux_config);
-
-        // Return boxed transport (to flatten complex type)
-        Ok(base_transport.boxed())
-    }
-}
-
-mod behaviour {
-    use crate::{alias, discovery};
-    use libp2p::swarm::NetworkBehaviour;
-    use libp2p::{gossipsub, identity};
-
-    /// Behavior of the Swarm which composes all desired behaviors:
-    /// Right now its just [`discovery::Behaviour`] and [`gossipsub::Behaviour`].
-    #[derive(NetworkBehaviour)]
-    pub struct Behaviour {
-        pub discovery: discovery::Behaviour,
-        pub gossipsub: gossipsub::Behaviour,
-    }
-
-    impl Behaviour {
-        pub fn new(keypair: &identity::Keypair) -> alias::AnyResult<Self> {
-            Ok(Self {
-                discovery: discovery::Behaviour::new(keypair)?,
-                gossipsub: gossipsub_behaviour(keypair),
-            })
-        }
-    }
-
-    fn gossipsub_behaviour(keypair: &identity::Keypair) -> gossipsub::Behaviour {
-        use gossipsub::{ConfigBuilder, MessageAuthenticity, ValidationMode};
-
-        // build a gossipsub network behaviour
-        //  => signed message authenticity + strict validation mode means the message-ID is
-        //     automatically provided by gossipsub w/out needing to provide custom message-ID function
-        gossipsub::Behaviour::new(
-            MessageAuthenticity::Signed(keypair.clone()),
-            ConfigBuilder::default()
-                .max_transmit_size(1024 * 1024)
-                .validation_mode(ValidationMode::Strict)
-                .build()
-                .expect("the configuration should always be valid"),
-        )
-        .expect("creating gossipsub behavior should always work")
-    }
-}

rust/util/Cargo.toml

@@ -1,15 +0,0 @@
-[package]
-name = "util"
-version = { workspace = true }
-edition = { workspace = true }
-publish = false
-
-[lib]
-doctest = false
-name = "util"
-path = "src/lib.rs"
-
-[lints]
-workspace = true
-
-[dependencies]

rust/util/src/lib.rs

@@ -1 +0,0 @@
-pub mod wakerdeque;

rust/util/src/wakerdeque.rs

@@ -1,55 +0,0 @@
-use std::collections::VecDeque;
-use std::fmt::{Debug, Formatter};
-use std::task::{Context, Waker};
-
-/// A wrapper around [`VecDeque`] which wakes (if it can) on any `push_*` methods,
-/// and updates the internally stored waker by consuming [`Context`] on any `pop_*` methods.
-pub struct WakerDeque<T> {
-    waker: Option<Waker>,
-    deque: VecDeque<T>,
-}
-
-impl<T: Debug> Debug for WakerDeque<T> {
-    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
-        self.deque.fmt(f)
-    }
-}
-
-impl<T> WakerDeque<T> {
-    pub fn new() -> Self {
-        Self {
-            waker: None,
-            deque: VecDeque::new(),
-        }
-    }
-
-    fn update(&mut self, cx: &mut Context<'_>) {
-        self.waker = Some(cx.waker().clone());
-    }
-
-    fn wake(&mut self) {
-        let Some(ref mut w) = self.waker else { return };
-        w.wake_by_ref();
-        self.waker = None;
-    }
-
-    pub fn pop_front(&mut self, cx: &mut Context<'_>) -> Option<T> {
-        self.update(cx);
-        self.deque.pop_front()
-    }
-
-    pub fn pop_back(&mut self, cx: &mut Context<'_>) -> Option<T> {
-        self.update(cx);
-        self.deque.pop_back()
-    }
-
-    pub fn push_front(&mut self, value: T) {
-        self.wake();
-        self.deque.push_front(value);
-    }
-
-    pub fn push_back(&mut self, value: T) {
-        self.wake();
-        self.deque.push_back(value);
-    }
-}

src/exo/master/adapters/chat_completions.py

@@ -26,7 +26,11 @@ from exo.shared.types.chunks import (
     ToolCallChunk,
 )
 from exo.shared.types.common import CommandId
-from exo.shared.types.text_generation import InputMessage, TextGenerationTaskParams
+from exo.shared.types.text_generation import (
+    InputMessage,
+    TextGenerationTaskParams,
+    resolve_reasoning_params,
+)
 
 
 def chat_request_to_text_generation(
@@ -75,6 +79,10 @@ def chat_request_to_text_generation(
             dumped: dict[str, Any] = msg_copy.model_dump(exclude_none=True)
             chat_template_messages.append(dumped)
 
+    resolved_effort, resolved_thinking = resolve_reasoning_params(
+        request.reasoning_effort, request.enable_thinking
+    )
+
     return TextGenerationTaskParams(
         model=request.model,
         input=input_messages
@@ -89,7 +97,8 @@ def chat_request_to_text_generation(
         seed=request.seed,
         stream=request.stream,
         tools=request.tools,
-        enable_thinking=request.enable_thinking,
+        reasoning_effort=resolved_effort,
+        enable_thinking=resolved_thinking,
         chat_template_messages=chat_template_messages
         if chat_template_messages
         else None,

src/exo/master/adapters/responses.py

@@ -42,7 +42,11 @@ from exo.shared.types.openai_responses import (
     ResponseTextDoneEvent,
     ResponseUsage,
 )
-from exo.shared.types.text_generation import InputMessage, TextGenerationTaskParams
+from exo.shared.types.text_generation import (
+    InputMessage,
+    TextGenerationTaskParams,
+    resolve_reasoning_params,
+)
 
 
 def _format_sse(event: ResponsesStreamEvent) -> str:
@@ -119,6 +123,11 @@ def responses_request_to_text_generation(
         )
         built_chat_template = chat_template_messages if chat_template_messages else None
 
+    effort_from_reasoning = request.reasoning.effort if request.reasoning else None
+    resolved_effort, resolved_thinking = resolve_reasoning_params(
+        effort_from_reasoning, request.enable_thinking
+    )
+
     return TextGenerationTaskParams(
         model=request.model,
         input=input_value,
@@ -132,6 +141,8 @@ def responses_request_to_text_generation(
         stop=request.stop,
         seed=request.seed,
         chat_template_messages=built_chat_template or request.chat_template_messages,
+        reasoning_effort=resolved_effort,
+        enable_thinking=resolved_thinking,
     )
 
 

src/exo/master/main.py

@@ -328,17 +328,22 @@ class Master:
                                         task_id=task_id,
                                     )
                                 )
-                        case TaskFinished():
-                            generated_events.append(
-                                TaskDeleted(
-                                    task_id=self.command_task_mapping[
-                                        command.finished_command_id
-                                    ]
+                            else:
+                                logger.warning(
+                                    f"Nonexistent command {command.cancelled_command_id} cancelled"
                                 )
-                            )
-                            self.command_task_mapping.pop(
-                                command.finished_command_id, None
-                            )
+                        case TaskFinished():
+                            if (
+                                task_id := self.command_task_mapping.pop(
+                                    command.finished_command_id, None
+                                )
+                            ) is not None:
+                                generated_events.append(TaskDeleted(task_id=task_id))
+                            else:
+                                logger.warning(
+                                    f"Finished command {command.finished_command_id} finished"
+                                )
+
                         case RequestEventLog():
                             # We should just be able to send everything, since other buffers will ignore old messages
                             # rate limit to 1000 at a time

src/exo/routing/router.py

@@ -258,6 +258,6 @@ def get_node_id_keypair(
 
         # if no valid credentials, create new ones and persist
         with open(path, "w+b") as f:
-            keypair = Keypair.generate_ed25519()
+            keypair = Keypair.generate()
             f.write(keypair.to_bytes())
             return keypair

src/exo/shared/models/model_cards.py

@@ -190,6 +190,8 @@ class ConfigData(BaseModel):
             ["DeepseekV3ForCausalLM"],
             ["Qwen3NextForCausalLM"],
             ["Qwen3MoeForCausalLM"],
+            ["Qwen3_5MoeForConditionalGeneration"],
+            ["Qwen3_5ForConditionalGeneration"],
             ["MiniMaxM2ForCausalLM"],
             ["LlamaForCausalLM"],
             ["GptOssForCausalLM"],

src/exo/shared/tests/test_resolve_reasoning_params.py

@@ -0,0 +1,30 @@
+import pytest
+
+from exo.shared.types.text_generation import resolve_reasoning_params
+
+
+def test_both_none_returns_none_none() -> None:
+    assert resolve_reasoning_params(None, None) == (None, None)
+
+
+def test_both_set_passes_through_unchanged() -> None:
+    assert resolve_reasoning_params("high", True) == ("high", True)
+    assert resolve_reasoning_params("none", True) == ("none", True)
+    assert resolve_reasoning_params("low", False) == ("low", False)
+
+
+def test_enable_thinking_true_derives_medium() -> None:
+    assert resolve_reasoning_params(None, True) == ("medium", True)
+
+
+def test_enable_thinking_false_derives_none() -> None:
+    assert resolve_reasoning_params(None, False) == ("none", False)
+
+
+def test_reasoning_effort_none_derives_thinking_false() -> None:
+    assert resolve_reasoning_params("none", None) == ("none", False)
+
+
+@pytest.mark.parametrize("effort", ["minimal", "low", "medium", "high", "xhigh"])
+def test_non_none_effort_derives_thinking_true(effort: str) -> None:
+    assert resolve_reasoning_params(effort, None) == (effort, True)  # pyright: ignore[reportArgumentType]

src/exo/shared/types/api.py

@@ -8,6 +8,7 @@ from pydantic import BaseModel, Field, field_validator
 from exo.shared.models.model_cards import ModelCard, ModelId
 from exo.shared.types.common import CommandId, NodeId
 from exo.shared.types.memory import Memory
+from exo.shared.types.text_generation import ReasoningEffort
 from exo.shared.types.worker.instances import Instance, InstanceId, InstanceMeta
 from exo.shared.types.worker.shards import Sharding, ShardMetadata
 from exo.utils.pydantic_ext import CamelCaseModel
@@ -198,6 +199,7 @@ class ChatCompletionRequest(BaseModel):
     top_p: float | None = None
     top_k: int | None = None
     tools: list[dict[str, Any]] | None = None
+    reasoning_effort: ReasoningEffort | None = None
     enable_thinking: bool | None = None
     tool_choice: str | dict[str, Any] | None = None
     parallel_tool_calls: bool | None = None

src/exo/shared/types/openai_responses.py

@@ -12,6 +12,7 @@ from typing import Any, Literal
 from pydantic import BaseModel, Field
 
 from exo.shared.types.common import ModelId
+from exo.shared.types.text_generation import ReasoningEffort
 
 # Type aliases
 ResponseStatus = Literal["completed", "failed", "in_progress", "incomplete"]
@@ -71,6 +72,13 @@ ResponseInputItem = (
 )
 
 
+class Reasoning(BaseModel, frozen=True):
+    """Reasoning configuration for OpenAI Responses API."""
+
+    effort: ReasoningEffort | None = None
+    summary: Literal["auto", "concise", "detailed"] | None = None
+
+
 class ResponsesRequest(BaseModel, frozen=True):
     """Request body for OpenAI Responses API.
 
@@ -89,8 +97,15 @@ class ResponsesRequest(BaseModel, frozen=True):
     stream: bool = False
     tools: list[dict[str, Any]] | None = None
     metadata: dict[str, str] | None = None
+    reasoning: Reasoning | None = None
 
     # --- exo extensions (not in OpenAI Responses API spec) ---
+    enable_thinking: bool | None = Field(
+        default=None,
+        description="[exo extension] Boolean thinking toggle. Not part of the OpenAI Responses API.",
+        json_schema_extra={"x-exo-extension": True},
+    )
+
     top_k: int | None = Field(
         default=None,
         description="[exo extension] Top-k sampling parameter. Not part of the OpenAI Responses API.",

src/exo/shared/types/text_generation.py

@@ -11,6 +11,29 @@ from pydantic import BaseModel
 from exo.shared.types.common import ModelId
 
 MessageRole = Literal["user", "assistant", "system", "developer"]
+ReasoningEffort = Literal["none", "minimal", "low", "medium", "high", "xhigh"]
+
+
+def resolve_reasoning_params(
+    reasoning_effort: ReasoningEffort | None,
+    enable_thinking: bool | None,
+) -> tuple[ReasoningEffort | None, bool | None]:
+    """
+    enable_thinking=True  -> reasoning_effort="medium"
+    enable_thinking=False -> reasoning_effort="none"
+    reasoning_effort="none" -> enable_thinking=False
+    reasoning_effort=<anything else> -> enable_thinking=True
+    """
+    resolved_effort: ReasoningEffort | None = reasoning_effort
+    resolved_thinking: bool | None = enable_thinking
+
+    if reasoning_effort is None and enable_thinking is not None:
+        resolved_effort = "medium" if enable_thinking else "none"
+
+    if enable_thinking is None and reasoning_effort is not None:
+        resolved_thinking = reasoning_effort != "none"
+
+    return resolved_effort, resolved_thinking
 
 
 class InputMessage(BaseModel, frozen=True):
@@ -40,6 +63,7 @@ class TextGenerationTaskParams(BaseModel, frozen=True):
     stop: str | list[str] | None = None
     seed: int | None = None
     chat_template_messages: list[dict[str, Any]] | None = None
+    reasoning_effort: ReasoningEffort | None = None
     enable_thinking: bool | None = None
     logprobs: bool = False
     top_logprobs: int | None = None

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

@@ -16,6 +16,7 @@ from mlx.nn.layers.distributed import (
 from mlx_lm.models.base import (
     scaled_dot_product_attention,  # pyright: ignore[reportUnknownVariableType]
 )
+from mlx_lm.models.cache import ArraysCache, KVCache
 from mlx_lm.models.deepseek_v3 import DeepseekV3MLP
 from mlx_lm.models.deepseek_v3 import Model as DeepseekV3Model
 from mlx_lm.models.deepseek_v32 import DeepseekV32MLP
@@ -31,10 +32,19 @@ from mlx_lm.models.llama import Model as LlamaModel
 from mlx_lm.models.minimax import MiniMaxAttention
 from mlx_lm.models.minimax import Model as MiniMaxModel
 from mlx_lm.models.ministral3 import Model as Ministral3Model
+from mlx_lm.models.qwen3_5 import DecoderLayer as Qwen3_5DecoderLayer
+from mlx_lm.models.qwen3_5 import Model as Qwen3_5TextModel
+from mlx_lm.models.qwen3_5 import Qwen3_5TextModel as Qwen3_5TextModelInner
+from mlx_lm.models.qwen3_5 import SparseMoeBlock as Qwen3_5SparseMoeBlock
+from mlx_lm.models.qwen3_5_moe import Model as Qwen3_5MoeModel
 from mlx_lm.models.qwen3_moe import Model as Qwen3MoeModel
 from mlx_lm.models.qwen3_moe import Qwen3MoeDecoderLayer, Qwen3MoeSparseMoeBlock
 from mlx_lm.models.qwen3_next import Model as Qwen3NextModel
-from mlx_lm.models.qwen3_next import Qwen3NextDecoderLayer, Qwen3NextSparseMoeBlock
+from mlx_lm.models.qwen3_next import (
+    Qwen3NextDecoderLayer,
+    Qwen3NextGatedDeltaNet,
+    Qwen3NextSparseMoeBlock,
+)
 from mlx_lm.models.step3p5 import Model as Step35Model
 from mlx_lm.models.step3p5 import Step3p5MLP as Step35MLP
 from mlx_lm.models.step3p5 import Step3p5Model as Step35InnerModel
@@ -191,9 +201,10 @@ class PipelineLastLayer(CustomMlxLayer):
                 # CacheList (used by MLA models like DeepSeekV32, GLM MoE DSA)
                 # doesn't have .keys directly; access via first sub-cache.
                 _cache = cache[0] if hasattr(cache, "caches") else cache  # type: ignore
-                _cache.keys = mx.depends(_cache.keys, output)  # type: ignore
+                if hasattr(_cache, "keys"):  # pyright: ignore[reportAny]
+                    _cache.keys = mx.depends(_cache.keys, output)  # type: ignore
             mx.eval(output)
-            if cache is not None:
+            if cache is not None and hasattr(_cache, "keys"):  # type: ignore
                 mx.eval(_cache.keys)  # type: ignore
 
         if not self.is_prefill:
@@ -248,6 +259,32 @@ def get_layers(inner_model_instance: nn.Module) -> list[_LayerCallable]:
     return layers
 
 
+def _patch_qwen35_cache(
+    model: Qwen3_5TextModel,
+    fa_idx: int,
+    has_full_attn: bool,
+    ssm_idx: int,
+    has_linear: bool,
+) -> None:
+    # Hacks to make make_mask happy.
+    original = model.make_cache
+
+    def patched() -> list[ArraysCache | KVCache]:
+        cache: list[ArraysCache | KVCache] = original()
+        if not has_full_attn:
+            entry = cache[fa_idx]
+            orig_make_mask = entry.make_mask
+            entry.make_mask = lambda n, **_kw: orig_make_mask(n)  # type: ignore
+        if not has_linear:
+            orig_ssm_make_mask = cache[ssm_idx].make_mask
+            cache[ssm_idx].make_mask = (  # type: ignore
+                lambda n, **kw: orig_ssm_make_mask(n, **kw) if kw else None  # type: ignore
+            )
+        return cache
+
+    model.make_cache = patched
+
+
 def pipeline_auto_parallel(
     model: nn.Module,
     group: mx.distributed.Group,
@@ -318,6 +355,24 @@ def pipeline_auto_parallel(
         inner_model_instance._swa_idx = 0 if not sliding_layers else sliding_layers[0]
         inner_model_instance._full_idx = 0 if not full_layers else full_layers[0]
 
+    if isinstance(inner_model_instance, Qwen3_5TextModelInner):
+        full_attn_layers = [
+            i for i, layer in enumerate(layers) if not getattr(layer, "is_linear", True)
+        ]
+        linear_layers = [
+            i for i, layer in enumerate(layers) if getattr(layer, "is_linear", False)
+        ]
+        inner_model_instance.fa_idx = full_attn_layers[0] if full_attn_layers else 0
+        inner_model_instance.ssm_idx = linear_layers[0] if linear_layers else 0
+        if not full_attn_layers or not linear_layers:
+            _patch_qwen35_cache(
+                cast(Qwen3_5TextModel, model),
+                fa_idx=inner_model_instance.fa_idx,
+                has_full_attn=bool(full_attn_layers),
+                ssm_idx=inner_model_instance.ssm_idx,
+                has_linear=bool(linear_layers),
+            )
+
     _set_layers(model, layers)
 
     assert isinstance(layers, list), (
@@ -347,7 +402,8 @@ def patch_pipeline_model[T](model: T, group: mx.distributed.Group) -> T:
         if cache is not None:
             last = cache[-1]  # type: ignore
             dep_cache = last[0] if hasattr(last, "caches") else last  # type: ignore
-            dep_cache.keys = mx.depends(dep_cache.keys, logits)  # type: ignore
+            if hasattr(dep_cache, "keys") and dep_cache.keys is not None:  # type: ignore
+                dep_cache.keys = mx.depends(dep_cache.keys, logits)  # type: ignore
 
         return logits
 
@@ -470,7 +526,9 @@ def tensor_auto_parallel(
             all_to_sharded_linear_in_place,
             sharded_to_all_linear_in_place,
         )
-    elif isinstance(model, (Qwen3MoeModel, Qwen3NextModel)):
+    elif isinstance(
+        model, (Qwen3MoeModel, Qwen3NextModel, Qwen3_5TextModel, Qwen3_5MoeModel)
+    ):
         tensor_parallel_sharding_strategy = QwenShardingStrategy(
             group,
             all_to_sharded_linear,
@@ -865,7 +923,9 @@ class QwenShardingStrategy(TensorParallelShardingStrategy):
         on_timeout: TimeoutCallback | None,
         on_layer_loaded: LayerLoadedCallback | None,
     ) -> nn.Module:
-        model = cast(Qwen3MoeModel | Qwen3NextModel, model)
+        model = cast(
+            Qwen3MoeModel | Qwen3NextModel | Qwen3_5TextModel | Qwen3_5MoeModel, model
+        )
         total = len(model.layers)
         for i, layer in enumerate(model.layers):
             eval_with_timeout(layer.parameters(), timeout_seconds / total, on_timeout)
@@ -886,16 +946,39 @@ class QwenShardingStrategy(TensorParallelShardingStrategy):
                 layer.self_attn.n_heads //= self.N
                 layer.self_attn.n_kv_heads //= self.N
             else:
-                assert isinstance(layer, Qwen3NextDecoderLayer)
+                assert isinstance(layer, (Qwen3NextDecoderLayer, Qwen3_5DecoderLayer))
                 if hasattr(layer, "linear_attn"):
                     linear_attn = layer.linear_attn
 
-                    linear_attn.in_proj_qkvz = self.all_to_sharded_linear(
-                        linear_attn.in_proj_qkvz
-                    )
-                    linear_attn.in_proj_ba = self.all_to_sharded_linear(
-                        linear_attn.in_proj_ba
-                    )
+                    if isinstance(linear_attn, Qwen3NextGatedDeltaNet):
+                        # Qwen3-Next: combined projections
+                        linear_attn.in_proj_qkvz = self.all_to_sharded_linear(
+                            linear_attn.in_proj_qkvz
+                        )
+                        linear_attn.in_proj_ba = self.all_to_sharded_linear(
+                            linear_attn.in_proj_ba
+                        )
+                    else:
+                        # Qwen3.5: separate projections
+                        # in_proj_qkv has sections [q(key_dim), k(key_dim), v(value_dim)]
+                        # that must be split section-aware, not as a contiguous block
+                        key_dim = linear_attn.key_dim
+                        value_dim = linear_attn.value_dim
+                        linear_attn.in_proj_qkv = shard_linear(
+                            linear_attn.in_proj_qkv,
+                            "all-to-sharded",
+                            segments=[key_dim, key_dim + key_dim],
+                            group=self.group,
+                        )
+                        linear_attn.in_proj_z = self.all_to_sharded_linear(
+                            linear_attn.in_proj_z
+                        )
+                        linear_attn.in_proj_b = self.all_to_sharded_linear(
+                            linear_attn.in_proj_b
+                        )
+                        linear_attn.in_proj_a = self.all_to_sharded_linear(
+                            linear_attn.in_proj_a
+                        )
                     linear_attn.out_proj = self.sharded_to_all_linear(
                         linear_attn.out_proj
                     )
@@ -957,11 +1040,20 @@ class QwenShardingStrategy(TensorParallelShardingStrategy):
                     layer.self_attn.num_key_value_heads //= self.N
 
             # Shard the MoE.
-            if isinstance(layer.mlp, (Qwen3MoeSparseMoeBlock, Qwen3NextSparseMoeBlock)):
+            if isinstance(
+                layer.mlp,
+                (
+                    Qwen3MoeSparseMoeBlock,
+                    Qwen3NextSparseMoeBlock,
+                    Qwen3_5SparseMoeBlock,
+                ),
+            ):
                 self.all_to_sharded_linear_in_place(layer.mlp.switch_mlp.gate_proj)
                 self.sharded_to_all_linear_in_place(layer.mlp.switch_mlp.down_proj)
                 self.all_to_sharded_linear_in_place(layer.mlp.switch_mlp.up_proj)
-                if isinstance(layer.mlp, Qwen3NextSparseMoeBlock):
+                if isinstance(
+                    layer.mlp, (Qwen3NextSparseMoeBlock, Qwen3_5SparseMoeBlock)
+                ):
                     self.all_to_sharded_linear_in_place(
                         layer.mlp.shared_expert.gate_proj
                     )

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

@@ -437,6 +437,7 @@ def mlx_generate(
     group: mx.distributed.Group | None,
     on_prefill_progress: Callable[[int, int], None] | None = None,
     distributed_prompt_progress_callback: Callable[[], None] | None = None,
+    on_generation_token: Callable[[], None] | None = None,
 ) -> Generator[GenerationResponse]:
     # Ensure that generation stats only contains peak memory for this generation
     mx.reset_peak_memory()
@@ -644,6 +645,9 @@ def mlx_generate(
                         full_prompt_tokens, caches, cache_snapshots
                     )
 
+        if on_generation_token is not None:
+            on_generation_token()
+
         yield GenerationResponse(
             text=text,
             token=out.token,

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

@@ -318,6 +318,9 @@ def get_eos_token_ids_for_model(model_id: ModelId) -> list[int] | None:
         return [151336, 151329, 151338]
     elif "gpt-oss" in model_id_lower:
         return [200002, 200012]
+    elif "qwen3.5" in model_id_lower or "qwen-3.5" in model_id_lower:
+        # For Qwen3.5: 248046 (<|im_end|>), 248044 (<|endoftext|>)
+        return [248046, 248044]
     return None
 
 
@@ -551,6 +554,8 @@ def apply_chat_template(
         # Jinja ignores unknown variables, so passing both is safe.
         extra_kwargs["enable_thinking"] = task_params.enable_thinking
         extra_kwargs["thinking"] = task_params.enable_thinking
+    if task_params.reasoning_effort is not None:
+        extra_kwargs["reasoning_effort"] = task_params.reasoning_effort
 
     patched_template: str | None = None
     if task_params.tools:

src/exo/worker/plan.py

@@ -297,10 +297,10 @@ def _pending_tasks(
             # the task status _should_ be set to completed by the LAST runner
             # it is currently set by the first
             # this is definitely a hack
-            if task.task_id in runner.completed:
+            if task.task_id in runner.completed or task.task_id in runner.in_progress:
                 continue
 
-            if isinstance(runner.status, RunnerReady) and all(
+            if isinstance(runner.status, (RunnerReady, RunnerRunning)) and all(
                 isinstance(all_runners[global_runner_id], (RunnerReady, RunnerRunning))
                 for global_runner_id in runner.bound_instance.instance.shard_assignments.runner_to_shard
             ):

src/exo/worker/runner/bootstrap.py

@@ -32,11 +32,19 @@ def entrypoint(
     # Import main after setting global logger - this lets us just import logger from this module
     try:
         if bound_instance.is_image_model:
-            from exo.worker.runner.image_models.runner import main
-        else:
-            from exo.worker.runner.llm_inference.runner import main
+            from exo.worker.runner.image_models.runner import Runner as ImageRunner
 
-        main(bound_instance, event_sender, task_receiver, cancel_receiver)
+            runner = ImageRunner(
+                bound_instance, event_sender, task_receiver, cancel_receiver
+            )
+            runner.main()
+        else:
+            from exo.worker.runner.llm_inference.runner import Runner
+
+            runner = Runner(
+                bound_instance, event_sender, task_receiver, cancel_receiver
+            )
+            runner.main()
 
     except ClosedResourceError:
         logger.warning("Runner communication closed unexpectedly")

src/exo/worker/runner/image_models/runner.py

@@ -182,272 +182,266 @@ def _send_image_chunk(
         )
 
 
-def main(
-    bound_instance: BoundInstance,
-    event_sender: MpSender[Event],
-    task_receiver: MpReceiver[Task],
-    cancel_receiver: MpReceiver[TaskId],
-):
-    soft, hard = resource.getrlimit(resource.RLIMIT_NOFILE)
-    resource.setrlimit(resource.RLIMIT_NOFILE, (min(max(soft, 2048), hard), hard))
+class Runner:
+    def __init__(
+        self,
+        bound_instance: BoundInstance,
+        event_sender: MpSender[Event],
+        task_receiver: MpReceiver[Task],
+        cancel_receiver: MpReceiver[TaskId],
+    ):
+        self.event_sender = event_sender
+        self.task_receiver = task_receiver
+        self.cancel_receiver = cancel_receiver
+        self.bound_instance = bound_instance
 
-    instance, runner_id, shard_metadata = (
-        bound_instance.instance,
-        bound_instance.bound_runner_id,
-        bound_instance.bound_shard,
-    )
-    device_rank = shard_metadata.device_rank
-    logger.info("hello from the runner")
-    if getattr(shard_metadata, "immediate_exception", False):
-        raise Exception("Fake exception - runner failed to spin up.")
-    if timeout := getattr(shard_metadata, "should_timeout", 0):
-        time.sleep(timeout)
+        soft, hard = resource.getrlimit(resource.RLIMIT_NOFILE)
+        resource.setrlimit(resource.RLIMIT_NOFILE, (min(max(soft, 2048), hard), hard))
 
-    setup_start_time = time.time()
-    cancelled_tasks = set[TaskId]()
+        self.instance, self.runner_id, self.shard_metadata = (
+            bound_instance.instance,
+            bound_instance.bound_runner_id,
+            bound_instance.bound_shard,
+        )
+        self.device_rank = self.shard_metadata.device_rank
 
-    image_model: DistributedImageModel | None = None
-    group = None
+        logger.info("hello from the runner")
+        if getattr(self.shard_metadata, "immediate_exception", False):
+            raise Exception("Fake exception - runner failed to spin up.")
+        if timeout := getattr(self.shard_metadata, "should_timeout", 0):
+            time.sleep(timeout)
 
-    current_status: RunnerStatus = RunnerIdle()
-    logger.info("runner created")
-    event_sender.send(
-        RunnerStatusUpdated(runner_id=runner_id, runner_status=current_status)
-    )
-    seen = set[TaskId]()
-    with task_receiver as tasks:
-        for task in tasks:
-            if task.task_id in seen:
-                logger.warning("repeat task - potential error")
-            seen.add(task.task_id)
-            cancelled_tasks.discard(TaskId("CANCEL_CURRENT_TASK"))
-            event_sender.send(
-                TaskStatusUpdated(task_id=task.task_id, task_status=TaskStatus.Running)
+        self.setup_start_time = time.time()
+        self.cancelled_tasks = set[TaskId]()
+
+        self.image_model: DistributedImageModel | None = None
+        self.group = None
+
+        self.current_status: RunnerStatus = RunnerIdle()
+        logger.info("runner created")
+        self.update_status(RunnerIdle())
+        self.seen = set[TaskId]()
+
+    def update_status(self, status: RunnerStatus):
+        self.current_status = status
+        self.event_sender.send(
+            RunnerStatusUpdated(
+                runner_id=self.runner_id, runner_status=self.current_status
             )
-            match task:
-                case ConnectToGroup() if isinstance(
-                    current_status, (RunnerIdle, RunnerFailed)
-                ):
-                    logger.info("runner connecting")
-                    current_status = RunnerConnecting()
-                    event_sender.send(
-                        RunnerStatusUpdated(
-                            runner_id=runner_id, runner_status=current_status
-                        )
-                    )
-                    event_sender.send(TaskAcknowledged(task_id=task.task_id))
-                    group = initialize_mlx(bound_instance)
+        )
 
-                    logger.info("runner connected")
-                    current_status = RunnerConnected()
+    def send_task_status(self, task: Task, status: TaskStatus):
+        self.event_sender.send(
+            TaskStatusUpdated(task_id=task.task_id, task_status=status)
+        )
 
-                # we load the model if it's connected with a group, or idle without a group. we should never tell a model to connect if it doesn't need to
-                case LoadModel() if (
-                    isinstance(current_status, RunnerConnected) and group is not None
-                ) or (isinstance(current_status, RunnerIdle) and group is None):
-                    current_status = RunnerLoading()
-                    logger.info("runner loading")
-                    event_sender.send(
-                        RunnerStatusUpdated(
-                            runner_id=runner_id, runner_status=current_status
-                        )
-                    )
-                    event_sender.send(TaskAcknowledged(task_id=task.task_id))
+    def acknowledge_task(self, task: Task):
+        self.event_sender.send(TaskAcknowledged(task_id=task.task_id))
 
-                    assert (
-                        ModelTask.TextToImage in shard_metadata.model_card.tasks
-                        or ModelTask.ImageToImage in shard_metadata.model_card.tasks
-                    ), f"Incorrect model task(s): {shard_metadata.model_card.tasks}"
-
-                    image_model = initialize_image_model(bound_instance)
-                    current_status = RunnerLoaded()
-                    logger.info("runner loaded")
-
-                case StartWarmup() if isinstance(current_status, RunnerLoaded):
-                    current_status = RunnerWarmingUp()
-                    logger.info("runner warming up")
-                    event_sender.send(
-                        RunnerStatusUpdated(
-                            runner_id=runner_id, runner_status=current_status
-                        )
-                    )
-                    event_sender.send(TaskAcknowledged(task_id=task.task_id))
-
-                    logger.info(f"warming up inference for instance: {instance}")
-
-                    assert image_model
-                    image = warmup_image_generator(model=image_model)
-                    if image is not None:
-                        logger.info(f"warmed up by generating {image.size} image")
-                    else:
-                        logger.info("warmup completed (non-primary node)")
-
-                    logger.info(
-                        f"runner initialized in {time.time() - setup_start_time} seconds"
-                    )
-
-                    current_status = RunnerReady()
-                    logger.info("runner ready")
-
-                case ImageGeneration(
-                    task_params=task_params, command_id=command_id
-                ) if isinstance(current_status, RunnerReady):
-                    assert image_model
-                    logger.info(f"received image generation request: {str(task)[:500]}")
-                    current_status = RunnerRunning()
-                    logger.info("runner running")
-                    event_sender.send(
-                        RunnerStatusUpdated(
-                            runner_id=runner_id, runner_status=current_status
-                        )
-                    )
-                    event_sender.send(TaskAcknowledged(task_id=task.task_id))
-
-                    try:
-                        image_index = 0
-                        for response in generate_image(
-                            model=image_model, task=task_params
-                        ):
-                            is_primary_output = _is_primary_output_node(shard_metadata)
-
-                            if is_primary_output:
-                                match response:
-                                    case PartialImageResponse():
-                                        logger.info(
-                                            f"sending partial ImageChunk {response.partial_index}/{response.total_partials}"
-                                        )
-                                        _process_image_response(
-                                            response,
-                                            command_id,
-                                            shard_metadata,
-                                            event_sender,
-                                            image_index,
-                                        )
-                                    case ImageGenerationResponse():
-                                        logger.info("sending final ImageChunk")
-                                        _process_image_response(
-                                            response,
-                                            command_id,
-                                            shard_metadata,
-                                            event_sender,
-                                            image_index,
-                                        )
-                                        image_index += 1
-                    # can we make this more explicit?
-                    except Exception as e:
-                        if _is_primary_output_node(shard_metadata):
-                            event_sender.send(
-                                ChunkGenerated(
-                                    command_id=command_id,
-                                    chunk=ErrorChunk(
-                                        model=shard_metadata.model_card.model_id,
-                                        finish_reason="error",
-                                        error_message=str(e),
-                                    ),
-                                )
-                            )
-                        raise
-                    finally:
-                        _send_traces_if_enabled(event_sender, task.task_id, device_rank)
-
-                    current_status = RunnerReady()
-                    logger.info("runner ready")
-
-                case ImageEdits(task_params=task_params, command_id=command_id) if (
-                    isinstance(current_status, RunnerReady)
-                ):
-                    assert image_model
-                    logger.info(f"received image edits request: {str(task)[:500]}")
-                    current_status = RunnerRunning()
-                    logger.info("runner running")
-                    event_sender.send(
-                        RunnerStatusUpdated(
-                            runner_id=runner_id, runner_status=current_status
-                        )
-                    )
-                    event_sender.send(TaskAcknowledged(task_id=task.task_id))
-
-                    try:
-                        image_index = 0
-                        for response in generate_image(
-                            model=image_model, task=task_params
-                        ):
-                            if _is_primary_output_node(shard_metadata):
-                                match response:
-                                    case PartialImageResponse():
-                                        logger.info(
-                                            f"sending partial ImageChunk {response.partial_index}/{response.total_partials}"
-                                        )
-                                        _process_image_response(
-                                            response,
-                                            command_id,
-                                            shard_metadata,
-                                            event_sender,
-                                            image_index,
-                                        )
-                                    case ImageGenerationResponse():
-                                        logger.info("sending final ImageChunk")
-                                        _process_image_response(
-                                            response,
-                                            command_id,
-                                            shard_metadata,
-                                            event_sender,
-                                            image_index,
-                                        )
-                                        image_index += 1
-                    except Exception as e:
-                        if _is_primary_output_node(shard_metadata):
-                            event_sender.send(
-                                ChunkGenerated(
-                                    command_id=command_id,
-                                    chunk=ErrorChunk(
-                                        model=shard_metadata.model_card.model_id,
-                                        finish_reason="error",
-                                        error_message=str(e),
-                                    ),
-                                )
-                            )
-                        raise
-                    finally:
-                        _send_traces_if_enabled(event_sender, task.task_id, device_rank)
-
-                    current_status = RunnerReady()
-                    logger.info("runner ready")
-
-                case Shutdown():
-                    current_status = RunnerShuttingDown()
-                    logger.info("runner shutting down")
-                    if not TYPE_CHECKING:
-                        del image_model, group
-                        mx.clear_cache()
-                        import gc
-
-                        gc.collect()
-
-                    event_sender.send(
-                        RunnerStatusUpdated(
-                            runner_id=runner_id, runner_status=current_status
-                        )
-                    )
-                    event_sender.send(TaskAcknowledged(task_id=task.task_id))
-
-                    current_status = RunnerShutdown()
-                case _:
-                    raise ValueError(
-                        f"Received {task.__class__.__name__} outside of state machine in {current_status=}"
-                    )
-            was_cancelled = (task.task_id in cancelled_tasks) or (
-                TaskId("CANCEL_CURRENT_TASK") in cancelled_tasks
-            )
-            if not was_cancelled:
-                event_sender.send(
-                    TaskStatusUpdated(
-                        task_id=task.task_id, task_status=TaskStatus.Complete
-                    )
+    def main(self):
+        with self.task_receiver as tasks:
+            for task in tasks:
+                if task.task_id in self.seen:
+                    logger.warning("repeat task - potential error")
+                self.seen.add(task.task_id)
+                self.cancelled_tasks.discard(TaskId("CANCEL_CURRENT_TASK"))
+                self.send_task_status(task, TaskStatus.Running)
+                self.handle_task(task)
+                was_cancelled = (task.task_id in self.cancelled_tasks) or (
+                    TaskId("CANCEL_CURRENT_TASK") in self.cancelled_tasks
                 )
-            event_sender.send(
-                RunnerStatusUpdated(runner_id=runner_id, runner_status=current_status)
-            )
+                if not was_cancelled:
+                    self.send_task_status(task, TaskStatus.Complete)
+                self.update_status(self.current_status)
 
-            if isinstance(current_status, RunnerShutdown):
-                break
+                if isinstance(self.current_status, RunnerShutdown):
+                    break
+
+    def handle_task(self, task: Task):
+        match task:
+            case ConnectToGroup() if isinstance(
+                self.current_status, (RunnerIdle, RunnerFailed)
+            ):
+                logger.info("runner connecting")
+                self.update_status(RunnerConnecting())
+                self.acknowledge_task(task)
+                self.group = initialize_mlx(self.bound_instance)
+
+                logger.info("runner connected")
+                self.current_status = RunnerConnected()
+
+            # we load the model if it's connected with a group, or idle without a group. we should never tell a model to connect if it doesn't need to
+            case LoadModel() if (
+                isinstance(self.current_status, RunnerConnected)
+                and self.group is not None
+            ) or (isinstance(self.current_status, RunnerIdle) and self.group is None):
+                logger.info("runner loading")
+                self.update_status(RunnerLoading())
+                self.acknowledge_task(task)
+
+                assert (
+                    ModelTask.TextToImage in self.shard_metadata.model_card.tasks
+                    or ModelTask.ImageToImage in self.shard_metadata.model_card.tasks
+                ), f"Incorrect model task(s): {self.shard_metadata.model_card.tasks}"
+
+                self.image_model = initialize_image_model(self.bound_instance)
+                self.current_status = RunnerLoaded()
+                logger.info("runner loaded")
+
+            case StartWarmup() if isinstance(self.current_status, RunnerLoaded):
+                logger.info("runner warming up")
+                self.update_status(RunnerWarmingUp())
+                self.acknowledge_task(task)
+
+                logger.info(f"warming up inference for instance: {self.instance}")
+
+                assert self.image_model
+                image = warmup_image_generator(model=self.image_model)
+                if image is not None:
+                    logger.info(f"warmed up by generating {image.size} image")
+                else:
+                    logger.info("warmup completed (non-primary node)")
+
+                logger.info(
+                    f"runner initialized in {time.time() - self.setup_start_time} seconds"
+                )
+
+                self.current_status = RunnerReady()
+                logger.info("runner ready")
+
+            case ImageGeneration(task_params=task_params, command_id=command_id) if (
+                isinstance(self.current_status, RunnerReady)
+            ):
+                assert self.image_model
+                logger.info(f"received image generation request: {str(task)[:500]}")
+                logger.info("runner running")
+                self.update_status(RunnerRunning())
+                self.acknowledge_task(task)
+
+                try:
+                    image_index = 0
+                    for response in generate_image(
+                        model=self.image_model, task=task_params
+                    ):
+                        is_primary_output = _is_primary_output_node(self.shard_metadata)
+
+                        if is_primary_output:
+                            match response:
+                                case PartialImageResponse():
+                                    logger.info(
+                                        f"sending partial ImageChunk {response.partial_index}/{response.total_partials}"
+                                    )
+                                    _process_image_response(
+                                        response,
+                                        command_id,
+                                        self.shard_metadata,
+                                        self.event_sender,
+                                        image_index,
+                                    )
+                                case ImageGenerationResponse():
+                                    logger.info("sending final ImageChunk")
+                                    _process_image_response(
+                                        response,
+                                        command_id,
+                                        self.shard_metadata,
+                                        self.event_sender,
+                                        image_index,
+                                    )
+                                    image_index += 1
+                # can we make this more explicit?
+                except Exception as e:
+                    if _is_primary_output_node(self.shard_metadata):
+                        self.event_sender.send(
+                            ChunkGenerated(
+                                command_id=command_id,
+                                chunk=ErrorChunk(
+                                    model=self.shard_metadata.model_card.model_id,
+                                    finish_reason="error",
+                                    error_message=str(e),
+                                ),
+                            )
+                        )
+                    raise
+                finally:
+                    _send_traces_if_enabled(
+                        self.event_sender, task.task_id, self.device_rank
+                    )
+
+                self.current_status = RunnerReady()
+                logger.info("runner ready")
+
+            case ImageEdits(task_params=task_params, command_id=command_id) if (
+                isinstance(self.current_status, RunnerReady)
+            ):
+                assert self.image_model
+                logger.info(f"received image edits request: {str(task)[:500]}")
+                logger.info("runner running")
+                self.update_status(RunnerRunning())
+                self.acknowledge_task(task)
+
+                try:
+                    image_index = 0
+                    for response in generate_image(
+                        model=self.image_model, task=task_params
+                    ):
+                        if _is_primary_output_node(self.shard_metadata):
+                            match response:
+                                case PartialImageResponse():
+                                    logger.info(
+                                        f"sending partial ImageChunk {response.partial_index}/{response.total_partials}"
+                                    )
+                                    _process_image_response(
+                                        response,
+                                        command_id,
+                                        self.shard_metadata,
+                                        self.event_sender,
+                                        image_index,
+                                    )
+                                case ImageGenerationResponse():
+                                    logger.info("sending final ImageChunk")
+                                    _process_image_response(
+                                        response,
+                                        command_id,
+                                        self.shard_metadata,
+                                        self.event_sender,
+                                        image_index,
+                                    )
+                                    image_index += 1
+                except Exception as e:
+                    if _is_primary_output_node(self.shard_metadata):
+                        self.event_sender.send(
+                            ChunkGenerated(
+                                command_id=command_id,
+                                chunk=ErrorChunk(
+                                    model=self.shard_metadata.model_card.model_id,
+                                    finish_reason="error",
+                                    error_message=str(e),
+                                ),
+                            )
+                        )
+                    raise
+                finally:
+                    _send_traces_if_enabled(
+                        self.event_sender, task.task_id, self.device_rank
+                    )
+
+                self.current_status = RunnerReady()
+                logger.info("runner ready")
+
+            case Shutdown():
+                logger.info("runner shutting down")
+                if not TYPE_CHECKING:
+                    del self.image_model, self.group
+                    mx.clear_cache()
+                    import gc
+
+                    gc.collect()
+
+                self.update_status(RunnerShuttingDown())
+                self.acknowledge_task(task)
+
+                self.current_status = RunnerShutdown()
+            case _:
+                raise ValueError(
+                    f"Received {task.__class__.__name__} outside of state machine in {self.current_status=}"
+                )

src/exo/worker/runner/llm_inference/batch_generator.py

@@ -0,0 +1,346 @@
+import itertools
+import math
+import time
+from abc import ABC, abstractmethod
+from collections import deque
+from collections.abc import Generator, Iterable
+from dataclasses import dataclass, field
+from typing import cast
+
+import mlx.core as mx
+from mlx_lm.tokenizer_utils import TokenizerWrapper
+
+from exo.shared.types.chunks import ErrorChunk, PrefillProgressChunk
+from exo.shared.types.common import ModelId
+from exo.shared.types.events import ChunkGenerated, Event
+from exo.shared.types.mlx import Model
+from exo.shared.types.tasks import TaskId, TextGeneration
+from exo.shared.types.text_generation import TextGenerationTaskParams
+from exo.shared.types.worker.runner_response import GenerationResponse, ToolCallResponse
+from exo.utils.channels import MpReceiver, MpSender
+from exo.worker.engines.mlx.cache import KVPrefixCache
+from exo.worker.engines.mlx.generator.generate import (
+    PrefillCancelled,
+    mlx_generate,
+    warmup_inference,
+)
+from exo.worker.engines.mlx.utils_mlx import (
+    apply_chat_template,
+    mx_any,
+)
+from exo.worker.runner.bootstrap import logger
+
+from .model_output_parsers import apply_all_parsers
+from .tool_parsers import ToolParser
+
+
+class Cancelled:
+    pass
+
+
+class Finished:
+    pass
+
+
+class GeneratorQueue[T]:
+    def __init__(self):
+        self._q = deque[T]()
+
+    def push(self, t: T):
+        self._q.append(t)
+
+    def gen(self) -> Generator[T | None]:
+        while True:
+            if len(self._q) == 0:
+                yield None
+            else:
+                yield self._q.popleft()
+
+
+class InferenceGenerator(ABC):
+    @abstractmethod
+    def warmup(self) -> None: ...
+
+    @abstractmethod
+    def submit(
+        self,
+        task: TextGeneration,
+    ) -> None: ...
+
+    @abstractmethod
+    def step(
+        self,
+    ) -> Iterable[
+        tuple[TaskId, ToolCallResponse | GenerationResponse | Cancelled | Finished]
+    ]: ...
+
+    @abstractmethod
+    def close(self) -> None: ...
+
+
+EXO_RUNNER_MUST_FAIL = "EXO RUNNER MUST FAIL"
+EXO_RUNNER_MUST_OOM = "EXO RUNNER MUST OOM"
+EXO_RUNNER_MUST_TIMEOUT = "EXO RUNNER MUST TIMEOUT"
+
+
+def _check_for_debug_prompts(task_params: TextGenerationTaskParams) -> None:
+    """Check for debug prompt triggers in the input."""
+    from exo.worker.engines.mlx.utils_mlx import mlx_force_oom
+
+    if len(task_params.input) == 0:
+        return
+    prompt = task_params.input[0].content
+    if not prompt:
+        return
+    if EXO_RUNNER_MUST_FAIL in prompt:
+        raise Exception("Artificial runner exception - for testing purposes only.")
+    if EXO_RUNNER_MUST_OOM in prompt:
+        mlx_force_oom()
+    if EXO_RUNNER_MUST_TIMEOUT in prompt:
+        time.sleep(100)
+
+
+@dataclass(eq=False)
+class SequentialGenerator(InferenceGenerator):
+    model: Model
+    tokenizer: TokenizerWrapper
+    group: mx.distributed.Group | None
+    kv_prefix_cache: KVPrefixCache | None
+    tool_parser: ToolParser | None
+    model_id: ModelId
+    device_rank: int
+    cancel_receiver: MpReceiver[TaskId]
+    event_sender: MpSender[Event]
+    check_for_cancel_every: int = 50
+
+    _cancelled_tasks: set[TaskId] = field(default_factory=set, init=False)
+    _maybe_queue: list[TextGeneration] = field(default_factory=list, init=False)
+    _queue: deque[TextGeneration] = field(default_factory=deque, init=False)
+    _active: (
+        tuple[
+            TextGeneration,
+            # mlx generator that does work
+            Generator[GenerationResponse],
+            # queue that the 1st generator should push to and 3rd generator should pull from
+            GeneratorQueue[GenerationResponse],
+            # generator to get parsed outputs
+            Generator[GenerationResponse | ToolCallResponse | None],
+        ]
+        | None
+    ) = field(default=None, init=False)
+
+    def warmup(self):
+        logger.info(f"warming up inference for instance: {self.model_id}")
+
+        t = time.monotonic()
+        toks = warmup_inference(
+            model=self.model,
+            tokenizer=self.tokenizer,
+            group=self.group,
+        )
+        logger.info(f"warmed up by generating {toks} tokens")
+        check_for_cancel_every = min(
+            math.ceil(toks / min(time.monotonic() - t, 0.001)), 100
+        )
+        if self.group is not None:
+            self.check_for_cancel_every = int(
+                mx.max(
+                    mx.distributed.all_gather(
+                        mx.array([check_for_cancel_every]),
+                        group=self.group,
+                    )
+                ).item()
+            )
+
+        logger.info(
+            f"runner checking for cancellation every {check_for_cancel_every} tokens"
+        )
+
+    def submit(
+        self,
+        task: TextGeneration,
+    ) -> None:
+        self._cancelled_tasks.discard(TaskId("CANCEL_CURRENT_TASK"))
+        self._maybe_queue.append(task)
+
+    def agree_on_tasks(self) -> None:
+        """Agree between all ranks about the task ordering (some may have received in different order or not at all)."""
+        agreed, different = mx_all_gather_tasks(self._maybe_queue, self.group)
+        self._queue.extend(task for task in self._maybe_queue if task in agreed)
+        self._maybe_queue = [task for task in self._maybe_queue if task in different]
+
+    def step(
+        self,
+    ) -> Iterable[
+        tuple[TaskId, GenerationResponse | ToolCallResponse | Cancelled | Finished]
+    ]:
+        if self._active is None:
+            self.agree_on_tasks()
+
+            if self._queue:
+                self._start_next()
+            else:
+                return map(lambda task: (task, Cancelled()), self._cancelled_tasks)
+
+        assert self._active is not None
+
+        task, mlx_gen, queue, output_generator = self._active
+        response = None
+        try:
+            queue.push(next(mlx_gen))
+            response = next(output_generator)
+        except (StopIteration, PrefillCancelled):
+            response = Finished()
+            self._active = None
+            if self._queue:
+                self._start_next()
+        except Exception as e:
+            self._send_error(task, e)
+            self._active = None
+            raise
+        return itertools.chain(
+            [] if response is None else [(task.task_id, response)],
+            map(lambda task: (task, Cancelled()), self._cancelled_tasks),
+        )
+
+    def _start_next(self) -> None:
+        task = self._queue.popleft()
+        try:
+            mlx_gen = self._build_generator(task)
+        except Exception as e:
+            self._send_error(task, e)
+            raise
+        queue = GeneratorQueue[GenerationResponse]()
+        output_generator = apply_all_parsers(
+            queue.gen(),
+            apply_chat_template(self.tokenizer, task.task_params),
+            self.tool_parser,
+            self.tokenizer,
+            type(self.model),
+            self.model_id,
+            task.task_params.tools,
+        )
+        self._active = (task, mlx_gen, queue, output_generator)
+
+    def _send_error(self, task: TextGeneration, e: Exception) -> None:
+        if self.device_rank == 0:
+            self.event_sender.send(
+                ChunkGenerated(
+                    command_id=task.command_id,
+                    chunk=ErrorChunk(
+                        model=self.model_id,
+                        finish_reason="error",
+                        error_message=str(e),
+                    ),
+                )
+            )
+
+    def _build_generator(self, task: TextGeneration) -> Generator[GenerationResponse]:
+        _check_for_debug_prompts(task.task_params)
+        prompt = apply_chat_template(self.tokenizer, task.task_params)
+
+        def on_prefill_progress(processed: int, total: int) -> None:
+            if self.device_rank == 0:
+                self.event_sender.send(
+                    ChunkGenerated(
+                        command_id=task.command_id,
+                        chunk=PrefillProgressChunk(
+                            model=self.model_id,
+                            processed_tokens=processed,
+                            total_tokens=total,
+                        ),
+                    )
+                )
+
+        def distributed_prompt_progress_callback() -> None:
+            self._cancelled_tasks.update(self.cancel_receiver.collect())
+            want_to_cancel = (task.task_id in self._cancelled_tasks) or (
+                TaskId("CANCEL_CURRENT_TASK") in self._cancelled_tasks
+            )
+            if mx_any(want_to_cancel, self.group):
+                raise PrefillCancelled()
+
+            self.agree_on_tasks()
+
+        tokens_since_cancel_check = self.check_for_cancel_every
+
+        def on_generation_token() -> None:
+            nonlocal tokens_since_cancel_check
+            tokens_since_cancel_check += 1
+            if tokens_since_cancel_check >= self.check_for_cancel_every:
+                tokens_since_cancel_check = 0
+                self._cancelled_tasks.update(self.cancel_receiver.collect())
+                want_to_cancel = (task.task_id in self._cancelled_tasks) or (
+                    TaskId("CANCEL_CURRENT_TASK") in self._cancelled_tasks
+                )
+                if mx_any(want_to_cancel, self.group):
+                    raise PrefillCancelled()
+
+                self.agree_on_tasks()
+
+        return mlx_generate(
+            model=self.model,
+            tokenizer=self.tokenizer,
+            task=task.task_params,
+            prompt=prompt,
+            kv_prefix_cache=self.kv_prefix_cache,
+            on_prefill_progress=on_prefill_progress,
+            distributed_prompt_progress_callback=distributed_prompt_progress_callback,
+            on_generation_token=on_generation_token,
+            group=self.group,
+        )
+
+    def close(self) -> None:
+        del self.model, self.tokenizer, self.group
+
+
+def mx_all_gather_tasks(
+    tasks: list[TextGeneration],
+    group: mx.distributed.Group | None,
+) -> tuple[list[TextGeneration], list[TextGeneration]]:
+    def encode_task_id(task_id: TaskId) -> list[int]:
+        utf8_task_id = task_id.encode()
+        return [
+            int.from_bytes(utf8_task_id[i : i + 1]) for i in range(len(utf8_task_id))
+        ]
+
+    def decode_task_id(encoded_task_id: list[int]) -> TaskId:
+        return TaskId(
+            bytes.decode(b"".join((x).to_bytes(length=1) for x in encoded_task_id))
+        )
+
+    uuid_byte_length = 36
+
+    n_tasks = len(tasks)
+    all_counts = cast(
+        list[int],
+        mx.distributed.all_gather(mx.array([n_tasks]), group=group).tolist(),
+    )
+    max_tasks = max(all_counts)
+    world_size: int = 1 if group is None else group.size()
+
+    if max_tasks == 0:
+        return [], []
+
+    padded = [encode_task_id(task.task_id) for task in tasks] + [
+        [0] * uuid_byte_length
+    ] * (max_tasks - n_tasks)
+    gathered = cast(
+        list[list[list[int]]],
+        mx.distributed.all_gather(mx.array(padded), group=group)
+        .reshape(world_size, max_tasks, -1)
+        .tolist(),
+    )
+    all_task_ids: list[list[TaskId]] = [
+        [decode_task_id(encoded_task_id) for encoded_task_id in rank_tasks[:count]]
+        for rank_tasks, count in zip(gathered, all_counts, strict=True)
+    ]
+
+    agreed_ids: set[TaskId] = set(all_task_ids[0])
+    for rank_tasks in all_task_ids[1:]:
+        agreed_ids &= set(rank_tasks)
+
+    local_tasks = {task.task_id: task for task in tasks}
+    agreed = [local_tasks[tid] for tid in sorted(agreed_ids)]
+    different = [task for task in tasks if task.task_id not in agreed_ids]
+    return agreed, different

src/exo/worker/runner/llm_inference/model_output_parsers.py

@@ -0,0 +1,378 @@
+from collections.abc import Generator
+from functools import cache
+from typing import Any
+
+from mlx_lm.models.deepseek_v32 import Model as DeepseekV32Model
+from mlx_lm.models.gpt_oss import Model as GptOssModel
+from mlx_lm.tokenizer_utils import TokenizerWrapper
+from openai_harmony import (  # pyright: ignore[reportMissingTypeStubs]
+    HarmonyEncodingName,
+    HarmonyError,  # pyright: ignore[reportUnknownVariableType]
+    Role,
+    StreamableParser,
+    load_harmony_encoding,
+)
+
+from exo.shared.types.api import ToolCallItem
+from exo.shared.types.common import ModelId
+from exo.shared.types.mlx import Model
+from exo.shared.types.worker.runner_response import GenerationResponse, ToolCallResponse
+from exo.worker.engines.mlx.utils_mlx import (
+    detect_thinking_prompt_suffix,
+)
+from exo.worker.runner.bootstrap import logger
+
+from .tool_parsers import ToolParser
+
+
+@cache
+def get_gpt_oss_encoding():
+    encoding = load_harmony_encoding(HarmonyEncodingName.HARMONY_GPT_OSS)
+    return encoding
+
+
+def apply_all_parsers(
+    receiver: Generator[GenerationResponse | None],
+    prompt: str,
+    tool_parser: ToolParser | None,
+    tokenizer: TokenizerWrapper,
+    model_type: type[Model],
+    model_id: ModelId,
+    tools: list[dict[str, Any]] | None,
+) -> Generator[GenerationResponse | ToolCallResponse | None]:
+    mlx_generator = receiver
+
+    if tokenizer.has_thinking:
+        mlx_generator = parse_thinking_models(
+            mlx_generator,
+            tokenizer,
+            starts_in_thinking=detect_thinking_prompt_suffix(prompt, tokenizer),
+        )
+
+    if issubclass(model_type, GptOssModel):
+        mlx_generator = parse_gpt_oss(mlx_generator)
+    elif (
+        issubclass(model_type, DeepseekV32Model)
+        and "deepseek" in model_id.normalize().lower()
+    ):
+        mlx_generator = parse_deepseek_v32(mlx_generator)
+    elif tool_parser:
+        mlx_generator = parse_tool_calls(mlx_generator, tool_parser, tools)
+
+    return mlx_generator
+
+
+def parse_gpt_oss(
+    responses: Generator[GenerationResponse | None],
+) -> Generator[GenerationResponse | ToolCallResponse | None]:
+    encoding = get_gpt_oss_encoding()
+    stream = StreamableParser(encoding, role=Role.ASSISTANT)
+    thinking = False
+    current_tool_name: str | None = None
+    tool_arg_parts: list[str] = []
+
+    for response in responses:
+        if response is None:
+            yield None
+            continue
+        try:
+            stream.process(response.token)
+        except HarmonyError:
+            logger.error("Encountered critical Harmony Error, returning early")
+            return
+
+        delta = stream.last_content_delta
+        ch = stream.current_channel
+        recipient = stream.current_recipient
+
+        # Debug: log every token with state
+        logger.debug(
+            f"parse_gpt_oss token={response.token} text={response.text!r} "
+            f"recipient={recipient!r} ch={ch!r} delta={delta!r} "
+            f"state={stream.state} current_tool={current_tool_name!r}"
+        )
+
+        if recipient != current_tool_name:
+            if current_tool_name is not None:
+                prefix = "functions."
+                if current_tool_name.startswith(prefix):
+                    current_tool_name = current_tool_name[len(prefix) :]
+                logger.info(
+                    f"parse_gpt_oss yielding tool call: name={current_tool_name!r}"
+                )
+                yield ToolCallResponse(
+                    tool_calls=[
+                        ToolCallItem(
+                            name=current_tool_name,
+                            arguments="".join(tool_arg_parts).strip(),
+                        )
+                    ],
+                    usage=response.usage,
+                )
+                tool_arg_parts = []
+            current_tool_name = recipient
+
+        # If inside a tool call, accumulate arguments
+        if current_tool_name is not None:
+            if delta:
+                tool_arg_parts.append(delta)
+            continue
+
+        if ch == "analysis" and not thinking:
+            thinking = True
+
+        if ch != "analysis" and thinking:
+            thinking = False
+
+        if delta:
+            yield response.model_copy(update={"text": delta, "is_thinking": thinking})
+
+        if response.finish_reason is not None:
+            yield response
+
+
+def parse_deepseek_v32(
+    responses: Generator[GenerationResponse | None],
+) -> Generator[GenerationResponse | ToolCallResponse | None]:
+    """Parse DeepSeek V3.2 DSML tool calls from the generation stream.
+
+    Uses accumulated-text matching (not per-token marker checks) because
+    DSML markers like <｜DSML｜function_calls> may span multiple tokens.
+    Also handles <think>...</think> blocks for thinking mode.
+    """
+    from exo.worker.engines.mlx.dsml_encoding import (
+        THINKING_END,
+        THINKING_START,
+        TOOL_CALLS_END,
+        TOOL_CALLS_START,
+        parse_dsml_output,
+    )
+
+    accumulated = ""
+    in_tool_call = False
+    thinking = False
+    # Tokens buffered while we detect the start of a DSML block
+    pending_buffer: list[GenerationResponse] = []
+    # Text accumulated during a tool call block
+    tool_call_text = ""
+
+    for response in responses:
+        if response is None:
+            yield None
+            continue
+
+        # ── Handle thinking tags ──
+        if not thinking and THINKING_START in response.text:
+            thinking = True
+            # Yield any text before the <think> tag
+            before = response.text[: response.text.index(THINKING_START)]
+            if before:
+                yield response.model_copy(update={"text": before})
+            continue
+
+        if thinking and THINKING_END in response.text:
+            thinking = False
+            # Yield any text after the </think> tag
+            after = response.text[
+                response.text.index(THINKING_END) + len(THINKING_END) :
+            ]
+            if after:
+                yield response.model_copy(update={"text": after, "is_thinking": False})
+            continue
+
+        if thinking:
+            yield response.model_copy(update={"is_thinking": True})
+            continue
+
+        # ── Handle tool call accumulation ──
+        if in_tool_call:
+            tool_call_text += response.text
+            if TOOL_CALLS_END in tool_call_text:
+                # Parse the accumulated DSML block
+                parsed = parse_dsml_output(tool_call_text)
+                if parsed is not None:
+                    logger.info(f"parsed DSML tool calls: {parsed}")
+                    yield ToolCallResponse(
+                        tool_calls=parsed,
+                        usage=response.usage,
+                        stats=response.stats,
+                    )
+                else:
+                    logger.warning(
+                        f"DSML tool call parsing failed for: {tool_call_text}"
+                    )
+                    yield response.model_copy(update={"text": tool_call_text})
+                in_tool_call = False
+                tool_call_text = ""
+                continue
+
+            # EOS reached before end marker — yield buffered text as-is
+            if response.finish_reason is not None:
+                logger.info("DSML tool call parsing interrupted by EOS")
+                yield response.model_copy(update={"text": tool_call_text})
+                in_tool_call = False
+                tool_call_text = ""
+            continue
+
+        # ── Detect start of tool call block ──
+        accumulated += response.text
+
+        if TOOL_CALLS_START in accumulated:
+            # The start marker might be split across pending_buffer + current token
+            start_idx = accumulated.index(TOOL_CALLS_START)
+            # Yield any pending tokens that are purely before the marker
+            pre_text = accumulated[:start_idx]
+            if pre_text:
+                # Flush pending buffer tokens that contributed text before the marker
+                for buf_resp in pending_buffer:
+                    if pre_text:
+                        chunk = buf_resp.text
+                        if len(chunk) <= len(pre_text):
+                            yield buf_resp
+                            pre_text = pre_text[len(chunk) :]
+                        else:
+                            yield buf_resp.model_copy(update={"text": pre_text})
+                            pre_text = ""
+            pending_buffer = []
+            tool_call_text = accumulated[start_idx:]
+            accumulated = ""
+
+            # Check if the end marker is already present (entire tool call in one token)
+            if TOOL_CALLS_END in tool_call_text:
+                parsed = parse_dsml_output(tool_call_text)
+                if parsed is not None:
+                    logger.info(f"parsed DSML tool calls: {parsed}")
+                    yield ToolCallResponse(
+                        tool_calls=parsed,
+                        usage=response.usage,
+                        stats=response.stats,
+                    )
+                else:
+                    logger.warning(
+                        f"DSML tool call parsing failed for: {tool_call_text}"
+                    )
+                    yield response.model_copy(update={"text": tool_call_text})
+                tool_call_text = ""
+            else:
+                in_tool_call = True
+            continue
+
+        # Check if accumulated text might be the start of a DSML marker
+        # Buffer tokens if we see a partial match at the end
+        if _could_be_dsml_prefix(accumulated):
+            pending_buffer.append(response)
+            continue
+
+        # No partial match — flush all pending tokens and the current one
+        for buf_resp in pending_buffer:
+            yield buf_resp
+        pending_buffer = []
+        accumulated = ""
+        yield response
+
+    # Flush any remaining pending buffer at generator end
+    for buf_resp in pending_buffer:
+        yield buf_resp
+
+
+def _could_be_dsml_prefix(text: str) -> bool:
+    """Check if the end of text could be the start of a DSML function_calls marker.
+
+    We look for suffixes of text that are prefixes of the TOOL_CALLS_START pattern.
+    This allows us to buffer tokens until we can determine if a tool call is starting.
+    """
+    from exo.worker.engines.mlx.dsml_encoding import TOOL_CALLS_START
+
+    # Only check the last portion of text that could overlap with the marker
+    max_check = len(TOOL_CALLS_START)
+    tail = text[-max_check:] if len(text) > max_check else text
+
+    # Check if any suffix of tail is a prefix of TOOL_CALLS_START
+    for i in range(len(tail)):
+        suffix = tail[i:]
+        if TOOL_CALLS_START.startswith(suffix):
+            return True
+    return False
+
+
+def parse_thinking_models(
+    responses: Generator[GenerationResponse | None],
+    tokenizer: TokenizerWrapper,
+    starts_in_thinking: bool = True,
+) -> Generator[GenerationResponse | None]:
+    """Route thinking tokens via is_thinking flag.
+
+    Swallows think tag tokens, sets is_thinking on all others.
+    Always yields tokens with finish_reason to avoid hanging the chunk stream.
+    """
+    in_thinking = starts_in_thinking
+    for response in responses:
+        if response is None:
+            yield None
+            continue
+
+        is_think_tag = (
+            tokenizer.think_end is not None and response.text == tokenizer.think_end
+        ) or (
+            tokenizer.think_start is not None and response.text == tokenizer.think_start
+        )
+
+        if is_think_tag:
+            in_thinking = response.text != tokenizer.think_end
+            # Never swallow finish_reason — the chunk stream needs it to terminate.
+            if response.finish_reason is not None:
+                yield response.model_copy(update={"text": "", "is_thinking": False})
+            continue
+        yield response.model_copy(update={"is_thinking": in_thinking})
+
+
+def parse_tool_calls(
+    responses: Generator[GenerationResponse | None],
+    tool_parser: ToolParser,
+    tools: list[dict[str, Any]] | None,
+) -> Generator[GenerationResponse | ToolCallResponse | None]:
+    in_tool_call = False
+    tool_call_text_parts: list[str] = []
+    for response in responses:
+        if response is None:
+            yield None
+            continue
+        if not in_tool_call and response.text.startswith(tool_parser.start_parsing):
+            in_tool_call = True
+
+        if in_tool_call:
+            tool_call_text_parts.append(response.text)
+            if response.text.endswith(tool_parser.end_parsing):
+                # parse the actual tool calls from the tool call text
+                parsed = tool_parser.parse("".join(tool_call_text_parts).strip(), tools)
+                logger.info(f"parsed {tool_call_text_parts=} into {parsed=}")
+                if parsed is not None:
+                    yield ToolCallResponse(
+                        tool_calls=parsed, usage=response.usage, stats=response.stats
+                    )
+                else:
+                    logger.warning(
+                        f"tool call parsing failed for text {''.join(tool_call_text_parts)}"
+                    )
+                    response.text = "".join(tool_call_text_parts)
+                    yield response
+
+                in_tool_call = False
+                tool_call_text_parts = []
+                continue
+
+            if response.finish_reason is not None:
+                logger.info(
+                    "tool call parsing interrupted, yield partial tool call as text"
+                )
+                response = response.model_copy(
+                    update={
+                        "text": "".join(tool_call_text_parts),
+                        "token": 0,
+                    }
+                )
+                yield response
+
+        else:
+            # fallthrough
+            yield response

src/exo/worker/runner/llm_inference/tool_parsers.py

@@ -1,4 +1,5 @@
 import json
+import math
 from dataclasses import dataclass
 from typing import Any, Callable
 
@@ -9,7 +10,177 @@ from exo.shared.types.api import ToolCallItem
 class ToolParser:
     start_parsing: str
     end_parsing: str
-    parse_tool_calls: Callable[[str], list[ToolCallItem] | None]
+    _inner_parser: Callable[[str], list[ToolCallItem] | None]
+
+    def parse(
+        self, text: str, tools: list[dict[str, Any]] | None
+    ) -> list[ToolCallItem] | None:
+        parsed = self._inner_parser(text)
+        if parsed is None:
+            return None
+        if tools is not None:
+            parsed = _coerce_tool_calls_to_schema(parsed, tools)
+        return parsed
+
+
+def _json_type_matches(value: Any, expected_type: str) -> bool:  # pyright: ignore[reportAny]
+    if expected_type == "object":
+        return isinstance(value, dict)
+    if expected_type == "array":
+        return isinstance(value, list)
+    if expected_type == "string":
+        return isinstance(value, str)
+    if expected_type == "integer":
+        return isinstance(value, int) and not isinstance(value, bool)
+    if expected_type == "number":
+        return (isinstance(value, int) and not isinstance(value, bool)) or isinstance(
+            value, float
+        )
+    if expected_type == "boolean":
+        return isinstance(value, bool)
+    if expected_type == "null":
+        return value is None
+    return False
+
+
+def _coerce_tool_arg_with_schema(value: Any, schema: dict[str, Any]) -> Any:  # pyright: ignore[reportAny]
+    schema_type = schema.get("type")
+
+    if isinstance(schema_type, list):
+        for candidate in schema_type:  # pyright: ignore[reportUnknownVariableType]
+            if not isinstance(candidate, str):
+                continue
+            if candidate == "null" and value is None:
+                return None
+            candidate_schema = {**schema, "type": candidate}
+            coerced = _coerce_tool_arg_with_schema(value, candidate_schema)  # pyright: ignore[reportAny]
+            if _json_type_matches(coerced, candidate):
+                return coerced  # pyright: ignore[reportAny]
+        return value  # pyright: ignore[reportAny]
+
+    if not isinstance(schema_type, str):
+        return value  # pyright: ignore[reportAny]
+
+    if schema_type == "object":
+        parsed = value  # pyright: ignore[reportAny]
+        if isinstance(parsed, str):
+            try:
+                parsed = json.loads(parsed)  # pyright: ignore[reportAny]
+            except Exception:
+                return value  # pyright: ignore[reportAny]
+        if not isinstance(parsed, dict):
+            return value  # pyright: ignore[reportAny]
+        properties = schema.get("properties")
+        if not isinstance(properties, dict):
+            return parsed  # pyright: ignore[reportUnknownVariableType]
+        return {
+            key: (
+                _coerce_tool_arg_with_schema(prop_value, prop_schema)  # pyright: ignore[reportUnknownArgumentType]
+                if isinstance(prop_schema, dict)
+                else prop_value
+            )
+            for key, prop_value in parsed.items()  # pyright: ignore[reportUnknownVariableType]
+            for prop_schema in [properties.get(key)]  # type: ignore
+        }
+
+    if schema_type == "array":
+        parsed = value  # pyright: ignore[reportAny]
+        if isinstance(parsed, str):
+            try:
+                parsed = json.loads(parsed)  # pyright: ignore[reportAny]
+            except Exception:
+                return value  # pyright: ignore[reportAny]
+        if not isinstance(parsed, list):
+            return value  # pyright: ignore[reportAny]
+        item_schema = schema.get("items")
+        if not isinstance(item_schema, dict):
+            return parsed  # pyright: ignore[reportUnknownVariableType]
+        return [_coerce_tool_arg_with_schema(item, item_schema) for item in parsed]  # type: ignore
+
+    if schema_type == "integer":
+        if isinstance(value, bool):
+            return value
+        if isinstance(value, int):
+            return value
+        if isinstance(value, float) and value.is_integer():
+            return int(value)
+        if isinstance(value, str):
+            try:
+                return int(value.strip())
+            except ValueError:
+                return value
+        return value
+
+    if schema_type == "number":
+        if isinstance(value, bool):
+            return value
+        if isinstance(value, (int, float)):
+            return value
+        if isinstance(value, str):
+            try:
+                num = float(value.strip())
+                if math.isfinite(num):
+                    return num
+            except ValueError:
+                return value
+        return value
+
+    if schema_type == "boolean":
+        if isinstance(value, bool):
+            return value
+        if isinstance(value, str):
+            lowered = value.strip().lower()
+            if lowered == "true":
+                return True
+            if lowered == "false":
+                return False
+        return value
+
+    return value  # pyright: ignore[reportAny]
+
+
+def _coerce_tool_calls_to_schema(
+    tool_calls: list[ToolCallItem], tools: list[dict[str, Any]]
+) -> list[ToolCallItem]:
+    schema_by_name: dict[str, dict[str, Any]] = {}
+    for tool in tools:
+        function = tool.get("function")
+        if not isinstance(function, dict):
+            continue
+        name = function.get("name")  # type: ignore
+        parameters = function.get("parameters")  # type: ignore
+        if isinstance(name, str) and isinstance(parameters, dict):
+            schema_by_name[name] = parameters
+
+    if not schema_by_name:
+        return tool_calls
+
+    coerced_calls: list[ToolCallItem] = []
+    for tool_call in tool_calls:
+        schema = schema_by_name.get(tool_call.name)
+        if schema is None:
+            coerced_calls.append(tool_call)
+            continue
+
+        try:
+            parsed_args = json.loads(tool_call.arguments)  # pyright: ignore[reportAny]
+        except Exception:
+            coerced_calls.append(tool_call)
+            continue
+
+        if not isinstance(parsed_args, dict):
+            coerced_calls.append(tool_call)
+            continue
+
+        coerced_args = _coerce_tool_arg_with_schema(parsed_args, schema)  # pyright: ignore[reportAny]
+        if not isinstance(coerced_args, dict):
+            coerced_calls.append(tool_call)
+            continue
+
+        coerced_calls.append(
+            tool_call.model_copy(update={"arguments": json.dumps(coerced_args)})
+        )
+    return coerced_calls
 
 
 def make_mlx_parser(
@@ -33,7 +204,7 @@ def make_mlx_parser(
     return ToolParser(
         start_parsing=tool_call_start,
         end_parsing=tool_call_end,
-        parse_tool_calls=parse_tool_calls,
+        _inner_parser=parse_tool_calls,
     )
 
 
@@ -62,7 +233,7 @@ def make_json_parser() -> ToolParser:
     return ToolParser(
         start_parsing="<tool_call>",
         end_parsing="</tool_call>",
-        parse_tool_calls=_parse_json_calls,
+        _inner_parser=_parse_json_calls,
     )
 
 

src/exo/worker/runner/runner_supervisor.py

@@ -12,13 +12,22 @@ from anyio import (
 )
 from loguru import logger
 
+from exo.shared.types.chunks import ErrorChunk
 from exo.shared.types.events import (
+    ChunkGenerated,
     Event,
     RunnerStatusUpdated,
     TaskAcknowledged,
     TaskStatusUpdated,
 )
-from exo.shared.types.tasks import Task, TaskId, TaskStatus
+from exo.shared.types.tasks import (
+    ImageEdits,
+    ImageGeneration,
+    Task,
+    TaskId,
+    TaskStatus,
+    TextGeneration,
+)
 from exo.shared.types.worker.instances import BoundInstance
 from exo.shared.types.worker.runners import (
     RunnerConnecting,
@@ -52,6 +61,7 @@ class RunnerSupervisor:
     _tg: TaskGroup = field(default_factory=TaskGroup, init=False)
     status: RunnerStatus = field(default_factory=RunnerIdle, init=False)
     pending: dict[TaskId, anyio.Event] = field(default_factory=dict, init=False)
+    in_progress: dict[TaskId, Task] = field(default_factory=dict, init=False)
     completed: set[TaskId] = field(default_factory=set, init=False)
     cancelled: set[TaskId] = field(default_factory=set, init=False)
     _cancel_watch_runner: anyio.CancelScope = field(
@@ -147,9 +157,11 @@ class RunnerSupervisor:
         logger.info(f"Starting task {task}")
         event = anyio.Event()
         self.pending[task.task_id] = event
+        self.in_progress[task.task_id] = task
         try:
             await self._task_sender.send_async(task)
         except ClosedResourceError:
+            self.in_progress.pop(task.task_id, None)
             logger.warning(f"Task {task} dropped, runner closed communication.")
             return
         await event.wait()
@@ -157,10 +169,17 @@ class RunnerSupervisor:
     async def cancel_task(self, task_id: TaskId):
         if task_id in self.completed:
             logger.info(f"Unable to cancel {task_id} as it has been completed")
+            self.cancelled.add(task_id)
             return
         self.cancelled.add(task_id)
         with anyio.move_on_after(0.5) as scope:
-            await self._cancel_sender.send_async(task_id)
+            try:
+                await self._cancel_sender.send_async(task_id)
+            except ClosedResourceError:
+                # typically occurs when trying to shut down a failed instance
+                logger.warning(
+                    f"Cancelling task {task_id} failed, runner closed communication"
+                )
         if scope.cancel_called:
             logger.error("RunnerSupervisor cancel pipe blocked")
             await self._check_runner(TimeoutError("cancel pipe blocked"))
@@ -189,6 +208,7 @@ class RunnerSupervisor:
                                 RunnerShuttingDown,
                             ),
                         )
+                        self.in_progress.pop(event.task_id, None)
                         self.completed.add(event.task_id)
                     await self._event_sender.send(event)
         except (ClosedResourceError, BrokenResourceError) as e:
@@ -233,6 +253,22 @@ class RunnerSupervisor:
 
         logger.opt(exception=e).error(f"Runner terminated with {cause}")
 
+        for task in self.in_progress.values():
+            if isinstance(task, (TextGeneration, ImageGeneration, ImageEdits)):
+                with anyio.CancelScope(shield=True):
+                    await self._event_sender.send(
+                        ChunkGenerated(
+                            command_id=task.command_id,
+                            chunk=ErrorChunk(
+                                model=self.shard_metadata.model_card.model_id,
+                                error_message=(
+                                    "Runner shutdown before completing command "
+                                    f"({cause})"
+                                ),
+                            ),
+                        )
+                    )
+
         try:
             self.status = RunnerFailed(error_message=f"Terminated ({cause})")
             with anyio.CancelScope(shield=True):

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

@@ -20,6 +20,8 @@ class FakeRunnerSupervisor:
     bound_instance: BoundInstance
     status: RunnerStatus
     completed: set[TaskId] = field(default_factory=set)
+    in_progress: set[TaskId] = field(default_factory=set)
+    pending: dict[TaskId, object] = field(default_factory=dict)
 
 
 class OtherTask(BaseTask):

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

@@ -19,7 +19,7 @@ from exo.worker.engines.mlx.dsml_encoding import (
     encode_messages,
     parse_dsml_output,
 )
-from exo.worker.runner.llm_inference.runner import parse_deepseek_v32
+from exo.worker.runner.llm_inference.model_output_parsers import parse_deepseek_v32
 
 # ── Shared fixtures ──────────────────────────────────────────────
 

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

@@ -6,6 +6,8 @@ from typing import Callable
 import mlx.core as mx
 import pytest
 
+import exo.worker.runner.llm_inference.batch_generator as mlx_batch_generator
+import exo.worker.runner.llm_inference.model_output_parsers as mlx_model_output_parsers
 import exo.worker.runner.llm_inference.runner as mlx_runner
 from exo.shared.types.chunks import TokenChunk
 from exo.shared.types.events import (
@@ -114,27 +116,41 @@ def patch_out_mlx(monkeypatch: pytest.MonkeyPatch):
     # initialize_mlx returns a mock group
     monkeypatch.setattr(mlx_runner, "initialize_mlx", make_nothin(MockGroup()))
     monkeypatch.setattr(mlx_runner, "load_mlx_items", make_nothin((1, MockTokenizer)))
-    monkeypatch.setattr(mlx_runner, "warmup_inference", make_nothin(1))
-    monkeypatch.setattr(mlx_runner, "_check_for_debug_prompts", nothin)
-    monkeypatch.setattr(mlx_runner, "mx_any", make_nothin(False))
+    monkeypatch.setattr(mlx_batch_generator, "warmup_inference", make_nothin(1))
+    monkeypatch.setattr(mlx_batch_generator, "_check_for_debug_prompts", nothin)
+    monkeypatch.setattr(mlx_batch_generator, "mx_any", make_nothin(False))
+
+    def fake_all_gather(
+        tasks: list[TextGeneration], group: object
+    ) -> tuple[list[TextGeneration], list[TextGeneration]]:
+        return (tasks, [])
+
+    monkeypatch.setattr(mlx_batch_generator, "mx_all_gather_tasks", fake_all_gather)
     # Mock apply_chat_template since we're using a fake tokenizer (integer 1).
     # Returns a prompt without thinking tag so detect_thinking_prompt_suffix returns None.
-    monkeypatch.setattr(mlx_runner, "apply_chat_template", make_nothin("test prompt"))
-    monkeypatch.setattr(mlx_runner, "detect_thinking_prompt_suffix", make_nothin(False))
+    monkeypatch.setattr(
+        mlx_batch_generator, "apply_chat_template", make_nothin("test prompt")
+    )
+    monkeypatch.setattr(
+        mlx_model_output_parsers, "detect_thinking_prompt_suffix", make_nothin(False)
+    )
 
     def fake_generate(*_1: object, **_2: object):
         yield GenerationResponse(token=0, text="hi", finish_reason="stop", usage=None)
 
-    monkeypatch.setattr(mlx_runner, "mlx_generate", fake_generate)
+    monkeypatch.setattr(mlx_batch_generator, "mlx_generate", fake_generate)
 
 
 # Use a fake event_sender to remove test flakiness.
 class EventCollector:
-    def __init__(self) -> None:
+    def __init__(self, on_event: Callable[[Event], None] | None = None) -> None:
         self.events: list[Event] = []
+        self._on_event = on_event
 
     def send(self, event: Event) -> None:
         self.events.append(event)
+        if self._on_event:
+            self._on_event(event)
 
     def close(self) -> None:
         pass
@@ -159,7 +175,7 @@ class MockGroup:
         return 1
 
 
-def _run(tasks: Iterable[Task]):
+def _run(tasks: Iterable[Task], send_after_ready: list[Task] | None = None):
     bound_instance = get_bound_mlx_ring_instance(
         instance_id=INSTANCE_1_ID,
         model_id=MODEL_A_ID,
@@ -169,7 +185,23 @@ def _run(tasks: Iterable[Task]):
 
     task_sender, task_receiver = mp_channel[Task]()
     _cancel_sender, cancel_receiver = mp_channel[TaskId]()
-    event_sender = EventCollector()
+
+    on_event: Callable[[Event], None] | None = None
+    if send_after_ready:
+        _saw_running = False
+
+        def _on_event(event: Event) -> None:
+            nonlocal _saw_running
+            if isinstance(event, RunnerStatusUpdated):
+                if isinstance(event.runner_status, RunnerRunning):
+                    _saw_running = True
+                elif _saw_running and isinstance(event.runner_status, RunnerReady):
+                    for t in send_after_ready:
+                        task_sender.send(t)
+
+        on_event = _on_event
+
+    event_sender = EventCollector(on_event=on_event)
 
     with task_sender:
         for t in tasks:
@@ -183,18 +215,22 @@ def _run(tasks: Iterable[Task]):
             "exo.worker.runner.llm_inference.runner.mx.distributed.all_gather",
             make_nothin(mx.array([1])),
         ):
-            mlx_runner.main(
+            runner = mlx_runner.Runner(
                 bound_instance,
                 event_sender,  # pyright: ignore[reportArgumentType]
                 task_receiver,
                 cancel_receiver,
             )
+            runner.main()
 
         return event_sender.events
 
 
 def test_events_processed_in_correct_order(patch_out_mlx: pytest.MonkeyPatch):
-    events = _run([INIT_TASK, LOAD_TASK, WARMUP_TASK, CHAT_TASK, SHUTDOWN_TASK])
+    events = _run(
+        [INIT_TASK, LOAD_TASK, WARMUP_TASK, CHAT_TASK],
+        send_after_ready=[SHUTDOWN_TASK],
+    )
 
     expected_chunk = ChunkGenerated(
         command_id=COMMAND_1_ID,

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

@@ -4,7 +4,7 @@ from exo.shared.types.worker.runner_response import (
     GenerationResponse,
     ToolCallResponse,
 )
-from exo.worker.runner.llm_inference.runner import parse_gpt_oss
+from exo.worker.runner.llm_inference.model_output_parsers import parse_gpt_oss
 
 # Token IDs from mlx-community/gpt-oss-20b-MXFP4-Q8 tokenizer.
 # These are stable since they come from the model's vocabulary.
@@ -107,7 +107,7 @@ def _collect(
     def _gen() -> Generator[GenerationResponse, None, None]:
         yield from _make_gen_responses(tokens)
 
-    return list(parse_gpt_oss(_gen()))
+    return list(x for x in parse_gpt_oss(_gen()) if x is not None)
 
 
 def _get_tool_call(

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

@@ -1,10 +1,11 @@
 """Tests for parse_tool_calls generator, especially unclosed tool call handling."""
 
+import json
 from collections.abc import Generator
 from typing import Any
 
 from exo.shared.types.worker.runner_response import GenerationResponse, ToolCallResponse
-from exo.worker.runner.llm_inference.runner import parse_tool_calls
+from exo.worker.runner.llm_inference.model_output_parsers import parse_tool_calls
 from exo.worker.runner.llm_inference.tool_parsers import make_mlx_parser
 
 
@@ -40,6 +41,7 @@ class TestParseToolCalls:
             parse_tool_calls(
                 _make_responses(texts, finish_on_last=False),
                 _dummy_parser,
+                tools=None,
             )
         )
 
@@ -53,6 +55,7 @@ class TestParseToolCalls:
             parse_tool_calls(
                 _make_responses(texts),
                 _dummy_parser,
+                tools=None,
             )
         )
 
@@ -77,9 +80,101 @@ class TestParseToolCalls:
             parse_tool_calls(
                 _make_responses(texts, finish_on_last=False),
                 make_mlx_parser("<tool_call>", "</tool_call>", _failing_parser),
+                tools=None,
             )
         )
 
         assert len(results) == 1
         assert isinstance(results[0], GenerationResponse)
         assert results[0].text == "<tool_call>bad content</tool_call>"
+
+    def test_tool_schema_coerces_string_arguments_to_expected_types(self):
+        """Tool argument values should be coerced using provided JSON schema."""
+
+        def _parser_with_string_args(_text: str) -> dict[str, Any]:
+            return {
+                "name": "process",
+                "arguments": {
+                    "action": "output",
+                    "id": "0",
+                    "verbose": "true",
+                    "temperature": "0.75",
+                },
+            }
+
+        tools = [
+            {
+                "type": "function",
+                "function": {
+                    "name": "process",
+                    "description": "Manage background processes",
+                    "parameters": {
+                        "type": "object",
+                        "properties": {
+                            "action": {"type": "string"},
+                            "id": {"type": "integer"},
+                            "verbose": {"type": "boolean"},
+                            "temperature": {"type": "number"},
+                        },
+                        "required": ["action"],
+                    },
+                },
+            }
+        ]
+
+        results = list(
+            parse_tool_calls(
+                _make_responses(["<tool_call>", "process", "</tool_call>"]),
+                make_mlx_parser(
+                    "<tool_call>", "</tool_call>", _parser_with_string_args
+                ),
+                tools,
+            )
+        )
+
+        assert len(results) == 1
+        assert isinstance(results[0], ToolCallResponse)
+
+        args = json.loads(results[0].tool_calls[0].arguments)  # pyright: ignore[reportAny]
+        assert args == {
+            "action": "output",
+            "id": 0,
+            "verbose": True,
+            "temperature": 0.75,
+        }
+
+    def test_schema_coercion_skips_unknown_tools(self):
+        """If no matching tool schema exists, arguments should remain unchanged."""
+
+        def _parser_with_string_id(_text: str) -> dict[str, Any]:
+            return {
+                "name": "process",
+                "arguments": {"action": "output", "id": "0"},
+            }
+
+        tools = [
+            {
+                "type": "function",
+                "function": {
+                    "name": "different_tool",
+                    "parameters": {
+                        "type": "object",
+                        "properties": {"id": {"type": "integer"}},
+                    },
+                },
+            }
+        ]
+
+        results = list(
+            parse_tool_calls(
+                _make_responses(["<tool_call>", "process", "</tool_call>"]),
+                make_mlx_parser("<tool_call>", "</tool_call>", _parser_with_string_id),
+                tools,
+            )
+        )
+
+        assert len(results) == 1
+        assert isinstance(results[0], ToolCallResponse)
+
+        args = json.loads(results[0].tool_calls[0].arguments)  # pyright: ignore[reportAny]
+        assert args == {"action": "output", "id": "0"}

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

@@ -1 +1,93 @@
-# TODO:
+import multiprocessing as mp
+from typing import cast
+
+import anyio
+import pytest
+
+from exo.shared.models.model_cards import ModelId
+from exo.shared.types.chunks import ErrorChunk
+from exo.shared.types.common import CommandId, NodeId
+from exo.shared.types.events import ChunkGenerated, Event, RunnerStatusUpdated
+from exo.shared.types.tasks import Task, TaskId, TextGeneration
+from exo.shared.types.text_generation import InputMessage, TextGenerationTaskParams
+from exo.shared.types.worker.instances import BoundInstance, InstanceId
+from exo.shared.types.worker.runners import RunnerFailed, RunnerId
+from exo.utils.channels import channel, mp_channel
+from exo.worker.runner.runner_supervisor import RunnerSupervisor
+from exo.worker.tests.unittests.conftest import get_bound_mlx_ring_instance
+
+
+class _DeadProcess:
+    exitcode = -6
+
+    def start(self) -> None:
+        return None
+
+    def is_alive(self) -> bool:
+        return False
+
+    def join(self, _timeout: float | None = None) -> None:
+        return None
+
+    def terminate(self) -> None:
+        return None
+
+    def kill(self) -> None:
+        return None
+
+
+@pytest.mark.asyncio
+async def test_check_runner_emits_error_chunk_for_inflight_text_generation() -> None:
+    event_sender, event_receiver = channel[Event]()
+    task_sender, _ = mp_channel[Task]()
+    cancel_sender, _ = mp_channel[TaskId]()
+    _, ev_recv = mp_channel[Event]()
+
+    bound_instance: BoundInstance = get_bound_mlx_ring_instance(
+        instance_id=InstanceId("instance-a"),
+        model_id=ModelId("mlx-community/Llama-3.2-1B-Instruct-4bit"),
+        runner_id=RunnerId("runner-a"),
+        node_id=NodeId("node-a"),
+    )
+
+    supervisor = RunnerSupervisor(
+        shard_metadata=bound_instance.bound_shard,
+        bound_instance=bound_instance,
+        runner_process=cast("mp.Process", cast(object, _DeadProcess())),
+        initialize_timeout=400,
+        _ev_recv=ev_recv,
+        _task_sender=task_sender,
+        _event_sender=event_sender,
+        _cancel_sender=cancel_sender,
+    )
+
+    command_id = CommandId("cmd-a")
+    task = TextGeneration(
+        task_id=TaskId("task-a"),
+        instance_id=bound_instance.instance.instance_id,
+        command_id=command_id,
+        task_params=TextGenerationTaskParams(
+            model=bound_instance.bound_shard.model_card.model_id,
+            input=[InputMessage(role="user", content="hi")],
+            stream=True,
+        ),
+    )
+    supervisor.in_progress[task.task_id] = task
+    supervisor.shutdown = lambda: None
+
+    await supervisor._check_runner(RuntimeError("boom"))  # pyright: ignore[reportPrivateUsage]
+
+    got_chunk = await event_receiver.receive()
+    got_status = await event_receiver.receive()
+
+    assert isinstance(got_chunk, ChunkGenerated)
+    assert got_chunk.command_id == command_id
+    assert isinstance(got_chunk.chunk, ErrorChunk)
+    assert "Runner shutdown before completing command" in got_chunk.chunk.error_message
+
+    assert isinstance(got_status, RunnerStatusUpdated)
+    assert isinstance(got_status.runner_status, RunnerFailed)
+
+    event_sender.close()
+    with anyio.move_on_after(0.1):
+        await event_receiver.aclose()

tmp/gen_card.py

@@ -0,0 +1,33 @@
+"""
+Generates inference model cards for EXO.
+Usage:
+    uv run tmp/gen_card.py mlx-community/my_cool_model-8bit [repo-id/model-id-2] [...]
+
+Model Cards require cleanup for family & quantization data
+"""
+
+import sys
+
+import anyio
+
+from exo.shared.models.model_cards import ModelCard, ModelId
+
+
+async def main():
+    if len(sys.argv) == 1:
+        print(f"USAGE: {sys.argv[0]} repo-id/model-id-1 [repo-id/model-id-2] [...]")
+        quit(1)
+    print("Remember! Model Cards require cleanup for family & quantization data")
+    for arg in sys.argv[1:]:
+        mid = ModelId(arg)
+        mc = await ModelCard.fetch_from_hf(mid)
+        await mc.save(
+            anyio.Path(__file__).parent.parent
+            / "resources"
+            / "inference_model_cards"
+            / (mid.normalize() + ".toml")
+        )
+
+
+if __name__ == "__main__":
+    anyio.run(main)

uv.lock

@@ -418,7 +418,7 @@ requires-dist = [
     { name = "mflux", specifier = "==0.15.5" },
     { name = "mlx", marker = "sys_platform == 'darwin'", git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks" },
     { name = "mlx", extras = ["cpu"], marker = "sys_platform == 'linux'", specifier = "==0.30.6" },
-    { name = "mlx-lm", specifier = "==0.30.7" },
+    { name = "mlx-lm", git = "https://github.com/ml-explore/mlx-lm?rev=834fac934c4e04de9b3d723e2b9287a2c60cfd4a" },
     { name = "msgspec", specifier = ">=0.19.0" },
     { name = "openai-harmony", specifier = ">=0.0.8" },
     { name = "pillow", specifier = ">=11.0,<12.0" },
@@ -1104,8 +1104,8 @@ wheels = [
 
 [[package]]
 name = "mlx-lm"
-version = "0.30.7"
-source = { registry = "https://pypi.org/simple" }
+version = "0.30.8"
+source = { git = "https://github.com/ml-explore/mlx-lm?rev=834fac934c4e04de9b3d723e2b9287a2c60cfd4a#834fac934c4e04de9b3d723e2b9287a2c60cfd4a" }
 dependencies = [
     { name = "jinja2", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
     { name = "mlx", version = "0.30.7.dev20260225+257d5692", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#257d5692fc7af6bba3b8afaeb63c549b7d1e43d5" }, marker = "sys_platform == 'darwin'" },
@@ -1115,10 +1115,6 @@ dependencies = [
     { name = "sentencepiece", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
     { name = "transformers", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
 ]
-sdist = { url = "https://files.pythonhosted.org/packages/66/0d/56542e2ae13ec6f542d3977d7cff89a205d4f6c5122e0ce23f33265f61c9/mlx_lm-0.30.7.tar.gz", hash = "sha256:e5f31ac58d9f2381f28e1ba639ff903e64f7cff1bdc245c0bc97f72264be329c", size = 275764, upload-time = "2026-02-12T18:41:11.86Z" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/1e/17/a41c798a3d9cbdc47f39c6db5bba4c2cd199203ead26bf911cb03b644070/mlx_lm-0.30.7-py3-none-any.whl", hash = "sha256:17442a4bf01c4c2d3bca1e647712fe44f19890c3f1eadc8589d389e57b44b9bf", size = 386591, upload-time = "2026-02-12T18:41:10.236Z" },
-]
 
 [[package]]
 name = "more-itertools"