feat: auto-discover models from HuggingFace Hub cache and custom paths

Searches ~/.cache/huggingface/hub by default so models downloaded by
other tools (HF CLI, transformers, LM Studio, etc.) are used without
re-downloading. Adds EXO_MODELS_PATH env var support and a Model Search
Paths setting in the macOS app.

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

.mlx_typings/mlx/core/__init__.pyi

@@ -215,22 +215,6 @@ class StreamContext:
         traceback: object | None = ...,
     ) -> None: ...
 
-def device_info() -> dict[str, str | int]:
-    """
-    Get information about the GPU device and system settings.
-
-    Currently returns:
-
-    * ``architecture``
-    * ``max_buffer_size``
-    * ``max_recommended_working_set_size``
-    * ``memory_size``
-    * ``resource_limit``
-
-    Returns:
-        dict: A dictionary with string keys and string or integer values.
-    """
-
 def abs(a: array, /, *, stream: Stream | Device | None = ...) -> array:
     """
     Element-wise absolute value.

.mlx_typings/mlx_lm/generate.pyi

@@ -73,11 +73,9 @@ class GenerationResponse:
     finish_reason: Optional[str] = ...
 
 def maybe_quantize_kv_cache(
-    prompt_cache: Any,
-    quantized_kv_start: int | None,
-    kv_group_size: int | None,
-    kv_bits: int | None,
-) -> None: ...
+    prompt_cache, quantized_kv_start, kv_group_size, kv_bits
+):  # -> None:
+    ...
 def generate_step(
     prompt: mx.array,
     model: nn.Module,

.mlx_typings/mlx_lm/models/cache.pyi

@@ -16,7 +16,7 @@ class Cache(Protocol):
         self, keys: mx.array, values: mx.array
     ) -> tuple[mx.array, mx.array]: ...
     @property
-    def state(self) -> tuple[mx.array | None, mx.array | None]: ...
+    def state(self) -> tuple[mx.array, mx.array]: ...
     @state.setter
     def state(self, v) -> None: ...
 
@@ -92,14 +92,13 @@ class _BaseCache(Cache):
     values: mx.array
     offset: int
     @property
-    def state(self) -> tuple[mx.array | None, mx.array | None]: ...
+    def state(self) -> tuple[mx.array, mx.array]: ...
     @state.setter
     def state(self, v) -> None: ...
     @property
     def meta_state(self) -> Literal[""]: ...
     @meta_state.setter
     def meta_state(self, v) -> None: ...
-    def trim(self, n: int) -> int: ...
     def is_trimmable(self) -> Literal[False]: ...
     @classmethod
     def from_state(cls, state, meta_state) -> Self: ...
@@ -115,13 +114,15 @@ class ConcatenateKVCache(_BaseCache):
     def update_and_fetch(self, keys, values):  # -> tuple[Any | array, Any | array]:
         ...
     @property
-    def state(self) -> tuple[mx.array | None, mx.array | None]: ...
+    def state(self):  # -> tuple[Any | array | None, Any | array | None]:
+        ...
     @state.setter
     def state(self, v):  # -> None:
         ...
     def is_trimmable(self):  # -> Literal[True]:
         ...
-    def trim(self, n: int) -> int: ...
+    def trim(self, n):  # -> int:
+        ...
     def make_mask(self, *args, **kwargs):  # -> array | Literal['causal'] | None:
         ...
 
@@ -131,7 +132,10 @@ class QuantizedKVCache(_BaseCache):
     def update_and_fetch(self, keys, values):  # -> Any:
         ...
     @property
-    def state(self) -> tuple[mx.array | None, mx.array | None]: ...
+    def state(
+        self,
+    ):  # -> tuple[Any | tuple[array, array, array] | None, Any | tuple[array, array, array] | None] | Any:
+        ...
     @state.setter
     def state(self, v):  # -> None:
         ...
@@ -143,7 +147,8 @@ class QuantizedKVCache(_BaseCache):
         ...
     def is_trimmable(self):  # -> Literal[True]:
         ...
-    def trim(self, n: int) -> int: ...
+    def trim(self, n):  # -> int:
+        ...
     def make_mask(self, *args, **kwargs):  # -> array | Literal['causal'] | None:
         ...
 
@@ -155,18 +160,18 @@ class KVCache(_BaseCache):
     @property
     def state(
         self,
-    ) -> tuple[mx.array | None, mx.array | None]: ...
+    ) -> tuple[array, array]: ...
     @state.setter
     def state(self, v) -> None: ...
     def is_trimmable(self):  # -> Literal[True]:
         ...
-    def trim(self, n: int) -> int: ...
+    def trim(self, n):  # -> int:
+        ...
     def to_quantized(
         self, group_size: int = ..., bits: int = ...
     ) -> QuantizedKVCache: ...
-    def make_mask(
-        self, *args: Any, **kwargs: Any
-    ) -> mx.array | Literal["causal"] | None: ...
+    def make_mask(self, *args, **kwargs):  # -> array | Literal['causal'] | None:
+        ...
 
 class RotatingKVCache(_BaseCache):
     step = ...
@@ -178,7 +183,8 @@ class RotatingKVCache(_BaseCache):
     @property
     def state(
         self,
-    ) -> tuple[mx.array | None, mx.array | None]: ...
+    ):  # -> tuple[Any | array, Any | array] | tuple[Any | array | None, Any | array | None]:
+        ...
     @state.setter
     def state(self, v):  # -> None:
         ...
@@ -190,7 +196,8 @@ class RotatingKVCache(_BaseCache):
         ...
     def is_trimmable(self):  # -> bool:
         ...
-    def trim(self, n: int) -> int: ...
+    def trim(self, n):  # -> int:
+        ...
     def to_quantized(
         self, group_size: int = ..., bits: int = ...
     ) -> QuantizedKVCache: ...
@@ -205,7 +212,8 @@ class ArraysCache(_BaseCache):
         ...
     def __getitem__(self, idx): ...
     @property
-    def state(self) -> tuple[mx.array | None, mx.array | None]: ...
+    def state(self):  # -> list[Any | array] | list[array]:
+        ...
     @state.setter
     def state(self, v):  # -> None:
         ...
@@ -219,7 +227,8 @@ class ArraysCache(_BaseCache):
         In-place extend this cache with the other cache.
         """
 
-    def make_mask(self, N: int) -> mx.array | None: ...
+    def make_mask(self, N: int):  # -> array | None:
+        ...
 
 class MambaCache(ArraysCache):
     def __init__(self, left_padding: Optional[List[int]] = ...) -> None: ...
@@ -230,7 +239,8 @@ class ChunkedKVCache(KVCache):
         ...
     def update_and_fetch(self, keys, values):  # -> tuple[array, array]:
         ...
-    def trim(self, n: int) -> int: ...
+    def trim(self, n):  # -> int:
+        ...
     @property
     def meta_state(self):  # -> tuple[str, ...]:
         ...
@@ -243,9 +253,10 @@ class CacheList(_BaseCache):
     def __getitem__(self, idx): ...
     def is_trimmable(self):  # -> bool:
         ...
-    def trim(self, n: int) -> int: ...
+    def trim(self, n): ...
     @property
-    def state(self) -> list[tuple[mx.array | None, mx.array | None]]: ...
+    def state(self):  # -> list[Any]:
+        ...
     @state.setter
     def state(self, v):  # -> None:
         ...

.mlx_typings/mlx_lm/models/qwen3_5.pyi

@@ -1,153 +0,0 @@
-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

@@ -1,19 +0,0 @@
-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,15 +7,6 @@ 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.toml

@@ -3,6 +3,7 @@ resolver = "3"
 members = [
     "rust/networking",
     "rust/exo_pyo3_bindings",
+    "rust/util",
 ]
 
 [workspace.package]
@@ -23,6 +24,7 @@ opt-level = 3
 [workspace.dependencies]
 ## Crate members as common dependencies
 networking = { path = "rust/networking" }
+util = { path = "rust/util" }
 
 # Macro dependecies
 extend = "1.2"
@@ -32,7 +34,6 @@ delegate = "0.13"
 keccak-const = "0.2"
 
 # Async dependencies
-async-stream = "0.3"
 tokio = "1.46"
 futures-lite = "2.6.1"
 futures-timer = "3.0"

app/EXO/EXO/ContentView.swift

@@ -21,21 +21,6 @@ struct ContentView: View {
     @State private var showAllNodes = false
     @State private var showAllInstances = false
     @State private var baseURLCopied = false
-    @State private var showAdvanced = false
-    @State private var showDebugInfo = false
-    private enum BugReportPhase: Equatable {
-        case idle
-        case prompting
-        case sending(String)
-        case success(String)
-        case failure(String)
-    }
-    @State private var bugReportPhase: BugReportPhase = .idle
-    @State private var bugReportUserDescription: String = ""
-    @State private var uninstallInProgress = false
-    @State private var pendingNamespace: String = ""
-    @State private var pendingHFToken: String = ""
-    @State private var pendingEnableImageModels = false
 
     var body: some View {
         VStack(alignment: .leading, spacing: 12) {
@@ -394,283 +379,6 @@ struct ContentView: View {
         }
     }
 
-    private var thunderboltStatusText: String {
-        switch networkStatusService.status.thunderboltBridgeState {
-        case .some(.disabled):
-            return "Thunderbolt Bridge: Disabled"
-        case .some(.deleted):
-            return "Thunderbolt Bridge: Deleted"
-        case .some(.enabled):
-            return "Thunderbolt Bridge: Enabled"
-        case nil:
-            return "Thunderbolt Bridge: Unknown"
-        }
-    }
-
-    private var thunderboltStatusColor: Color {
-        switch networkStatusService.status.thunderboltBridgeState {
-        case .some(.disabled), .some(.deleted):
-            return .green
-        case .some(.enabled):
-            return .red
-        case nil:
-            return .secondary
-        }
-    }
-
-    /// Shows TB bridge status for all nodes from exo cluster state
-    private var clusterThunderboltBridgeView: some View {
-        let bridgeStatuses = stateService.latestSnapshot?.nodeThunderboltBridge ?? [:]
-        let localNodeId = stateService.localNodeId
-        let nodeProfiles = stateService.latestSnapshot?.nodeProfiles ?? [:]
-
-        return VStack(alignment: .leading, spacing: 1) {
-            if bridgeStatuses.isEmpty {
-                Text("Cluster TB Bridge: No data")
-                    .font(.caption2)
-                    .foregroundColor(.secondary)
-            } else {
-                Text("Cluster TB Bridge Status:")
-                    .font(.caption2)
-                    .foregroundColor(.secondary)
-                ForEach(Array(bridgeStatuses.keys.sorted()), id: \.self) { nodeId in
-                    if let status = bridgeStatuses[nodeId] {
-                        let nodeName =
-                            nodeProfiles[nodeId]?.friendlyName ?? String(nodeId.prefix(8))
-                        let isLocal = nodeId == localNodeId
-                        let prefix = isLocal ? "  \(nodeName) (local):" : "  \(nodeName):"
-                        let statusText =
-                            !status.exists
-                            ? "N/A"
-                            : (status.enabled ? "Enabled" : "Disabled")
-                        let color: Color =
-                            !status.exists
-                            ? .secondary
-                            : (status.enabled ? .red : .green)
-                        Text("\(prefix) \(statusText)")
-                            .font(.caption2)
-                            .foregroundColor(color)
-                    }
-                }
-            }
-        }
-    }
-
-    private var interfaceIpList: some View {
-        let statuses = networkStatusService.status.interfaceStatuses
-        return VStack(alignment: .leading, spacing: 1) {
-            Text("Interfaces (en0–en7):")
-                .font(.caption2)
-                .foregroundColor(.secondary)
-            if statuses.isEmpty {
-                Text("  Unknown")
-                    .font(.caption2)
-                    .foregroundColor(.secondary)
-            } else {
-                ForEach(statuses, id: \.interfaceName) { status in
-                    let ipText = status.ipAddress ?? "No IP"
-                    Text("  \(status.interfaceName): \(ipText)")
-                        .font(.caption2)
-                        .foregroundColor(status.ipAddress == nil ? .red : .green)
-                }
-            }
-        }
-    }
-
-    private var debugSection: some View {
-        VStack(alignment: .leading, spacing: 4) {
-            HoverButton(
-                title: "Debug Info",
-                tint: .primary,
-                trailingSystemImage: showDebugInfo ? "chevron.up" : "chevron.down",
-                small: true
-            ) {
-                showDebugInfo.toggle()
-            }
-            if showDebugInfo {
-                VStack(alignment: .leading, spacing: 4) {
-                    Text("Version: \(buildTag)")
-                        .font(.caption2)
-                        .foregroundColor(.secondary)
-                    Text("Commit: \(buildCommit)")
-                        .font(.caption2)
-                        .foregroundColor(.secondary)
-                    Text(thunderboltStatusText)
-                        .font(.caption2)
-                        .foregroundColor(thunderboltStatusColor)
-                    clusterThunderboltBridgeView
-                    interfaceIpList
-                    rdmaStatusView
-                    sendBugReportButton
-                        .padding(.top, 6)
-                }
-                .padding(.leading, 8)
-                .transition(.opacity)
-            }
-        }
-        .animation(.easeInOut(duration: 0.25), value: showDebugInfo)
-    }
-
-    private var rdmaStatusView: some View {
-        let rdmaStatuses = stateService.latestSnapshot?.nodeRdmaCtl ?? [:]
-        let localNodeId = stateService.localNodeId
-        let nodeProfiles = stateService.latestSnapshot?.nodeProfiles ?? [:]
-        let localDevices = networkStatusService.status.localRdmaDevices
-        let localPorts = networkStatusService.status.localRdmaActivePorts
-
-        return VStack(alignment: .leading, spacing: 1) {
-            if rdmaStatuses.isEmpty {
-                Text("Cluster RDMA: No data")
-                    .font(.caption2)
-                    .foregroundColor(.secondary)
-            } else {
-                Text("Cluster RDMA Status:")
-                    .font(.caption2)
-                    .foregroundColor(.secondary)
-                ForEach(Array(rdmaStatuses.keys.sorted()), id: \.self) { nodeId in
-                    if let status = rdmaStatuses[nodeId] {
-                        let nodeName =
-                            nodeProfiles[nodeId]?.friendlyName ?? String(nodeId.prefix(8))
-                        let isLocal = nodeId == localNodeId
-                        let prefix = isLocal ? "  \(nodeName) (local):" : "  \(nodeName):"
-                        let statusText = status.enabled ? "Enabled" : "Disabled"
-                        let color: Color = status.enabled ? .green : .orange
-                        Text("\(prefix) \(statusText)")
-                            .font(.caption2)
-                            .foregroundColor(color)
-                    }
-                }
-            }
-            if !localDevices.isEmpty {
-                Text("  Local Devices: \(localDevices.joined(separator: ", "))")
-                    .font(.caption2)
-                    .foregroundColor(.secondary)
-            }
-            if !localPorts.isEmpty {
-                Text("  Local Active Ports:")
-                    .font(.caption2)
-                    .foregroundColor(.secondary)
-                ForEach(localPorts, id: \.device) { port in
-                    Text("    \(port.device) port \(port.port): \(port.state)")
-                        .font(.caption2)
-                        .foregroundColor(.green)
-                }
-            }
-        }
-    }
-
-    private var sendBugReportButton: some View {
-        VStack(alignment: .leading, spacing: 6) {
-            switch bugReportPhase {
-            case .idle:
-                Button {
-                    bugReportPhase = .prompting
-                    bugReportUserDescription = ""
-                } label: {
-                    HStack {
-                        Text("Send Bug Report")
-                            .font(.caption)
-                            .fontWeight(.semibold)
-                        Spacer()
-                    }
-                    .padding(.vertical, 6)
-                    .padding(.horizontal, 8)
-                    .background(
-                        RoundedRectangle(cornerRadius: 6)
-                            .fill(Color.accentColor.opacity(0.12))
-                    )
-                }
-                .buttonStyle(.plain)
-
-            case .prompting:
-                VStack(alignment: .leading, spacing: 6) {
-                    Text("What's the issue? (optional)")
-                        .font(.caption2)
-                        .foregroundColor(.secondary)
-                    TextEditor(text: $bugReportUserDescription)
-                        .font(.caption2)
-                        .frame(height: 60)
-                        .overlay(
-                            RoundedRectangle(cornerRadius: 4)
-                                .stroke(Color.secondary.opacity(0.3), lineWidth: 1)
-                        )
-                    HStack(spacing: 8) {
-                        Button("Send") {
-                            Task {
-                                await sendBugReport()
-                            }
-                        }
-                        .font(.caption2)
-                        .buttonStyle(.borderedProminent)
-                        .controlSize(.small)
-                        Button("Cancel") {
-                            bugReportPhase = .idle
-                        }
-                        .font(.caption2)
-                        .buttonStyle(.bordered)
-                        .controlSize(.small)
-                    }
-                }
-                .padding(8)
-                .background(
-                    RoundedRectangle(cornerRadius: 6)
-                        .fill(Color.accentColor.opacity(0.06))
-                )
-
-            case .sending(let message):
-                HStack(spacing: 6) {
-                    ProgressView()
-                        .scaleEffect(0.6)
-                    Text(message)
-                        .font(.caption2)
-                        .foregroundColor(.secondary)
-                }
-
-            case .success(let message):
-                VStack(alignment: .leading, spacing: 6) {
-                    Text(message)
-                        .font(.caption2)
-                        .foregroundColor(.secondary)
-                        .fixedSize(horizontal: false, vertical: true)
-                    Button {
-                        openGitHubIssue()
-                    } label: {
-                        HStack(spacing: 4) {
-                            Image(systemName: "arrow.up.right.square")
-                                .imageScale(.small)
-                            Text("Create GitHub Issue")
-                                .font(.caption2)
-                        }
-                    }
-                    .buttonStyle(.bordered)
-                    .controlSize(.small)
-                    Button("Done") {
-                        bugReportPhase = .idle
-                        bugReportUserDescription = ""
-                    }
-                    .font(.caption2)
-                    .buttonStyle(.plain)
-                    .foregroundColor(.secondary)
-                }
-
-            case .failure(let message):
-                VStack(alignment: .leading, spacing: 4) {
-                    Text(message)
-                        .font(.caption2)
-                        .foregroundColor(.red)
-                        .fixedSize(horizontal: false, vertical: true)
-                    Button("Dismiss") {
-                        bugReportPhase = .idle
-                    }
-                    .font(.caption2)
-                    .buttonStyle(.plain)
-                    .foregroundColor(.secondary)
-                }
-            }
-        }
-        .animation(.easeInOut(duration: 0.2), value: bugReportPhase)
-    }
-
     private var processToggleBinding: Binding<Bool> {
         Binding(
             get: {
@@ -711,143 +419,6 @@ struct ContentView: View {
         )
     }
 
-    private func sendBugReport() async {
-        bugReportPhase = .sending("Collecting logs...")
-        let service = BugReportService()
-        let description = bugReportUserDescription.trimmingCharacters(in: .whitespacesAndNewlines)
-        do {
-            let outcome = try await service.sendReport(
-                isManual: true,
-                userDescription: description.isEmpty ? nil : description
-            )
-            if outcome.success {
-                bugReportPhase = .success(outcome.message)
-            } else {
-                bugReportPhase = .failure(outcome.message)
-            }
-        } catch {
-            bugReportPhase = .failure(error.localizedDescription)
-        }
-    }
-
-    private func openGitHubIssue() {
-        let description = bugReportUserDescription.trimmingCharacters(in: .whitespacesAndNewlines)
-
-        var bodyParts: [String] = []
-        bodyParts.append("## Describe the bug")
-        bodyParts.append("")
-        if !description.isEmpty {
-            bodyParts.append(description)
-        } else {
-            bodyParts.append("A clear and concise description of what the bug is.")
-        }
-        bodyParts.append("")
-        bodyParts.append("## Environment")
-        bodyParts.append("")
-        bodyParts.append("- macOS Version: \(ProcessInfo.processInfo.operatingSystemVersionString)")
-        bodyParts.append("- EXO Version: \(buildTag) (\(buildCommit))")
-        bodyParts.append("")
-        bodyParts.append("## Additional context")
-        bodyParts.append("")
-        bodyParts.append("A bug report with diagnostic logs was submitted via the app.")
-
-        let body = bodyParts.joined(separator: "\n")
-
-        var components = URLComponents(string: "https://github.com/exo-explore/exo/issues/new")!
-        components.queryItems = [
-            URLQueryItem(name: "template", value: "bug_report.md"),
-            URLQueryItem(name: "title", value: "[BUG] "),
-            URLQueryItem(name: "body", value: body),
-            URLQueryItem(name: "labels", value: "bug"),
-        ]
-
-        if let url = components.url {
-            NSWorkspace.shared.open(url)
-        }
-    }
-
-    private func showUninstallConfirmationAlert() {
-        let alert = NSAlert()
-        alert.messageText = "Uninstall EXO"
-        alert.informativeText = """
-            This will remove EXO and all its system components:
-
-            • Network configuration daemon
-            • Launch at login registration
-            • EXO network location
-
-            The app will be moved to Trash.
-            """
-        alert.alertStyle = .warning
-        alert.addButton(withTitle: "Uninstall")
-        alert.addButton(withTitle: "Cancel")
-
-        // Style the Uninstall button as destructive
-        if let uninstallButton = alert.buttons.first {
-            uninstallButton.hasDestructiveAction = true
-        }
-
-        let response = alert.runModal()
-        if response == .alertFirstButtonReturn {
-            performUninstall()
-        }
-    }
-
-    private func performUninstall() {
-        uninstallInProgress = true
-
-        // Stop EXO process first
-        controller.cancelPendingLaunch()
-        controller.stop()
-        stateService.stopPolling()
-
-        // Run the privileged uninstall on a background thread
-        // Using .utility QoS to avoid priority inversion with NSAppleScript's subprocess
-        DispatchQueue.global(qos: .utility).async {
-            do {
-                // Remove network setup daemon and components (requires admin privileges)
-                try NetworkSetupHelper.uninstall()
-
-                DispatchQueue.main.async {
-                    // Unregister from launch at login
-                    LaunchAtLoginHelper.disable()
-
-                    // Move app to trash
-                    self.moveAppToTrash()
-
-                    // Quit the app
-                    DispatchQueue.main.asyncAfter(deadline: .now() + 0.5) {
-                        NSApplication.shared.terminate(nil)
-                    }
-                }
-            } catch {
-                DispatchQueue.main.async {
-                    self.showErrorAlert(message: error.localizedDescription)
-                    self.uninstallInProgress = false
-                }
-            }
-        }
-    }
-
-    private func showErrorAlert(message: String) {
-        let alert = NSAlert()
-        alert.messageText = "Uninstall Failed"
-        alert.informativeText = message
-        alert.alertStyle = .critical
-        alert.addButton(withTitle: "OK")
-        alert.runModal()
-    }
-
-    private func moveAppToTrash() {
-        guard let appURL = Bundle.main.bundleURL as URL? else { return }
-        do {
-            try FileManager.default.trashItem(at: appURL, resultingItemURL: nil)
-        } catch {
-            // If we can't trash the app, that's OK - user can do it manually
-            // The important system components have already been cleaned up
-        }
-    }
-
     private var buildTag: String {
         Bundle.main.infoDictionary?["EXOBuildTag"] as? String ?? "unknown"
     }

app/EXO/EXO/ExoProcessController.swift

@@ -5,8 +5,8 @@ import Foundation
 private let customNamespaceKey = "EXOCustomNamespace"
 private let hfTokenKey = "EXOHFToken"
 private let enableImageModelsKey = "EXOEnableImageModels"
-private let offlineModeKey = "EXOOfflineMode"
 private let onboardingCompletedKey = "EXOOnboardingCompleted"
+private let modelSearchPathsKey = "EXOModelSearchPaths"
 
 @MainActor
 final class ExoProcessController: ObservableObject {
@@ -61,12 +61,12 @@ final class ExoProcessController: ObservableObject {
             UserDefaults.standard.set(enableImageModels, forKey: enableImageModelsKey)
         }
     }
-    @Published var offlineMode: Bool = {
-        return UserDefaults.standard.bool(forKey: offlineModeKey)
+    @Published var modelSearchPaths: String = {
+        return UserDefaults.standard.string(forKey: modelSearchPathsKey) ?? ""
     }()
     {
         didSet {
-            UserDefaults.standard.set(offlineMode, forKey: offlineModeKey)
+            UserDefaults.standard.set(modelSearchPaths, forKey: modelSearchPathsKey)
         }
     }
 
@@ -276,8 +276,8 @@ final class ExoProcessController: ObservableObject {
         if enableImageModels {
             environment["EXO_ENABLE_IMAGE_MODELS"] = "true"
         }
-        if offlineMode {
-            environment["EXO_OFFLINE"] = "true"
+        if !modelSearchPaths.isEmpty {
+            environment["EXO_MODELS_PATH"] = modelSearchPaths
         }
 
         var paths: [String] = []

app/EXO/EXO/Services/BugReportService.swift

@@ -38,8 +38,7 @@ struct BugReportService {
     func sendReport(
         baseURL: URL = URL(string: "http://127.0.0.1:52415")!,
         now: Date = Date(),
-        isManual: Bool = false,
-        userDescription: String? = nil
+        isManual: Bool = false
     ) async throws -> BugReportOutcome {
         let timestamp = Self.runTimestampString(now)
         let dayPrefix = Self.dayPrefixString(now)
@@ -61,8 +60,7 @@ struct BugReportService {
             ifconfig: ifconfigText,
             debugInfo: debugInfo,
             isManual: isManual,
-            clusterTbBridgeStatus: clusterTbBridgeStatus,
-            userDescription: userDescription
+            clusterTbBridgeStatus: clusterTbBridgeStatus
         )
 
         let eventLogFiles = readAllEventLogs()
@@ -308,8 +306,7 @@ struct BugReportService {
         ifconfig: String,
         debugInfo: DebugInfo,
         isManual: Bool,
-        clusterTbBridgeStatus: [[String: Any]]?,
-        userDescription: String? = nil
+        clusterTbBridgeStatus: [[String: Any]]?
     ) -> Data? {
         let system = readSystemMetadata()
         let exo = readExoMetadata()
@@ -326,9 +323,6 @@ struct BugReportService {
         if let tbStatus = clusterTbBridgeStatus {
             payload["cluster_thunderbolt_bridge"] = tbStatus
         }
-        if let desc = userDescription, !desc.isEmpty {
-            payload["user_description"] = desc
-        }
         return try? JSONSerialization.data(withJSONObject: payload, options: [.prettyPrinted])
     }
 

app/EXO/EXO/Views/SettingsView.swift

@@ -13,7 +13,7 @@ struct SettingsView: View {
     @State private var pendingNamespace: String = ""
     @State private var pendingHFToken: String = ""
     @State private var pendingEnableImageModels = false
-    @State private var pendingOfflineMode = false
+    @State private var pendingModelSearchPaths: String = ""
     @State private var needsRestart = false
     @State private var bugReportInFlight = false
     @State private var bugReportMessage: String?
@@ -43,7 +43,7 @@ struct SettingsView: View {
             pendingNamespace = controller.customNamespace
             pendingHFToken = controller.hfToken
             pendingEnableImageModels = controller.enableImageModels
-            pendingOfflineMode = controller.offlineMode
+            pendingModelSearchPaths = controller.modelSearchPaths
             needsRestart = false
         }
     }
@@ -74,13 +74,6 @@ struct SettingsView: View {
                     .foregroundColor(.secondary)
             }
 
-            Section {
-                Toggle("Offline Mode", isOn: $pendingOfflineMode)
-                Text("Skip internet checks and use only locally available models.")
-                    .font(.caption)
-                    .foregroundColor(.secondary)
-            }
-
             Section {
                 HStack {
                     Spacer()
@@ -106,6 +99,19 @@ struct SettingsView: View {
                     .foregroundColor(.secondary)
             }
 
+            Section {
+                LabeledContent("Model Search Paths") {
+                    TextField("/path/one:/path/two", text: $pendingModelSearchPaths)
+                        .textFieldStyle(.roundedBorder)
+                        .frame(width: 200)
+                }
+                Text(
+                    "Extra directories to search for pre-downloaded models (colon-separated). HuggingFace cache (~/.cache/huggingface/hub) is checked automatically."
+                )
+                .font(.caption)
+                .foregroundColor(.secondary)
+            }
+
             Section {
                 HStack {
                     Spacer()
@@ -454,22 +460,22 @@ struct SettingsView: View {
 
     private var hasGeneralChanges: Bool {
         pendingNamespace != controller.customNamespace || pendingHFToken != controller.hfToken
-            || pendingOfflineMode != controller.offlineMode
     }
 
     private var hasModelChanges: Bool {
         pendingEnableImageModels != controller.enableImageModels
+            || pendingModelSearchPaths != controller.modelSearchPaths
     }
 
     private func applyGeneralSettings() {
         controller.customNamespace = pendingNamespace
         controller.hfToken = pendingHFToken
-        controller.offlineMode = pendingOfflineMode
         restartIfRunning()
     }
 
     private func applyModelSettings() {
         controller.enableImageModels = pendingEnableImageModels
+        controller.modelSearchPaths = pendingModelSearchPaths
         restartIfRunning()
     }
 

bench/bench.toml

@@ -2,4 +2,6 @@
 #
 # Lists the suite files to include. Each file defines benchmarks
 # with shared constraints, topology, and default args.
-include = ["single-m3-ultra.toml"]
+include = [
+    "single-m3-ultra.toml",
+]

bench/exo_bench.py

@@ -75,7 +75,7 @@ def load_tokenizer_for_bench(model_id: str) -> Any:
         model_path = Path(
             snapshot_download(
                 model_id,
-                allow_patterns=["*.json", "*.py", "*.tiktoken", "*.model"],
+                allow_patterns=["*.json", "*.py", "*.tiktoken"],
             )
         )
 

bench/harness.py

@@ -371,7 +371,7 @@ def run_planning_phase(
                 unwrap_instance(p["DownloadCompleted"]["shardMetadata"])["modelCard"][
                     "modelId"
                 ],
-                p["DownloadCompleted"]["total"]["inBytes"],
+                p["DownloadCompleted"]["totalBytes"]["inBytes"],
             )
             for p in node_downloads
             if "DownloadCompleted" in p

bench/pyproject.toml

@@ -4,13 +4,12 @@ version = "0.1.0"
 description = "Benchmarking tool for exo distributed inference"
 requires-python = ">=3.13"
 dependencies = [
-  "httpx>=0.27.0",
-  "loguru>=0.7.3",
-  "transformers>=5.0.0",
-  "huggingface-hub>=0.33.4",
-  "tiktoken>=0.12.0",
-  "jinja2>=3.1.0",
-  "protobuf>=5.29.0",
+    "httpx>=0.27.0",
+    "loguru>=0.7.3",
+    "transformers>=5.0.0",
+    "huggingface-hub>=0.33.4",
+    "tiktoken>=0.12.0",
+    "jinja2>=3.1.0",
 ]
 
 [build-system]

bench/single-m3-ultra.toml

@@ -2,10 +2,10 @@
 #
 # Shared constraints applied to ALL benchmarks in this file.
 constraints = [
-  "All(MacOsBuild(=25D125))",
-  "Hosts(=1)",
-  "All(Chip(m3_ultra))",
-  "All(GpuCores(=80))",
+    "All(MacOsBuild(=25D125))",
+    "Hosts(=1)",
+    "All(Chip(m3_ultra))",
+    "All(GpuCores(=80))",
 ]
 
 [topology]

dashboard/src/lib/components/ChatForm.svelte

@@ -7,6 +7,8 @@
     editingImage,
     clearEditingImage,
     selectedChatModel,
+    setSelectedChatModel,
+    instances,
     ttftMs,
     tps,
     totalTokens,
@@ -28,18 +30,6 @@
     modelTasks?: Record<string, string[]>;
     modelCapabilities?: Record<string, string[]>;
     onSend?: () => void;
-    onAutoSend?: (
-      content: string,
-      files?: {
-        id: string;
-        name: string;
-        type: string;
-        textContent?: string;
-        preview?: string;
-      }[],
-    ) => void;
-    onOpenModelPicker?: () => void;
-    modelDisplayOverride?: string;
   }
 
   let {
@@ -51,9 +41,6 @@
     modelTasks = {},
     modelCapabilities = {},
     onSend,
-    onAutoSend,
-    onOpenModelPicker,
-    modelDisplayOverride,
   }: Props = $props();
 
   let message = $state("");
@@ -64,12 +51,27 @@
   const thinkingEnabled = $derived(thinkingEnabledStore());
   let loading = $derived(isLoading());
   const currentModel = $derived(selectedChatModel());
+  const instanceData = $derived(instances());
   const currentTtft = $derived(ttftMs());
   const currentTps = $derived(tps());
   const currentTokens = $derived(totalTokens());
   const currentEditingImage = $derived(editingImage());
   const isEditMode = $derived(currentEditingImage !== null);
 
+  // Custom dropdown state
+  let isModelDropdownOpen = $state(false);
+  let dropdownButtonRef: HTMLButtonElement | undefined = $state();
+  let dropdownPosition = $derived(() => {
+    if (!dropdownButtonRef || !isModelDropdownOpen)
+      return { top: 0, left: 0, width: 0 };
+    const rect = dropdownButtonRef.getBoundingClientRect();
+    return {
+      top: rect.top,
+      left: rect.left,
+      width: rect.width,
+    };
+  });
+
   // Accept all supported file types
   const acceptString = getAcceptString(["image", "text", "pdf"]);
 
@@ -122,14 +124,73 @@
         uploadedFiles.length > 0),
   );
 
-  // Short label for the currently selected model
-  const currentModelLabel = $derived(
-    currentModel
-      ? currentModel.split("/").pop() || currentModel
-      : modelDisplayOverride
-        ? modelDisplayOverride.split("/").pop() || modelDisplayOverride
-        : "",
-  );
+  // Extract available models from running instances
+  const availableModels = $derived(() => {
+    const models: Array<{ id: string; label: string; isImageModel: boolean }> =
+      [];
+    for (const [, instance] of Object.entries(instanceData)) {
+      const modelId = getInstanceModelId(instance);
+      if (
+        modelId &&
+        modelId !== "Unknown" &&
+        !models.some((m) => m.id === modelId)
+      ) {
+        models.push({
+          id: modelId,
+          label: modelId.split("/").pop() || modelId,
+          isImageModel: modelSupportsImageGeneration(modelId),
+        });
+      }
+    }
+    return models;
+  });
+
+  // Track previous model IDs to detect newly added models (plain variable to avoid reactive loop)
+  let previousModelIds: Set<string> = new Set();
+
+  // Auto-select the first available model if none is selected, if current selection is stale, or if a new model is added
+  $effect(() => {
+    const models = availableModels();
+    const currentModelIds = new Set(models.map((m) => m.id));
+
+    if (models.length > 0) {
+      // Find newly added models (in current but not in previous)
+      const newModels = models.filter((m) => !previousModelIds.has(m.id));
+
+      // If no model selected, select the first available
+      if (!currentModel) {
+        setSelectedChatModel(models[0].id);
+      }
+      // If current model is stale (no longer has a running instance), reset to first available
+      else if (!models.some((m) => m.id === currentModel)) {
+        setSelectedChatModel(models[0].id);
+      }
+      // If a new model was just added, select it
+      else if (newModels.length > 0 && previousModelIds.size > 0) {
+        setSelectedChatModel(newModels[0].id);
+      }
+    } else {
+      // No instances running - clear the selected model
+      if (currentModel) {
+        setSelectedChatModel("");
+      }
+    }
+
+    // Update previous model IDs for next comparison
+    previousModelIds = currentModelIds;
+  });
+
+  function getInstanceModelId(instanceWrapped: unknown): string {
+    if (!instanceWrapped || typeof instanceWrapped !== "object") return "";
+    const keys = Object.keys(instanceWrapped as Record<string, unknown>);
+    if (keys.length === 1) {
+      const instance = (instanceWrapped as Record<string, unknown>)[
+        keys[0]
+      ] as { shardAssignments?: { modelId?: string } };
+      return instance?.shardAssignments?.modelId || "";
+    }
+    return "";
+  }
 
   async function handleFiles(files: File[]) {
     if (files.length === 0) return;
@@ -216,15 +277,6 @@
     uploadedFiles = [];
     resetTextareaHeight();
 
-    // When onAutoSend is provided, the parent controls all send logic
-    // (including launching non-running models before sending)
-    if (onAutoSend) {
-      onAutoSend(content, files);
-      onSend?.();
-      setTimeout(() => textareaRef?.focus(), 10);
-      return;
-    }
-
     // Use image editing if in edit mode
     if (isEditMode && currentEditingImage && content) {
       editImage(content, currentEditingImage.imageDataUrl);
@@ -368,7 +420,7 @@
     {/if}
 
     <!-- Model selector (when enabled) -->
-    {#if showModelSelector}
+    {#if showModelSelector && availableModels().length > 0}
       <div
         class="flex items-center justify-between gap-2 px-3 py-2 border-b border-exo-medium-gray/30"
       >
@@ -377,22 +429,33 @@
             class="text-xs text-exo-light-gray uppercase tracking-wider flex-shrink-0"
             >MODEL:</span
           >
-          <!-- Model button — opens the full model picker -->
+          <!-- Custom dropdown -->
           <div class="relative flex-1 max-w-xs">
             <button
+              bind:this={dropdownButtonRef}
               type="button"
-              onclick={() => onOpenModelPicker?.()}
-              class="w-full bg-exo-medium-gray/50 border border-exo-yellow/30 rounded pl-3 pr-8 py-1.5 text-xs font-mono text-left tracking-wide cursor-pointer transition-all duration-200 hover:border-exo-yellow/50 focus:outline-none focus:border-exo-yellow/70"
+              onclick={() => (isModelDropdownOpen = !isModelDropdownOpen)}
+              class="w-full bg-exo-medium-gray/50 border border-exo-yellow/30 rounded pl-3 pr-8 py-1.5 text-xs font-mono text-left tracking-wide cursor-pointer transition-all duration-200 hover:border-exo-yellow/50 focus:outline-none focus:border-exo-yellow/70 {isModelDropdownOpen
+                ? 'border-exo-yellow/70'
+                : ''}"
             >
-              {#if currentModelLabel}
-                <span class="text-exo-yellow truncate">{currentModelLabel}</span
+              {#if availableModels().find((m) => m.id === currentModel)}
+                <span class="text-exo-yellow truncate"
+                  >{availableModels().find((m) => m.id === currentModel)
+                    ?.label}</span
+                >
+              {:else if availableModels().length > 0}
+                <span class="text-exo-yellow truncate"
+                  >{availableModels()[0].label}</span
                 >
               {:else}
                 <span class="text-exo-light-gray/50">— SELECT MODEL —</span>
               {/if}
             </button>
             <div
-              class="absolute right-2 top-1/2 -translate-y-1/2 pointer-events-none"
+              class="absolute right-2 top-1/2 -translate-y-1/2 pointer-events-none transition-transform duration-200 {isModelDropdownOpen
+                ? 'rotate-180'
+                : ''}"
             >
               <svg
                 class="w-3 h-3 text-exo-yellow/60"
@@ -409,6 +472,78 @@
               </svg>
             </div>
           </div>
+
+          {#if isModelDropdownOpen}
+            <!-- Backdrop to close dropdown -->
+            <button
+              type="button"
+              class="fixed inset-0 z-[9998] cursor-default"
+              onclick={() => (isModelDropdownOpen = false)}
+              aria-label="Close dropdown"
+            ></button>
+
+            <!-- Dropdown Panel - fixed positioning to escape overflow:hidden -->
+            <div
+              class="fixed bg-exo-dark-gray border border-exo-yellow/30 rounded shadow-lg shadow-black/50 z-[9999] max-h-48 overflow-y-auto"
+              style="bottom: calc(100vh - {dropdownPosition()
+                .top}px + 4px); left: {dropdownPosition()
+                .left}px; width: {dropdownPosition().width}px;"
+            >
+              <div class="py-1">
+                {#each availableModels() as model}
+                  <button
+                    type="button"
+                    onclick={() => {
+                      setSelectedChatModel(model.id);
+                      isModelDropdownOpen = false;
+                    }}
+                    class="w-full px-3 py-2 text-left text-xs font-mono tracking-wide transition-colors duration-100 flex items-center gap-2 {currentModel ===
+                    model.id
+                      ? 'bg-transparent text-exo-yellow'
+                      : 'text-exo-light-gray hover:text-exo-yellow'}"
+                  >
+                    {#if currentModel === model.id}
+                      <svg
+                        class="w-3 h-3 flex-shrink-0"
+                        fill="currentColor"
+                        viewBox="0 0 20 20"
+                      >
+                        <path
+                          fill-rule="evenodd"
+                          d="M16.707 5.293a1 1 0 010 1.414l-8 8a1 1 0 01-1.414 0l-4-4a1 1 0 011.414-1.414L8 12.586l7.293-7.293a1 1 0 011.414 0z"
+                          clip-rule="evenodd"
+                        />
+                      </svg>
+                    {:else}
+                      <span class="w-3"></span>
+                    {/if}
+                    {#if model.isImageModel}
+                      <svg
+                        class="w-3.5 h-3.5 flex-shrink-0 text-exo-yellow"
+                        fill="none"
+                        viewBox="0 0 24 24"
+                        stroke="currentColor"
+                        stroke-width="2"
+                        aria-label="Image generation model"
+                      >
+                        <rect
+                          x="3"
+                          y="3"
+                          width="18"
+                          height="18"
+                          rx="2"
+                          ry="2"
+                        />
+                        <circle cx="8.5" cy="8.5" r="1.5" />
+                        <polyline points="21 15 16 10 5 21" />
+                      </svg>
+                    {/if}
+                    <span class="truncate flex-1">{model.label}</span>
+                  </button>
+                {/each}
+              </div>
+            </div>
+          {/if}
         </div>
         <!-- Thinking toggle -->
         {#if modelSupportsThinking()}

dashboard/src/lib/components/ChatModelSelector.svelte

@@ -1,401 +0,0 @@
-<script lang="ts" module>
-  export interface ChatModelInfo {
-    id: string;
-    name: string;
-    base_model: string;
-    storage_size_megabytes: number;
-    capabilities: string[];
-    family: string;
-    quantization: string;
-  }
-
-  // Auto mode tier list (for when user just starts typing)
-  export const AUTO_TIERS: string[][] = [
-    // Tier 1 (frontier)
-    ["DeepSeek V3.1", "GLM-5", "Kimi K2.5", "Qwen3 Coder Next"],
-    // Tier 2 (excellent)
-    [
-      "Kimi K2",
-      "Qwen3 235B",
-      "MiniMax M2.5",
-      "Step 3.5 Flash",
-      "Qwen3 Next 80B",
-    ],
-    // Tier 3 (great)
-    [
-      "GLM 4.7",
-      "MiniMax M2.1",
-      "Qwen3 Coder 480B",
-      "GLM 4.5 Air",
-      "Llama 3.3 70B",
-    ],
-    // Tier 4 (good)
-    ["GPT-OSS 120B", "Qwen3 30B", "Llama 3.1 70B", "GLM 4.7 Flash"],
-    // Tier 5 (small/fast)
-    [
-      "Llama 3.1 8B",
-      "GPT-OSS 20B",
-      "Llama 3.2 3B",
-      "Qwen3 0.6B",
-      "Llama 3.2 1B",
-    ],
-  ];
-
-  /** Return the tier index (0 = best) for a base_model name. */
-  export function getAutoTierIndex(baseModel: string): number {
-    for (let i = 0; i < AUTO_TIERS.length; i++) {
-      if (AUTO_TIERS[i].includes(baseModel)) return i;
-    }
-    return AUTO_TIERS.length; // not in any tier → lowest priority
-  }
-
-  /** Auto mode: walk tiers top-down, pick biggest fitting variant from highest tier. */
-  export function pickAutoModel(
-    modelList: ChatModelInfo[],
-    memoryGB: number,
-  ): ChatModelInfo | null {
-    for (const tier of AUTO_TIERS) {
-      const candidates: ChatModelInfo[] = [];
-      for (const baseModel of tier) {
-        const variants = modelList
-          .filter(
-            (m) =>
-              m.base_model === baseModel &&
-              (m.storage_size_megabytes || 0) / 1024 <= memoryGB &&
-              (m.storage_size_megabytes || 0) > 0,
-          )
-          .sort(
-            (a, b) =>
-              (b.storage_size_megabytes || 0) - (a.storage_size_megabytes || 0),
-          );
-        if (variants[0]) candidates.push(variants[0]);
-      }
-      if (candidates.length > 0) {
-        candidates.sort(
-          (a, b) =>
-            (b.storage_size_megabytes || 0) - (a.storage_size_megabytes || 0),
-        );
-        return candidates[0];
-      }
-    }
-    return null;
-  }
-</script>
-
-<script lang="ts">
-  interface CategoryRecommendation {
-    category: "coding" | "writing" | "agentic" | "biggest";
-    label: string;
-    model: ChatModelInfo | null;
-    tooltip: string;
-  }
-
-  interface Props {
-    models: ChatModelInfo[];
-    clusterLabel: string;
-    totalMemoryGB: number;
-    onSelect: (modelId: string, category: string) => void;
-    onAddModel: () => void;
-    class?: string;
-  }
-
-  let {
-    models,
-    clusterLabel,
-    totalMemoryGB,
-    onSelect,
-    onAddModel,
-    class: className = "",
-  }: Props = $props();
-
-  // --- Hardcoded Rankings ---
-  const CODING_RANKING = [
-    "Qwen3 Coder Next",
-    "Qwen3 Coder 480B",
-    "Qwen3 30B",
-    "GPT-OSS 20B",
-    "Llama 3.1 8B",
-    "Llama 3.2 3B",
-    "Qwen3 0.6B",
-  ];
-
-  const WRITING_RANKING = [
-    "Kimi K2.5",
-    "Kimi K2",
-    "Qwen3 Next 80B",
-    "Llama 3.3 70B",
-    "MiniMax M2.5",
-    "GLM 4.5 Air",
-    "GLM 4.7 Flash",
-    "GPT-OSS 20B",
-    "Llama 3.1 8B",
-    "Llama 3.2 3B",
-    "Qwen3 0.6B",
-  ];
-
-  const AGENTIC_RANKING = [
-    "DeepSeek V3.1",
-    "GLM-5",
-    "Qwen3 235B",
-    "Step 3.5 Flash",
-    "GLM 4.7",
-    "MiniMax M2.1",
-    "GPT-OSS 120B",
-    "Llama 3.3 70B",
-    "Llama 3.1 70B",
-    "GLM 4.7 Flash",
-    "GPT-OSS 20B",
-    "Qwen3 30B",
-    "Llama 3.1 8B",
-    "Llama 3.2 3B",
-    "Qwen3 0.6B",
-  ];
-
-  function getModelSizeGB(m: ChatModelInfo): number {
-    return (m.storage_size_megabytes || 0) / 1024;
-  }
-
-  function fitsInMemory(m: ChatModelInfo): boolean {
-    return getModelSizeGB(m) <= totalMemoryGB && getModelSizeGB(m) > 0;
-  }
-
-  /** For a given base_model name, find the biggest quant variant that fits in memory. */
-  function pickBestVariant(baseModel: string): ChatModelInfo | null {
-    const variants = models
-      .filter((m) => m.base_model === baseModel && fitsInMemory(m))
-      .sort((a, b) => getModelSizeGB(b) - getModelSizeGB(a));
-    return variants[0] ?? null;
-  }
-
-  /** Walk a ranked list of base_model names, return the first that has a fitting variant. */
-  function pickFromRanking(ranking: string[]): ChatModelInfo | null {
-    for (const baseModel of ranking) {
-      const pick = pickBestVariant(baseModel);
-      if (pick) return pick;
-    }
-    return null;
-  }
-
-  /** Pick the single biggest model that fits. */
-  function pickBiggest(): ChatModelInfo | null {
-    const fitting = models
-      .filter((m) => fitsInMemory(m))
-      .sort((a, b) => getModelSizeGB(b) - getModelSizeGB(a));
-    return fitting[0] ?? null;
-  }
-
-  const recommendations = $derived.by((): CategoryRecommendation[] => {
-    return [
-      {
-        category: "coding",
-        label: "Best for Coding",
-        model: pickFromRanking(CODING_RANKING),
-        tooltip:
-          "Ranked by coding benchmark performance (LiveCodeBench, SWE-bench)",
-      },
-      {
-        category: "writing",
-        label: "Best for Writing",
-        model: pickFromRanking(WRITING_RANKING),
-        tooltip: "Ranked by creative writing quality and instruction following",
-      },
-      {
-        category: "agentic",
-        label: "Best Agentic",
-        model: pickFromRanking(AGENTIC_RANKING),
-        tooltip: "Ranked by reasoning, planning, and tool-use capability",
-      },
-      {
-        category: "biggest",
-        label: "Biggest",
-        model: pickBiggest(),
-        tooltip: "Largest model that fits in your available memory",
-      },
-    ];
-  });
-
-  function formatSize(mb: number): string {
-    const gb = mb / 1024;
-    if (gb >= 100) return `${Math.round(gb)} GB`;
-    return `${gb.toFixed(1)} GB`;
-  }
-
-  // Category icons (SVG paths)
-  const categoryIcons: Record<string, string> = {
-    coding:
-      "M8 9l3 3-3 3m5 0h3M5 20h14a2 2 0 002-2V6a2 2 0 00-2-2H5a2 2 0 00-2 2v12a2 2 0 002 2z",
-    writing:
-      "M11 5H6a2 2 0 00-2 2v11a2 2 0 002 2h11a2 2 0 002-2v-5m-1.414-9.414a2 2 0 112.828 2.828L11.828 15H9v-2.828l8.586-8.586z",
-    agentic:
-      "M9.663 17h4.673M12 3v1m6.364 1.636l-.707.707M21 12h-1M4 12H3m3.343-5.657l-.707-.707m2.828 9.9a5 5 0 117.072 0l-.548.547A3.374 3.374 0 0014 18.469V19a2 2 0 11-4 0v-.531c0-.895-.356-1.754-.988-2.386l-.548-.547z",
-    biggest:
-      "M19 11H5m14 0a2 2 0 012 2v6a2 2 0 01-2 2H5a2 2 0 01-2-2v-6a2 2 0 012-2m14 0V9a2 2 0 00-2-2M5 11V9a2 2 0 012-2m0 0V5a2 2 0 012-2h6a2 2 0 012 2v2M7 7h10",
-  };
-
-  let hoveredTooltip = $state<string | null>(null);
-  let tooltipAnchor = $state<{ x: number; y: number } | null>(null);
-
-  function showTooltip(category: string, e: MouseEvent | FocusEvent) {
-    hoveredTooltip = category;
-    const target = e.currentTarget as HTMLElement;
-    const rect = target.getBoundingClientRect();
-    tooltipAnchor = { x: rect.left + rect.width / 2, y: rect.top };
-  }
-
-  function hideTooltip() {
-    hoveredTooltip = null;
-    tooltipAnchor = null;
-  }
-</script>
-
-<div class="flex flex-col items-center justify-center gap-6 {className}">
-  <!-- Header -->
-  <div class="text-center">
-    <p class="text-xs text-exo-light-gray uppercase tracking-[0.2em] mb-1">
-      Recommended for your
-    </p>
-    <p class="text-sm text-white font-mono tracking-wide">{clusterLabel}</p>
-  </div>
-
-  <!-- Category Cards Grid -->
-  <div class="grid grid-cols-2 gap-3 w-full max-w-md">
-    {#each recommendations as rec}
-      {#if rec.model}
-        <button
-          type="button"
-          onclick={() => rec.model && onSelect(rec.model.id, rec.category)}
-          class="group relative flex flex-col items-start gap-2 p-4 rounded-lg border border-exo-medium-gray/50 bg-exo-dark-gray/50 hover:border-exo-yellow/40 hover:bg-exo-dark-gray transition-all duration-200 cursor-pointer text-left"
-        >
-          <!-- Category icon + label -->
-          <div class="flex items-center gap-2 w-full">
-            <svg
-              class="w-4 h-4 text-exo-yellow/70 group-hover:text-exo-yellow transition-colors flex-shrink-0"
-              fill="none"
-              viewBox="0 0 24 24"
-              stroke="currentColor"
-              stroke-width="1.5"
-            >
-              <path
-                stroke-linecap="round"
-                stroke-linejoin="round"
-                d={categoryIcons[rec.category]}
-              />
-            </svg>
-            <span
-              class="text-xs font-mono uppercase tracking-wider text-exo-light-gray group-hover:text-white transition-colors"
-            >
-              {rec.label}
-            </span>
-            <!-- Info tooltip -->
-            <div class="ml-auto flex-shrink-0">
-              <span
-                role="button"
-                tabindex="-1"
-                class="text-exo-light-gray/40 hover:text-exo-light-gray transition-colors cursor-help inline-flex"
-                onmouseenter={(e: MouseEvent) => showTooltip(rec.category, e)}
-                onmouseleave={() => hideTooltip()}
-                onclick={(e: MouseEvent) => {
-                  e.stopPropagation();
-                  if (hoveredTooltip === rec.category) {
-                    hideTooltip();
-                  } else {
-                    showTooltip(rec.category, e);
-                  }
-                }}
-              >
-                <svg
-                  class="w-3.5 h-3.5"
-                  fill="none"
-                  viewBox="0 0 24 24"
-                  stroke="currentColor"
-                  stroke-width="2"
-                >
-                  <circle cx="12" cy="12" r="10" />
-                  <path d="M12 16v-4m0-4h.01" />
-                </svg>
-              </span>
-            </div>
-          </div>
-
-          <!-- Model name + size -->
-          <div class="w-full">
-            <p class="text-sm text-white font-mono truncate">
-              {rec.model.base_model}
-            </p>
-            <p class="text-xs text-exo-light-gray/60 font-mono mt-0.5">
-              {formatSize(rec.model.storage_size_megabytes)}
-              {#if rec.model.quantization}
-                <span class="text-exo-light-gray/40"
-                  >&middot; {rec.model.quantization}</span
-                >
-              {/if}
-            </p>
-          </div>
-        </button>
-      {:else}
-        <!-- No model fits for this category -->
-        <div
-          class="flex flex-col items-start gap-2 p-4 rounded-lg border border-exo-medium-gray/30 bg-exo-dark-gray/30 opacity-50"
-        >
-          <div class="flex items-center gap-2">
-            <svg
-              class="w-4 h-4 text-exo-light-gray/40 flex-shrink-0"
-              fill="none"
-              viewBox="0 0 24 24"
-              stroke="currentColor"
-              stroke-width="1.5"
-            >
-              <path
-                stroke-linecap="round"
-                stroke-linejoin="round"
-                d={categoryIcons[rec.category]}
-              />
-            </svg>
-            <span
-              class="text-xs font-mono uppercase tracking-wider text-exo-light-gray/50"
-              >{rec.label}</span
-            >
-          </div>
-          <p class="text-xs text-exo-light-gray/40 font-mono">No model fits</p>
-        </div>
-      {/if}
-    {/each}
-  </div>
-
-  <!-- Add Model Button -->
-  <button
-    type="button"
-    onclick={onAddModel}
-    class="flex items-center gap-2 px-4 py-2 text-xs font-mono uppercase tracking-wider text-exo-light-gray hover:text-exo-yellow border border-exo-medium-gray/30 hover:border-exo-yellow/30 rounded-lg transition-all duration-200 cursor-pointer"
-  >
-    <svg
-      class="w-3.5 h-3.5"
-      fill="none"
-      viewBox="0 0 24 24"
-      stroke="currentColor"
-      stroke-width="2"
-    >
-      <path stroke-linecap="round" stroke-linejoin="round" d="M12 4v16m8-8H4" />
-    </svg>
-    Add Model
-  </button>
-
-  <!-- Auto hint -->
-  <p class="text-xs text-exo-light-gray/40 font-mono tracking-wide text-center">
-    Or just start typing &mdash; we'll pick the best model automatically
-  </p>
-</div>
-
-<!-- Fixed-position tooltip (escapes overflow-hidden ancestors) -->
-{#if hoveredTooltip && tooltipAnchor}
-  {@const rec = recommendations.find((r) => r.category === hoveredTooltip)}
-  {#if rec}
-    <div
-      class="fixed z-[9999] px-3 py-2 bg-exo-black border border-exo-medium-gray/50 rounded text-xs text-exo-light-gray whitespace-nowrap shadow-lg pointer-events-none"
-      style="left: {tooltipAnchor.x}px; top: {tooltipAnchor.y -
-        8}px; transform: translate(-50%, -100%);"
-    >
-      {rec.tooltip}
-    </div>
-  {/if}
-{/if}

dashboard/src/lib/components/ChatSidebar.svelte

@@ -2,6 +2,7 @@
   import {
     conversations,
     activeConversationId,
+    createConversation,
     loadConversation,
     deleteConversation,
     deleteAllConversations,
@@ -16,15 +17,9 @@
 
   interface Props {
     class?: string;
-    onNewChat?: () => void;
-    onSelectConversation?: () => void;
   }
 
-  let {
-    class: className = "",
-    onNewChat,
-    onSelectConversation,
-  }: Props = $props();
+  let { class: className = "" }: Props = $props();
 
   const conversationList = $derived(conversations());
   const activeId = $derived(activeConversationId());
@@ -47,11 +42,10 @@
   );
 
   function handleNewChat() {
-    onNewChat?.();
+    createConversation();
   }
 
   function handleSelectConversation(id: string) {
-    onSelectConversation?.();
     loadConversation(id);
   }
 

dashboard/src/lib/components/DeviceIcon.svelte

@@ -67,8 +67,8 @@
   const studioMemH = $derived((ramPercent / 100) * studioMemTotalH);
 
   // ── MacBook dimensions (same ratios as TopologyGraph) ──
-  const mbW = $derived((size * 1.6 * 0.85) / 1.15);
-  const mbH = $derived(size * 0.85);
+  const mbW = $derived(size * 1.6);
+  const mbH = $derived(size * 1.15);
   const mbX = $derived(cx - mbW / 2);
   const mbY = $derived(cy - mbH / 2);
 

dashboard/src/lib/components/HeaderNav.svelte

@@ -95,7 +95,7 @@
     {#if showHome}
       <button
         onclick={handleHome}
-        class="text-sm text-white/70 hover:text-exo-yellow transition-colors tracking-wider uppercase flex items-center gap-2 cursor-pointer"
+        class="text-sm text-exo-light-gray hover:text-exo-yellow transition-colors tracking-wider uppercase flex items-center gap-2 cursor-pointer"
         title="Back to topology view"
       >
         <svg
@@ -116,7 +116,7 @@
     {/if}
     <a
       href="/#/downloads"
-      class="text-sm text-white/70 hover:text-exo-yellow transition-colors tracking-wider uppercase flex items-center gap-2 cursor-pointer"
+      class="text-sm text-exo-light-gray hover:text-exo-yellow transition-colors tracking-wider uppercase flex items-center gap-2 cursor-pointer"
       title="View downloads overview"
     >
       {#if downloadProgress}

dashboard/src/lib/components/ModelFilterPopover.svelte

@@ -6,7 +6,6 @@
     capabilities: string[];
     sizeRange: { min: number; max: number } | null;
     downloadedOnly: boolean;
-    readyOnly: boolean;
   }
 
   type ModelFilterPopoverProps = {
@@ -191,58 +190,34 @@
       </div>
     </div>
 
-    <!-- Availability filters -->
+    <!-- Downloaded only -->
     <div>
       <h4 class="text-xs font-mono text-white/50 mb-2">Availability</h4>
-      <div class="flex flex-wrap gap-1.5">
-        <button
-          type="button"
-          class="px-2 py-1 text-xs font-mono rounded transition-colors {filters.downloadedOnly
-            ? 'bg-green-500/20 text-green-400 border border-green-500/30'
-            : 'bg-white/5 text-white/60 hover:bg-white/10 border border-transparent'}"
-          onclick={() =>
-            onChange({ ...filters, downloadedOnly: !filters.downloadedOnly })}
+      <button
+        type="button"
+        class="px-2 py-1 text-xs font-mono rounded transition-colors {filters.downloadedOnly
+          ? 'bg-green-500/20 text-green-400 border border-green-500/30'
+          : 'bg-white/5 text-white/60 hover:bg-white/10 border border-transparent'}"
+        onclick={() =>
+          onChange({ ...filters, downloadedOnly: !filters.downloadedOnly })}
+      >
+        <svg
+          class="w-3.5 h-3.5 inline-block"
+          viewBox="0 0 24 24"
+          fill="none"
+          stroke="currentColor"
+          stroke-width="2"
+          stroke-linecap="round"
+          stroke-linejoin="round"
         >
-          <svg
-            class="w-3.5 h-3.5 inline-block"
-            viewBox="0 0 24 24"
-            fill="none"
-            stroke="currentColor"
-            stroke-width="2"
-            stroke-linecap="round"
-            stroke-linejoin="round"
-          >
-            <path
-              class="text-white/40"
-              d="M20 20a2 2 0 0 0 2-2V8a2 2 0 0 0-2-2h-7.9a2 2 0 0 1-1.69-.9L9.6 3.9A2 2 0 0 0 7.93 3H4a2 2 0 0 0-2 2v13a2 2 0 0 0 2 2Z"
-            />
-            <path class="text-green-400" d="m9 13 2 2 4-4" />
-          </svg>
-          <span class="ml-1">Downloaded</span>
-        </button>
-        <button
-          type="button"
-          class="px-2 py-1 text-xs font-mono rounded transition-colors {filters.readyOnly
-            ? 'bg-green-500/20 text-green-400 border border-green-500/30'
-            : 'bg-white/5 text-white/60 hover:bg-white/10 border border-transparent'}"
-          onclick={() =>
-            onChange({ ...filters, readyOnly: !filters.readyOnly })}
-        >
-          <svg
-            class="w-3.5 h-3.5 inline-block"
-            viewBox="0 0 24 24"
-            fill="none"
-            stroke="currentColor"
-            stroke-width="2"
-            stroke-linecap="round"
-            stroke-linejoin="round"
-          >
-            <circle cx="12" cy="12" r="10" />
-            <path d="m9 12 2 2 4-4" />
-          </svg>
-          <span class="ml-1">Ready</span>
-        </button>
-      </div>
+          <path
+            class="text-white/40"
+            d="M20 20a2 2 0 0 0 2-2V8a2 2 0 0 0-2-2h-7.9a2 2 0 0 1-1.69-.9L9.6 3.9A2 2 0 0 0 7.93 3H4a2 2 0 0 0-2 2v13a2 2 0 0 0 2 2Z"
+          />
+          <path class="text-green-400" d="m9 13 2 2 4-4" />
+        </svg>
+        <span class="ml-1">Downloaded</span>
+      </button>
     </div>
 
     <!-- Size range -->

dashboard/src/lib/components/ModelPickerGroup.svelte

@@ -41,7 +41,6 @@
     onShowInfo: (group: ModelGroup) => void;
     downloadStatusMap?: Map<string, DownloadAvailability>;
     launchedAt?: number;
-    instanceStatuses?: Record<string, { status: string; statusClass: string }>;
   };
 
   let {
@@ -57,7 +56,6 @@
     onShowInfo,
     downloadStatusMap,
     launchedAt,
-    instanceStatuses = {},
   }: ModelPickerGroupProps = $props();
 
   // Group-level download status: show if any variant is downloaded
@@ -94,12 +92,6 @@
   const anyVariantFits = $derived(
     group.variants.some((v) => canModelFit(v.id)),
   );
-  // Check if any variant has an active instance (ready, loading, downloading)
-  const anyVariantHasInstance = $derived(
-    instanceStatuses
-      ? group.variants.some((v) => instanceStatuses[v.id] != null)
-      : false,
-  );
   const groupFitStatus = $derived.by((): ModelFitStatus => {
     let hasClusterCapacityOnly = false;
     for (const variant of group.variants) {
@@ -130,35 +122,16 @@
     !group.hasMultipleVariants &&
       group.variants.some((v) => v.id === selectedModelId),
   );
-
-  // Group-level instance status: show the "best" status across all variants
-  const groupInstanceStatus = $derived.by(() => {
-    if (!instanceStatuses) return null;
-    const readyStatuses = ["READY", "LOADED", "RUNNING"];
-    const loadingStatuses = ["LOADING", "WARMING UP"];
-    let bestStatus: { status: string; statusClass: string } | null = null;
-    for (const variant of group.variants) {
-      const s = instanceStatuses[variant.id];
-      if (!s) continue;
-      if (readyStatuses.includes(s.status)) return s; // Ready is best
-      if (loadingStatuses.includes(s.status) || s.status === "DOWNLOADING") {
-        bestStatus = s;
-      }
-    }
-    return bestStatus;
-  });
 </script>
 
 <div
-  class="border-b border-white/5 last:border-b-0 {!anyVariantFits &&
-  !anyVariantHasInstance
+  class="border-b border-white/5 last:border-b-0 {!anyVariantFits
     ? 'opacity-50'
     : ''}"
 >
   <!-- Main row -->
   <div
-    class="flex items-center gap-2 px-3 py-2.5 transition-colors {anyVariantFits ||
-    anyVariantHasInstance
+    class="flex items-center gap-2 px-3 py-2.5 transition-colors {anyVariantFits
       ? 'hover:bg-white/5 cursor-pointer'
       : 'cursor-not-allowed'} {isMainSelected
       ? 'bg-exo-yellow/10 border-l-2 border-exo-yellow'
@@ -168,7 +141,7 @@
         onToggleExpand();
       } else {
         const modelId = group.variants[0]?.id;
-        if (modelId && (canModelFit(modelId) || instanceStatuses[modelId])) {
+        if (modelId && canModelFit(modelId)) {
           onSelectModel(modelId);
         }
       }
@@ -182,7 +155,7 @@
           onToggleExpand();
         } else {
           const modelId = group.variants[0]?.id;
-          if (modelId && (canModelFit(modelId) || instanceStatuses[modelId])) {
+          if (modelId && canModelFit(modelId)) {
             onSelectModel(modelId);
           }
         }
@@ -373,47 +346,6 @@
       </span>
     {/if}
 
-    <!-- Instance status badge -->
-    {#if groupInstanceStatus}
-      {#if groupInstanceStatus.status === "READY" || groupInstanceStatus.status === "LOADED" || groupInstanceStatus.status === "RUNNING"}
-        <span class="flex-shrink-0" title="Running">
-          <svg
-            class="w-3 h-3 text-green-400"
-            viewBox="0 0 12 12"
-            fill="currentColor"
-          >
-            <circle cx="6" cy="6" r="5" />
-          </svg>
-        </span>
-      {:else if groupInstanceStatus.status === "DOWNLOADING"}
-        <span class="flex-shrink-0 animate-pulse" title="Downloading">
-          <svg
-            class="w-3.5 h-3.5 text-blue-400"
-            viewBox="0 0 24 24"
-            fill="none"
-            stroke="currentColor"
-            stroke-width="2"
-            stroke-linecap="round"
-            stroke-linejoin="round"
-          >
-            <path d="M21 15v4a2 2 0 01-2 2H5a2 2 0 01-2-2v-4" />
-            <polyline points="7 10 12 15 17 10" />
-            <line x1="12" y1="15" x2="12" y2="3" />
-          </svg>
-        </span>
-      {:else if groupInstanceStatus.status === "LOADING" || groupInstanceStatus.status === "WARMING UP"}
-        <span class="flex-shrink-0 animate-pulse" title="Loading">
-          <svg
-            class="w-3 h-3 text-yellow-400"
-            viewBox="0 0 12 12"
-            fill="currentColor"
-          >
-            <circle cx="6" cy="6" r="5" />
-          </svg>
-        </span>
-      {/if}
-    {/if}
-
     <!-- Check mark if selected (single-variant) -->
     {#if isMainSelected}
       <svg
@@ -488,30 +420,20 @@
       {#each group.variants as variant}
         {@const fitStatus = getModelFitStatus(variant.id)}
         {@const modelCanFit = canModelFit(variant.id)}
-        {@const variantHasInstance = instanceStatuses[variant.id] != null}
         {@const isSelected = selectedModelId === variant.id}
-        <div
-          class="w-full flex items-center gap-3 px-3 py-2 pl-10 hover:bg-white/5 transition-colors text-left {!modelCanFit &&
-          !variantHasInstance
+        <button
+          type="button"
+          class="w-full flex items-center gap-3 px-3 py-2 pl-10 hover:bg-white/5 transition-colors text-left {!modelCanFit
             ? 'opacity-50 cursor-not-allowed'
             : 'cursor-pointer'} {isSelected
             ? 'bg-exo-yellow/10 border-l-2 border-exo-yellow'
             : 'border-l-2 border-transparent'}"
-          role="button"
-          tabindex="0"
+          disabled={!modelCanFit}
           onclick={() => {
-            if (modelCanFit || variantHasInstance) {
+            if (modelCanFit) {
               onSelectModel(variant.id);
             }
           }}
-          onkeydown={(e) => {
-            if (e.key === "Enter" || e.key === " ") {
-              e.preventDefault();
-              if (modelCanFit) {
-                onSelectModel(variant.id);
-              }
-            }
-          }}
         >
           <!-- Quantization badge -->
           <span
@@ -556,48 +478,6 @@
             {/if}
           {/if}
 
-          <!-- Instance status badge -->
-          {#if instanceStatuses[variant.id]}
-            {@const instStatus = instanceStatuses[variant.id]}
-            {#if instStatus.status === "READY" || instStatus.status === "LOADED" || instStatus.status === "RUNNING"}
-              <span class="flex-shrink-0" title="Running">
-                <svg
-                  class="w-3 h-3 text-green-400"
-                  viewBox="0 0 12 12"
-                  fill="currentColor"
-                >
-                  <circle cx="6" cy="6" r="5" />
-                </svg>
-              </span>
-            {:else if instStatus.status === "DOWNLOADING"}
-              <span class="flex-shrink-0 animate-pulse" title="Downloading">
-                <svg
-                  class="w-3.5 h-3.5 text-blue-400"
-                  viewBox="0 0 24 24"
-                  fill="none"
-                  stroke="currentColor"
-                  stroke-width="2"
-                  stroke-linecap="round"
-                  stroke-linejoin="round"
-                >
-                  <path d="M21 15v4a2 2 0 01-2 2H5a2 2 0 01-2-2v-4" />
-                  <polyline points="7 10 12 15 17 10" />
-                  <line x1="12" y1="15" x2="12" y2="3" />
-                </svg>
-              </span>
-            {:else if instStatus.status === "LOADING" || instStatus.status === "WARMING UP"}
-              <span class="flex-shrink-0 animate-pulse" title="Loading">
-                <svg
-                  class="w-3 h-3 text-yellow-400"
-                  viewBox="0 0 12 12"
-                  fill="currentColor"
-                >
-                  <circle cx="6" cy="6" r="5" />
-                </svg>
-              </span>
-            {/if}
-          {/if}
-
           <!-- Check mark if selected -->
           {#if isSelected}
             <svg
@@ -610,36 +490,7 @@
               />
             </svg>
           {/if}
-
-          <!-- Info button -->
-          <button
-            type="button"
-            class="p-1 rounded hover:bg-white/10 transition-colors flex-shrink-0"
-            onclick={(e) => {
-              e.stopPropagation();
-              onShowInfo({
-                id: variant.id,
-                name: variant.name || variant.id,
-                capabilities: group.capabilities,
-                family: group.family,
-                variants: [variant],
-                smallestVariant: variant,
-                hasMultipleVariants: false,
-              });
-            }}
-            title="View variant details"
-          >
-            <svg
-              class="w-4 h-4 text-white/30 hover:text-white/50"
-              viewBox="0 0 24 24"
-              fill="currentColor"
-            >
-              <path
-                d="M12 2C6.48 2 2 6.48 2 12s4.48 10 10 10 10-4.48 10-10S17.52 2 12 2zm1 15h-2v-6h2v6zm0-8h-2V7h2v2z"
-              />
-            </svg>
-          </button>
-        </div>
+        </button>
       {/each}
     </div>
   {/if}

dashboard/src/lib/components/ModelPickerModal.svelte

@@ -36,7 +36,6 @@
     capabilities: string[];
     sizeRange: { min: number; max: number } | null;
     downloadedOnly: boolean;
-    readyOnly: boolean;
   }
 
   interface HuggingFaceModel {
@@ -50,11 +49,6 @@
 
   type ModelFitStatus = "fits_now" | "fits_cluster_capacity" | "too_large";
 
-  export type InstanceStatus = {
-    status: string;
-    statusClass: string;
-  };
-
   type ModelPickerModalProps = {
     isOpen: boolean;
     models: ModelInfo[];
@@ -81,7 +75,6 @@
         macmon_info?: { memory?: { ram_total?: number } };
       }
     >;
-    instanceStatuses?: Record<string, InstanceStatus>;
   };
 
   let {
@@ -103,7 +96,6 @@
     usedMemoryGB,
     downloadsData,
     topologyNodes,
-    instanceStatuses = {},
   }: ModelPickerModalProps = $props();
 
   // Local state
@@ -115,7 +107,6 @@
     capabilities: [],
     sizeRange: null,
     downloadedOnly: false,
-    readyOnly: false,
   });
   let infoGroup = $state<ModelGroup | null>(null);
 
@@ -449,16 +440,6 @@
       );
     }
 
-    // Filter to ready/running models only
-    if (filters.readyOnly) {
-      result = result.filter((g) =>
-        g.variants.some((v) => {
-          const s = instanceStatuses[v.id];
-          return s && s.statusClass === "ready";
-        }),
-      );
-    }
-
     // Sort: fits-now first, then fits-cluster-capacity, then too-large
     result.sort((a, b) => {
       const getGroupFitRank = (group: ModelGroup): number => {
@@ -569,19 +550,13 @@
   }
 
   function clearFilters() {
-    filters = {
-      capabilities: [],
-      sizeRange: null,
-      downloadedOnly: false,
-      readyOnly: false,
-    };
+    filters = { capabilities: [], sizeRange: null, downloadedOnly: false };
   }
 
   const hasActiveFilters = $derived(
     filters.capabilities.length > 0 ||
       filters.sizeRange !== null ||
-      filters.downloadedOnly ||
-      filters.readyOnly,
+      filters.downloadedOnly,
   );
 </script>
 
@@ -850,7 +825,6 @@
                 onShowInfo={(g) => (infoGroup = g)}
                 downloadStatusMap={getVariantDownloadMap(group)}
                 launchedAt={recentTimestamps.get(group.variants[0]?.id ?? "")}
-                {instanceStatuses}
               />
             {/each}
           {/if}
@@ -918,7 +892,6 @@
               {onToggleFavorite}
               onShowInfo={(g) => (infoGroup = g)}
               downloadStatusMap={getVariantDownloadMap(group)}
-              {instanceStatuses}
             />
           {/each}
           <!-- Other models -->
@@ -945,7 +918,6 @@
               {onToggleFavorite}
               onShowInfo={(g) => (infoGroup = g)}
               downloadStatusMap={getVariantDownloadMap(group)}
-              {instanceStatuses}
             />
           {/each}
         {/if}
@@ -968,11 +940,6 @@
             >Downloaded</span
           >
         {/if}
-        {#if filters.readyOnly}
-          <span class="px-1.5 py-0.5 bg-green-500/20 text-green-400 rounded"
-            >Ready</span
-          >
-        {/if}
         {#if filters.sizeRange}
           <span class="px-1.5 py-0.5 bg-exo-yellow/20 text-exo-yellow rounded">
             {filters.sizeRange.min}GB - {filters.sizeRange.max}GB

dashboard/src/lib/components/index.ts

@@ -12,4 +12,3 @@ export { default as HuggingFaceResultItem } from "./HuggingFaceResultItem.svelte
 export { default as ModelFilterPopover } from "./ModelFilterPopover.svelte";
 export { default as ModelPickerGroup } from "./ModelPickerGroup.svelte";
 export { default as ModelPickerModal } from "./ModelPickerModal.svelte";
-export { default as ChatModelSelector } from "./ChatModelSelector.svelte";

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

@@ -847,10 +847,6 @@ class AppStore {
     this.thinkingEnabled = conversation.enableThinking ?? true;
     this.refreshConversationModelFromInstances();
 
-    // Sync global selection to the loaded conversation's model so reactive
-    // effects in +page.svelte can determine the correct chat launch state.
-    this.selectedChatModel = conversation.modelId || "";
-
     return true;
   }
 
@@ -3158,23 +3154,6 @@ class AppStore {
     return (await response.json()) as TraceStatsResponse;
   }
 
-  /**
-   * Delete traces by task IDs
-   */
-  async deleteTraces(
-    taskIds: string[],
-  ): Promise<{ deleted: string[]; notFound: string[] }> {
-    const response = await fetch("/v1/traces/delete", {
-      method: "POST",
-      headers: { "Content-Type": "application/json" },
-      body: JSON.stringify({ taskIds }),
-    });
-    if (!response.ok) {
-      throw new Error(`Failed to delete traces: ${response.status}`);
-    }
-    return await response.json();
-  }
-
   /**
    * Get the URL for the raw trace file (for Perfetto)
    */
@@ -3318,5 +3297,3 @@ export const fetchTraceStats = (taskId: string) =>
   appStore.fetchTraceStats(taskId);
 export const getTraceRawUrl = (taskId: string) =>
   appStore.getTraceRawUrl(taskId);
-export const deleteTraces = (taskIds: string[]) =>
-  appStore.deleteTraces(taskIds);

dashboard/src/routes/downloads/+page.svelte

@@ -412,7 +412,7 @@
                   <div>{col.label}</div>
                   {#if col.diskAvailable != null}
                     <div
-                      class="text-[9px] text-white/70 normal-case tracking-normal mt-0.5"
+                      class="text-[9px] text-exo-light-gray/60 normal-case tracking-normal mt-0.5"
                     >
                       {formatBytes(col.diskAvailable)} free
                     </div>
@@ -436,7 +436,7 @@
                       </div>
                       {#if row.prettyName}
                         <div
-                          class="text-[10px] text-white/60"
+                          class="text-[10px] text-exo-light-gray/60"
                           title={row.modelId}
                         >
                           {row.modelId}
@@ -450,7 +450,7 @@
                       title="View model details"
                     >
                       <svg
-                        class="w-4 h-4 text-white/60 hover:text-white/80"
+                        class="w-4 h-4 text-white/30 hover:text-white/60"
                         viewBox="0 0 24 24"
                         fill="currentColor"
                       >
@@ -469,11 +469,11 @@
                   <td class="px-4 py-3 text-center align-middle">
                     {#if cell.kind === "completed"}
                       <div
-                        class="flex flex-col items-center gap-1"
+                        class="flex flex-col items-center gap-0.5"
                         title="Completed ({formatBytes(cell.totalBytes)})"
                       >
                         <svg
-                          class="w-7 h-7 text-green-400"
+                          class="w-5 h-5 text-green-400"
                           viewBox="0 0 20 20"
                           fill="currentColor"
                         >
@@ -483,18 +483,18 @@
                             clip-rule="evenodd"
                           ></path>
                         </svg>
-                        <span class="text-xs text-white/70"
+                        <span class="text-[10px] text-exo-light-gray/70"
                           >{formatBytes(cell.totalBytes)}</span
                         >
                         <button
                           type="button"
-                          class="text-white/50 hover:text-red-400 transition-colors mt-0.5 cursor-pointer"
+                          class="text-exo-light-gray/40 hover:text-red-400 transition-colors mt-0.5"
                           onclick={() =>
                             deleteDownload(col.nodeId, row.modelId)}
                           title="Delete from this node"
                         >
                           <svg
-                            class="w-5 h-5"
+                            class="w-3.5 h-3.5"
                             viewBox="0 0 20 20"
                             fill="none"
                             stroke="currentColor"
@@ -517,11 +517,11 @@
                           cell.speed,
                         )} - ETA {formatEta(cell.etaMs)}"
                       >
-                        <span class="text-exo-yellow text-sm font-medium"
+                        <span class="text-exo-yellow text-xs font-medium"
                           >{clampPercent(cell.percentage).toFixed(1)}%</span
                         >
                         <div
-                          class="w-16 h-2 bg-exo-black/60 rounded-sm overflow-hidden"
+                          class="w-14 h-1.5 bg-exo-black/60 rounded-sm overflow-hidden"
                         >
                           <div
                             class="h-full bg-gradient-to-r from-exo-yellow to-exo-yellow/70 transition-all duration-300"
@@ -530,25 +530,25 @@
                             ).toFixed(1)}%"
                           ></div>
                         </div>
-                        <span class="text-[10px] text-white/70"
+                        <span class="text-[9px] text-exo-light-gray/60"
                           >{formatSpeed(cell.speed)}</span
                         >
                       </div>
                     {:else if cell.kind === "pending"}
                       <div
-                        class="flex flex-col items-center gap-1"
+                        class="flex flex-col items-center gap-0.5"
                         title={cell.downloaded > 0
-                          ? `${formatBytes(cell.downloaded)} / ${formatBytes(cell.total)} downloaded (paused)`
+                          ? `${formatBytes(cell.downloaded)} / ${formatBytes(cell.total)} downloaded`
                           : "Download pending"}
                       >
                         {#if cell.downloaded > 0 && cell.total > 0}
-                          <span class="text-white/70 text-xs"
+                          <span class="text-exo-light-gray/70 text-[10px]"
                             >{formatBytes(cell.downloaded)} / {formatBytes(
                               cell.total,
                             )}</span
                           >
                           <div
-                            class="w-full h-1.5 bg-white/10 rounded-full overflow-hidden"
+                            class="w-full h-1 bg-white/10 rounded-full overflow-hidden"
                           >
                             <div
                               class="h-full bg-exo-light-gray/40 rounded-full"
@@ -558,65 +558,21 @@
                               ).toFixed(1)}%"
                             ></div>
                           </div>
-                          {#if row.shardMetadata}
-                            <button
-                              type="button"
-                              class="text-white/50 hover:text-exo-yellow transition-colors cursor-pointer"
-                              onclick={() =>
-                                startDownload(col.nodeId, row.shardMetadata!)}
-                              title="Resume download on this node"
-                            >
-                              <svg
-                                class="w-5 h-5"
-                                viewBox="0 0 20 20"
-                                fill="none"
-                                stroke="currentColor"
-                                stroke-width="2"
-                              >
-                                <path
-                                  d="M10 3v10m0 0l-3-3m3 3l3-3M3 17h14"
-                                  stroke-linecap="round"
-                                  stroke-linejoin="round"
-                                ></path>
-                              </svg>
-                            </button>
-                          {:else}
-                            <span class="text-white/50 text-[10px]">paused</span
-                            >
-                          {/if}
-                        {:else if row.shardMetadata}
-                          <button
-                            type="button"
-                            class="text-white/50 hover:text-exo-yellow transition-colors cursor-pointer"
-                            onclick={() =>
-                              startDownload(col.nodeId, row.shardMetadata!)}
-                            title="Start download on this node"
+                          <span class="text-exo-light-gray/40 text-[9px]"
+                            >paused</span
                           >
-                            <svg
-                              class="w-6 h-6"
-                              viewBox="0 0 20 20"
-                              fill="none"
-                              stroke="currentColor"
-                              stroke-width="2"
-                            >
-                              <path
-                                d="M10 3v10m0 0l-3-3m3 3l3-3M3 17h14"
-                                stroke-linecap="round"
-                                stroke-linejoin="round"
-                              ></path>
-                            </svg>
-                          </button>
                         {:else}
-                          <span class="text-white/40 text-sm">...</span>
+                          <span class="text-exo-light-gray/50 text-sm">...</span
+                          >
                         {/if}
                       </div>
                     {:else if cell.kind === "failed"}
                       <div
-                        class="flex flex-col items-center gap-1"
+                        class="flex flex-col items-center gap-0.5"
                         title="Download failed"
                       >
                         <svg
-                          class="w-7 h-7 text-red-400"
+                          class="w-5 h-5 text-red-400"
                           viewBox="0 0 20 20"
                           fill="currentColor"
                         >
@@ -629,13 +585,13 @@
                         {#if row.shardMetadata}
                           <button
                             type="button"
-                            class="text-white/50 hover:text-exo-yellow transition-colors cursor-pointer"
+                            class="text-exo-light-gray/40 hover:text-exo-yellow transition-colors"
                             onclick={() =>
                               startDownload(col.nodeId, row.shardMetadata!)}
                             title="Retry download on this node"
                           >
                             <svg
-                              class="w-5 h-5"
+                              class="w-3.5 h-3.5"
                               viewBox="0 0 20 20"
                               fill="none"
                               stroke="currentColor"
@@ -661,13 +617,13 @@
                         {#if row.shardMetadata}
                           <button
                             type="button"
-                            class="text-white/50 hover:text-exo-yellow transition-colors mt-0.5 opacity-0 group-hover:opacity-100 cursor-pointer"
+                            class="text-exo-light-gray/30 hover:text-exo-yellow transition-colors mt-0.5 opacity-0 group-hover:opacity-100"
                             onclick={() =>
                               startDownload(col.nodeId, row.shardMetadata!)}
                             title="Download to this node"
                           >
                             <svg
-                              class="w-5 h-5"
+                              class="w-3.5 h-3.5"
                               viewBox="0 0 20 20"
                               fill="none"
                               stroke="currentColor"

dashboard/src/routes/traces/+page.svelte

@@ -3,7 +3,6 @@
   import {
     listTraces,
     getTraceRawUrl,
-    deleteTraces,
     type TraceListItem,
   } from "$lib/stores/app.svelte";
   import HeaderNav from "$lib/components/HeaderNav.svelte";
@@ -11,51 +10,6 @@
   let traces = $state<TraceListItem[]>([]);
   let loading = $state(true);
   let error = $state<string | null>(null);
-  let selectedIds = $state<Set<string>>(new Set());
-  let deleting = $state(false);
-
-  let allSelected = $derived(
-    traces.length > 0 && selectedIds.size === traces.length,
-  );
-
-  function toggleSelect(taskId: string) {
-    const next = new Set(selectedIds);
-    if (next.has(taskId)) {
-      next.delete(taskId);
-    } else {
-      next.add(taskId);
-    }
-    selectedIds = next;
-  }
-
-  function toggleSelectAll() {
-    if (allSelected) {
-      selectedIds = new Set();
-    } else {
-      selectedIds = new Set(traces.map((t) => t.taskId));
-    }
-  }
-
-  async function handleDelete() {
-    if (selectedIds.size === 0) return;
-    const count = selectedIds.size;
-    if (
-      !confirm(
-        `Delete ${count} trace${count === 1 ? "" : "s"}? This cannot be undone.`,
-      )
-    )
-      return;
-    deleting = true;
-    try {
-      await deleteTraces([...selectedIds]);
-      selectedIds = new Set();
-      await refresh();
-    } catch (e) {
-      error = e instanceof Error ? e.message : "Failed to delete traces";
-    } finally {
-      deleting = false;
-    }
-  }
 
   function formatBytes(bytes: number): string {
     if (!bytes || bytes <= 0) return "0B";
@@ -155,16 +109,6 @@
         </h1>
       </div>
       <div class="flex items-center gap-3">
-        {#if selectedIds.size > 0}
-          <button
-            type="button"
-            class="text-xs font-mono text-red-400 hover:text-red-300 transition-colors uppercase border border-red-500/40 px-2 py-1 rounded"
-            onclick={handleDelete}
-            disabled={deleting}
-          >
-            {deleting ? "Deleting..." : `Delete (${selectedIds.size})`}
-          </button>
-        {/if}
         <button
           type="button"
           class="text-xs font-mono text-exo-light-gray hover:text-exo-yellow transition-colors uppercase border border-exo-medium-gray/40 px-2 py-1 rounded"
@@ -199,41 +143,14 @@
       </div>
     {:else}
       <div class="space-y-3">
-        <div class="flex items-center gap-2 px-1">
-          <button
-            type="button"
-            class="text-xs font-mono uppercase transition-colors {allSelected
-              ? 'text-exo-yellow'
-              : 'text-exo-light-gray hover:text-exo-yellow'}"
-            onclick={toggleSelectAll}
-          >
-            {allSelected ? "Deselect all" : "Select all"}
-          </button>
-        </div>
         {#each traces as trace}
-          {@const isSelected = selectedIds.has(trace.taskId)}
-          <!-- svelte-ignore a11y_no_static_element_interactions -->
           <div
-            role="button"
-            tabindex="0"
-            class="w-full text-left rounded border-l-2 border-r border-t border-b transition-all p-4 flex items-center justify-between gap-4 cursor-pointer {isSelected
-              ? 'bg-exo-yellow/10 border-l-exo-yellow border-r-exo-medium-gray/30 border-t-exo-medium-gray/30 border-b-exo-medium-gray/30'
-              : 'bg-exo-black/30 border-l-transparent border-r-exo-medium-gray/30 border-t-exo-medium-gray/30 border-b-exo-medium-gray/30 hover:bg-white/[0.03]'}"
-            onclick={() => toggleSelect(trace.taskId)}
-            onkeydown={(e) => {
-              if (e.key === "Enter" || e.key === " ") {
-                e.preventDefault();
-                toggleSelect(trace.taskId);
-              }
-            }}
+            class="rounded border border-exo-medium-gray/30 bg-exo-black/30 p-4 flex items-center justify-between gap-4"
           >
             <div class="min-w-0 flex-1">
               <a
                 href="#/traces/{trace.taskId}"
-                class="text-sm font-mono transition-colors truncate block {isSelected
-                  ? 'text-exo-yellow'
-                  : 'text-white hover:text-exo-yellow'}"
-                onclick={(e) => e.stopPropagation()}
+                class="text-sm font-mono text-white hover:text-exo-yellow transition-colors truncate block"
               >
                 {trace.taskId}
               </a>
@@ -243,11 +160,7 @@
                 )}
               </div>
             </div>
-            <!-- svelte-ignore a11y_click_events_have_key_events -->
-            <div
-              class="flex items-center gap-2 shrink-0"
-              onclick={(e) => e.stopPropagation()}
-            >
+            <div class="flex items-center gap-2 shrink-0">
               <a
                 href="#/traces/{trace.taskId}"
                 class="text-xs font-mono text-exo-light-gray hover:text-exo-yellow transition-colors uppercase border border-exo-medium-gray/40 px-2 py-1 rounded"

docs/api.md

@@ -183,27 +183,6 @@ Same schema as `/v1/chat/completions`.
 **Response:**
 Chat completion plus benchmarking metrics.
 
-### Cancel Command
-
-**POST** `/v1/cancel/{command_id}`
-
-Cancels an active generation command (text or image). Notifies workers and closes the stream.
-
-**Path parameters:**
-
-* `command_id`: string, ID of the command to cancel
-
-**Response (example):**
-
-```json
-{
-  "message": "Command cancelled.",
-  "command_id": "cmd-abc-123"
-}
-```
-
-Returns 404 if the command is not found or already completed.
-
 ## 5. Image Generation & Editing
 
 ### Image Generation
@@ -287,7 +266,6 @@ GET     /v1/models
 
 POST    /v1/chat/completions
 POST    /bench/chat/completions
-POST    /v1/cancel/{command_id}
 
 POST    /v1/images/generations
 POST    /bench/images/generations

flake.nix

@@ -108,7 +108,6 @@
                 package = pkgsSwift.swiftPackages.swift-format;
               };
               shfmt.enable = true;
-              taplo.enable = true;
             };
           };
 

nix/mlx.nix

@@ -41,7 +41,7 @@ let
 
   mlx = stdenv.mkDerivation rec {
     pname = "mlx";
-    version = let v = "0.30.7.dev20260225+257d5692"; in
+    version = let v = "0.30.7.dev20260220+13998a05"; in
       assert v == uvLockMlxVersion || throw "MLX version mismatch: nix/mlx.nix has ${v} but uv.lock has ${uvLockMlxVersion}. Update both the version and hash in nix/mlx.nix.";
       v;
     pyproject = true;
@@ -49,8 +49,8 @@ let
     src = fetchFromGitHub {
       owner = "rltakashige";
       repo = "mlx-jaccl-fix-small-recv";
-      rev = "257d5692fc7af6bba3b8afaeb63c549b7d1e43d5";
-      hash = "sha256-GosFIWxIB48Egb1MqJrR3xhsUsQeWdRk5rV93USY6wQ=";
+      rev = "13998a054715edcdc93618fb1496c79c7c25ff7c";
+      hash = "sha256-fAqA3hFwNBx7FcoGnhQsIFpAIRbC2EerACm4Fvne0Cc=";
     };
 
     patches = [

pyproject.toml

@@ -1,35 +1,35 @@
 [project]
 name = "exo"
-version = "0.3.68"
+version = "0.3.0"
 description = "Exo"
 readme = "README.md"
 requires-python = ">=3.13"
 dependencies = [
-  "aiofiles>=24.1.0",
-  "aiohttp>=3.12.14",
-  "types-aiofiles>=24.1.0.20250708",
-  "pydantic>=2.11.7",
-  "fastapi>=0.116.1",
-  "filelock>=3.18.0",
-  "rustworkx>=0.17.1",
-  "huggingface-hub>=0.33.4",
-  "psutil>=7.0.0",
-  "loguru>=0.7.3",
-  "exo_pyo3_bindings",                         # rust bindings
-  "anyio==4.11.0",
-  "mlx; sys_platform == 'darwin'",
-  "mlx[cpu]==0.30.6; sys_platform == 'linux'",
-  "mlx-lm",
-  "tiktoken>=0.12.0",                          # required for kimi k2 tokenizer
-  "hypercorn>=0.18.0",
-  "openai-harmony>=0.0.8",
-  "httpx>=0.28.1",
-  "tomlkit>=0.14.0",
-  "pillow>=11.0,<12.0",                        # compatibility with mflux
-  "mflux==0.15.5",
-  "python-multipart>=0.0.21",
-  "msgspec>=0.19.0",
-  "zstandard>=0.23.0",
+    "aiofiles>=24.1.0",
+    "aiohttp>=3.12.14",
+    "types-aiofiles>=24.1.0.20250708",
+    "pydantic>=2.11.7",
+    "fastapi>=0.116.1",
+    "filelock>=3.18.0",
+    "rustworkx>=0.17.1",
+    "huggingface-hub>=0.33.4",
+    "psutil>=7.0.0",
+    "loguru>=0.7.3",
+    "exo_pyo3_bindings", # rust bindings
+    "anyio==4.11.0",
+    "mlx; sys_platform == 'darwin'",
+    "mlx[cpu]==0.30.6; sys_platform == 'linux'",
+    "mlx-lm==0.30.7",
+    "tiktoken>=0.12.0", # required for kimi k2 tokenizer
+    "hypercorn>=0.18.0",
+    "openai-harmony>=0.0.8",
+    "httpx>=0.28.1",
+    "tomlkit>=0.14.0",
+    "pillow>=11.0,<12.0", # compatibility with mflux
+    "mflux==0.15.5",
+    "python-multipart>=0.0.21",
+    "msgspec>=0.19.0",
+    "zstandard>=0.23.0",
 ]
 
 [project.scripts]
@@ -38,12 +38,12 @@ exo = "exo.main:main"
 # dependencies only required for development
 [dependency-groups]
 dev = [
-  "basedpyright>=1.29.0",
-  "pyinstaller>=6.17.0",
-  "pytest>=8.4.0",
-  "pytest-asyncio>=1.0.0",
-  "pytest-env",
-  "ruff>=0.11.13",
+    "basedpyright>=1.29.0",
+    "pyinstaller>=6.17.0",
+    "pytest>=8.4.0",
+    "pytest-asyncio>=1.0.0",
+    "pytest-env",
+    "ruff>=0.11.13",
 ]
 
 # mlx[cuda] requires a newer version of mlx. the ideal on linux is: default to mlx[cpu] unless[cuda] specified.
@@ -57,12 +57,15 @@ dev = [
 ###
 
 [tool.uv.workspace]
-members = ["rust/exo_pyo3_bindings", "bench"]
+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/ml-explore/mlx-lm", rev = "834fac934c4e04de9b3d723e2b9287a2c60cfd4a" }
+#mlx-lm = { git = "https://github.com/davidmcc73/mlx-lm", branch = "stable" }
 # 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 }
@@ -92,15 +95,7 @@ reportUnnecessaryTypeIgnoreComment = "error"
 pythonVersion = "3.13"
 pythonPlatform = "Darwin"
 
-exclude = [
-  "**/.venv",
-  "**/venv",
-  "**/__pycache__",
-  "**/exo_scripts",
-  "**/.direnv",
-  "**/rust",
-  "**/.github",
-]
+exclude = ["**/.venv", "**/venv", "**/__pycache__", "**/exo_scripts", "**/.direnv", "**/rust", "**/.github"]
 stubPath = ".mlx_typings"
 
 [[tool.basedpyright.executionEnvironments]]
@@ -114,18 +109,17 @@ root = "src"
 [tool.uv]
 required-version = ">=0.8.6"
 prerelease = "allow"
-environments = ["sys_platform == 'darwin'", "sys_platform == 'linux'"]
+environments = [
+    "sys_platform == 'darwin'",
+    "sys_platform == 'linux'",
+]
 
 ###
 # ruff configuration
 ###
 
 [tool.ruff]
-extend-exclude = [
-  "shared/protobufs/**",
-  "*mlx_typings/**",
-  "rust/exo_pyo3_bindings/**",
-]
+extend-exclude = ["shared/protobufs/**", "*mlx_typings/**", "rust/exo_pyo3_bindings/**"]
 
 [tool.ruff.lint]
 extend-select = ["I", "N", "B", "A", "PIE", "SIM"]
@@ -133,7 +127,13 @@ extend-select = ["I", "N", "B", "A", "PIE", "SIM"]
 [tool.pytest.ini_options]
 pythonpath = "."
 asyncio_mode = "auto"
-markers = ["slow: marks tests as slow (deselected by default)"]
-env = ["EXO_TESTS=1"]
+markers = [
+    "slow: marks tests as slow (deselected by default)"
+]
+env = [
+  "EXO_TESTS=1"
+]
 addopts = "-m 'not slow' --ignore=tests/start_distributed_test.py"
-filterwarnings = ["ignore:builtin type Swig:DeprecationWarning"]
+filterwarnings = [
+    "ignore:builtin type Swig:DeprecationWarning",
+]

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

@@ -1,12 +0,0 @@
-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

@@ -1,12 +0,0 @@
-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

@@ -1,12 +0,0 @@
-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

@@ -1,12 +0,0 @@
-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

@@ -1,12 +0,0 @@
-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

@@ -1,12 +0,0 @@
-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

@@ -1,12 +0,0 @@
-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

@@ -1,12 +0,0 @@
-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

@@ -1,12 +0,0 @@
-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

@@ -1,12 +0,0 @@
-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

@@ -1,12 +0,0 @@
-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

@@ -1,12 +0,0 @@
-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

@@ -1,12 +0,0 @@
-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

@@ -1,12 +0,0 @@
-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

@@ -5,6 +5,7 @@ edition = { workspace = true }
 publish = false
 
 [lib]
+doctest = false
 path = "src/lib.rs"
 name = "exo_pyo3_bindings"
 
@@ -25,33 +26,37 @@ networking = { workspace = true }
 
 # interop
 pyo3 = { version = "0.27.2", features = [
-  # "abi3-py313", # tells pyo3 (and maturin) to build using the stable ABI with minimum Python version 3.13
-  # "nightly", # enables better-supported GIL integration
-  "experimental-async", # async support in #[pyfunction] & #[pymethods]
-  #"experimental-inspect", # inspection of generated binary => easier to automate type-hint generation
-  #"py-clone", # adding Clone-ing of `Py<T>` without GIL (may cause panics - remove if panics happen)
-  # "multiple-pymethods", # allows multiple #[pymethods] sections per class
+    # "abi3-py313", # tells pyo3 (and maturin) to build using the stable ABI with minimum Python version 3.13
+    # "nightly", # enables better-supported GIL integration
+    "experimental-async", # async support in #[pyfunction] & #[pymethods]
+    #"experimental-inspect", # inspection of generated binary => easier to automate type-hint generation
+    #"py-clone", # adding Clone-ing of `Py<T>` without GIL (may cause panics - remove if panics happen)
+    # "multiple-pymethods", # allows multiple #[pymethods] sections per class
 
-  # integrations with other libraries
-  # "arc_lock", "bigdecimal", "either", "hashbrown", "indexmap", "num-bigint", "num-complex", "num-rational",
-  # "ordered-float", "rust_decimal", "smallvec",
-  # "anyhow", "chrono", "chrono-local", "chrono-tz", "eyre", "jiff-02", "lock_api", "parking-lot", "time",  "serde",
+    # integrations with other libraries
+    # "arc_lock", "bigdecimal", "either", "hashbrown", "indexmap", "num-bigint", "num-complex", "num-rational",
+    # "ordered-float", "rust_decimal", "smallvec",
+    # "anyhow", "chrono", "chrono-local", "chrono-tz", "eyre", "jiff-02", "lock_api", "parking-lot", "time",  "serde",
 ] }
 pyo3-stub-gen = { version = "0.17.2" }
-pyo3-async-runtimes = { version = "0.27.0", features = [
-  "attributes",
-  "tokio-runtime",
-  "testing",
-] }
+pyo3-async-runtimes = { version = "0.27.0", features = ["attributes", "tokio-runtime", "testing"] }
 pyo3-log = "0.13.2"
 
 # macro dependencies
 extend = { workspace = true }
+delegate = { workspace = true }
 
 # async runtime
 tokio = { workspace = true, features = ["full", "tracing"] }
 futures-lite = { workspace = true }
 
+# utility dependencies
+util = { workspace = true }
+
 # Tracing
 log = { workspace = true }
 env_logger = "0.11"
+
+# Networking
+libp2p = { workspace = true, features = ["full"] }
+pin-project = "1.1.10"

rust/exo_pyo3_bindings/README.md

@@ -0,0 +1 @@
+TODO: do something here....

rust/exo_pyo3_bindings/exo_pyo3_bindings.pyi

@@ -2,6 +2,7 @@
 # ruff: noqa: E501, F401
 
 import builtins
+import enum
 import typing
 
 @typing.final
@@ -10,6 +11,29 @@ class AllQueuesFullError(builtins.Exception):
     def __repr__(self) -> builtins.str: ...
     def __str__(self) -> builtins.str: ...
 
+@typing.final
+class ConnectionUpdate:
+    @property
+    def update_type(self) -> ConnectionUpdateType:
+        r"""
+        Whether this is a connection or disconnection event
+        """
+    @property
+    def peer_id(self) -> builtins.str:
+        r"""
+        Identity of the peer that we have connected to or disconnected from.
+        """
+    @property
+    def remote_ipv4(self) -> builtins.str:
+        r"""
+        Remote connection's IPv4 address.
+        """
+    @property
+    def remote_tcp_port(self) -> builtins.int:
+        r"""
+        Remote connection's TCP port.
+        """
+
 @typing.final
 class Keypair:
     r"""
@@ -34,15 +58,21 @@ class Keypair:
         Convert the `Keypair` into the corresponding `PeerId` string, which we use as our `NodeId`.
         """
 
-@typing.final
-class MessageTooLargeError(builtins.Exception):
-    def __new__(cls, *args: typing.Any) -> MessageTooLargeError: ...
-    def __repr__(self) -> builtins.str: ...
-    def __str__(self) -> builtins.str: ...
-
 @typing.final
 class NetworkingHandle:
     def __new__(cls, identity: Keypair) -> NetworkingHandle: ...
+    async def connection_update_recv(self) -> ConnectionUpdate:
+        r"""
+        Receives the next `ConnectionUpdate` from networking.
+        """
+    async def connection_update_recv_many(self, limit: builtins.int) -> builtins.list[ConnectionUpdate]:
+        r"""
+        Receives at most `limit` `ConnectionUpdate`s from networking and returns them.
+        
+        For `limit = 0`, an empty collection of `ConnectionUpdate`s will be returned immediately.
+        For `limit > 0`, if there are no `ConnectionUpdate`s in the channel's queue this method
+        will sleep until a `ConnectionUpdate`s is sent.
+        """
     async def gossipsub_subscribe(self, topic: builtins.str) -> builtins.bool:
         r"""
         Subscribe to a `GossipSub` topic.
@@ -61,7 +91,18 @@ class NetworkingHandle:
         
         If no peers are found that subscribe to this topic, throws `NoPeersSubscribedToTopicError` exception.
         """
-    async def recv(self) -> PyFromSwarm: ...
+    async def gossipsub_recv(self) -> tuple[builtins.str, bytes]:
+        r"""
+        Receives the next message from the `GossipSub` network.
+        """
+    async def gossipsub_recv_many(self, limit: builtins.int) -> builtins.list[tuple[builtins.str, bytes]]:
+        r"""
+        Receives at most `limit` messages from the `GossipSub` network and returns them.
+        
+        For `limit = 0`, an empty collection of messages will be returned immediately.
+        For `limit > 0`, if there are no messages in the channel's queue this method
+        will sleep until a message is sent.
+        """
 
 @typing.final
 class NoPeersSubscribedToTopicError(builtins.Exception):
@@ -69,26 +110,11 @@ class NoPeersSubscribedToTopicError(builtins.Exception):
     def __repr__(self) -> builtins.str: ...
     def __str__(self) -> builtins.str: ...
 
-class PyFromSwarm:
-    @typing.final
-    class Connection(PyFromSwarm):
-        __match_args__ = ("peer_id", "connected",)
-        @property
-        def peer_id(self) -> builtins.str: ...
-        @property
-        def connected(self) -> builtins.bool: ...
-        def __new__(cls, peer_id: builtins.str, connected: builtins.bool) -> PyFromSwarm.Connection: ...
-    
-    @typing.final
-    class Message(PyFromSwarm):
-        __match_args__ = ("origin", "topic", "data",)
-        @property
-        def origin(self) -> builtins.str: ...
-        @property
-        def topic(self) -> builtins.str: ...
-        @property
-        def data(self) -> bytes: ...
-        def __new__(cls, origin: builtins.str, topic: builtins.str, data: bytes) -> PyFromSwarm.Message: ...
-    
-    ...
+@typing.final
+class ConnectionUpdateType(enum.Enum):
+    r"""
+    Connection or disconnection event discriminant type.
+    """
+    Connected = ...
+    Disconnected = ...
 

rust/exo_pyo3_bindings/pyproject.toml

@@ -4,18 +4,21 @@ build-backend = "maturin"
 
 [project]
 name = "exo_pyo3_bindings"
-version = "0.2.0"
+version = "0.1.0"
 description = "Add your description here"
 readme = "README.md"
 authors = [
-  { name = "Andrei Cravtov", email = "the.andrei.cravtov@gmail.com" },
-  { name = "Evan Quiney", email = "evanev7@gmail.com" },
+    { name = "Andrei Cravtov", email = "the.andrei.cravtov@gmail.com" }
 ]
 requires-python = ">=3.13"
 dependencies = []
 
 [dependency-groups]
-dev = ["exo_pyo3_bindings", "pytest>=8.4.0", "pytest-asyncio>=1.0.0"]
+dev = [
+    "exo_pyo3_bindings",
+    "pytest>=8.4.0",
+    "pytest-asyncio>=1.0.0",
+]
 
 [tool.maturin]
 #purelib = true

rust/exo_pyo3_bindings/src/allow_threading.rs

@@ -1,22 +1,37 @@
-//! See: <https://pyo3.rs/v0.27.2/async-await.html#detaching-from-the-interpreter-across-await>
+//! SEE: https://pyo3.rs/v0.26.0/async-await.html#detaching-from-the-interpreter-across-await
+//!
+
+use pin_project::pin_project;
 use pyo3::prelude::*;
 use std::{
     future::Future,
-    pin::{Pin, pin},
+    pin::Pin,
     task::{Context, Poll},
 };
 
-pub struct AllowThreads<F>(pub(crate) F);
+/// SEE: https://pyo3.rs/v0.26.0/async-await.html#detaching-from-the-interpreter-across-await
+#[pin_project]
+#[repr(transparent)]
+pub(crate) struct AllowThreads<F>(#[pin] F);
+
+impl<F> AllowThreads<F>
+where
+    Self: Future,
+{
+    pub fn new(f: F) -> Self {
+        Self(f)
+    }
+}
 
 impl<F> Future for AllowThreads<F>
 where
-    F: Future + Unpin + Send,
+    F: Future + Send,
     F::Output: Send,
 {
     type Output = F::Output;
 
-    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
+    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
         let waker = cx.waker();
-        Python::attach(|py| py.detach(|| pin!(&mut self.0).poll(&mut Context::from_waker(waker))))
+        Python::attach(|py| py.detach(|| self.project().0.poll(&mut Context::from_waker(waker))))
     }
 }

rust/exo_pyo3_bindings/src/lib.rs

@@ -1,5 +1,171 @@
-use pyo3_stub_gen::define_stub_info_gatherer;
+//! TODO: crate documentation
+//!
+//! this is here as a placeholder documentation
+//!
+//!
 
 mod allow_threading;
+mod ident;
+mod networking;
+
+use crate::ident::PyKeypair;
+use crate::networking::networking_submodule;
+use pyo3::prelude::PyModule;
+use pyo3::types::PyModuleMethods;
+use pyo3::{Bound, PyResult, pyclass, pymodule};
+use pyo3_stub_gen::define_stub_info_gatherer;
+
+/// Namespace for all the constants used by this crate.
+pub(crate) mod r#const {
+    pub const MPSC_CHANNEL_SIZE: usize = 1024;
+}
+
+/// Namespace for crate-wide extension traits/methods
+pub(crate) mod ext {
+    use crate::allow_threading::AllowThreads;
+    use extend::ext;
+    use pyo3::exceptions::{PyConnectionError, PyRuntimeError};
+    use pyo3::types::PyBytes;
+    use pyo3::{Py, PyErr, PyResult, Python};
+    use tokio::runtime::Runtime;
+    use tokio::sync::mpsc;
+    use tokio::sync::mpsc::error::TryRecvError;
+    use tokio::task::JoinHandle;
+
+    #[ext(pub, name = ByteArrayExt)]
+    impl [u8] {
+        fn pybytes(&self) -> Py<PyBytes> {
+            Python::attach(|py| PyBytes::new(py, self).unbind())
+        }
+    }
+
+    #[ext(pub, name = ResultExt)]
+    impl<T, E> Result<T, E>
+    where
+        E: ToString,
+    {
+        fn pyerr(self) -> PyResult<T> {
+            self.map_err(|e| PyRuntimeError::new_err(e.to_string()))
+        }
+    }
+
+    pub trait FutureExt: Future + Sized {
+        /// SEE: https://pyo3.rs/v0.26.0/async-await.html#detaching-from-the-interpreter-across-await
+        fn allow_threads_py(self) -> AllowThreads<Self>
+        where
+            AllowThreads<Self>: Future,
+        {
+            AllowThreads::new(self)
+        }
+    }
+
+    impl<T: Future> FutureExt for T {}
+
+    #[ext(pub, name = PyErrExt)]
+    impl PyErr {
+        fn receiver_channel_closed() -> Self {
+            PyConnectionError::new_err("Receiver channel closed unexpectedly")
+        }
+    }
+
+    #[ext(pub, name = PyResultExt)]
+    impl<T> PyResult<T> {
+        fn write_unraisable(self) -> Option<T> {
+            Python::attach(|py| self.write_unraisable_with(py))
+        }
+
+        fn write_unraisable_with(self, py: Python<'_>) -> Option<T> {
+            match self {
+                Ok(v) => Some(v),
+                Err(e) => {
+                    // write error back to python
+                    e.write_unraisable(py, None);
+                    None
+                }
+            }
+        }
+    }
+
+    #[ext(pub, name = TokioRuntimeExt)]
+    impl Runtime {
+        fn spawn_with_scope<F>(&self, py: Python<'_>, future: F) -> PyResult<JoinHandle<F::Output>>
+        where
+            F: Future + Send + 'static,
+            F::Output: Send + 'static,
+        {
+            let locals = pyo3_async_runtimes::tokio::get_current_locals(py)?;
+            Ok(self.spawn(pyo3_async_runtimes::tokio::scope(locals, future)))
+        }
+    }
+
+    #[ext(pub, name = TokioMpscSenderExt)]
+    impl<T> mpsc::Sender<T> {
+        /// Sends a value, waiting until there is capacity.
+        ///
+        /// A successful send occurs when it is determined that the other end of the
+        /// channel has not hung up already. An unsuccessful send would be one where
+        /// the corresponding receiver has already been closed.
+        async fn send_py(&self, value: T) -> PyResult<()> {
+            self.send(value)
+                .await
+                .map_err(|_| PyErr::receiver_channel_closed())
+        }
+    }
+
+    #[ext(pub, name = TokioMpscReceiverExt)]
+    impl<T> mpsc::Receiver<T> {
+        /// Receives the next value for this receiver.
+        async fn recv_py(&mut self) -> PyResult<T> {
+            self.recv().await.ok_or_else(PyErr::receiver_channel_closed)
+        }
+
+        /// Receives at most `limit` values for this receiver and returns them.
+        ///
+        /// For `limit = 0`, an empty collection of messages will be returned immediately.
+        /// For `limit > 0`, if there are no messages in the channel's queue this method
+        /// will sleep until a message is sent.
+        async fn recv_many_py(&mut self, limit: usize) -> PyResult<Vec<T>> {
+            // get updates from receiver channel
+            let mut updates = Vec::with_capacity(limit);
+            let received = self.recv_many(&mut updates, limit).await;
+
+            // if we received zero items, then the channel was unexpectedly closed
+            if limit != 0 && received == 0 {
+                return Err(PyErr::receiver_channel_closed());
+            }
+
+            Ok(updates)
+        }
+
+        /// Tries to receive the next value for this receiver.
+        fn try_recv_py(&mut self) -> PyResult<Option<T>> {
+            match self.try_recv() {
+                Ok(v) => Ok(Some(v)),
+                Err(TryRecvError::Empty) => Ok(None),
+                Err(TryRecvError::Disconnected) => Err(PyErr::receiver_channel_closed()),
+            }
+        }
+    }
+}
+
+/// A Python module implemented in Rust. The name of this function must match
+/// the `lib.name` setting in the `Cargo.toml`, else Python will not be able to
+/// import the module.
+#[pymodule(name = "exo_pyo3_bindings")]
+fn main_module(m: &Bound<'_, PyModule>) -> PyResult<()> {
+    // install logger
+    pyo3_log::init();
+
+    // TODO: for now this is all NOT a submodule, but figure out how to make the submodule system
+    //       work with maturin, where the types generate correctly, in the right folder, without
+    //       too many importing issues...
+    m.add_class::<PyKeypair>()?;
+    networking_submodule(m)?;
+
+    // top-level constructs
+    // TODO: ...
+
+    Ok(())
+}
 
 define_stub_info_gatherer!(stub_info);

rust/exo_pyo3_bindings/src/networking.rs

@@ -0,0 +1,571 @@
+#![allow(
+    clippy::multiple_inherent_impl,
+    clippy::unnecessary_wraps,
+    clippy::unused_self,
+    clippy::needless_pass_by_value
+)]
+
+use crate::r#const::MPSC_CHANNEL_SIZE;
+use crate::ext::{ByteArrayExt as _, FutureExt, PyErrExt as _};
+use crate::ext::{ResultExt as _, TokioMpscReceiverExt as _, TokioMpscSenderExt as _};
+use crate::ident::PyKeypair;
+use crate::pyclass;
+use libp2p::futures::StreamExt as _;
+use libp2p::gossipsub;
+use libp2p::gossipsub::{IdentTopic, Message, MessageId, PublishError};
+use libp2p::swarm::SwarmEvent;
+use networking::discovery;
+use networking::swarm::create_swarm;
+use pyo3::prelude::{PyModule, PyModuleMethods as _};
+use pyo3::types::PyBytes;
+use pyo3::{Bound, Py, PyErr, PyResult, PyTraverseError, PyVisit, Python, pymethods};
+use pyo3_stub_gen::derive::{gen_stub_pyclass, gen_stub_pyclass_enum, gen_stub_pymethods};
+use std::net::IpAddr;
+use tokio::sync::{Mutex, mpsc, oneshot};
+
+mod exception {
+    use pyo3::types::PyTuple;
+    use pyo3::{exceptions::PyException, prelude::*};
+    use pyo3_stub_gen::derive::*;
+
+    #[gen_stub_pyclass]
+    #[pyclass(frozen, extends=PyException, name="NoPeersSubscribedToTopicError")]
+    pub struct PyNoPeersSubscribedToTopicError {}
+
+    impl PyNoPeersSubscribedToTopicError {
+        const MSG: &'static str = "\
+        No peers are currently subscribed to receive messages on this topic. \
+        Wait for peers to subscribe or check your network connectivity.";
+
+        ///   Creates a new  [ `PyErr` ]  of this type.
+        ///
+        ///   [`PyErr`] :  https://docs.rs/pyo3/latest/pyo3/struct.PyErr.html   "PyErr in pyo3"
+        pub(crate) fn new_err() -> PyErr {
+            PyErr::new::<Self, _>(()) // TODO: check if this needs to be replaced???
+        }
+    }
+
+    #[gen_stub_pymethods]
+    #[pymethods]
+    impl PyNoPeersSubscribedToTopicError {
+        #[new]
+        #[pyo3(signature = (*args))]
+        #[allow(unused_variables)]
+        pub(crate) fn new(args: &Bound<'_, PyTuple>) -> Self {
+            Self {}
+        }
+
+        fn __repr__(&self) -> String {
+            format!("PeerId(\"{}\")", Self::MSG)
+        }
+
+        fn __str__(&self) -> String {
+            Self::MSG.to_string()
+        }
+    }
+
+    #[gen_stub_pyclass]
+    #[pyclass(frozen, extends=PyException, name="AllQueuesFullError")]
+    pub struct PyAllQueuesFullError {}
+
+    impl PyAllQueuesFullError {
+        const MSG: &'static str =
+            "All libp2p peers are unresponsive, resend the message or reconnect.";
+
+        ///   Creates a new  [ `PyErr` ]  of this type.
+        ///
+        ///   [`PyErr`] :  https://docs.rs/pyo3/latest/pyo3/struct.PyErr.html   "PyErr in pyo3"
+        pub(crate) fn new_err() -> PyErr {
+            PyErr::new::<Self, _>(()) // TODO: check if this needs to be replaced???
+        }
+    }
+
+    #[gen_stub_pymethods]
+    #[pymethods]
+    impl PyAllQueuesFullError {
+        #[new]
+        #[pyo3(signature = (*args))]
+        #[allow(unused_variables)]
+        pub(crate) fn new(args: &Bound<'_, PyTuple>) -> Self {
+            Self {}
+        }
+
+        fn __repr__(&self) -> String {
+            format!("PeerId(\"{}\")", Self::MSG)
+        }
+
+        fn __str__(&self) -> String {
+            Self::MSG.to_string()
+        }
+    }
+}
+
+/// Connection or disconnection event discriminant type.
+#[gen_stub_pyclass_enum]
+#[pyclass(eq, eq_int, name = "ConnectionUpdateType")]
+#[derive(Debug, Clone, PartialEq)]
+enum PyConnectionUpdateType {
+    Connected = 0,
+    Disconnected,
+}
+
+#[gen_stub_pyclass]
+#[pyclass(frozen, name = "ConnectionUpdate")]
+#[derive(Debug, Clone)]
+struct PyConnectionUpdate {
+    /// Whether this is a connection or disconnection event
+    #[pyo3(get)]
+    update_type: PyConnectionUpdateType,
+
+    /// Identity of the peer that we have connected to or disconnected from.
+    #[pyo3(get)]
+    peer_id: String,
+
+    /// Remote connection's IPv4 address.
+    #[pyo3(get)]
+    remote_ipv4: String,
+
+    /// Remote connection's TCP port.
+    #[pyo3(get)]
+    remote_tcp_port: u16,
+}
+
+enum ToTask {
+    GossipsubSubscribe {
+        topic: String,
+        result_tx: oneshot::Sender<PyResult<bool>>,
+    },
+    GossipsubUnsubscribe {
+        topic: String,
+        result_tx: oneshot::Sender<bool>,
+    },
+    GossipsubPublish {
+        topic: String,
+        data: Vec<u8>,
+        result_tx: oneshot::Sender<PyResult<MessageId>>,
+    },
+}
+
+#[allow(clippy::enum_glob_use)]
+async fn networking_task(
+    mut swarm: networking::swarm::Swarm,
+    mut to_task_rx: mpsc::Receiver<ToTask>,
+    connection_update_tx: mpsc::Sender<PyConnectionUpdate>,
+    gossipsub_message_tx: mpsc::Sender<(String, Vec<u8>)>,
+) {
+    use SwarmEvent::*;
+    use ToTask::*;
+    use networking::swarm::BehaviourEvent::*;
+
+    log::info!("RUST: networking task started");
+
+    loop {
+        tokio::select! {
+            message = to_task_rx.recv() => {
+                // handle closed channel
+                let Some(message) = message else {
+                    log::info!("RUST: channel closed");
+                    break;
+                };
+
+                // dispatch incoming messages
+                match message {
+                    GossipsubSubscribe { topic, result_tx } => {
+                        // try to subscribe
+                        let result = swarm.behaviour_mut()
+                            .gossipsub.subscribe(&IdentTopic::new(topic));
+
+                        // send response oneshot
+                        if let Err(e) = result_tx.send(result.pyerr()) {
+                            log::error!("RUST: could not subscribe to gossipsub topic since channel already closed: {e:?}");
+                            continue;
+                        }
+                    }
+                    GossipsubUnsubscribe { topic, result_tx } => {
+                        // try to unsubscribe from the topic
+                        let result = swarm.behaviour_mut()
+                            .gossipsub.unsubscribe(&IdentTopic::new(topic));
+
+                        // send response oneshot (or exit if connection closed)
+                        if let Err(e) = result_tx.send(result) {
+                            log::error!("RUST: could not unsubscribe from gossipsub topic since channel already closed: {e:?}");
+                            continue;
+                        }
+                    }
+                    GossipsubPublish { topic, data, result_tx } => {
+                        // try to publish the data -> catch NoPeersSubscribedToTopic error & convert to correct exception
+                        let result = swarm.behaviour_mut().gossipsub.publish(
+                            IdentTopic::new(topic), data);
+                        let pyresult: PyResult<MessageId> = if let Err(PublishError::NoPeersSubscribedToTopic) = result {
+                            Err(exception::PyNoPeersSubscribedToTopicError::new_err())
+                        } else if let Err(PublishError::AllQueuesFull(_)) = result {
+                            Err(exception::PyAllQueuesFullError::new_err())
+                        } else {
+                            result.pyerr()
+                        };
+
+                        // send response oneshot (or exit if connection closed)
+                        if let Err(e) = result_tx.send(pyresult) {
+                            log::error!("RUST: could not publish gossipsub message since channel already closed: {e:?}");
+                            continue;
+                        }
+                    }
+                }
+            }
+
+            // architectural solution to this problem:
+            // create keep_alive behavior who's job it is to dial peers discovered by mDNS (and drop when expired)
+            //   -> it will emmit TRUE connected/disconnected events consumable elsewhere
+            //
+            // gossipsub will feed off-of dial attempts created by networking, and that will bootstrap its' peers list
+            // then for actual communication it will dial those peers if need-be
+            swarm_event = swarm.select_next_some() => {
+                match swarm_event {
+                    Behaviour(Gossipsub(gossipsub::Event::Message {
+                        message: Message {
+                            topic,
+                            data,
+                            ..
+                        },
+                        ..
+                    })) => {
+                        // topic-ID is just the topic hash!!! (since we used identity hasher)
+                        let message = (topic.into_string(), data);
+
+                        // send incoming message to channel (or exit if connection closed)
+                        if let Err(e) = gossipsub_message_tx.send(message).await {
+                            log::error!("RUST: could not send incoming gossipsub message since channel already closed: {e}");
+                            continue;
+                        }
+                    },
+                    Behaviour(Discovery(discovery::Event::ConnectionEstablished { peer_id, remote_ip, remote_tcp_port, .. })) => {
+                        // grab IPv4 string
+                        let remote_ipv4 = match remote_ip {
+                            IpAddr::V4(ip) => ip.to_string(),
+                            IpAddr::V6(ip) => {
+                                log::warn!("RUST: ignoring connection to IPv6 address: {ip}");
+                                continue;
+                            }
+                        };
+
+                        // send connection event to channel (or exit if connection closed)
+                        if let Err(e) = connection_update_tx.send(PyConnectionUpdate {
+                            update_type: PyConnectionUpdateType::Connected,
+                            peer_id: peer_id.to_base58(),
+                            remote_ipv4,
+                            remote_tcp_port,
+                        }).await {
+                            log::error!("RUST: could not send connection update since channel already closed: {e}");
+                            continue;
+                        }
+                    },
+                    Behaviour(Discovery(discovery::Event::ConnectionClosed { peer_id, remote_ip, remote_tcp_port, .. })) => {
+                        // grab IPv4 string
+                        let remote_ipv4 = match remote_ip {
+                            IpAddr::V4(ip) => ip.to_string(),
+                            IpAddr::V6(ip) => {
+                                log::warn!("RUST: ignoring disconnection from IPv6 address: {ip}");
+                                continue;
+                            }
+                        };
+
+                        // send disconnection event to channel (or exit if connection closed)
+                        if let Err(e) = connection_update_tx.send(PyConnectionUpdate {
+                            update_type: PyConnectionUpdateType::Disconnected,
+                            peer_id: peer_id.to_base58(),
+                            remote_ipv4,
+                            remote_tcp_port,
+                        }).await {
+                            log::error!("RUST: could not send connection update since channel already closed: {e}");
+                            continue;
+                        }
+                    },
+                    e => {
+                        log::info!("RUST: other event {e:?}");
+                    }
+                }
+            }
+        }
+    }
+
+    log::info!("RUST: networking task stopped");
+}
+
+#[gen_stub_pyclass]
+#[pyclass(name = "NetworkingHandle")]
+#[derive(Debug)]
+struct PyNetworkingHandle {
+    // channels
+    to_task_tx: Option<mpsc::Sender<ToTask>>,
+    connection_update_rx: Mutex<mpsc::Receiver<PyConnectionUpdate>>,
+    gossipsub_message_rx: Mutex<mpsc::Receiver<(String, Vec<u8>)>>,
+}
+
+impl Drop for PyNetworkingHandle {
+    fn drop(&mut self) {
+        // TODO: may or may not need to await a "kill-signal" oneshot channel message,
+        //       to ensure that the networking task is done BEFORE exiting the clear function...
+        //       but this may require GIL?? and it may not be safe to call GIL here??
+        self.to_task_tx = None; // Using Option<T> as a trick to force channel to be dropped
+    }
+}
+
+#[allow(clippy::expect_used)]
+impl PyNetworkingHandle {
+    fn new(
+        to_task_tx: mpsc::Sender<ToTask>,
+        connection_update_rx: mpsc::Receiver<PyConnectionUpdate>,
+        gossipsub_message_rx: mpsc::Receiver<(String, Vec<u8>)>,
+    ) -> Self {
+        Self {
+            to_task_tx: Some(to_task_tx),
+            connection_update_rx: Mutex::new(connection_update_rx),
+            gossipsub_message_rx: Mutex::new(gossipsub_message_rx),
+        }
+    }
+
+    const fn to_task_tx(&self) -> &mpsc::Sender<ToTask> {
+        self.to_task_tx
+            .as_ref()
+            .expect("The sender should only be None after de-initialization.")
+    }
+}
+
+#[gen_stub_pymethods]
+#[pymethods]
+impl PyNetworkingHandle {
+    // NOTE: `async fn`s here that use `.await` will wrap the future in `.allow_threads_py()`
+    //       immediately beforehand to release the interpreter.
+    //       SEE: https://pyo3.rs/v0.26.0/async-await.html#detaching-from-the-interpreter-across-await
+
+    // ---- Lifecycle management methods ----
+
+    #[new]
+    fn py_new(identity: Bound<'_, PyKeypair>) -> PyResult<Self> {
+        use pyo3_async_runtimes::tokio::get_runtime;
+
+        // create communication channels
+        let (to_task_tx, to_task_rx) = mpsc::channel(MPSC_CHANNEL_SIZE);
+        let (connection_update_tx, connection_update_rx) = mpsc::channel(MPSC_CHANNEL_SIZE);
+        let (gossipsub_message_tx, gossipsub_message_rx) = mpsc::channel(MPSC_CHANNEL_SIZE);
+
+        // get identity
+        let identity = identity.borrow().0.clone();
+
+        // create networking swarm (within tokio context!! or it crashes)
+        let swarm = get_runtime()
+            .block_on(async { create_swarm(identity) })
+            .pyerr()?;
+
+        // spawn tokio task running the networking logic
+        get_runtime().spawn(async move {
+            networking_task(
+                swarm,
+                to_task_rx,
+                connection_update_tx,
+                gossipsub_message_tx,
+            )
+            .await;
+        });
+        Ok(Self::new(
+            to_task_tx,
+            connection_update_rx,
+            gossipsub_message_rx,
+        ))
+    }
+
+    #[gen_stub(skip)]
+    const fn __traverse__(&self, _visit: PyVisit<'_>) -> Result<(), PyTraverseError> {
+        Ok(()) // This is needed purely so `__clear__` can work
+    }
+
+    #[gen_stub(skip)]
+    fn __clear__(&mut self) {
+        // TODO: may or may not need to await a "kill-signal" oneshot channel message,
+        //       to ensure that the networking task is done BEFORE exiting the clear function...
+        //       but this may require GIL?? and it may not be safe to call GIL here??
+        self.to_task_tx = None; // Using Option<T> as a trick to force channel to be dropped
+    }
+
+    // ---- Connection update receiver methods ----
+
+    /// Receives the next `ConnectionUpdate` from networking.
+    async fn connection_update_recv(&self) -> PyResult<PyConnectionUpdate> {
+        self.connection_update_rx
+            .lock()
+            .allow_threads_py() // allow-threads-aware async call
+            .await
+            .recv_py()
+            .allow_threads_py() // allow-threads-aware async call
+            .await
+    }
+
+    /// Receives at most `limit` `ConnectionUpdate`s from networking and returns them.
+    ///
+    /// For `limit = 0`, an empty collection of `ConnectionUpdate`s will be returned immediately.
+    /// For `limit > 0`, if there are no `ConnectionUpdate`s in the channel's queue this method
+    /// will sleep until a `ConnectionUpdate`s is sent.
+    async fn connection_update_recv_many(&self, limit: usize) -> PyResult<Vec<PyConnectionUpdate>> {
+        self.connection_update_rx
+            .lock()
+            .allow_threads_py() // allow-threads-aware async call
+            .await
+            .recv_many_py(limit)
+            .allow_threads_py() // allow-threads-aware async call
+            .await
+    }
+
+    // TODO: rn this blocks main thread if anything else is awaiting the channel (bc its a mutex)
+    //       so its too dangerous to expose just yet. figure out a better semantics for handling this,
+    //       so things don't randomly block
+    // /// Tries to receive the next `ConnectionUpdate` from networking.
+    // fn connection_update_try_recv(&self) -> PyResult<Option<PyConnectionUpdate>> {
+    //     self.connection_update_rx.blocking_lock().try_recv_py()
+    // }
+    //
+    // /// Checks if the `ConnectionUpdate` channel is empty.
+    // fn connection_update_is_empty(&self) -> bool {
+    //     self.connection_update_rx.blocking_lock().is_empty()
+    // }
+    //
+    // /// Returns the number of `ConnectionUpdate`s in the channel.
+    // fn connection_update_len(&self) -> usize {
+    //     self.connection_update_rx.blocking_lock().len()
+    // }
+
+    // ---- Gossipsub management methods ----
+
+    /// Subscribe to a `GossipSub` topic.
+    ///
+    /// Returns `True` if the subscription worked. Returns `False` if we were already subscribed.
+    async fn gossipsub_subscribe(&self, topic: String) -> PyResult<bool> {
+        let (tx, rx) = oneshot::channel();
+
+        // send off request to subscribe
+        self.to_task_tx()
+            .send_py(ToTask::GossipsubSubscribe {
+                topic,
+                result_tx: tx,
+            })
+            .allow_threads_py() // allow-threads-aware async call
+            .await?;
+
+        // wait for response & return any errors
+        rx.allow_threads_py() // allow-threads-aware async call
+            .await
+            .map_err(|_| PyErr::receiver_channel_closed())?
+    }
+
+    /// Unsubscribes from a `GossipSub` topic.
+    ///
+    /// Returns `True` if we were subscribed to this topic. Returns `False` if we were not subscribed.
+    async fn gossipsub_unsubscribe(&self, topic: String) -> PyResult<bool> {
+        let (tx, rx) = oneshot::channel();
+
+        // send off request to unsubscribe
+        self.to_task_tx()
+            .send_py(ToTask::GossipsubUnsubscribe {
+                topic,
+                result_tx: tx,
+            })
+            .allow_threads_py() // allow-threads-aware async call
+            .await?;
+
+        // wait for response & convert any errors
+        rx.allow_threads_py() // allow-threads-aware async call
+            .await
+            .map_err(|_| PyErr::receiver_channel_closed())
+    }
+
+    /// Publishes a message with multiple topics to the `GossipSub` network.
+    ///
+    /// If no peers are found that subscribe to this topic, throws `NoPeersSubscribedToTopicError` exception.
+    async fn gossipsub_publish(&self, topic: String, data: Py<PyBytes>) -> PyResult<()> {
+        let (tx, rx) = oneshot::channel();
+
+        // send off request to subscribe
+        let data = Python::attach(|py| Vec::from(data.as_bytes(py)));
+        self.to_task_tx()
+            .send_py(ToTask::GossipsubPublish {
+                topic,
+                data,
+                result_tx: tx,
+            })
+            .allow_threads_py() // allow-threads-aware async call
+            .await?;
+
+        // wait for response & return any errors => ignore messageID for now!!!
+        let _ = rx
+            .allow_threads_py() // allow-threads-aware async call
+            .await
+            .map_err(|_| PyErr::receiver_channel_closed())??;
+        Ok(())
+    }
+
+    // ---- Gossipsub message receiver methods ----
+
+    /// Receives the next message from the `GossipSub` network.
+    async fn gossipsub_recv(&self) -> PyResult<(String, Py<PyBytes>)> {
+        self.gossipsub_message_rx
+            .lock()
+            .allow_threads_py() // allow-threads-aware async call
+            .await
+            .recv_py()
+            .allow_threads_py() // allow-threads-aware async call
+            .await
+            .map(|(t, d)| (t, d.pybytes()))
+    }
+
+    /// Receives at most `limit` messages from the `GossipSub` network and returns them.
+    ///
+    /// For `limit = 0`, an empty collection of messages will be returned immediately.
+    /// For `limit > 0`, if there are no messages in the channel's queue this method
+    /// will sleep until a message is sent.
+    async fn gossipsub_recv_many(&self, limit: usize) -> PyResult<Vec<(String, Py<PyBytes>)>> {
+        Ok(self
+            .gossipsub_message_rx
+            .lock()
+            .allow_threads_py() // allow-threads-aware async call
+            .await
+            .recv_many_py(limit)
+            .allow_threads_py() // allow-threads-aware async call
+            .await?
+            .into_iter()
+            .map(|(t, d)| (t, d.pybytes()))
+            .collect())
+    }
+
+    // TODO: rn this blocks main thread if anything else is awaiting the channel (bc its a mutex)
+    //       so its too dangerous to expose just yet. figure out a better semantics for handling this,
+    //       so things don't randomly block
+    // /// Tries to receive the next message from the `GossipSub` network.
+    // fn gossipsub_try_recv(&self) -> PyResult<Option<(String, Py<PyBytes>)>> {
+    //     Ok(self
+    //         .gossipsub_message_rx
+    //         .blocking_lock()
+    //         .try_recv_py()?
+    //         .map(|(t, d)| (t, d.pybytes())))
+    // }
+    //
+    // /// Checks if the `GossipSub` message channel is empty.
+    // fn gossipsub_is_empty(&self) -> bool {
+    //     self.gossipsub_message_rx.blocking_lock().is_empty()
+    // }
+    //
+    // /// Returns the number of `GossipSub` messages in the channel.
+    // fn gossipsub_len(&self) -> usize {
+    //     self.gossipsub_message_rx.blocking_lock().len()
+    // }
+}
+
+pub fn networking_submodule(m: &Bound<'_, PyModule>) -> PyResult<()> {
+    m.add_class::<exception::PyNoPeersSubscribedToTopicError>()?;
+    m.add_class::<exception::PyAllQueuesFullError>()?;
+
+    m.add_class::<PyConnectionUpdateType>()?;
+    m.add_class::<PyConnectionUpdate>()?;
+    m.add_class::<PyConnectionUpdateType>()?;
+    m.add_class::<PyNetworkingHandle>()?;
+
+    Ok(())
+}

rust/exo_pyo3_bindings/tests/dummy.rs

@@ -0,0 +1,54 @@
+#[cfg(test)]
+mod tests {
+    use core::mem::drop;
+    use core::option::Option::Some;
+    use core::time::Duration;
+    use tokio;
+    use tokio::sync::mpsc;
+
+    #[tokio::test]
+    async fn test_drop_channel() {
+        struct Ping;
+
+        let (tx, mut rx) = mpsc::channel::<Ping>(10);
+
+        let _ = tokio::spawn(async move {
+            println!("TASK: entered");
+
+            loop {
+                tokio::select! {
+                    result = rx.recv() => {
+                        match result {
+                            Some(_) => {
+                                println!("TASK: pinged");
+                            }
+                            None => {
+                                println!("TASK: closing channel");
+                                break;
+                            }
+                        }
+                    }
+                    _ = tokio::time::sleep(Duration::from_secs_f32(0.1)) => {
+                        println!("TASK: heartbeat");
+                    }
+                }
+            }
+
+            println!("TASK: exited");
+        });
+
+        let tx2 = tx.clone();
+
+        tokio::time::sleep(Duration::from_secs_f32(0.11)).await;
+
+        tx.send(Ping).await.expect("Should not fail");
+        drop(tx);
+
+        tokio::time::sleep(Duration::from_secs_f32(0.11)).await;
+
+        tx2.send(Ping).await.expect("Should not fail");
+        drop(tx2);
+
+        tokio::time::sleep(Duration::from_secs_f32(0.11)).await;
+    }
+}

rust/exo_pyo3_bindings/tests/test_python.py

@@ -1,20 +1,16 @@
 import asyncio
 
 import pytest
-from exo_pyo3_bindings import (
-    Keypair,
-    NetworkingHandle,
-    NoPeersSubscribedToTopicError,
-    PyFromSwarm,
-)
+from exo_pyo3_bindings import Keypair, NetworkingHandle, NoPeersSubscribedToTopicError
 
 
 @pytest.mark.asyncio
 async def test_sleep_on_multiple_items() -> None:
     print("PYTHON: starting handle")
-    h = NetworkingHandle(Keypair.generate())
+    h = NetworkingHandle(Keypair.generate_ed25519())
 
-    rt = asyncio.create_task(_await_recv(h))
+    ct = asyncio.create_task(_await_cons(h))
+    mt = asyncio.create_task(_await_msg(h))
 
     # sleep for 4 ticks
     for i in range(4):
@@ -26,11 +22,13 @@ async def test_sleep_on_multiple_items() -> None:
             print("caught it", e)
 
 
-async def _await_recv(h: NetworkingHandle):
+async def _await_cons(h: NetworkingHandle):
     while True:
-        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}")
+        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}")

rust/networking/Cargo.toml

@@ -5,6 +5,7 @@ edition = { workspace = true }
 publish = false
 
 [lib]
+doctest = false
 name = "networking"
 path = "src/lib.rs"
 
@@ -12,22 +13,25 @@ path = "src/lib.rs"
 workspace = true
 
 [dependencies]
+# datastructures
+either = { workspace = true }
+
 # macro dependencies
 extend = { workspace = true }
+delegate = { workspace = true }
 
 # async
 tokio = { workspace = true, features = ["full"] }
 futures-lite = { workspace = true }
+futures-timer = { workspace = true }
 
 # utility dependencies
+util = { workspace = true }
 tracing-subscriber = { version = "0.3.19", features = ["default", "env-filter"] }
-pin-project = "1.1.10"
+keccak-const = { workspace = true }
 
 # tracing/logging
 log = { workspace = true }
-iroh = "0.96.1"
-iroh-gossip = "0.96.0"
-iroh-blobs = "0.98.0"
-iroh-docs = "0.96.0"
 
 # networking
+libp2p = { workspace = true, features = ["full"] }

rust/networking/examples/chatroom.rs

@@ -1,3 +1,7 @@
+use futures_lite::StreamExt;
+use libp2p::{gossipsub, identity, swarm::SwarmEvent};
+use networking::{discovery, swarm};
+use tokio::{io, io::AsyncBufReadExt as _, select};
 use tracing_subscriber::EnvFilter;
 use tracing_subscriber::filter::LevelFilter;
 
@@ -6,4 +10,65 @@ async fn main() {
     let _ = tracing_subscriber::fmt()
         .with_env_filter(EnvFilter::from_default_env().add_directive(LevelFilter::INFO.into()))
         .try_init();
+
+    // Configure swarm
+    let mut swarm =
+        swarm::create_swarm(identity::Keypair::generate_ed25519()).expect("Swarm creation failed");
+
+    // Create a Gossipsub topic & subscribe
+    let topic = gossipsub::IdentTopic::new("test-net");
+    swarm
+        .behaviour_mut()
+        .gossipsub
+        .subscribe(&topic)
+        .expect("Subscribing to topic failed");
+
+    // Read full lines from stdin
+    let mut stdin = io::BufReader::new(io::stdin()).lines();
+    println!("Enter messages via STDIN and they will be sent to connected peers using Gossipsub");
+
+    // Kick it off
+    loop {
+        select! {
+            // on gossipsub outgoing
+            Ok(Some(line)) = stdin.next_line() => {
+                if let Err(e) = swarm
+                    .behaviour_mut().gossipsub
+                    .publish(topic.clone(), line.as_bytes()) {
+                    println!("Publish error: {e:?}");
+                }
+            }
+            event = swarm.next() => match event {
+                // on gossipsub incoming
+                Some(SwarmEvent::Behaviour(swarm::BehaviourEvent::Gossipsub(gossipsub::Event::Message {
+                    propagation_source: peer_id,
+                    message_id: id,
+                    message,
+                }))) => println!(
+                        "\n\nGot message: '{}' with id: {id} from peer: {peer_id}\n\n",
+                        String::from_utf8_lossy(&message.data),
+                    ),
+
+                // on discovery
+                Some(SwarmEvent::Behaviour(swarm::BehaviourEvent::Discovery(e)) )=> match e {
+                    discovery::Event::ConnectionEstablished {
+                        peer_id, connection_id, remote_ip, remote_tcp_port
+                    } => {
+                        println!("\n\nConnected to: {peer_id}; connection ID: {connection_id}; remote IP: {remote_ip}; remote TCP port: {remote_tcp_port}\n\n");
+                    }
+                    discovery::Event::ConnectionClosed {
+                        peer_id, connection_id, remote_ip, remote_tcp_port
+                    } => {
+                        eprintln!("\n\nDisconnected from: {peer_id}; connection ID: {connection_id}; remote IP: {remote_ip}; remote TCP port: {remote_tcp_port}\n\n");
+                    }
+                }
+
+                // ignore outgoing errors: those are normal
+                e@Some(SwarmEvent::OutgoingConnectionError { .. }) => { log::debug!("Outgoing connection error: {e:?}"); }
+
+                // otherwise log any other event
+                e => { log::info!("Other event {e:?}"); }
+            }
+        }
+    }
 }

rust/networking/src/RESEARCH_NOTES.txt

@@ -0,0 +1,44 @@
+https://github.com/ml-explore/mlx/commit/3fe98bacc7640d857acf3539f1d21b47a32e5609
+^raw sockets distributed -> `<net/ndrv.h>` -> https://newosxbook.com/code/xnu-3247.1.106/bsd/net/ndrv.h.auto.html
+--> header file for a networking component found in the macOS kernel (XNU) that defines structures for network device driver registration, specifically the ndrv_demux_desc and ndrv_protocol_desc structures used for demultiplexing protocol data at the network interface level. It specifies how to describe protocol data, such as an Ethernet type or a SNAP header, and how to associate these descriptions with a specific protocol family to receive matching packets.
+--> Used to bind an NDRV socket so that packets that match given protocol demux descriptions can be received.
+--> An NDRV socket is a special kind of socket in the Darwin/macOS operating system's XNU kernel, used for low-level network packet manipulation and binding to specific protocols for packet processing. It allows user-space applications or drivers to directly write Layer 2 (L2) network packets or interact with the network stack at a lower level, often by binding to protocol descriptors like the ndrv_protocol_desc. This type of socket is used for functions such as capturing and injecting packets, especially in network infrastructure software like routers or for kernel-level network monitoring and security tools.
+--> also called PF_NDRV sockets --> https://newosxbook.com/bonus/vol1ch16.html
+----> they are conceptually similar to https://scapy.disruptivelabs.in/networking/socket-interface PF_RAW or PF_PACKET
+
+https://stackoverflow.com/questions/17169298/af-packet-on-osx
+^AF_PACKET duplicates the packets as soon as it receives them from the physical layer (for incoming packets) or just before sending them out to the physical layer (for outgoing packets). -> this is on Linux only
+^it doesn't exist on OS X so you can use /dev/bpfX (Berkeley Packet Filter) for sniffing
+
+https://www.unix.com/man_page/mojave/4/ip/
+^OS X manpages for IP
+
+https://developer.apple.com/documentation/kernel/implementing_drivers_system_extensions_and_kexts
+^driver kit, system extensions & kexts for macOS
+
+----
+
+To set up a Linux system to use a Thunderbolt connection as a network device, connect the two computers with a Thunderbolt cable, load the thunderbolt-net kernel module (usually automatic but modprobe is an option for manual loading), and then the operating system will create virtual Ethernet interfaces (e.g., thunderbolt0) for networking. You can then use standard tools like ifconfig or your desktop environment's network manager to configure these new interfaces for a link-local network.
+--> https://gist.github.com/geosp/80fbd39e617b7d1d9421683df4ea224a
+----> here is a guide on how to set up thunderbolt-ethernet on linux
+----> I may be able to steal the thunderbolt-net code ideas to implement a kernel module for MacOS
+
+https://chatgpt.com/s/t_68af8e41a8548191993281a014f846a7
+^GPT discussion about making socket interface
+
+https://chatgpt.com/s/t_68afb798a85c8191973c02a0fa7a48a3 --> link-local address,,??
+https://chatgpt.com/s/t_68afb02987e08191b2b0044d3667ece2
+^GPT discussion about accessing TB on MacOS low level interactions
+
+--------------------------------
+
+https://www.intel.com/content/www/us/en/support/articles/000098893/software.html
+^Thunderbolt Share & Thunderbolt Networking Mode => intel's equivalent of thunderbolt bridge
+
+
+---------------------------------
+
+https://www.zerotier.com/blog/how-zerotier-eliminated-kernel-extensions-on-macos/
+-->fake ethernet devices on MacOS -> omg??? we can detect thunderbolt bridge, then bind to it, then re-expose it as fake ethernet??
+-->ps: https://chatgpt.com/s/t_68afb2b25fb881919526763fb5d7359c, AF/PF_NDRV are one and the same!!!
+-->https://github.com/zerotier/ZeroTierOne/blob/dev/osdep/MacEthernetTapAgent.c

rust/networking/src/discovery.rs

@@ -0,0 +1,382 @@
+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

@@ -0,0 +1,44 @@
+//! TODO: crate documentation
+//!
+//! 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 {
+    use std::error::Error;
+
+    pub type AnyError = Box<dyn Error + Send + Sync + 'static>;
+    pub type AnyResult<T> = Result<T, AnyError>;
+}
+
+/// 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;
+
+    #[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,
+                }
+            } else {
+                return None;
+            };
+            let Some(Protocol::Tcp(port)) = ps.next() else {
+                return None;
+            };
+            Some((ip, port))
+        }
+    }
+}

rust/networking/src/swarm.rs

@@ -0,0 +1,143 @@
+use crate::alias;
+use crate::swarm::transport::tcp_transport;
+pub use behaviour::{Behaviour, BehaviourEvent};
+use libp2p::{SwarmBuilder, identity};
+
+pub type Swarm = libp2p::Swarm<Behaviour>;
+
+/// 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";
+
+/// Create and configure a swarm which listens to all ports on OS
+pub fn create_swarm(keypair: identity::Keypair) -> 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)
+}
+
+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/networking/tests/dummy.rs

@@ -0,0 +1,7 @@
+// maybe this will hold test in the future...??
+
+#[cfg(test)]
+mod tests {
+    #[test]
+    fn does_nothing() {}
+}

rust/util/Cargo.toml

@@ -0,0 +1,15 @@
+[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

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

rust/util/src/wakerdeque.rs

@@ -0,0 +1,55 @@
+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/download/coordinator.py

@@ -1,8 +1,11 @@
 import asyncio
+import socket
 from dataclasses import dataclass, field
+from random import random
 
 import anyio
 from anyio import current_time
+from anyio.abc import TaskGroup
 from loguru import logger
 
 from exo.download.download_utils import (
@@ -20,9 +23,13 @@ from exo.shared.types.commands import (
     ForwarderDownloadCommand,
     StartDownload,
 )
-from exo.shared.types.common import NodeId
+from exo.shared.types.common import NodeId, SessionId, SystemId
 from exo.shared.types.events import (
     Event,
+    EventId,
+    # TODO(evan): just for acks, should delete this ASAP
+    GlobalForwarderEvent,
+    LocalForwarderEvent,
     NodeDownloadProgress,
 )
 from exo.shared.types.worker.downloads import (
@@ -33,28 +40,41 @@ from exo.shared.types.worker.downloads import (
     DownloadProgress,
 )
 from exo.shared.types.worker.shards import PipelineShardMetadata, ShardMetadata
-from exo.utils.channels import Receiver, Sender
-from exo.utils.task_group import TaskGroup
+from exo.utils.channels import Receiver, Sender, channel
 
 
 @dataclass
 class DownloadCoordinator:
     node_id: NodeId
+    session_id: SessionId
     shard_downloader: ShardDownloader
     download_command_receiver: Receiver[ForwarderDownloadCommand]
-    event_sender: Sender[Event]
+    local_event_sender: Sender[LocalForwarderEvent]
+
+    # ack stuff
+    _global_event_receiver: Receiver[GlobalForwarderEvent]
+    _out_for_delivery: dict[EventId, LocalForwarderEvent] = field(default_factory=dict)
+
     offline: bool = False
 
+    _system_id: SystemId = field(default_factory=SystemId)
+
     # Local state
     download_status: dict[ModelId, DownloadProgress] = field(default_factory=dict)
     active_downloads: dict[ModelId, asyncio.Task[None]] = field(default_factory=dict)
 
-    _tg: TaskGroup = field(init=False, default_factory=TaskGroup)
+    # Internal event channel for forwarding (initialized in __post_init__)
+    event_sender: Sender[Event] = field(init=False)
+    event_receiver: Receiver[Event] = field(init=False)
+    _tg: TaskGroup = field(init=False, default_factory=anyio.create_task_group)
 
     # Per-model throttle for download progress events
     _last_progress_time: dict[ModelId, float] = field(default_factory=dict)
 
     def __post_init__(self) -> None:
+        self.event_sender, self.event_receiver = channel[Event]()
+        if self.offline:
+            self.shard_downloader.set_internet_connection(False)
         self.shard_downloader.on_progress(self._download_progress_callback)
 
     def _model_dir(self, model_id: ModelId) -> str:
@@ -103,16 +123,65 @@ class DownloadCoordinator:
         logger.info(
             f"Starting DownloadCoordinator{' (offline mode)' if self.offline else ''}"
         )
+        if not self.offline:
+            self._test_internet_connection()
         try:
             async with self._tg as tg:
                 tg.start_soon(self._command_processor)
+                tg.start_soon(self._forward_events)
                 tg.start_soon(self._emit_existing_download_progress)
+                tg.start_soon(self._resend_out_for_delivery)
+                tg.start_soon(self._clear_ofd)
+                if not self.offline:
+                    tg.start_soon(self._check_internet_connection)
         finally:
             for task in self.active_downloads.values():
                 task.cancel()
 
+    def _test_internet_connection(self) -> None:
+        # Try multiple endpoints since some ISPs/networks block specific IPs
+        for host in ("1.1.1.1", "8.8.8.8", "1.0.0.1"):
+            try:
+                socket.create_connection((host, 443), timeout=3).close()
+                self.shard_downloader.set_internet_connection(True)
+                logger.debug(f"Internet connectivity: True (via {host})")
+                return
+            except OSError:
+                continue
+        self.shard_downloader.set_internet_connection(False)
+        logger.debug("Internet connectivity: False")
+
+    async def _check_internet_connection(self) -> None:
+        first_connection = True
+        while True:
+            await asyncio.sleep(10)
+
+            # Assume that internet connection is set to False on 443 errors.
+            if self.shard_downloader.internet_connection:
+                continue
+
+            self._test_internet_connection()
+
+            if first_connection and self.shard_downloader.internet_connection:
+                first_connection = False
+                self._tg.start_soon(self._emit_existing_download_progress)
+
     def shutdown(self) -> None:
-        self._tg.cancel_tasks()
+        self._tg.cancel_scope.cancel()
+
+    # directly copied from worker
+    async def _resend_out_for_delivery(self) -> None:
+        # This can also be massively tightened, we should check events are at least a certain age before resending.
+        # Exponential backoff would also certainly help here.
+        while True:
+            await anyio.sleep(1 + random())
+            for event in self._out_for_delivery.copy().values():
+                await self.local_event_sender.send(event)
+
+    async def _clear_ofd(self) -> None:
+        with self._global_event_receiver as events:
+            async for event in events:
+                self._out_for_delivery.pop(event.event.event_id, None)
 
     async def _command_processor(self) -> None:
         with self.download_command_receiver as commands:
@@ -286,6 +355,23 @@ class DownloadCoordinator:
             )
             del self.download_status[model_id]
 
+    async def _forward_events(self) -> None:
+        idx = 0
+        with self.event_receiver as events:
+            async for event in events:
+                fe = LocalForwarderEvent(
+                    origin_idx=idx,
+                    origin=self._system_id,
+                    session=self.session_id,
+                    event=event,
+                )
+                idx += 1
+                logger.debug(
+                    f"DownloadCoordinator published event {idx}: {str(event)[:100]}"
+                )
+                await self.local_event_sender.send(fe)
+                self._out_for_delivery[event.event_id] = fe
+
     async def _emit_existing_download_progress(self) -> None:
         try:
             while True:

src/exo/download/download_utils.py

@@ -110,20 +110,54 @@ def map_repo_download_progress_to_download_progress_data(
     )
 
 
+def _resolve_hf_hub_model(search_dir: Path, normalized: str) -> Path | None:
+    """Try to find a model in HuggingFace Hub cache format.
+
+    HF Hub stores models as ``models--<org>--<name>/snapshots/<commit>/``
+    with symlinks to ``../../blobs/``.  The active commit is read from
+    ``refs/main``.
+    """
+    hf_model_dir = search_dir / f"models--{normalized}"
+    if not hf_model_dir.is_dir():
+        return None
+    # Resolve ref → snapshot
+    ref_file = hf_model_dir / "refs" / "main"
+    if ref_file.is_file():
+        commit_hash = ref_file.read_text().strip()
+        snapshot = hf_model_dir / "snapshots" / commit_hash
+        if snapshot.is_dir():
+            return snapshot
+    # Fallback: use latest snapshot by mtime
+    snapshots_dir = hf_model_dir / "snapshots"
+    if snapshots_dir.is_dir():
+        snapshots = sorted(
+            snapshots_dir.iterdir(), key=lambda p: p.stat().st_mtime, reverse=True
+        )
+        if snapshots:
+            return snapshots[0]
+    return None
+
+
 def resolve_model_in_path(model_id: ModelId) -> Path | None:
     """Search EXO_MODELS_PATH directories for a pre-existing model.
 
-    Checks each directory for the normalized name (org--model).  A candidate
-    is only returned if ``is_model_directory_complete`` confirms all weight
-    files are present.
+    Checks each directory for the normalized name (org--model) and the
+    HuggingFace Hub cache format (models--org--model/snapshots/<ref>/).
+    A candidate is only returned if ``is_model_directory_complete``
+    confirms all weight files are present.
     """
     if EXO_MODELS_PATH is None:
         return None
     normalized = model_id.normalize()
     for search_dir in EXO_MODELS_PATH:
+        # Try direct format: <dir>/<org--name>/
         candidate = search_dir / normalized
         if candidate.is_dir() and is_model_directory_complete(candidate):
             return candidate
+        # Try HF Hub cache format: <dir>/models--<org--name>/snapshots/<ref>/
+        hf_candidate = _resolve_hf_hub_model(search_dir, normalized)
+        if hf_candidate is not None and is_model_directory_complete(hf_candidate):
+            return hf_candidate
     return None
 
 
@@ -314,13 +348,9 @@ async def fetch_file_list_with_cache(
         _fetched_file_lists_this_session.add(cache_key)
         return file_list
     except Exception as e:
-        logger.opt(exception=e).warning(
-            "Ran into exception when fetching file list from HF."
-        )
-
         if await aios.path.exists(cache_file):
             logger.warning(
-                f"No cached file list for {model_id} - using local file list"
+                f"No internet and no cached file list for {model_id} - using local file list"
             )
             async with aiofiles.open(cache_file, "r") as f:
                 return TypeAdapter(list[FileListEntry]).validate_json(await f.read())
@@ -827,7 +857,6 @@ async def download_shard(
 
     for file in filtered_file_list:
         downloaded_bytes = await get_downloaded_size(target_dir / file.path)
-        final_file_exists = await aios.path.exists(target_dir / file.path)
         file_progress[file.path] = RepoFileDownloadProgress(
             repo_id=shard.model_card.model_id,
             repo_revision=revision,
@@ -837,9 +866,7 @@ async def download_shard(
             total=Memory.from_bytes(file.size or 0),
             speed=0,
             eta=timedelta(0),
-            status="complete"
-            if final_file_exists and downloaded_bytes == file.size
-            else "not_started",
+            status="complete" if downloaded_bytes == file.size else "not_started",
             start_time=time.time(),
         )
 

src/exo/download/impl_shard_downloader.py

@@ -15,13 +15,9 @@ from exo.shared.types.worker.shards import (
 )
 
 
-def exo_shard_downloader(
-    max_parallel_downloads: int = 8, offline: bool = False
-) -> ShardDownloader:
+def exo_shard_downloader(max_parallel_downloads: int = 8) -> ShardDownloader:
     return SingletonShardDownloader(
-        CachedShardDownloader(
-            ResumableShardDownloader(max_parallel_downloads, offline=offline)
-        )
+        CachedShardDownloader(ResumableShardDownloader(max_parallel_downloads))
     )
 
 
@@ -54,6 +50,10 @@ class SingletonShardDownloader(ShardDownloader):
         self.shard_downloader = shard_downloader
         self.active_downloads: dict[ShardMetadata, asyncio.Task[Path]] = {}
 
+    def set_internet_connection(self, value: bool) -> None:
+        self.internet_connection = value
+        self.shard_downloader.set_internet_connection(value)
+
     def on_progress(
         self,
         callback: Callable[[ShardMetadata, RepoDownloadProgress], Awaitable[None]],
@@ -90,6 +90,10 @@ class CachedShardDownloader(ShardDownloader):
         self.shard_downloader = shard_downloader
         self.cache: dict[tuple[str, ShardMetadata], Path] = {}
 
+    def set_internet_connection(self, value: bool) -> None:
+        self.internet_connection = value
+        self.shard_downloader.set_internet_connection(value)
+
     def on_progress(
         self,
         callback: Callable[[ShardMetadata, RepoDownloadProgress], Awaitable[None]],
@@ -119,9 +123,8 @@ class CachedShardDownloader(ShardDownloader):
 
 
 class ResumableShardDownloader(ShardDownloader):
-    def __init__(self, max_parallel_downloads: int = 8, offline: bool = False):
+    def __init__(self, max_parallel_downloads: int = 8):
         self.max_parallel_downloads = max_parallel_downloads
-        self.offline = offline
         self.on_progress_callbacks: list[
             Callable[[ShardMetadata, RepoDownloadProgress], Awaitable[None]]
         ] = []
@@ -148,7 +151,8 @@ class ResumableShardDownloader(ShardDownloader):
             self.on_progress_wrapper,
             max_parallel_downloads=self.max_parallel_downloads,
             allow_patterns=allow_patterns,
-            skip_internet=self.offline,
+            skip_internet=not self.internet_connection,
+            on_connection_lost=lambda: self.set_internet_connection(False),
         )
         return target_dir
 
@@ -164,7 +168,8 @@ class ResumableShardDownloader(ShardDownloader):
                 shard,
                 self.on_progress_wrapper,
                 skip_download=True,
-                skip_internet=self.offline,
+                skip_internet=not self.internet_connection,
+                on_connection_lost=lambda: self.set_internet_connection(False),
             )
 
         semaphore = asyncio.Semaphore(self.max_parallel_downloads)
@@ -193,6 +198,7 @@ class ResumableShardDownloader(ShardDownloader):
             shard,
             self.on_progress_wrapper,
             skip_download=True,
-            skip_internet=self.offline,
+            skip_internet=not self.internet_connection,
+            on_connection_lost=lambda: self.set_internet_connection(False),
         )
         return progress

src/exo/download/shard_downloader.py

@@ -16,6 +16,11 @@ from exo.shared.types.worker.shards import (
 
 # TODO: the PipelineShardMetadata getting reinstantiated is a bit messy. Should this be a classmethod?
 class ShardDownloader(ABC):
+    internet_connection: bool = False
+
+    def set_internet_connection(self, value: bool) -> None:
+        self.internet_connection = value
+
     @abstractmethod
     async def ensure_shard(
         self, shard: ShardMetadata, config_only: bool = False

src/exo/download/tests/test_coordinator_ack.py

@@ -0,0 +1,98 @@
+from typing import Any
+
+import anyio
+import pytest
+
+from exo.download.coordinator import DownloadCoordinator
+from exo.download.shard_downloader import NoopShardDownloader
+from exo.shared.models.model_cards import ModelCard, ModelTask
+from exo.shared.types.common import ModelId, NodeId, SessionId
+from exo.shared.types.events import (
+    GlobalForwarderEvent,
+    LocalForwarderEvent,
+    NodeDownloadProgress,
+)
+from exo.shared.types.memory import Memory
+from exo.shared.types.worker.downloads import (
+    DownloadPending,
+)
+from exo.shared.types.worker.shards import PipelineShardMetadata
+from exo.utils.channels import channel
+
+# Use the built‑in NoopShardDownloader directly – it already implements the required abstract interface.
+# No additional subclass is needed for this test.
+
+
+@pytest.mark.anyio
+async def test_ack_behaviour():
+    # Create channels (type Any for simplicity)
+    _, command_receiver = channel[Any]()
+    local_sender, _ = channel[Any]()
+    global_sender, global_receiver = channel[Any]()
+
+    # Minimal identifiers
+    node_id = NodeId()
+    session_id = SessionId(master_node_id=node_id, election_clock=0)
+
+    # Create a dummy model card and shard metadata
+    model_id = ModelId("test/model")
+    model_card = ModelCard(
+        model_id=model_id,
+        storage_size=Memory.from_bytes(0),
+        n_layers=1,
+        hidden_size=1,
+        supports_tensor=True,
+        tasks=[ModelTask.TextGeneration],
+    )
+    shard = PipelineShardMetadata(
+        model_card=model_card,
+        device_rank=0,
+        world_size=1,
+        start_layer=0,
+        end_layer=1,
+        n_layers=1,
+    )
+
+    # Instantiate the coordinator with the dummy downloader
+    coord = DownloadCoordinator(
+        node_id=node_id,
+        session_id=session_id,
+        shard_downloader=NoopShardDownloader(),
+        download_command_receiver=command_receiver,
+        local_event_sender=local_sender,
+        _global_event_receiver=global_receiver,
+    )
+
+    async with anyio.create_task_group() as tg:
+        # Start the forwarding and ack‑clearing loops
+        tg.start_soon(coord._forward_events)  # pyright: ignore[reportPrivateUsage]
+        tg.start_soon(coord._clear_ofd)  # pyright: ignore[reportPrivateUsage]
+
+        # Send a pending download progress event via the internal event sender
+        pending = DownloadPending(
+            node_id=node_id,
+            shard_metadata=shard,
+            model_directory="/tmp/model",
+        )
+        await coord.event_sender.send(NodeDownloadProgress(download_progress=pending))
+        # Allow the forwarder to process the event
+        await anyio.sleep(0.1)
+
+        # There should be exactly one entry awaiting ACK
+        assert len(coord._out_for_delivery) == 1  # pyright: ignore[reportPrivateUsage]
+        # Retrieve the stored LocalForwarderEvent
+        stored_fe: LocalForwarderEvent = next(iter(coord._out_for_delivery.values()))  # pyright: ignore[reportPrivateUsage]
+        # Simulate receiving a global ack for this event
+        ack = GlobalForwarderEvent(
+            origin_idx=0,
+            origin=node_id,
+            session=session_id,
+            event=stored_fe.event,
+        )
+        await global_sender.send(ack)
+        # Give the clear‑ofd task a moment to process the ack
+        await anyio.sleep(0.1)
+        # The out‑for‑delivery map should now be empty
+        assert len(coord._out_for_delivery) == 0  # pyright: ignore[reportPrivateUsage]
+        # Cancel background tasks
+        tg.cancel_scope.cancel()

src/exo/main.py

@@ -7,6 +7,7 @@ from dataclasses import dataclass, field
 from typing import Self
 
 import anyio
+from anyio.abc import TaskGroup
 from loguru import logger
 from pydantic import PositiveInt
 
@@ -15,7 +16,6 @@ from exo.download.coordinator import DownloadCoordinator
 from exo.download.impl_shard_downloader import exo_shard_downloader
 from exo.master.api import API  # TODO: should API be in master?
 from exo.master.main import Master
-from exo.routing.event_router import EventRouter
 from exo.routing.router import Router, get_node_id_keypair
 from exo.shared.constants import EXO_LOG
 from exo.shared.election import Election, ElectionResult
@@ -23,14 +23,12 @@ from exo.shared.logging import logger_cleanup, logger_setup
 from exo.shared.types.common import NodeId, SessionId
 from exo.utils.channels import Receiver, channel
 from exo.utils.pydantic_ext import CamelCaseModel
-from exo.utils.task_group import TaskGroup
 from exo.worker.main import Worker
 
 
 @dataclass
 class Node:
     router: Router
-    event_router: EventRouter
     download_coordinator: DownloadCoordinator | None
     worker: Worker | None
     election: Election  # Every node participates in election, as we do want a node to become master even if it isn't a master candidate if no master candidates are present.
@@ -40,7 +38,7 @@ class Node:
 
     node_id: NodeId
     offline: bool
-    _tg: TaskGroup = field(init=False, default_factory=TaskGroup)
+    _tg: TaskGroup = field(init=False, default_factory=anyio.create_task_group)
 
     @classmethod
     async def create(cls, args: "Args") -> Self:
@@ -54,12 +52,6 @@ class Node:
         await router.register_topic(topics.ELECTION_MESSAGES)
         await router.register_topic(topics.CONNECTION_MESSAGES)
         await router.register_topic(topics.DOWNLOAD_COMMANDS)
-        event_router = EventRouter(
-            session_id,
-            command_sender=router.sender(topics.COMMANDS),
-            external_outbound=router.sender(topics.LOCAL_EVENTS),
-            external_inbound=router.receiver(topics.GLOBAL_EVENTS),
-        )
 
         logger.info(f"Starting node {node_id}")
 
@@ -67,10 +59,13 @@ class Node:
         if not args.no_downloads:
             download_coordinator = DownloadCoordinator(
                 node_id,
-                exo_shard_downloader(offline=args.offline),
-                event_sender=event_router.sender(),
+                session_id,
+                exo_shard_downloader(),
                 download_command_receiver=router.receiver(topics.DOWNLOAD_COMMANDS),
+                local_event_sender=router.sender(topics.LOCAL_EVENTS),
                 offline=args.offline,
+                # TODO(evan): remove
+                _global_event_receiver=router.receiver(topics.GLOBAL_EVENTS),
             )
         else:
             download_coordinator = None
@@ -78,8 +73,9 @@ class Node:
         if args.spawn_api:
             api = API(
                 node_id,
+                session_id,
                 port=args.api_port,
-                event_receiver=event_router.receiver(),
+                global_event_receiver=router.receiver(topics.GLOBAL_EVENTS),
                 command_sender=router.sender(topics.COMMANDS),
                 download_command_sender=router.sender(topics.DOWNLOAD_COMMANDS),
                 election_receiver=router.receiver(topics.ELECTION_MESSAGES),
@@ -90,8 +86,9 @@ class Node:
         if not args.no_worker:
             worker = Worker(
                 node_id,
-                event_receiver=event_router.receiver(),
-                event_sender=event_router.sender(),
+                session_id,
+                global_event_receiver=router.receiver(topics.GLOBAL_EVENTS),
+                local_event_sender=router.sender(topics.LOCAL_EVENTS),
                 command_sender=router.sender(topics.COMMANDS),
                 download_command_sender=router.sender(topics.DOWNLOAD_COMMANDS),
             )
@@ -102,7 +99,6 @@ class Node:
         master = Master(
             node_id,
             session_id,
-            event_sender=event_router.sender(),
             global_event_sender=router.sender(topics.GLOBAL_EVENTS),
             local_event_receiver=router.receiver(topics.LOCAL_EVENTS),
             command_receiver=router.receiver(topics.COMMANDS),
@@ -125,7 +121,6 @@ class Node:
 
         return cls(
             router,
-            event_router,
             download_coordinator,
             worker,
             election,
@@ -141,7 +136,6 @@ class Node:
             signal.signal(signal.SIGINT, lambda _, __: self.shutdown())
             signal.signal(signal.SIGTERM, lambda _, __: self.shutdown())
             tg.start_soon(self.router.run)
-            tg.start_soon(self.event_router.run)
             tg.start_soon(self.election.run)
             if self.download_coordinator:
                 tg.start_soon(self.download_coordinator.run)
@@ -155,11 +149,11 @@ class Node:
 
     def shutdown(self):
         # if this is our second call to shutdown, just sys.exit
-        if self._tg.cancel_called():
+        if self._tg.cancel_scope.cancel_called:
             import sys
 
             sys.exit(1)
-        self._tg.cancel_tasks()
+        self._tg.cancel_scope.cancel()
 
     async def _elect_loop(self):
         with self.election_result_receiver as results:
@@ -176,17 +170,6 @@ class Node:
                 # - Shutdown and re-create the worker
                 # - Shut down and re-create the API
 
-                if result.is_new_master:
-                    await anyio.sleep(0)
-                    self.event_router.shutdown()
-                    self.event_router = EventRouter(
-                        result.session_id,
-                        self.router.sender(topics.COMMANDS),
-                        self.router.receiver(topics.GLOBAL_EVENTS),
-                        self.router.sender(topics.LOCAL_EVENTS),
-                    )
-                    self._tg.start_soon(self.event_router.run)
-
                 if (
                     result.session_id.master_node_id == self.node_id
                     and self.master is not None
@@ -200,7 +183,6 @@ class Node:
                     self.master = Master(
                         self.node_id,
                         result.session_id,
-                        event_sender=self.event_router.sender(),
                         global_event_sender=self.router.sender(topics.GLOBAL_EVENTS),
                         local_event_receiver=self.router.receiver(topics.LOCAL_EVENTS),
                         command_receiver=self.router.receiver(topics.COMMANDS),
@@ -223,16 +205,22 @@ class Node:
                         f"Node {result.session_id.master_node_id} elected master"
                     )
                 if result.is_new_master:
+                    await anyio.sleep(0)
                     if self.download_coordinator:
                         self.download_coordinator.shutdown()
                         self.download_coordinator = DownloadCoordinator(
                             self.node_id,
-                            exo_shard_downloader(offline=self.offline),
-                            event_sender=self.event_router.sender(),
+                            result.session_id,
+                            exo_shard_downloader(),
                             download_command_receiver=self.router.receiver(
                                 topics.DOWNLOAD_COMMANDS
                             ),
+                            local_event_sender=self.router.sender(topics.LOCAL_EVENTS),
                             offline=self.offline,
+                            # TODO(evan): remove
+                            _global_event_receiver=self.router.receiver(
+                                topics.GLOBAL_EVENTS
+                            ),
                         )
                         self._tg.start_soon(self.download_coordinator.run)
                     if self.worker:
@@ -240,8 +228,11 @@ class Node:
                         # TODO: add profiling etc to resource monitor
                         self.worker = Worker(
                             self.node_id,
-                            event_receiver=self.event_router.receiver(),
-                            event_sender=self.event_router.sender(),
+                            result.session_id,
+                            global_event_receiver=self.router.receiver(
+                                topics.GLOBAL_EVENTS
+                            ),
+                            local_event_sender=self.router.sender(topics.LOCAL_EVENTS),
                             command_sender=self.router.sender(topics.COMMANDS),
                             download_command_sender=self.router.sender(
                                 topics.DOWNLOAD_COMMANDS
@@ -249,7 +240,7 @@ class Node:
                         )
                         self._tg.start_soon(self.worker.run)
                     if self.api:
-                        self.api.reset(result.won_clock, self.event_router.receiver())
+                        self.api.reset(result.session_id, result.won_clock)
                 else:
                     if self.api:
                         self.api.unpause(result.won_clock)
@@ -261,7 +252,7 @@ def main():
     target = min(max(soft, 65535), hard)
     resource.setrlimit(resource.RLIMIT_NOFILE, (target, hard))
 
-    mp.set_start_method("spawn", force=True)
+    mp.set_start_method("spawn")
     # TODO: Refactor the current verbosity system
     logger_setup(EXO_LOG, args.verbosity)
     logger.info("Starting EXO")
@@ -279,16 +270,9 @@ def main():
         logger.info("FAST_SYNCH forced OFF")
 
     node = anyio.run(Node.create, args)
-    try:
-        anyio.run(node.run)
-    except BaseException as exception:
-        logger.opt(exception=exception).critical(
-            "EXO terminated due to unhandled exception"
-        )
-        raise
-    finally:
-        logger.info("EXO Shutdown complete")
-        logger_cleanup()
+    anyio.run(node.run)
+    logger.info("EXO Shutdown complete")
+    logger_cleanup()
 
 
 class Args(CamelCaseModel):
@@ -299,7 +283,7 @@ class Args(CamelCaseModel):
     tb_only: bool = False
     no_worker: bool = False
     no_downloads: bool = False
-    offline: bool = os.getenv("EXO_OFFLINE", "false").lower() == "true"
+    offline: bool = False
     fast_synch: bool | None = None  # None = auto, True = force on, False = force off
 
     @classmethod
@@ -350,7 +334,6 @@ class Args(CamelCaseModel):
         parser.add_argument(
             "--offline",
             action="store_true",
-            default=os.getenv("EXO_OFFLINE", "false").lower() == "true",
             help="Run in offline/air-gapped mode: skip internet checks, use only pre-staged local models",
         )
         fast_synch_group = parser.add_mutually_exclusive_group()

src/exo/master/adapters/chat_completions.py

@@ -26,11 +26,7 @@ from exo.shared.types.chunks import (
     ToolCallChunk,
 )
 from exo.shared.types.common import CommandId
-from exo.shared.types.text_generation import (
-    InputMessage,
-    TextGenerationTaskParams,
-    resolve_reasoning_params,
-)
+from exo.shared.types.text_generation import InputMessage, TextGenerationTaskParams
 
 
 def chat_request_to_text_generation(
@@ -79,10 +75,6 @@ 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
@@ -97,8 +89,7 @@ def chat_request_to_text_generation(
         seed=request.seed,
         stream=request.stream,
         tools=request.tools,
-        reasoning_effort=resolved_effort,
-        enable_thinking=resolved_thinking,
+        enable_thinking=request.enable_thinking,
         chat_template_messages=chat_template_messages
         if chat_template_messages
         else None,

src/exo/master/adapters/responses.py

@@ -42,11 +42,7 @@ from exo.shared.types.openai_responses import (
     ResponseTextDoneEvent,
     ResponseUsage,
 )
-from exo.shared.types.text_generation import (
-    InputMessage,
-    TextGenerationTaskParams,
-    resolve_reasoning_params,
-)
+from exo.shared.types.text_generation import InputMessage, TextGenerationTaskParams
 
 
 def _format_sse(event: ResponsesStreamEvent) -> str:
@@ -123,11 +119,6 @@ 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,
@@ -141,8 +132,6 @@ 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/api.py

@@ -11,7 +11,8 @@ from typing import Annotated, Literal, cast
 from uuid import uuid4
 
 import anyio
-from anyio import BrokenResourceError
+from anyio import BrokenResourceError, create_task_group
+from anyio.abc import TaskGroup
 from fastapi import FastAPI, File, Form, HTTPException, Query, Request, UploadFile
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.responses import FileResponse, JSONResponse, StreamingResponse
@@ -73,7 +74,6 @@ from exo.shared.types.api import (
     BenchChatCompletionResponse,
     BenchImageGenerationResponse,
     BenchImageGenerationTaskParams,
-    CancelCommandResponse,
     ChatCompletionChoice,
     ChatCompletionMessage,
     ChatCompletionRequest,
@@ -82,8 +82,6 @@ from exo.shared.types.api import (
     CreateInstanceResponse,
     DeleteDownloadResponse,
     DeleteInstanceResponse,
-    DeleteTracesRequest,
-    DeleteTracesResponse,
     ErrorInfo,
     ErrorResponse,
     FinishReason,
@@ -143,10 +141,11 @@ from exo.shared.types.commands import (
     TaskFinished,
     TextGeneration,
 )
-from exo.shared.types.common import CommandId, Id, NodeId, SystemId
+from exo.shared.types.common import CommandId, Id, NodeId, SessionId, SystemId
 from exo.shared.types.events import (
     ChunkGenerated,
     Event,
+    GlobalForwarderEvent,
     IndexedEvent,
     TracesMerged,
 )
@@ -174,7 +173,7 @@ from exo.shared.types.worker.instances import Instance, InstanceId, InstanceMeta
 from exo.shared.types.worker.shards import Sharding
 from exo.utils.banner import print_startup_banner
 from exo.utils.channels import Receiver, Sender, channel
-from exo.utils.task_group import TaskGroup
+from exo.utils.event_buffer import OrderedBuffer
 
 _API_EVENT_LOG_DIR = EXO_EVENT_LOG_DIR / "api"
 ONBOARDING_COMPLETE_FILE = EXO_CACHE_HOME / "onboarding_complete"
@@ -197,9 +196,10 @@ class API:
     def __init__(
         self,
         node_id: NodeId,
+        session_id: SessionId,
         *,
         port: int,
-        event_receiver: Receiver[IndexedEvent],
+        global_event_receiver: Receiver[GlobalForwarderEvent],
         command_sender: Sender[ForwarderCommand],
         download_command_sender: Sender[ForwarderDownloadCommand],
         # This lets us pause the API if an election is running
@@ -210,9 +210,11 @@ class API:
         self._system_id = SystemId()
         self.command_sender = command_sender
         self.download_command_sender = download_command_sender
-        self.event_receiver = event_receiver
+        self.global_event_receiver = global_event_receiver
         self.election_receiver = election_receiver
+        self.event_buffer: OrderedBuffer[Event] = OrderedBuffer[Event]()
         self.node_id: NodeId = node_id
+        self.session_id: SessionId = session_id
         self.last_completed_election: int = 0
         self.port = port
 
@@ -250,20 +252,19 @@ class API:
             CommandId, Sender[ImageChunk | ErrorChunk]
         ] = {}
         self._image_store = ImageStore(EXO_IMAGE_CACHE_DIR)
-        self._tg: TaskGroup = TaskGroup()
+        self._tg: TaskGroup | None = None
 
-    def reset(self, result_clock: int, event_receiver: Receiver[IndexedEvent]):
+    def reset(self, new_session_id: SessionId, result_clock: int):
         logger.info("Resetting API State")
         self._event_log.close()
         self._event_log = DiskEventLog(_API_EVENT_LOG_DIR)
         self.state = State()
         self._system_id = SystemId()
+        self.session_id = new_session_id
+        self.event_buffer = OrderedBuffer[Event]()
         self._text_generation_queues = {}
         self._image_generation_queues = {}
         self.unpause(result_clock)
-        self.event_receiver.close()
-        self.event_receiver = event_receiver
-        self._tg.start_soon(self._apply_state)
 
     def unpause(self, result_clock: int):
         logger.info("Unpausing API")
@@ -323,7 +324,6 @@ class API:
         self.app.get("/images/{image_id}")(self.get_image)
         self.app.post("/v1/messages", response_model=None)(self.claude_messages)
         self.app.post("/v1/responses", response_model=None)(self.openai_responses)
-        self.app.post("/v1/cancel/{command_id}")(self.cancel_command)
 
         # Ollama API
         self.app.head("/ollama/")(self.ollama_version)
@@ -344,7 +344,6 @@ class API:
         self.app.post("/download/start")(self.start_download)
         self.app.delete("/download/{node_id}/{model_id:path}")(self.delete_download)
         self.app.get("/v1/traces")(self.list_traces)
-        self.app.post("/v1/traces/delete")(self.delete_traces)
         self.app.get("/v1/traces/{task_id}")(self.get_trace)
         self.app.get("/v1/traces/{task_id}/stats")(self.get_trace_stats)
         self.app.get("/v1/traces/{task_id}/raw")(self.get_trace_raw)
@@ -569,25 +568,6 @@ class API:
             instance_id=instance_id,
         )
 
-    async def cancel_command(self, command_id: CommandId) -> CancelCommandResponse:
-        """Cancel an active command by closing its stream and notifying workers."""
-        sender = self._text_generation_queues.get(
-            command_id
-        ) or self._image_generation_queues.get(command_id)
-        if sender is None:
-            raise HTTPException(
-                status_code=404,
-                detail="Command not found or already completed",
-            )
-
-        await self._send(TaskCancelled(cancelled_command_id=command_id))
-        sender.close()
-
-        return CancelCommandResponse(
-            message="Command cancelled.",
-            command_id=command_id,
-        )
-
     async def _token_chunk_stream(
         self, command_id: CommandId
     ) -> AsyncGenerator[
@@ -1611,7 +1591,8 @@ class API:
         shutdown_ev = anyio.Event()
 
         try:
-            async with self._tg as tg:
+            async with create_task_group() as tg:
+                self._tg = tg
                 logger.info("Starting API")
                 tg.start_soon(self._apply_state)
                 tg.start_soon(self._pause_on_new_election)
@@ -1626,7 +1607,7 @@ class API:
         finally:
             self._event_log.close()
             self.command_sender.close()
-            self.event_receiver.close()
+            self.global_event_receiver.close()
 
     async def run_api(self, ev: anyio.Event):
         cfg = Config()
@@ -1643,31 +1624,38 @@ class API:
             )
 
     async def _apply_state(self):
-        with self.event_receiver as events:
-            async for i_event in events:
-                self._event_log.append(i_event.event)
-                self.state = apply(self.state, i_event)
-                event = i_event.event
+        with self.global_event_receiver as events:
+            async for f_event in events:
+                if f_event.session != self.session_id:
+                    continue
+                if f_event.origin != self.session_id.master_node_id:
+                    continue
+                self.event_buffer.ingest(f_event.origin_idx, f_event.event)
+                for idx, event in self.event_buffer.drain_indexed():
+                    self._event_log.append(event)
+                    self.state = apply(self.state, IndexedEvent(event=event, idx=idx))
 
-                if isinstance(event, ChunkGenerated):
-                    if queue := self._image_generation_queues.get(
-                        event.command_id, None
-                    ):
-                        assert isinstance(event.chunk, ImageChunk)
-                        try:
-                            await queue.send(event.chunk)
-                        except BrokenResourceError:
-                            self._image_generation_queues.pop(event.command_id, None)
-                    if queue := self._text_generation_queues.get(
-                        event.command_id, None
-                    ):
-                        assert not isinstance(event.chunk, ImageChunk)
-                        try:
-                            await queue.send(event.chunk)
-                        except BrokenResourceError:
-                            self._text_generation_queues.pop(event.command_id, None)
-                if isinstance(event, TracesMerged):
-                    self._save_merged_trace(event)
+                    if isinstance(event, ChunkGenerated):
+                        if queue := self._image_generation_queues.get(
+                            event.command_id, None
+                        ):
+                            assert isinstance(event.chunk, ImageChunk)
+                            try:
+                                await queue.send(event.chunk)
+                            except BrokenResourceError:
+                                self._image_generation_queues.pop(
+                                    event.command_id, None
+                                )
+                        if queue := self._text_generation_queues.get(
+                            event.command_id, None
+                        ):
+                            assert not isinstance(event.chunk, ImageChunk)
+                            try:
+                                await queue.send(event.chunk)
+                            except BrokenResourceError:
+                                self._text_generation_queues.pop(event.command_id, None)
+                    if isinstance(event, TracesMerged):
+                        self._save_merged_trace(event)
 
     def _save_merged_trace(self, event: TracesMerged) -> None:
         traces = [
@@ -1731,12 +1719,8 @@ class API:
         await self._send_download(command)
         return DeleteDownloadResponse(command_id=command.command_id)
 
-    @staticmethod
-    def _get_trace_path(task_id: str) -> Path:
-        trace_path = EXO_TRACING_CACHE_DIR / f"trace_{task_id}.json"
-        if not trace_path.resolve().is_relative_to(EXO_TRACING_CACHE_DIR.resolve()):
-            raise HTTPException(status_code=400, detail=f"Invalid task ID: {task_id}")
-        return trace_path
+    def _get_trace_path(self, task_id: str) -> Path:
+        return EXO_TRACING_CACHE_DIR / f"trace_{task_id}.json"
 
     async def list_traces(self) -> TraceListResponse:
         traces: list[TraceListItem] = []
@@ -1835,18 +1819,6 @@ class API:
             filename=f"trace_{task_id}.json",
         )
 
-    async def delete_traces(self, request: DeleteTracesRequest) -> DeleteTracesResponse:
-        deleted: list[str] = []
-        not_found: list[str] = []
-        for task_id in request.task_ids:
-            trace_path = self._get_trace_path(task_id)
-            if trace_path.exists():
-                trace_path.unlink()
-                deleted.append(task_id)
-            else:
-                not_found.append(task_id)
-        return DeleteTracesResponse(deleted=deleted, not_found=not_found)
-
     async def get_onboarding(self) -> JSONResponse:
         return JSONResponse({"completed": ONBOARDING_COMPLETE_FILE.exists()})
 

src/exo/master/main.py

@@ -1,6 +1,7 @@
 from datetime import datetime, timedelta, timezone
 
 import anyio
+from anyio.abc import TaskGroup
 from loguru import logger
 
 from exo.master.event_log import DiskEventLog
@@ -60,9 +61,8 @@ from exo.shared.types.tasks import (
     TextGeneration as TextGenerationTask,
 )
 from exo.shared.types.worker.instances import InstanceId
-from exo.utils.channels import Receiver, Sender
+from exo.utils.channels import Receiver, Sender, channel
 from exo.utils.event_buffer import MultiSourceBuffer
-from exo.utils.task_group import TaskGroup
 
 
 class Master:
@@ -72,21 +72,25 @@ class Master:
         session_id: SessionId,
         *,
         command_receiver: Receiver[ForwarderCommand],
-        event_sender: Sender[Event],
         local_event_receiver: Receiver[LocalForwarderEvent],
         global_event_sender: Sender[GlobalForwarderEvent],
         download_command_sender: Sender[ForwarderDownloadCommand],
     ):
+        self.state = State()
+        self._tg: TaskGroup = anyio.create_task_group()
         self.node_id = node_id
         self.session_id = session_id
-        self.state = State()
-        self._tg: TaskGroup = TaskGroup()
         self.command_task_mapping: dict[CommandId, TaskId] = {}
         self.command_receiver = command_receiver
         self.local_event_receiver = local_event_receiver
         self.global_event_sender = global_event_sender
         self.download_command_sender = download_command_sender
-        self.event_sender = event_sender
+        send, recv = channel[Event]()
+        self.event_sender: Sender[Event] = send
+        self._loopback_event_receiver: Receiver[Event] = recv
+        self._loopback_event_sender: Sender[LocalForwarderEvent] = (
+            local_event_receiver.clone_sender()
+        )
         self._system_id = SystemId()
         self._multi_buffer = MultiSourceBuffer[SystemId, Event]()
         self._event_log = DiskEventLog(EXO_EVENT_LOG_DIR / "master")
@@ -100,16 +104,19 @@ class Master:
             async with self._tg as tg:
                 tg.start_soon(self._event_processor)
                 tg.start_soon(self._command_processor)
+                tg.start_soon(self._loopback_processor)
                 tg.start_soon(self._plan)
         finally:
             self._event_log.close()
             self.global_event_sender.close()
             self.local_event_receiver.close()
             self.command_receiver.close()
+            self._loopback_event_sender.close()
+            self._loopback_event_receiver.close()
 
     async def shutdown(self):
         logger.info("Stopping Master")
-        self._tg.cancel_tasks()
+        self._tg.cancel_scope.cancel()
 
     async def _command_processor(self) -> None:
         with self.command_receiver as commands:
@@ -328,22 +335,17 @@ class Master:
                                         task_id=task_id,
                                     )
                                 )
-                            else:
-                                logger.warning(
-                                    f"Nonexistent command {command.cancelled_command_id} cancelled"
-                                )
                         case TaskFinished():
-                            if (
-                                task_id := self.command_task_mapping.pop(
-                                    command.finished_command_id, None
+                            generated_events.append(
+                                TaskDeleted(
+                                    task_id=self.command_task_mapping[
+                                        command.finished_command_id
+                                    ]
                                 )
-                            ) is not None:
-                                generated_events.append(TaskDeleted(task_id=task_id))
-                            else:
-                                logger.warning(
-                                    f"Finished command {command.finished_command_id} finished"
-                                )
-
+                            )
+                            self.command_task_mapping.pop(
+                                command.finished_command_id, None
+                            )
                         case RequestEventLog():
                             # We should just be able to send everything, since other buffers will ignore old messages
                             # rate limit to 1000 at a time
@@ -407,6 +409,22 @@ class Master:
                     self._event_log.append(event)
                     await self._send_event(indexed)
 
+    async def _loopback_processor(self) -> None:
+        # this would ideally not be necessary.
+        # this is WAY less hacky than how I was working around this before
+        local_index = 0
+        with self._loopback_event_receiver as events:
+            async for event in events:
+                await self._loopback_event_sender.send(
+                    LocalForwarderEvent(
+                        origin=self._system_id,
+                        origin_idx=local_index,
+                        session=self.session_id,
+                        event=event,
+                    )
+                )
+                local_index += 1
+
     # This function is re-entrant, take care!
     async def _send_event(self, event: IndexedEvent):
         # Convenience method since this line is ugly

src/exo/master/placement.py

@@ -112,11 +112,7 @@ def place_instance(
         cycle for cycle in smallest_cycles if topology.is_rdma_cycle(cycle)
     ]
 
-    if command.instance_meta == InstanceMeta.MlxJaccl:
-        if not smallest_rdma_cycles:
-            raise ValueError(
-                "Requested RDMA (MlxJaccl) but no RDMA-connected cycles available"
-            )
+    if command.instance_meta == InstanceMeta.MlxJaccl and smallest_rdma_cycles != []:
         smallest_cycles = smallest_rdma_cycles
 
     cycles_with_leaf_nodes: list[Cycle] = [
@@ -133,11 +129,6 @@ def place_instance(
         ),
     )
 
-    # Single-node: force Pipeline/Ring (Tensor and Jaccl require multi-node)
-    if len(selected_cycle) == 1:
-        command.instance_meta = InstanceMeta.MlxRing
-        command.sharding = Sharding.Pipeline
-
     shard_assignments = get_shard_assignments(
         command.model_card, selected_cycle, command.sharding, node_memory
     )
@@ -147,6 +138,9 @@ def place_instance(
     instance_id = InstanceId()
     target_instances = dict(deepcopy(current_instances))
 
+    if len(selected_cycle) == 1:
+        command.instance_meta = InstanceMeta.MlxRing
+
     match command.instance_meta:
         case InstanceMeta.MlxJaccl:
             # TODO(evan): shard assignments should contain information about ranks, this is ugly

src/exo/master/tests/test_cancel_command.py

@@ -1,77 +0,0 @@
-# pyright: reportUnusedFunction=false, reportAny=false
-from typing import Any
-from unittest.mock import AsyncMock, MagicMock
-
-from fastapi import FastAPI
-from fastapi.testclient import TestClient
-
-from exo.shared.types.common import CommandId
-
-
-def _make_api() -> Any:
-    """Create a minimal API instance with cancel route and error handler."""
-    from exo.master.api import API
-
-    app = FastAPI()
-    api = object.__new__(API)
-    api.app = app
-    api._text_generation_queues = {}  # pyright: ignore[reportPrivateUsage]
-    api._image_generation_queues = {}  # pyright: ignore[reportPrivateUsage]
-    api._send = AsyncMock()  # pyright: ignore[reportPrivateUsage]
-    api._setup_exception_handlers()  # pyright: ignore[reportPrivateUsage]
-    app.post("/v1/cancel/{command_id}")(api.cancel_command)
-    return api
-
-
-def test_cancel_nonexistent_command_returns_404() -> None:
-    """Cancel for an unknown command_id returns 404 in OpenAI error format."""
-    api = _make_api()
-    client = TestClient(api.app)
-
-    response = client.post("/v1/cancel/nonexistent-id")
-    assert response.status_code == 404
-    data: dict[str, Any] = response.json()
-    assert "error" in data
-    assert data["error"]["message"] == "Command not found or already completed"
-    assert data["error"]["type"] == "Not Found"
-    assert data["error"]["code"] == 404
-
-
-def test_cancel_active_text_generation() -> None:
-    """Cancel an active text generation command: returns 200, sender.close() called."""
-    api = _make_api()
-    client = TestClient(api.app)
-
-    cid = CommandId("text-cmd-123")
-    sender = MagicMock()
-    api._text_generation_queues[cid] = sender
-
-    response = client.post(f"/v1/cancel/{cid}")
-    assert response.status_code == 200
-    data: dict[str, Any] = response.json()
-    assert data["message"] == "Command cancelled."
-    assert data["command_id"] == str(cid)
-    sender.close.assert_called_once()
-    api._send.assert_called_once()
-    task_cancelled = api._send.call_args[0][0]
-    assert task_cancelled.cancelled_command_id == cid
-
-
-def test_cancel_active_image_generation() -> None:
-    """Cancel an active image generation command: returns 200, sender.close() called."""
-    api = _make_api()
-    client = TestClient(api.app)
-
-    cid = CommandId("img-cmd-456")
-    sender = MagicMock()
-    api._image_generation_queues[cid] = sender
-
-    response = client.post(f"/v1/cancel/{cid}")
-    assert response.status_code == 200
-    data: dict[str, Any] = response.json()
-    assert data["message"] == "Command cancelled."
-    assert data["command_id"] == str(cid)
-    sender.close.assert_called_once()
-    api._send.assert_called_once()
-    task_cancelled = api._send.call_args[0][0]
-    assert task_cancelled.cancelled_command_id == cid

src/exo/master/tests/test_master.py

@@ -17,7 +17,6 @@ from exo.shared.types.commands import (
 )
 from exo.shared.types.common import ModelId, NodeId, SessionId, SystemId
 from exo.shared.types.events import (
-    Event,
     GlobalForwarderEvent,
     IndexedEvent,
     InstanceCreated,
@@ -51,22 +50,6 @@ async def test_master():
     command_sender, co_receiver = channel[ForwarderCommand]()
     local_event_sender, le_receiver = channel[LocalForwarderEvent]()
     fcds, _fcdr = channel[ForwarderDownloadCommand]()
-    ev_send, ev_recv = channel[Event]()
-
-    async def mock_event_router():
-        idx = 0
-        sid = SystemId()
-        with ev_recv as master_events:
-            async for event in master_events:
-                await local_event_sender.send(
-                    LocalForwarderEvent(
-                        origin=sid,
-                        origin_idx=idx,
-                        session=session_id,
-                        event=event,
-                    )
-                )
-                idx += 1
 
     all_events: list[IndexedEvent] = []
 
@@ -84,7 +67,6 @@ async def test_master():
     master = Master(
         node_id,
         session_id,
-        event_sender=ev_send,
         global_event_sender=ge_sender,
         local_event_receiver=le_receiver,
         command_receiver=co_receiver,
@@ -93,7 +75,6 @@ async def test_master():
     logger.info("run the master")
     async with anyio.create_task_group() as tg:
         tg.start_soon(master.run)
-        tg.start_soon(mock_event_router)
 
         # inject a NodeGatheredInfo event
         logger.info("inject a NodeGatheredInfo event")
@@ -216,5 +197,4 @@ async def test_master():
             input=[InputMessage(role="user", content="Hello, how are you?")],
         )
 
-        ev_send.close()
         await master.shutdown()

src/exo/master/tests/test_placement.py

@@ -14,12 +14,10 @@ from exo.shared.models.model_cards import ModelCard, ModelId, ModelTask
 from exo.shared.topology import Topology
 from exo.shared.types.commands import PlaceInstance
 from exo.shared.types.common import CommandId, NodeId
-from exo.shared.types.events import InstanceCreated, InstanceDeleted, TaskStatusUpdated
+from exo.shared.types.events import InstanceCreated, InstanceDeleted
 from exo.shared.types.memory import Memory
 from exo.shared.types.multiaddr import Multiaddr
 from exo.shared.types.profiling import NetworkInterfaceInfo, NodeNetworkInfo
-from exo.shared.types.tasks import TaskId, TaskStatus, TextGeneration
-from exo.shared.types.text_generation import InputMessage, TextGenerationTaskParams
 from exo.shared.types.topology import Connection, SocketConnection
 from exo.shared.types.worker.instances import (
     Instance,
@@ -458,117 +456,3 @@ def test_tensor_rdma_backend_connectivity_matrix(
         else:
             ip_part = coordinator.split(":")[0]
             assert len(ip_part.split(".")) == 4
-
-
-def _make_task(
-    instance_id: InstanceId,
-    status: TaskStatus = TaskStatus.Running,
-) -> TextGeneration:
-    return TextGeneration(
-        task_id=TaskId(),
-        task_status=status,
-        instance_id=instance_id,
-        command_id=CommandId(),
-        task_params=TextGenerationTaskParams(
-            model=ModelId("test-model"),
-            input=[InputMessage(role="user", content="hello")],
-        ),
-    )
-
-
-def test_get_transition_events_delete_instance_cancels_running_tasks(
-    instance: Instance,
-):
-    # arrange
-    instance_id = InstanceId()
-    current_instances: dict[InstanceId, Instance] = {instance_id: instance}
-    target_instances: dict[InstanceId, Instance] = {}
-    task = _make_task(instance_id, TaskStatus.Running)
-    tasks = {task.task_id: task}
-
-    # act
-    events = get_transition_events(current_instances, target_instances, tasks)
-
-    # assert – cancellation event should come before the deletion event
-    assert len(events) == 2
-    assert isinstance(events[0], TaskStatusUpdated)
-    assert events[0].task_id == task.task_id
-    assert events[0].task_status == TaskStatus.Cancelled
-    assert isinstance(events[1], InstanceDeleted)
-    assert events[1].instance_id == instance_id
-
-
-def test_get_transition_events_delete_instance_cancels_pending_tasks(
-    instance: Instance,
-):
-    # arrange
-    instance_id = InstanceId()
-    current_instances: dict[InstanceId, Instance] = {instance_id: instance}
-    target_instances: dict[InstanceId, Instance] = {}
-    task = _make_task(instance_id, TaskStatus.Pending)
-    tasks = {task.task_id: task}
-
-    # act
-    events = get_transition_events(current_instances, target_instances, tasks)
-
-    # assert
-    assert len(events) == 2
-    assert isinstance(events[0], TaskStatusUpdated)
-    assert events[0].task_id == task.task_id
-    assert events[0].task_status == TaskStatus.Cancelled
-    assert isinstance(events[1], InstanceDeleted)
-
-
-def test_get_transition_events_delete_instance_ignores_completed_tasks(
-    instance: Instance,
-):
-    # arrange
-    instance_id = InstanceId()
-    current_instances: dict[InstanceId, Instance] = {instance_id: instance}
-    target_instances: dict[InstanceId, Instance] = {}
-    tasks = {
-        t.task_id: t
-        for t in [
-            _make_task(instance_id, TaskStatus.Complete),
-            _make_task(instance_id, TaskStatus.Failed),
-            _make_task(instance_id, TaskStatus.TimedOut),
-            _make_task(instance_id, TaskStatus.Cancelled),
-        ]
-    }
-
-    # act
-    events = get_transition_events(current_instances, target_instances, tasks)
-
-    # assert – only the InstanceDeleted event, no cancellations
-    assert len(events) == 1
-    assert isinstance(events[0], InstanceDeleted)
-
-
-def test_get_transition_events_delete_instance_cancels_only_matching_tasks(
-    instance: Instance,
-):
-    # arrange
-    instance_id_a = InstanceId()
-    instance_id_b = InstanceId()
-    current_instances: dict[InstanceId, Instance] = {
-        instance_id_a: instance,
-        instance_id_b: instance,
-    }
-    # only delete instance A, keep instance B
-    target_instances: dict[InstanceId, Instance] = {instance_id_b: instance}
-
-    task_a = _make_task(instance_id_a, TaskStatus.Running)
-    task_b = _make_task(instance_id_b, TaskStatus.Running)
-    tasks = {task_a.task_id: task_a, task_b.task_id: task_b}
-
-    # act
-    events = get_transition_events(current_instances, target_instances, tasks)
-
-    # assert – only task_a should be cancelled
-    cancel_events = [e for e in events if isinstance(e, TaskStatusUpdated)]
-    delete_events = [e for e in events if isinstance(e, InstanceDeleted)]
-    assert len(cancel_events) == 1
-    assert cancel_events[0].task_id == task_a.task_id
-    assert cancel_events[0].task_status == TaskStatus.Cancelled
-    assert len(delete_events) == 1
-    assert delete_events[0].instance_id == instance_id_a

src/exo/routing/connection_message.py

@@ -1,4 +1,6 @@
-from exo_pyo3_bindings import PyFromSwarm
+from enum import Enum
+
+from exo_pyo3_bindings import ConnectionUpdate, ConnectionUpdateType
 
 from exo.shared.types.common import NodeId
 from exo.utils.pydantic_ext import CamelCaseModel
@@ -6,10 +8,30 @@ from exo.utils.pydantic_ext import CamelCaseModel
 """Serialisable types for Connection Updates/Messages"""
 
 
+class ConnectionMessageType(Enum):
+    Connected = 0
+    Disconnected = 1
+
+    @staticmethod
+    def from_update_type(update_type: ConnectionUpdateType):
+        match update_type:
+            case ConnectionUpdateType.Connected:
+                return ConnectionMessageType.Connected
+            case ConnectionUpdateType.Disconnected:
+                return ConnectionMessageType.Disconnected
+
+
 class ConnectionMessage(CamelCaseModel):
     node_id: NodeId
-    connected: bool
+    connection_type: ConnectionMessageType
+    remote_ipv4: str
+    remote_tcp_port: int
 
     @classmethod
-    def from_update(cls, update: PyFromSwarm.Connection) -> "ConnectionMessage":
-        return cls(node_id=NodeId(update.peer_id), connected=update.connected)
+    def from_update(cls, update: ConnectionUpdate) -> "ConnectionMessage":
+        return cls(
+            node_id=NodeId(update.peer_id),
+            connection_type=ConnectionMessageType.from_update_type(update.update_type),
+            remote_ipv4=update.remote_ipv4,
+            remote_tcp_port=update.remote_tcp_port,
+        )

src/exo/routing/event_router.py

@@ -1,161 +0,0 @@
-from dataclasses import dataclass, field
-from random import random
-
-import anyio
-from anyio import BrokenResourceError, ClosedResourceError
-from anyio.abc import CancelScope
-from loguru import logger
-
-from exo.shared.types.commands import ForwarderCommand, RequestEventLog
-from exo.shared.types.common import SessionId, SystemId
-from exo.shared.types.events import (
-    Event,
-    EventId,
-    GlobalForwarderEvent,
-    IndexedEvent,
-    LocalForwarderEvent,
-)
-from exo.utils.channels import Receiver, Sender, channel
-from exo.utils.event_buffer import OrderedBuffer
-from exo.utils.task_group import TaskGroup
-
-
-@dataclass
-class EventRouter:
-    session_id: SessionId
-    command_sender: Sender[ForwarderCommand]
-    external_inbound: Receiver[GlobalForwarderEvent]
-    external_outbound: Sender[LocalForwarderEvent]
-    _system_id: SystemId = field(init=False, default_factory=SystemId)
-    internal_outbound: list[Sender[IndexedEvent]] = field(
-        init=False, default_factory=list
-    )
-    event_buffer: OrderedBuffer[Event] = field(
-        init=False, default_factory=OrderedBuffer
-    )
-    out_for_delivery: dict[EventId, tuple[float, LocalForwarderEvent]] = field(
-        init=False, default_factory=dict
-    )
-    _tg: TaskGroup = field(init=False, default_factory=TaskGroup)
-
-    _nack_cancel_scope: CancelScope | None = field(init=False, default=None)
-    _nack_attempts: int = field(init=False, default=0)
-    _nack_base_seconds: float = field(init=False, default=0.5)
-    _nack_cap_seconds: float = field(init=False, default=10.0)
-
-    async def run(self):
-        try:
-            async with self._tg as tg:
-                tg.start_soon(self._run_ext_in)
-                tg.start_soon(self._simple_retry)
-        finally:
-            self.external_outbound.close()
-            for send in self.internal_outbound:
-                send.close()
-
-    # can make this better in future
-    async def _simple_retry(self):
-        while True:
-            await anyio.sleep(1 + random())
-            # list here is a shallow clone for shared mutation
-            for e_id, (time, event) in list(self.out_for_delivery.items()):
-                if anyio.current_time() > time + 5:
-                    self.out_for_delivery[e_id] = (anyio.current_time(), event)
-                    await self.external_outbound.send(event)
-
-    def sender(self) -> Sender[Event]:
-        send, recv = channel[Event]()
-        if self._tg.is_running():
-            self._tg.start_soon(self._ingest, SystemId(), recv)
-        else:
-            self._tg.queue(self._ingest, SystemId(), recv)
-        return send
-
-    def receiver(self) -> Receiver[IndexedEvent]:
-        send, recv = channel[IndexedEvent]()
-        self.internal_outbound.append(send)
-        return recv
-
-    def shutdown(self) -> None:
-        self._tg.cancel_tasks()
-
-    async def _ingest(self, system_id: SystemId, recv: Receiver[Event]):
-        idx = 0
-        with recv as events:
-            async for event in events:
-                f_ev = LocalForwarderEvent(
-                    origin_idx=idx,
-                    origin=system_id,
-                    session=self.session_id,
-                    event=event,
-                )
-                idx += 1
-                await self.external_outbound.send(f_ev)
-                self.out_for_delivery[event.event_id] = (anyio.current_time(), f_ev)
-
-    async def _run_ext_in(self):
-        buf = OrderedBuffer[Event]()
-        with self.external_inbound as events:
-            async for event in events:
-                if event.session != self.session_id:
-                    continue
-                if event.origin != self.session_id.master_node_id:
-                    continue
-
-                buf.ingest(event.origin_idx, event.event)
-                event_id = event.event.event_id
-                if event_id in self.out_for_delivery:
-                    self.out_for_delivery.pop(event_id)
-
-                drained = buf.drain_indexed()
-                if drained:
-                    self._nack_attempts = 0
-                    if self._nack_cancel_scope:
-                        self._nack_cancel_scope.cancel()
-
-                if not drained and (
-                    self._nack_cancel_scope is None
-                    or self._nack_cancel_scope.cancel_called
-                ):
-                    # Request the next index.
-                    self._tg.start_soon(self._nack_request, buf.next_idx_to_release)
-                    continue
-
-                for idx, event in drained:
-                    to_clear = set[int]()
-                    for i, sender in enumerate(self.internal_outbound):
-                        try:
-                            await sender.send(IndexedEvent(idx=idx, event=event))
-                        except (ClosedResourceError, BrokenResourceError):
-                            to_clear.add(i)
-                    for i in sorted(to_clear, reverse=True):
-                        self.internal_outbound.pop(i)
-
-    async def _nack_request(self, since_idx: int) -> None:
-        # We request all events after (and including) the missing index.
-        # This function is started whenever we receive an event that is out of sequence.
-        # It is cancelled as soon as we receiver an event that is in sequence.
-
-        if since_idx < 0:
-            logger.warning(f"Negative value encountered for nack request {since_idx=}")
-            since_idx = 0
-
-        with CancelScope() as scope:
-            self._nack_cancel_scope = scope
-            delay: float = self._nack_base_seconds * (2.0**self._nack_attempts)
-            delay = min(self._nack_cap_seconds, delay)
-            self._nack_attempts += 1
-            try:
-                await anyio.sleep(delay)
-                logger.info(
-                    f"Nack attempt {self._nack_attempts}: Requesting Event Log from {since_idx}"
-                )
-                await self.command_sender.send(
-                    ForwarderCommand(
-                        origin=self._system_id,
-                        command=RequestEventLog(since_idx=since_idx),
-                    )
-                )
-            finally:
-                if self._nack_cancel_scope is scope:
-                    self._nack_cancel_scope = None

src/exo/routing/router.py

@@ -8,16 +8,16 @@ from typing import cast
 from anyio import (
     BrokenResourceError,
     ClosedResourceError,
+    create_task_group,
     move_on_after,
     sleep_forever,
 )
+from anyio.abc import TaskGroup
 from exo_pyo3_bindings import (
     AllQueuesFullError,
     Keypair,
-    MessageTooLargeError,
     NetworkingHandle,
     NoPeersSubscribedToTopicError,
-    PyFromSwarm,
 )
 from filelock import FileLock
 from loguru import logger
@@ -25,7 +25,6 @@ from loguru import logger
 from exo.shared.constants import EXO_NODE_ID_KEYPAIR
 from exo.utils.channels import Receiver, Sender, channel
 from exo.utils.pydantic_ext import CamelCaseModel
-from exo.utils.task_group import TaskGroup
 
 from .connection_message import ConnectionMessage
 from .topics import CONNECTION_MESSAGES, PublishPolicy, TypedTopic
@@ -112,9 +111,10 @@ class Router:
         self._net: NetworkingHandle = handle
         self._tmp_networking_sender: Sender[tuple[str, bytes]] | None = send
         self._id_count = count()
-        self._tg: TaskGroup = TaskGroup()
+        self._tg: TaskGroup | None = None
 
     async def register_topic[T: CamelCaseModel](self, topic: TypedTopic[T]):
+        assert self._tg is None, "Attempted to register topic after setup time"
         send = self._tmp_networking_sender
         if send:
             self._tmp_networking_sender = None
@@ -122,8 +122,7 @@ class Router:
             send = self.networking_receiver.clone_sender()
         router = TopicRouter[T](topic, send)
         self.topic_routers[topic.topic] = cast(TopicRouter[CamelCaseModel], router)
-        if self._tg.is_running():
-            await self._networking_subscribe(topic.topic)
+        await self._networking_subscribe(str(topic.topic))
 
     def sender[T: CamelCaseModel](self, topic: TypedTopic[T]) -> Sender[T]:
         router = self.topic_routers.get(topic.topic, None)
@@ -149,15 +148,14 @@ class Router:
     async def run(self):
         logger.debug("Starting Router")
         try:
-            async with self._tg as tg:
+            async with create_task_group() as tg:
+                self._tg = tg
                 for topic in self.topic_routers:
                     router = self.topic_routers[topic]
                     tg.start_soon(router.run)
                 tg.start_soon(self._networking_recv)
+                tg.start_soon(self._networking_recv_connection_messages)
                 tg.start_soon(self._networking_publish)
-                # subscribe to pending topics
-                for topic in self.topic_routers:
-                    await self._networking_subscribe(topic)
                 # Router only shuts down if you cancel it.
                 await sleep_forever()
         finally:
@@ -167,7 +165,9 @@ class Router:
 
     async def shutdown(self):
         logger.debug("Shutting down Router")
-        self._tg.cancel_tasks()
+        if not self._tg:
+            return
+        self._tg.cancel_scope.cancel()
 
     async def _networking_subscribe(self, topic: str):
         await self._net.gossipsub_subscribe(topic)
@@ -178,41 +178,28 @@ class Router:
         logger.info(f"Unsubscribed from {topic}")
 
     async def _networking_recv(self):
-        try:
-            while True:
-                from_swarm = await self._net.recv()
-                logger.debug(from_swarm)
-                match from_swarm:
-                    case PyFromSwarm.Message(origin, topic, data):
-                        logger.trace(
-                            f"Received message on {topic} from {origin} with payload {data}"
-                        )
-                        if topic not in self.topic_routers:
-                            logger.warning(
-                                f"Received message on unknown or inactive topic {topic}"
-                            )
-                            continue
-                        router = self.topic_routers[topic]
-                        await router.publish_bytes(data)
-                    case PyFromSwarm.Connection():
-                        message = ConnectionMessage.from_update(from_swarm)
-                        logger.trace(
-                            f"Received message on connection_messages with payload {message}"
-                        )
-                        if CONNECTION_MESSAGES.topic in self.topic_routers:
-                            router = self.topic_routers[CONNECTION_MESSAGES.topic]
-                            assert router.topic.model_type == ConnectionMessage
-                            router = cast(TopicRouter[ConnectionMessage], router)
-                            await router.publish(message)
-                    case _:
-                        logger.critical(
-                            "failed to exhaustively check FromSwarm messages - logic error"
-                        )
-        except Exception as exception:
-            logger.opt(exception=exception).error(
-                "Gossipsub receive loop terminated unexpectedly"
+        while True:
+            topic, data = await self._net.gossipsub_recv()
+            logger.trace(f"Received message on {topic} with payload {data}")
+            if topic not in self.topic_routers:
+                logger.warning(f"Received message on unknown or inactive topic {topic}")
+                continue
+
+            router = self.topic_routers[topic]
+            await router.publish_bytes(data)
+
+    async def _networking_recv_connection_messages(self):
+        while True:
+            update = await self._net.connection_update_recv()
+            message = ConnectionMessage.from_update(update)
+            logger.trace(
+                f"Received message on connection_messages with payload {message}"
             )
-            raise
+            if CONNECTION_MESSAGES.topic in self.topic_routers:
+                router = self.topic_routers[CONNECTION_MESSAGES.topic]
+                assert router.topic.model_type == ConnectionMessage
+                router = cast(TopicRouter[ConnectionMessage], router)
+                await router.publish(message)
 
     async def _networking_publish(self):
         with self.networking_receiver as networked_items:
@@ -224,10 +211,6 @@ class Router:
                     pass
                 except AllQueuesFullError:
                     logger.warning(f"All peer queues full, dropping message on {topic}")
-                except MessageTooLargeError:
-                    logger.warning(
-                        f"Message too large for gossipsub on {topic} ({len(data)} bytes), dropping"
-                    )
 
 
 def get_node_id_keypair(
@@ -258,6 +241,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()
+            keypair = Keypair.generate_ed25519()
             f.write(keypair.to_bytes())
             return keypair

src/exo/shared/constants.py

@@ -36,12 +36,26 @@ EXO_MODELS_DIR = (
 
 # Read-only search path for pre-downloaded models (colon-separated directories)
 _EXO_MODELS_PATH_ENV = os.environ.get("EXO_MODELS_PATH", None)
-EXO_MODELS_PATH: tuple[Path, ...] | None = (
-    tuple(Path(p).expanduser() for p in _EXO_MODELS_PATH_ENV.split(":") if p)
-    if _EXO_MODELS_PATH_ENV is not None
-    else None
+
+# Well-known model cache directories from other inference engines
+_WELL_KNOWN_MODEL_PATHS: tuple[Path, ...] = tuple(
+    p for p in (Path.home() / ".cache" / "huggingface" / "hub",) if p.is_dir()
 )
 
+
+def _build_models_path() -> tuple[Path, ...] | None:
+    if _EXO_MODELS_PATH_ENV is not None:
+        user_paths = tuple(
+            Path(p).expanduser() for p in _EXO_MODELS_PATH_ENV.split(":") if p
+        )
+        return user_paths + tuple(
+            p for p in _WELL_KNOWN_MODEL_PATHS if p not in user_paths
+        )
+    return _WELL_KNOWN_MODEL_PATHS if _WELL_KNOWN_MODEL_PATHS else None
+
+
+EXO_MODELS_PATH: tuple[Path, ...] | None = _build_models_path()
+
 _RESOURCES_DIR_ENV = os.environ.get("EXO_RESOURCES_DIR", None)
 RESOURCES_DIR = (
     find_resources() if _RESOURCES_DIR_ENV is None else Path.home() / _RESOURCES_DIR_ENV
@@ -78,6 +92,4 @@ EXO_ENABLE_IMAGE_MODELS = (
     os.getenv("EXO_ENABLE_IMAGE_MODELS", "false").lower() == "true"
 )
 
-EXO_OFFLINE = os.getenv("EXO_OFFLINE", "false").lower() == "true"
-
 EXO_TRACING_ENABLED = os.getenv("EXO_TRACING_ENABLED", "false").lower() == "true"

src/exo/shared/election.py

@@ -4,8 +4,10 @@ import anyio
 from anyio import (
     CancelScope,
     Event,
+    create_task_group,
     get_cancelled_exc_class,
 )
+from anyio.abc import TaskGroup
 from loguru import logger
 
 from exo.routing.connection_message import ConnectionMessage
@@ -13,7 +15,6 @@ from exo.shared.types.commands import ForwarderCommand
 from exo.shared.types.common import NodeId, SessionId
 from exo.utils.channels import Receiver, Sender
 from exo.utils.pydantic_ext import CamelCaseModel
-from exo.utils.task_group import TaskGroup
 
 DEFAULT_ELECTION_TIMEOUT = 3.0
 
@@ -81,12 +82,13 @@ class Election:
         self._candidates: list[ElectionMessage] = []
         self._campaign_cancel_scope: CancelScope | None = None
         self._campaign_done: Event | None = None
-        self._tg = TaskGroup()
+        self._tg: TaskGroup | None = None
 
     async def run(self):
         logger.info("Starting Election")
         try:
-            async with self._tg as tg:
+            async with create_task_group() as tg:
+                self._tg = tg
                 tg.start_soon(self._election_receiver)
                 tg.start_soon(self._connection_receiver)
                 tg.start_soon(self._command_counter)
@@ -122,7 +124,12 @@ class Election:
         )
 
     async def shutdown(self) -> None:
-        self._tg.cancel_tasks()
+        if not self._tg:
+            logger.warning(
+                "Attempted to shutdown election service that was not running"
+            )
+            return
+        self._tg.cancel_scope.cancel()
 
     async def _election_receiver(self) -> None:
         with self._em_receiver as election_messages:
@@ -136,6 +143,7 @@ class Election:
                 if message.clock > self.clock:
                     self.clock = message.clock
                     logger.debug(f"New clock: {self.clock}")
+                    assert self._tg is not None
                     logger.debug("Starting new campaign")
                     candidates: list[ElectionMessage] = [message]
                     logger.debug(f"Candidates: {candidates}")
@@ -170,6 +178,7 @@ class Election:
                 # These messages are strictly peer to peer
                 self.clock += 1
                 logger.debug(f"New clock: {self.clock}")
+                assert self._tg is not None
                 candidates: list[ElectionMessage] = []
                 self._candidates = candidates
                 logger.debug("Starting new campaign")

src/exo/shared/models/model_cards.py

@@ -90,7 +90,6 @@ class ModelCard(CamelCaseModel):
     base_model: str = ""
     capabilities: list[str] = []
     uses_cfg: bool = False
-    trust_remote_code: bool = True
 
     @field_validator("tasks", mode="before")
     @classmethod
@@ -138,7 +137,6 @@ class ModelCard(CamelCaseModel):
             hidden_size=config_data.hidden_size or 0,
             supports_tensor=config_data.supports_tensor,
             tasks=[ModelTask.TextGeneration],
-            trust_remote_code=False,
         )
         await mc.save_to_custom_dir()
         _card_cache[model_id] = mc
@@ -190,8 +188,6 @@ class ConfigData(BaseModel):
             ["DeepseekV3ForCausalLM"],
             ["Qwen3NextForCausalLM"],
             ["Qwen3MoeForCausalLM"],
-            ["Qwen3_5MoeForConditionalGeneration"],
-            ["Qwen3_5ForConditionalGeneration"],
             ["MiniMaxM2ForCausalLM"],
             ["LlamaForCausalLM"],
             ["GptOssForCausalLM"],

src/exo/shared/tests/test_election.py

@@ -1,7 +1,7 @@
 import pytest
 from anyio import create_task_group, fail_after, move_on_after
 
-from exo.routing.connection_message import ConnectionMessage
+from exo.routing.connection_message import ConnectionMessage, ConnectionMessageType
 from exo.shared.election import Election, ElectionMessage, ElectionResult
 from exo.shared.types.commands import ForwarderCommand, TestCommand
 from exo.shared.types.common import NodeId, SessionId, SystemId
@@ -327,7 +327,14 @@ async def test_connection_message_triggers_new_round_broadcast() -> None:
             tg.start_soon(election.run)
 
             # Send any connection message object; we close quickly to cancel before result creation
-            await cm_tx.send(ConnectionMessage(node_id=NodeId(), connected=True))
+            await cm_tx.send(
+                ConnectionMessage(
+                    node_id=NodeId(),
+                    connection_type=ConnectionMessageType.Connected,
+                    remote_ipv4="",
+                    remote_tcp_port=0,
+                )
+            )
 
             # Expect a broadcast for the new round at clock=1
             while True:

src/exo/shared/tests/test_resolve_reasoning_params.py

@@ -1,30 +0,0 @@
-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,7 +8,6 @@ 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
@@ -199,7 +198,6 @@ 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
@@ -264,11 +262,6 @@ class DeleteInstanceResponse(BaseModel):
     instance_id: InstanceId
 
 
-class CancelCommandResponse(BaseModel):
-    message: str
-    command_id: CommandId
-
-
 ImageSize = Literal[
     "auto",
     "512x512",
@@ -444,12 +437,3 @@ class TraceListItem(CamelCaseModel):
 
 class TraceListResponse(CamelCaseModel):
     traces: list[TraceListItem]
-
-
-class DeleteTracesRequest(CamelCaseModel):
-    task_ids: list[str]
-
-
-class DeleteTracesResponse(CamelCaseModel):
-    deleted: list[str]
-    not_found: list[str]

src/exo/shared/types/mlx.py

@@ -2,8 +2,6 @@
 
 from collections.abc import Sequence
 
-from mlx import core as mx
-from mlx import nn as nn
 from mlx_lm.models.cache import (
     ArraysCache,
     CacheList,
@@ -16,16 +14,3 @@ from mlx_lm.models.cache import (
 KVCacheType = Sequence[
     KVCache | RotatingKVCache | QuantizedKVCache | ArraysCache | CacheList
 ]
-
-
-# Model is a wrapper function to fix the fact that mlx is not strongly typed in the same way that EXO is.
-# For example - MLX has no guarantee of the interface that nn.Module will expose. But we need a guarantee that it has a __call__() function
-class Model(nn.Module):
-    layers: list[nn.Module]
-
-    def __call__(
-        self,
-        x: mx.array,
-        cache: KVCacheType | None,
-        input_embeddings: mx.array | None = None,
-    ) -> mx.array: ...

src/exo/shared/types/openai_responses.py

@@ -12,7 +12,6 @@ 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"]
@@ -72,13 +71,6 @@ 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.
 
@@ -97,15 +89,8 @@ 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,29 +11,6 @@ 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):
@@ -63,7 +40,6 @@ 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/utils/channels.py

@@ -5,7 +5,7 @@ from math import inf
 from multiprocessing.synchronize import Event
 from queue import Empty, Full
 from types import TracebackType
-from typing import Any, Self
+from typing import Self
 
 from anyio import (
     CapacityLimiter,
@@ -157,7 +157,7 @@ class MpSender[T]:
     ) -> None:
         self.close()
 
-    def __getstate__(self) -> dict[str, Any]:
+    def __getstate__(self):
         d = self.__dict__.copy()
         d.pop("__orig_class__", None)
         return d

src/exo/utils/info_gatherer/info_gatherer.py

@@ -8,7 +8,8 @@ from subprocess import CalledProcessError
 from typing import Self, cast
 
 import anyio
-from anyio import fail_after, open_process, to_thread
+from anyio import create_task_group, fail_after, open_process, to_thread
+from anyio.abc import TaskGroup
 from anyio.streams.buffered import BufferedByteReceiveStream
 from anyio.streams.text import TextReceiveStream
 from loguru import logger
@@ -29,7 +30,6 @@ from exo.shared.types.thunderbolt import (
 )
 from exo.utils.channels import Sender
 from exo.utils.pydantic_ext import TaggedModel
-from exo.utils.task_group import TaskGroup
 
 from .macmon import MacmonMetrics
 from .system_info import (
@@ -381,7 +381,7 @@ class InfoGatherer:
     static_info_poll_interval: float | None = 60
     rdma_ctl_poll_interval: float | None = 10 if IS_DARWIN else None
     disk_poll_interval: float | None = 30
-    _tg: TaskGroup = field(init=False, default_factory=TaskGroup)
+    _tg: TaskGroup = field(init=False, default_factory=create_task_group)
 
     async def run(self):
         async with self._tg as tg:
@@ -408,7 +408,7 @@ class InfoGatherer:
                 await self.info_sender.send(nc)
 
     def shutdown(self):
-        self._tg.cancel_tasks()
+        self._tg.cancel_scope.cancel()
 
     async def _monitor_static_info(self):
         if self.static_info_poll_interval is None:

src/exo/utils/task_group.py

@@ -1,65 +0,0 @@
-from collections.abc import Awaitable, Callable
-from dataclasses import dataclass, field
-from types import TracebackType
-from typing import Any, Unpack
-
-from anyio import create_task_group
-from anyio.abc import TaskGroup as TaskGroupABC
-
-
-@dataclass
-class TaskGroup:
-    _tg: TaskGroupABC | None = field(default=None, init=False)
-    _queued: list[tuple[Any, Any, Any]] | None = field(default_factory=list, init=False)
-
-    def is_running(self) -> bool:
-        return self._tg is not None
-
-    def cancel_tasks(self):
-        assert self._tg
-        self._tg.cancel_scope.cancel()
-
-    def cancel_called(self) -> bool:
-        assert self._tg
-        return self._tg.cancel_scope.cancel_called
-
-    def start_soon[*T](
-        self,
-        func: Callable[[Unpack[T]], Awaitable[Any]],
-        *args: Unpack[T],
-        name: object = None,
-    ) -> None:
-        assert self._tg is not None
-        assert self._queued is None
-        self._tg.start_soon(func, *args, name=name)
-
-    def queue[*T](
-        self,
-        func: Callable[[Unpack[T]], Awaitable[Any]],
-        *args: Unpack[T],
-        name: object = None,
-    ) -> None:
-        assert self._tg is None
-        assert self._queued is not None
-        self._queued.append((func, args, name))
-
-    async def __aenter__(self) -> TaskGroupABC:
-        assert self._tg is None
-        assert self._queued is not None
-        self._tg = create_task_group()
-        r = await self._tg.__aenter__()
-        for func, args, name in self._queued:  # pyright: ignore[reportAny]
-            self._tg.start_soon(func, *args, name=name)  # pyright: ignore[reportAny]
-        self._queued = None
-        return r
-
-    async def __aexit__(
-        self,
-        exc_type: type[BaseException] | None,
-        exc_val: BaseException | None,
-        exc_tb: TracebackType | None,
-    ) -> bool:
-        """Exit the task group context waiting for all tasks to finish."""
-        assert self._tg is not None, "aenter sets self.lazy, so it exists when we aexit"
-        assert self._queued is None
-        return await self._tg.__aexit__(exc_type, exc_val, exc_tb)

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

@@ -0,0 +1,17 @@
+import mlx.core as mx
+import mlx.nn as nn
+from mlx_lm.models.cache import KVCache
+
+# These are wrapper functions to fix the fact that mlx is not strongly typed in the same way that EXO is.
+# For example - MLX has no guarantee of the interface that nn.Module will expose. But we need a guarantee that it has a __call__() function
+
+
+class Model(nn.Module):
+    layers: list[nn.Module]
+
+    def __call__(
+        self,
+        x: mx.array,
+        cache: list[KVCache] | None,
+        input_embeddings: mx.array | None = None,
+    ) -> mx.array: ...

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

@@ -16,7 +16,6 @@ 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
@@ -32,22 +31,14 @@ 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_moe import Qwen3MoeSparseMoeBlock
 from mlx_lm.models.qwen3_next import Model as Qwen3NextModel
-from mlx_lm.models.qwen3_next import (
-    Qwen3NextDecoderLayer,
-    Qwen3NextGatedDeltaNet,
-    Qwen3NextSparseMoeBlock,
-)
+from mlx_lm.models.qwen3_next import Qwen3NextDecoderLayer, 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
+from transformers.models.qwen3.modeling_qwen3 import Qwen3DecoderLayer
 
 from exo.shared.logging import logger
 from exo.shared.types.worker.shards import PipelineShardMetadata
@@ -59,21 +50,6 @@ TimeoutCallback = Callable[[], None]
 LayerLoadedCallback = Callable[[int, int], None]  # (layers_loaded, total_layers)
 
 
-_pending_prefill_sends: list[tuple[mx.array, int, mx.distributed.Group]] = []
-
-
-def flush_prefill_sends() -> None:
-    for output, dst, group in _pending_prefill_sends:
-        sent = mx.distributed.send(output, dst, group=group)
-        mx.async_eval(sent)
-    _pending_prefill_sends.clear()
-
-
-def clear_prefill_sends() -> None:
-    # Discard pending sends (e.g. on cancellation).
-    _pending_prefill_sends.clear()
-
-
 def eval_with_timeout(
     mlx_item: Any,  # pyright: ignore[reportAny]
     timeout_seconds: float = 60.0,
@@ -153,11 +129,11 @@ class PipelineFirstLayer(CustomMlxLayer):
 
     def __call__(self, x: mx.array, *args: object, **kwargs: object) -> mx.array:
         if self.r != 0:
-            # We want to avoid GPU timeout errors by evalling the distributed operation
-            # so that it stays on CPU, which does not have a timeout.
-            mx.eval(x)
             x = mx.distributed.recv_like(x, (self.r - 1), group=self.group)
-            mx.eval(x)
+            if self.is_prefill:
+                # We want to avoid GPU timeout errors by evalling the distributed operation
+                # so that it stays on CPU, which does not have a timeout.
+                mx.eval(x)
         return self.original_layer(x, *args, **kwargs)
 
 
@@ -175,7 +151,6 @@ class PipelineLastLayer(CustomMlxLayer):
         self.group = group
         self.original_layer_signature = signature(self.original_layer.__call__)
         self.is_prefill: bool = False
-        self.queue_sends: bool = False
 
     def __call__(self, x: mx.array, *args: object, **kwargs: object) -> mx.array:
         cache = self.original_layer_signature.bind_partial(
@@ -184,34 +159,24 @@ class PipelineLastLayer(CustomMlxLayer):
 
         output: mx.array = self.original_layer(x, *args, **kwargs)
 
-        # Eval layer output to materialize it before send — this splits the graph
-        # so the send is isolated and the receiving rank's recv can complete.
-        mx.eval(output)
-
         if self.r != self.s - 1:
-            if self.queue_sends:
-                _pending_prefill_sends.append(
-                    (output, (self.r + 1) % self.s, self.group)
-                )
-            else:
-                output = mx.distributed.send(
-                    output, (self.r + 1) % self.s, group=self.group
-                )
+            output = mx.distributed.send(
+                output, (self.r + 1) % self.s, group=self.group
+            )
             if cache is not None:
                 # 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
-                if hasattr(_cache, "keys"):  # pyright: ignore[reportAny]
-                    _cache.keys = mx.depends(_cache.keys, output)  # type: ignore
-            mx.eval(output)
-            if cache is not None and hasattr(_cache, "keys"):  # type: ignore
-                mx.eval(_cache.keys)  # type: ignore
+                _cache.keys = mx.depends(_cache.keys, output)  # type: ignore
+            if self.is_prefill:
+                mx.eval(output)
+                if cache is not None:
+                    mx.eval(_cache.keys)  # type: ignore
 
         if not self.is_prefill:
             output = mx.distributed.all_gather(output, group=self.group)[
                 -output.shape[0] :
             ]
-            mx.eval(output)
 
         return output
 
@@ -222,12 +187,6 @@ def set_pipeline_prefill(model: nn.Module, is_prefill: bool) -> None:
             layer.is_prefill = is_prefill
 
 
-def set_pipeline_queue_sends(model: nn.Module, queue_sends: bool) -> None:
-    for layer in model.layers:  # type: ignore
-        if isinstance(layer, PipelineLastLayer):
-            layer.queue_sends = queue_sends
-
-
 def get_inner_model(model: nn.Module) -> nn.Module:
     inner = getattr(model, "model", None)
     if isinstance(inner, nn.Module):
@@ -259,32 +218,6 @@ 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,
@@ -355,24 +288,6 @@ 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), (
@@ -402,8 +317,7 @@ 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
-            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
+            dep_cache.keys = mx.depends(dep_cache.keys, logits)  # type: ignore
 
         return logits
 
@@ -526,9 +440,7 @@ def tensor_auto_parallel(
             all_to_sharded_linear_in_place,
             sharded_to_all_linear_in_place,
         )
-    elif isinstance(
-        model, (Qwen3MoeModel, Qwen3NextModel, Qwen3_5TextModel, Qwen3_5MoeModel)
-    ):
+    elif isinstance(model, (Qwen3MoeModel, Qwen3NextModel)):
         tensor_parallel_sharding_strategy = QwenShardingStrategy(
             group,
             all_to_sharded_linear,
@@ -923,14 +835,12 @@ class QwenShardingStrategy(TensorParallelShardingStrategy):
         on_timeout: TimeoutCallback | None,
         on_layer_loaded: LayerLoadedCallback | None,
     ) -> nn.Module:
-        model = cast(
-            Qwen3MoeModel | Qwen3NextModel | Qwen3_5TextModel | Qwen3_5MoeModel, model
-        )
+        model = cast(Qwen3MoeModel | Qwen3NextModel, model)
         total = len(model.layers)
         for i, layer in enumerate(model.layers):
             eval_with_timeout(layer.parameters(), timeout_seconds / total, on_timeout)
             # Shard the self attention
-            if isinstance(layer, Qwen3MoeDecoderLayer):
+            if isinstance(layer, Qwen3DecoderLayer):
                 layer.self_attn.q_proj = self.all_to_sharded_linear(
                     layer.self_attn.q_proj
                 )
@@ -943,42 +853,17 @@ class QwenShardingStrategy(TensorParallelShardingStrategy):
                 layer.self_attn.o_proj = self.sharded_to_all_linear(
                     layer.self_attn.o_proj
                 )
-                layer.self_attn.n_heads //= self.N
-                layer.self_attn.n_kv_heads //= self.N
             else:
-                assert isinstance(layer, (Qwen3NextDecoderLayer, Qwen3_5DecoderLayer))
+                assert isinstance(layer, Qwen3NextDecoderLayer)
                 if hasattr(layer, "linear_attn"):
                     linear_attn = layer.linear_attn
 
-                    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.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
+                    )
                     linear_attn.out_proj = self.sharded_to_all_linear(
                         linear_attn.out_proj
                     )
@@ -1040,20 +925,11 @@ class QwenShardingStrategy(TensorParallelShardingStrategy):
                     layer.self_attn.num_key_value_heads //= self.N
 
             # Shard the MoE.
-            if isinstance(
-                layer.mlp,
-                (
-                    Qwen3MoeSparseMoeBlock,
-                    Qwen3NextSparseMoeBlock,
-                    Qwen3_5SparseMoeBlock,
-                ),
-            ):
+            if isinstance(layer.mlp, (Qwen3MoeSparseMoeBlock, Qwen3NextSparseMoeBlock)):
                 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, Qwen3_5SparseMoeBlock)
-                ):
+                if isinstance(layer.mlp, Qwen3NextSparseMoeBlock):
                     self.all_to_sharded_linear_in_place(
                         layer.mlp.shared_expert.gate_proj
                     )

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

@@ -13,7 +13,8 @@ from mlx_lm.models.cache import (
 from mlx_lm.tokenizer_utils import TokenizerWrapper
 
 from exo.shared.types.memory import Memory
-from exo.shared.types.mlx import KVCacheType, Model
+from exo.shared.types.mlx import KVCacheType
+from exo.worker.engines.mlx import Model
 from exo.worker.engines.mlx.constants import CACHE_GROUP_SIZE, KV_CACHE_BITS
 from exo.worker.runner.bootstrap import logger
 
@@ -253,9 +254,9 @@ def trim_cache(
             if snapshot is not None and snapshot.states[i] is not None:
                 cache[i] = deepcopy(snapshot.states[i])  # type: ignore
             else:
-                c.state = [None] * len(c.state)
+                c.state = [None] * len(c.state)  # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType]
         else:
-            c.trim(num_tokens)
+            c.trim(num_tokens)  # pyright: ignore[reportUnknownMemberType]
 
 
 def encode_prompt(tokenizer: TokenizerWrapper, prompt: str) -> mx.array:

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

@@ -1,14 +1,9 @@
-import functools
-import math
 import time
 from copy import deepcopy
 from typing import Callable, Generator, cast, get_args
 
 import mlx.core as mx
-from mlx_lm.generate import (
-    maybe_quantize_kv_cache,
-    stream_generate,
-)
+from mlx_lm.generate import stream_generate
 from mlx_lm.models.cache import ArraysCache, RotatingKVCache
 from mlx_lm.sample_utils import make_sampler
 from mlx_lm.tokenizer_utils import TokenizerWrapper
@@ -23,19 +18,13 @@ from exo.shared.types.api import (
 )
 from exo.shared.types.common import ModelId
 from exo.shared.types.memory import Memory
-from exo.shared.types.mlx import KVCacheType, Model
+from exo.shared.types.mlx import KVCacheType
 from exo.shared.types.text_generation import InputMessage, TextGenerationTaskParams
 from exo.shared.types.worker.runner_response import (
     GenerationResponse,
 )
-from exo.worker.engines.mlx.auto_parallel import (
-    PipelineFirstLayer,
-    PipelineLastLayer,
-    clear_prefill_sends,
-    flush_prefill_sends,
-    set_pipeline_prefill,
-    set_pipeline_queue_sends,
-)
+from exo.worker.engines.mlx import Model
+from exo.worker.engines.mlx.auto_parallel import set_pipeline_prefill
 from exo.worker.engines.mlx.cache import (
     CacheSnapshot,
     KVPrefixCache,
@@ -66,130 +55,6 @@ class PrefillCancelled(BaseException):
     """Raised when prefill is cancelled via the progress callback."""
 
 
-def _has_pipeline_communication_layer(model: Model):
-    for layer in model.layers:
-        if isinstance(layer, (PipelineFirstLayer, PipelineLastLayer)):
-            return True
-    return False
-
-
-def pipeline_parallel_prefill(
-    model: Model,
-    prompt: mx.array,
-    prompt_cache: KVCacheType,
-    prefill_step_size: int,
-    kv_group_size: int | None,
-    kv_bits: int | None,
-    prompt_progress_callback: Callable[[int, int], None],
-    distributed_prompt_progress_callback: Callable[[], None] | None,
-    group: mx.distributed.Group,
-) -> None:
-    """Prefill the KV cache for pipeline parallel with overlapping stages.
-
-    Each rank processes the full prompt through its real cache, offset by leading
-    and trailing dummy iterations.
-
-    Total iterations per rank = N_real_chunks + world_size - 1:
-      - rank r leading dummies  (skip_pipeline_io, throwaway cache)
-      - N_real_chunks real      (pipeline IO active, real cache)
-      - (world_size-1-r) trailing dummies (skip_pipeline_io, throwaway cache)
-
-    e.g.
-    Timeline (2 ranks, 3 chunks of 10240 tokens @ step=4096):
-        iter 0: R0 real[0:4096]     R1 dummy
-        iter 1: R0 real[4096:8192]  R1 real[0:4096]
-        iter 2: R0 real[8192:10240] R1 real[4096:8192]
-        iter 3: R0 dummy            R1 real[8192:10240]
-
-    This function is designed to match mlx_lm's stream_generate exactly in terms of
-    side effects (given the same prefill step size)
-    """
-    prefill_step_size = prefill_step_size // min(4, group.size())
-
-    quantize_cache_fn: Callable[..., None] = functools.partial(
-        maybe_quantize_kv_cache,
-        quantized_kv_start=0,
-        kv_group_size=kv_group_size,
-        kv_bits=kv_bits,
-    )
-
-    _prompt_cache: KVCacheType = prompt_cache
-    rank = group.rank()
-    world_size = group.size()
-
-    # Build list of real prompt chunk sizes
-    total = len(prompt)
-    real_chunk_sizes: list[int] = []
-    remaining = total - 1
-    while remaining:
-        n = min(prefill_step_size, remaining)
-        real_chunk_sizes.append(n)
-        remaining -= n
-    n_real = len(real_chunk_sizes)
-
-    # Each rank does: [rank leading dummies] [N real chunks] [world_size-1-rank trailing dummies]
-    n_leading = rank
-    n_trailing = world_size - 1 - rank
-    n_total = n_leading + n_real + n_trailing
-
-    t_start = time.perf_counter()
-    processed = 0
-    logger.info(
-        f"[R{rank}] Pipeline prefill: {n_real} real + {n_leading} leading + {n_trailing} trailing = {n_total} iterations"
-    )
-    clear_prefill_sends()
-
-    # Initial callback matching generate_step
-    prompt_progress_callback(0, total)
-
-    try:
-        with mx.stream(generation_stream):
-            for _ in range(n_leading):
-                if distributed_prompt_progress_callback is not None:
-                    distributed_prompt_progress_callback()
-
-            for i in range(n_real):
-                chunk_size = real_chunk_sizes[i]
-                model(
-                    prompt[processed : processed + chunk_size][None],
-                    cache=_prompt_cache,
-                )
-                quantize_cache_fn(_prompt_cache)
-                processed += chunk_size
-
-                if distributed_prompt_progress_callback is not None:
-                    distributed_prompt_progress_callback()
-
-                flush_prefill_sends()
-
-                prompt_progress_callback(processed, total)
-
-            for _ in range(n_trailing):
-                if distributed_prompt_progress_callback is not None:
-                    distributed_prompt_progress_callback()
-
-    finally:
-        clear_prefill_sends()
-
-    # Post-loop: process remaining 1 token + add +1 entry to match stream_generate.
-    for _ in range(2):
-        with mx.stream(generation_stream):
-            model(prompt[-1:][None], cache=_prompt_cache)
-            quantize_cache_fn(_prompt_cache)
-        flush_prefill_sends()
-
-    assert _prompt_cache is not None
-    mx.eval([c.state for c in _prompt_cache])  # type: ignore
-
-    # Final callback matching generate_step
-    prompt_progress_callback(total, total)
-
-    logger.info(
-        f"[R{rank}] Prefill: {n_real} real + {n_leading}+{n_trailing} dummy iterations, "
-        f"Processed {processed} tokens in {(time.perf_counter() - t_start) * 1000:.1f}ms"
-    )
-
-
 def prefill(
     model: Model,
     tokenizer: TokenizerWrapper,
@@ -198,7 +63,6 @@ def prefill(
     cache: KVCacheType,
     group: mx.distributed.Group | None,
     on_prefill_progress: Callable[[int, int], None] | None,
-    distributed_prompt_progress_callback: Callable[[], None] | None,
 ) -> tuple[float, int, list[CacheSnapshot]]:
     """Prefill the KV cache with prompt tokens.
 
@@ -230,57 +94,31 @@ def prefill(
         if on_prefill_progress is not None:
             on_prefill_progress(processed, total)
 
-    def combined_progress_callback(processed: int, total: int) -> None:
-        if distributed_prompt_progress_callback is not None:
-            distributed_prompt_progress_callback()
-        progress_callback(processed, total)
-
     set_pipeline_prefill(model, is_prefill=True)
 
     mx_barrier(group)
     logger.info("Starting prefill")
 
-    is_pipeline = _has_pipeline_communication_layer(model)
-
-    prefill_step_size = 4096
-
+    # Use max_tokens=1 because max_tokens=0 does not work.
+    # We just throw away the generated token - we only care about filling the cache
     try:
-        if is_pipeline and num_tokens >= prefill_step_size:
-            set_pipeline_queue_sends(model, queue_sends=True)
-            assert group is not None, "Pipeline prefill requires a distributed group"
-            pipeline_parallel_prefill(
-                model=model,
-                prompt=prompt_tokens,
-                prompt_cache=cache,
-                prefill_step_size=prefill_step_size,
-                kv_group_size=KV_GROUP_SIZE,
-                kv_bits=KV_BITS,
-                prompt_progress_callback=progress_callback,
-                distributed_prompt_progress_callback=distributed_prompt_progress_callback,
-                group=group,
-            )
-        else:
-            # Use max_tokens=1 because max_tokens=0 does not work.
-            # We just throw away the generated token - we only care about filling the cache
-            for _ in stream_generate(
-                model=model,
-                tokenizer=tokenizer,
-                prompt=prompt_tokens,
-                max_tokens=1,
-                sampler=sampler,
-                prompt_cache=cache,
-                prefill_step_size=prefill_step_size,
-                kv_group_size=KV_GROUP_SIZE,
-                kv_bits=KV_BITS,
-                prompt_progress_callback=combined_progress_callback,
-            ):
-                break  # Stop after first iteration - cache is now filled
+        for _ in stream_generate(
+            model=model,
+            tokenizer=tokenizer,
+            prompt=prompt_tokens,
+            max_tokens=1,
+            sampler=sampler,
+            prompt_cache=cache,
+            prefill_step_size=4096,
+            kv_group_size=KV_GROUP_SIZE,
+            kv_bits=KV_BITS,
+            prompt_progress_callback=progress_callback,
+        ):
+            break  # Stop after first iteration - cache is now filled
     except PrefillCancelled:
-        set_pipeline_queue_sends(model, queue_sends=False)
         set_pipeline_prefill(model, is_prefill=False)
         raise
 
-    set_pipeline_queue_sends(model, queue_sends=False)
     set_pipeline_prefill(model, is_prefill=False)
 
     # stream_generate added 1 extra generated token to the cache, so we should trim it.
@@ -293,7 +131,7 @@ def prefill(
                 cache[i] = deepcopy(pre_gen.states[i])  # type: ignore
         else:
             assert not isinstance(c, (ArraysCache, RotatingKVCache))
-            c.trim(2)
+            c.trim(2)  # pyright: ignore[reportUnknownMemberType]
 
     elapsed = time.perf_counter() - start_time
     tokens_per_sec = num_tokens / elapsed if elapsed > 0 else 0.0
@@ -410,9 +248,6 @@ def extract_top_logprobs(
     for i in range(top_logprobs):
         token_id = int(top_indices[i].item())
         token_logprob = float(top_values[i].item())
-        if math.isnan(token_logprob):
-            continue
-
         # Decode token ID to string
         token_str = tokenizer.decode([token_id])
         # Get byte representation
@@ -436,8 +271,6 @@ def mlx_generate(
     kv_prefix_cache: KVPrefixCache | None,
     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()
@@ -499,7 +332,6 @@ def mlx_generate(
         caches,
         group,
         on_prefill_progress,
-        distributed_prompt_progress_callback,
     )
     cache_snapshots: list[CacheSnapshot] | None = ssm_snapshots_list or None
 
@@ -645,9 +477,6 @@ 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,