Claude-generated settings, no idea if it works

this pr broke macmon - revert it

.mlx_typings/mlx_lm/models/cache.pyi

@@ -249,7 +249,8 @@ class ChunkedKVCache(KVCache):
         ...
 
 class CacheList(_BaseCache):
-    def __init__(self, *caches) -> None: ...
+    caches: tuple[_BaseCache, ...]
+    def __init__(self, *caches: _BaseCache) -> None: ...
     def __getitem__(self, idx): ...
     def is_trimmable(self):  # -> bool:
         ...

app/EXO/EXO/ExoProcessController.swift

@@ -5,6 +5,7 @@ import Foundation
 private let customNamespaceKey = "EXOCustomNamespace"
 private let hfTokenKey = "EXOHFToken"
 private let enableImageModelsKey = "EXOEnableImageModels"
+private let offlineModeKey = "EXOOfflineMode"
 private let onboardingCompletedKey = "EXOOnboardingCompleted"
 
 @MainActor
@@ -60,6 +61,14 @@ final class ExoProcessController: ObservableObject {
             UserDefaults.standard.set(enableImageModels, forKey: enableImageModelsKey)
         }
     }
+    @Published var offlineMode: Bool = {
+        return UserDefaults.standard.bool(forKey: offlineModeKey)
+    }()
+    {
+        didSet {
+            UserDefaults.standard.set(offlineMode, forKey: offlineModeKey)
+        }
+    }
 
     /// Fires once when EXO transitions to `.running` for the very first time (fresh install).
     @Published private(set) var isFirstLaunchReady = false
@@ -267,6 +276,9 @@ final class ExoProcessController: ObservableObject {
         if enableImageModels {
             environment["EXO_ENABLE_IMAGE_MODELS"] = "true"
         }
+        if offlineMode {
+            environment["EXO_OFFLINE"] = "true"
+        }
 
         var paths: [String] = []
         if let existing = environment["PATH"], !existing.isEmpty {

app/EXO/EXO/Views/SettingsView.swift

@@ -13,6 +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 needsRestart = false
     @State private var bugReportInFlight = false
     @State private var bugReportMessage: String?
@@ -42,6 +43,7 @@ struct SettingsView: View {
             pendingNamespace = controller.customNamespace
             pendingHFToken = controller.hfToken
             pendingEnableImageModels = controller.enableImageModels
+            pendingOfflineMode = controller.offlineMode
             needsRestart = false
         }
     }
@@ -72,6 +74,13 @@ 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()
@@ -445,6 +454,7 @@ struct SettingsView: View {
 
     private var hasGeneralChanges: Bool {
         pendingNamespace != controller.customNamespace || pendingHFToken != controller.hfToken
+            || pendingOfflineMode != controller.offlineMode
     }
 
     private var hasModelChanges: Bool {
@@ -454,6 +464,7 @@ struct SettingsView: View {
     private func applyGeneralSettings() {
         controller.customNamespace = pendingNamespace
         controller.hfToken = pendingHFToken
+        controller.offlineMode = pendingOfflineMode
         restartIfRunning()
     }
 

dashboard/src/lib/components/HeaderNav.svelte

@@ -170,5 +170,30 @@
       {/if}
       Downloads
     </a>
+    <a
+      href="/#/settings"
+      class="text-sm text-white/70 hover:text-exo-yellow transition-colors tracking-wider uppercase flex items-center gap-2 cursor-pointer"
+      title="Settings"
+    >
+      <svg
+        class="w-4 h-4"
+        fill="none"
+        viewBox="0 0 24 24"
+        stroke="currentColor"
+        stroke-width="2"
+      >
+        <path
+          stroke-linecap="round"
+          stroke-linejoin="round"
+          d="M10.325 4.317c.426-1.756 2.924-1.756 3.35 0a1.724 1.724 0 002.573 1.066c1.543-.94 3.31.826 2.37 2.37a1.724 1.724 0 001.065 2.572c1.756.426 1.756 2.924 0 3.35a1.724 1.724 0 00-1.066 2.573c.94 1.543-.826 3.31-2.37 2.37a1.724 1.724 0 00-2.572 1.065c-.426 1.756-2.924 1.756-3.35 0a1.724 1.724 0 00-2.573-1.066c-1.543.94-3.31-.826-2.37-2.37a1.724 1.724 0 00-1.065-2.572c-1.756-.426-1.756-2.924 0-3.35a1.724 1.724 0 001.066-2.573c-.94-1.543.826-3.31 2.37-2.37.996.608 2.296.07 2.572-1.065z"
+        />
+        <path
+          stroke-linecap="round"
+          stroke-linejoin="round"
+          d="M15 12a3 3 0 11-6 0 3 3 0 016 0z"
+        />
+      </svg>
+      Settings
+    </a>
   </nav>
 </header>

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

@@ -0,0 +1,87 @@
+/**
+ * SettingsStore - Manages exo runtime settings via the /settings API.
+ */
+
+export interface MemorySettings {
+  oom_prevention: boolean;
+  memory_threshold: number;
+  memory_floor_gb: number;
+}
+
+export interface GenerationSettings {
+  prefill_step_size: number;
+  max_tokens: number;
+  kv_cache_bits: 4 | 8 | null;
+}
+
+export interface ExoSettings {
+  memory: MemorySettings;
+  generation: GenerationSettings;
+}
+
+function defaultSettings(): ExoSettings {
+  return {
+    memory: {
+      oom_prevention: false,
+      memory_threshold: 0.8,
+      memory_floor_gb: 5.0,
+    },
+    generation: {
+      prefill_step_size: 4096,
+      max_tokens: 32168,
+      kv_cache_bits: null,
+    },
+  };
+}
+
+class SettingsStore {
+  settings = $state<ExoSettings>(defaultSettings());
+  loading = $state(false);
+  error = $state<string | null>(null);
+
+  async load(): Promise<void> {
+    this.loading = true;
+    this.error = null;
+    try {
+      const response = await fetch("/settings");
+      if (!response.ok) {
+        throw new Error(`Failed to fetch settings: ${response.status}`);
+      }
+      this.settings = (await response.json()) as ExoSettings;
+    } catch (err) {
+      console.error("Failed to load settings:", err);
+      this.error = err instanceof Error ? err.message : "Unknown error";
+    } finally {
+      this.loading = false;
+    }
+  }
+
+  async save(updated: ExoSettings): Promise<boolean> {
+    this.loading = true;
+    this.error = null;
+    try {
+      const response = await fetch("/settings", {
+        method: "POST",
+        headers: { "Content-Type": "application/json" },
+        body: JSON.stringify(updated),
+      });
+      if (!response.ok) {
+        throw new Error(`Failed to save settings: ${response.status}`);
+      }
+      this.settings = (await response.json()) as ExoSettings;
+      return true;
+    } catch (err) {
+      console.error("Failed to save settings:", err);
+      this.error = err instanceof Error ? err.message : "Unknown error";
+      return false;
+    } finally {
+      this.loading = false;
+    }
+  }
+
+  resetToDefaults(): ExoSettings {
+    return defaultSettings();
+  }
+}
+
+export const settingsStore = new SettingsStore();

dashboard/src/routes/+page.svelte

@@ -42,6 +42,7 @@
     setSelectedChatModel,
     selectedChatModel,
     sendMessage,
+    messages,
     debugMode,
     toggleDebugMode,
     topologyOnlyMode,
@@ -889,6 +890,7 @@
       availableModels.some((m) => m.id === defaults.modelId)
     ) {
       selectPreviewModel(defaults.modelId);
+      setSelectedChatModel(defaults.modelId);
     }
   }
 
@@ -1324,6 +1326,7 @@
 
   function handleModelPickerSelect(modelId: string) {
     selectPreviewModel(modelId);
+    setSelectedChatModel(modelId);
     saveLaunchDefaults();
     isModelPickerOpen = false;
   }
@@ -2284,7 +2287,8 @@
     selectedChatCategory = null;
     pendingAutoMessage = null;
     userForcedIdle = true;
-    setSelectedChatModel("");
+    // Restore chat model from the sidebar preview selection so both selectors stay in sync
+    setSelectedChatModel(selectedModelId ?? "");
     clearChat();
   }
 
@@ -2514,6 +2518,11 @@
         return;
       }
     }
+
+    // Fallthrough: model exists but has no active instance/download/loading state
+    chatLaunchState = "idle";
+    pendingChatModelId = null;
+    selectedChatCategory = null;
   });
 
   // Suggested prompts per category
@@ -2637,7 +2646,11 @@
   }
 
   // Launch a model for seamless chat
-  async function launchModelForChat(modelId: string, category: string) {
+  async function launchModelForChat(
+    modelId: string,
+    category: string,
+    skipCreate = false,
+  ) {
     userForcedIdle = false;
     pendingChatModelId = modelId;
     selectedChatCategory = category;
@@ -2645,7 +2658,7 @@
     // Check if already running — skip straight to chat
     if (hasRunningInstance(modelId)) {
       setSelectedChatModel(modelId);
-      createConversation();
+      if (!skipCreate) createConversation();
       chatLaunchState = "ready";
       return;
     }
@@ -2654,7 +2667,7 @@
     if (hasExistingInstance(modelId)) {
       setSelectedChatModel(modelId);
       pendingChatModelId = modelId;
-      createConversation();
+      if (!skipCreate) createConversation();
       const dlStatus = getModelDownloadStatus(modelId);
       if (dlStatus.isDownloading) {
         chatLaunchState = "downloading";
@@ -2707,7 +2720,7 @@
 
       setSelectedChatModel(modelId);
       recordRecentLaunch(modelId);
-      createConversation();
+      if (!skipCreate) createConversation();
       chatLaunchState = "downloading";
     } catch (error) {
       addToast({ type: "error", message: `Network error: ${error}` });
@@ -2993,6 +3006,7 @@
   // Handle model selection from the picker when opened from chat context
   function handleChatPickerSelect(modelId: string) {
     setSelectedChatModel(modelId);
+    selectPreviewModel(modelId);
     userForcedIdle = false;
     isModelPickerOpen = false;
   }
@@ -3012,6 +3026,7 @@
 
     // Model is selected and running — send directly
     if (model && hasRunningInstance(model)) {
+      chatLaunchState = "ready";
       sendMessage(content, files, null);
       return;
     }
@@ -3020,7 +3035,7 @@
     if (model) {
       pendingAutoMessage = { content, files };
       userForcedIdle = false;
-      launchModelForChat(model, "picker");
+      launchModelForChat(model, "picker", messages().length > 0);
       return;
     }
 
@@ -5804,43 +5819,7 @@
           class="flex-1 flex flex-col min-w-0 overflow-hidden"
           in:fade={{ duration: 300, delay: 100 }}
         >
-          {#if chatLaunchState === "idle"}
-            <!-- No running instance: show model selector -->
-            <div
-              class="flex-1 overflow-y-auto flex items-center justify-center px-8 py-6"
-            >
-              <ChatModelSelector
-                models={models.map((m) => ({
-                  id: m.id,
-                  name: m.name ?? "",
-                  base_model: m.base_model ?? "",
-                  storage_size_megabytes: m.storage_size_megabytes ?? 0,
-                  capabilities: m.capabilities ?? [],
-                  family: m.family ?? "",
-                  quantization: m.quantization ?? "",
-                }))}
-                clusterLabel={chatClusterLabel}
-                totalMemoryGB={availableMemoryGB()}
-                onSelect={handleChatModelSelect}
-                onAddModel={handleChatAddModel}
-              />
-            </div>
-            <div
-              class="flex-shrink-0 px-8 pb-6 pt-4 bg-gradient-to-t from-exo-black via-exo-black to-transparent"
-            >
-              <div class="max-w-7xl mx-auto">
-                <ChatForm
-                  placeholder="Ask anything — we'll pick the best model automatically"
-                  showModelSelector={!!bestRunningModelId}
-                  modelDisplayOverride={bestRunningModelId ?? undefined}
-                  modelTasks={modelTasks()}
-                  modelCapabilities={modelCapabilities()}
-                  onAutoSend={handleAutoSend}
-                  onOpenModelPicker={openChatModelPicker}
-                />
-              </div>
-            </div>
-          {:else if chatLaunchState !== "idle" && chatLaunchState !== "ready"}
+          {#if chatLaunchState !== "idle" && chatLaunchState !== "ready"}
             <!-- Model launching/downloading/loading: show progress -->
             <div class="flex-1 flex items-center justify-center px-8 py-6">
               <div class="flex flex-col items-center gap-6 max-w-md w-full">
@@ -5947,8 +5926,8 @@
                 />
               </div>
             </div>
-          {:else}
-            <!-- Normal chat: model is running -->
+          {:else if messages().length > 0 || chatLaunchState === "ready"}
+            <!-- Normal chat: show messages -->
             <div
               class="flex-1 overflow-y-auto px-8 py-6"
               bind:this={chatScrollRef}
@@ -6004,6 +5983,42 @@
                 />
               </div>
             </div>
+          {:else}
+            <!-- No running instance, no messages: show model selector -->
+            <div
+              class="flex-1 overflow-y-auto flex items-center justify-center px-8 py-6"
+            >
+              <ChatModelSelector
+                models={models.map((m) => ({
+                  id: m.id,
+                  name: m.name ?? "",
+                  base_model: m.base_model ?? "",
+                  storage_size_megabytes: m.storage_size_megabytes ?? 0,
+                  capabilities: m.capabilities ?? [],
+                  family: m.family ?? "",
+                  quantization: m.quantization ?? "",
+                }))}
+                clusterLabel={chatClusterLabel}
+                totalMemoryGB={availableMemoryGB()}
+                onSelect={handleChatModelSelect}
+                onAddModel={handleChatAddModel}
+              />
+            </div>
+            <div
+              class="flex-shrink-0 px-8 pb-6 pt-4 bg-gradient-to-t from-exo-black via-exo-black to-transparent"
+            >
+              <div class="max-w-7xl mx-auto">
+                <ChatForm
+                  placeholder="Ask anything — we'll pick the best model automatically"
+                  showModelSelector={!!bestRunningModelId}
+                  modelDisplayOverride={bestRunningModelId ?? undefined}
+                  modelTasks={modelTasks()}
+                  modelCapabilities={modelCapabilities()}
+                  onAutoSend={handleAutoSend}
+                  onOpenModelPicker={openChatModelPicker}
+                />
+              </div>
+            </div>
           {/if}
         </div>
 

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

@@ -0,0 +1,193 @@
+<script lang="ts">
+  import { onMount } from "svelte";
+  import { fade } from "svelte/transition";
+  import HeaderNav from "$lib/components/HeaderNav.svelte";
+  import { settingsStore, type ExoSettings } from "$lib/stores/settings.svelte";
+  import { addToast } from "$lib/stores/toast.svelte";
+
+  let draft = $state<ExoSettings | null>(null);
+  const loading = $derived(settingsStore.loading);
+
+  onMount(async () => {
+    await settingsStore.load();
+    draft = structuredClone(settingsStore.settings);
+  });
+
+  async function handleSave() {
+    if (!draft) return;
+    const ok = await settingsStore.save(draft);
+    if (ok) {
+      addToast({ type: "success", message: "Settings saved" });
+    } else {
+      addToast({ type: "error", message: settingsStore.error ?? "Failed to save settings" });
+    }
+  }
+
+  function handleReset() {
+    draft = settingsStore.resetToDefaults();
+  }
+
+  const KV_OPTIONS: { label: string; value: 4 | 8 | null }[] = [
+    { label: "None (full precision)", value: null },
+    { label: "4-bit", value: 4 },
+    { label: "8-bit", value: 8 },
+  ];
+</script>
+
+<HeaderNav showHome={true} />
+
+{#if draft}
+  <div class="min-h-screen bg-background text-foreground" in:fade={{ duration: 200 }}>
+    <div class="max-w-2xl mx-auto px-6 py-8">
+      <h1 class="text-2xl font-bold text-exo-yellow tracking-wider uppercase mb-8">Settings</h1>
+
+      <!-- Memory / Safety -->
+      <section class="mb-10">
+        <h2 class="text-sm font-semibold text-white/50 tracking-widest uppercase mb-4">Memory / Safety</h2>
+        <div class="space-y-5">
+          <!-- OOM Prevention Toggle -->
+          <div class="flex items-center justify-between">
+            <div>
+              <div class="text-sm text-white/90">OOM Prevention</div>
+              <div class="text-xs text-white/40 mt-0.5">Stop generation when memory is low</div>
+            </div>
+            <button
+              onclick={() => { if (draft) draft.memory.oom_prevention = !draft.memory.oom_prevention; }}
+              class="relative w-11 h-6 rounded-full transition-colors duration-200 cursor-pointer {draft.memory.oom_prevention ? 'bg-exo-yellow' : 'bg-exo-medium-gray'}"
+              role="switch"
+              aria-checked={draft.memory.oom_prevention}
+            >
+              <span
+                class="absolute top-0.5 left-0.5 w-5 h-5 rounded-full bg-white shadow transition-transform duration-200 {draft.memory.oom_prevention ? 'translate-x-5' : 'translate-x-0'}"
+              ></span>
+            </button>
+          </div>
+
+          <!-- Memory Threshold Slider -->
+          <div>
+            <div class="flex items-center justify-between mb-1.5">
+              <div>
+                <div class="text-sm text-white/90">Memory Threshold</div>
+                <div class="text-xs text-white/40 mt-0.5">KV cache eviction triggers above this level</div>
+              </div>
+              <span class="text-sm font-mono text-exo-yellow">{(draft.memory.memory_threshold * 100).toFixed(0)}%</span>
+            </div>
+            <input
+              type="range"
+              min="0.5"
+              max="0.99"
+              step="0.01"
+              bind:value={draft.memory.memory_threshold}
+              class="w-full h-1.5 rounded-full appearance-none cursor-pointer bg-exo-medium-gray accent-exo-yellow"
+            />
+          </div>
+
+          <!-- Memory Floor -->
+          <div>
+            <div class="flex items-center justify-between mb-1.5">
+              <div>
+                <div class="text-sm text-white/90">Memory Floor</div>
+                <div class="text-xs text-white/40 mt-0.5">Minimum free memory to reserve (GB)</div>
+              </div>
+              <span class="text-sm font-mono text-exo-yellow">{draft.memory.memory_floor_gb.toFixed(1)} GB</span>
+            </div>
+            <input
+              type="number"
+              min="0"
+              max="64"
+              step="0.5"
+              bind:value={draft.memory.memory_floor_gb}
+              class="w-full bg-exo-medium-gray border border-exo-light-gray/20 rounded px-3 py-1.5 text-sm text-white/90 font-mono focus:outline-none focus:border-exo-yellow/50"
+            />
+          </div>
+        </div>
+      </section>
+
+      <!-- Generation / Performance -->
+      <section class="mb-10">
+        <h2 class="text-sm font-semibold text-white/50 tracking-widest uppercase mb-4">Generation / Performance</h2>
+        <div class="space-y-5">
+          <!-- Prefill Step Size -->
+          <div>
+            <div class="flex items-center justify-between mb-1.5">
+              <div>
+                <div class="text-sm text-white/90">Prefill Step Size</div>
+                <div class="text-xs text-white/40 mt-0.5">Token chunk size during prompt processing</div>
+              </div>
+              <span class="text-sm font-mono text-exo-yellow">{draft.generation.prefill_step_size.toLocaleString()}</span>
+            </div>
+            <input
+              type="number"
+              min="128"
+              max="32768"
+              step="128"
+              bind:value={draft.generation.prefill_step_size}
+              class="w-full bg-exo-medium-gray border border-exo-light-gray/20 rounded px-3 py-1.5 text-sm text-white/90 font-mono focus:outline-none focus:border-exo-yellow/50"
+            />
+          </div>
+
+          <!-- Max Tokens -->
+          <div>
+            <div class="flex items-center justify-between mb-1.5">
+              <div>
+                <div class="text-sm text-white/90">Max Tokens</div>
+                <div class="text-xs text-white/40 mt-0.5">Maximum generation length per response</div>
+              </div>
+              <span class="text-sm font-mono text-exo-yellow">{draft.generation.max_tokens.toLocaleString()}</span>
+            </div>
+            <input
+              type="number"
+              min="1"
+              max="131072"
+              step="1024"
+              bind:value={draft.generation.max_tokens}
+              class="w-full bg-exo-medium-gray border border-exo-light-gray/20 rounded px-3 py-1.5 text-sm text-white/90 font-mono focus:outline-none focus:border-exo-yellow/50"
+            />
+          </div>
+
+          <!-- KV Cache Bits -->
+          <div>
+            <div class="mb-1.5">
+              <div class="text-sm text-white/90">KV Cache Quantization</div>
+              <div class="text-xs text-white/40 mt-0.5">Lower bits save memory at slight quality cost</div>
+            </div>
+            <select
+              bind:value={draft.generation.kv_cache_bits}
+              class="w-full bg-exo-medium-gray border border-exo-light-gray/20 rounded px-3 py-1.5 text-sm text-white/90 font-mono focus:outline-none focus:border-exo-yellow/50 cursor-pointer"
+            >
+              {#each KV_OPTIONS as opt}
+                <option value={opt.value}>{opt.label}</option>
+              {/each}
+            </select>
+          </div>
+        </div>
+      </section>
+
+      <!-- Action Buttons -->
+      <div class="flex items-center gap-3">
+        <button
+          onclick={handleSave}
+          disabled={loading}
+          class="px-5 py-2 rounded text-sm font-semibold tracking-wider uppercase transition-colors cursor-pointer
+            bg-exo-yellow text-exo-black hover:bg-exo-yellow-darker
+            disabled:opacity-50 disabled:cursor-not-allowed"
+        >
+          {loading ? "Saving..." : "Save"}
+        </button>
+        <button
+          onclick={handleReset}
+          disabled={loading}
+          class="px-5 py-2 rounded text-sm font-semibold tracking-wider uppercase transition-colors cursor-pointer
+            border border-exo-light-gray/30 text-white/70 hover:border-exo-yellow/50 hover:text-exo-yellow
+            disabled:opacity-50 disabled:cursor-not-allowed"
+        >
+          Reset to Defaults
+        </button>
+      </div>
+    </div>
+  </div>
+{:else}
+  <div class="min-h-screen bg-background flex items-center justify-center">
+    <div class="text-white/40 text-sm">Loading settings...</div>
+  </div>
+{/if}

nix/mlx.nix

@@ -41,7 +41,7 @@ let
 
   mlx = stdenv.mkDerivation rec {
     pname = "mlx";
-    version = let v = "0.30.7.dev20260220+13998a05"; in
+    version = let v = "0.30.7.dev20260225+257d5692"; 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 = "13998a054715edcdc93618fb1496c79c7c25ff7c";
-      hash = "sha256-fAqA3hFwNBx7FcoGnhQsIFpAIRbC2EerACm4Fvne0Cc=";
+      rev = "257d5692fc7af6bba3b8afaeb63c549b7d1e43d5";
+      hash = "sha256-GosFIWxIB48Egb1MqJrR3xhsUsQeWdRk5rV93USY6wQ=";
     };
 
     patches = [

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "exo"
-version = "0.3.0"
+version = "0.3.68"
 description = "Exo"
 readme = "README.md"
 requires-python = ">=3.13"

src/exo/download/coordinator.py

@@ -1,5 +1,4 @@
 import asyncio
-import socket
 from dataclasses import dataclass, field
 from random import random
 
@@ -73,8 +72,6 @@ class DownloadCoordinator:
 
     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:
@@ -123,8 +120,6 @@ 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)
@@ -132,40 +127,10 @@ class DownloadCoordinator:
                 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()
 

src/exo/download/impl_shard_downloader.py

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

src/exo/download/shard_downloader.py

@@ -16,11 +16,6 @@ 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/main.py

@@ -60,7 +60,7 @@ class Node:
             download_coordinator = DownloadCoordinator(
                 node_id,
                 session_id,
-                exo_shard_downloader(),
+                exo_shard_downloader(offline=args.offline),
                 download_command_receiver=router.receiver(topics.DOWNLOAD_COMMANDS),
                 local_event_sender=router.sender(topics.LOCAL_EVENTS),
                 offline=args.offline,
@@ -211,7 +211,7 @@ class Node:
                         self.download_coordinator = DownloadCoordinator(
                             self.node_id,
                             result.session_id,
-                            exo_shard_downloader(),
+                            exo_shard_downloader(offline=self.offline),
                             download_command_receiver=self.router.receiver(
                                 topics.DOWNLOAD_COMMANDS
                             ),
@@ -261,13 +261,6 @@ def main():
     if args.offline:
         logger.info("Running in OFFLINE mode — no internet checks, local models only")
 
-    # Set trust_remote_code override env var for runner subprocesses
-    if args.trust_remote_code:
-        os.environ["EXO_TRUST_REMOTE_CODE"] = "1"
-        logger.warning(
-            "--trust-remote-code enabled: models may execute arbitrary code during loading"
-        )
-
     # Set FAST_SYNCH override env var for runner subprocesses
     if args.fast_synch is True:
         os.environ["EXO_FAST_SYNCH"] = "on"
@@ -277,9 +270,16 @@ def main():
         logger.info("FAST_SYNCH forced OFF")
 
     node = anyio.run(Node.create, args)
-    anyio.run(node.run)
-    logger.info("EXO Shutdown complete")
-    logger_cleanup()
+    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()
 
 
 class Args(CamelCaseModel):
@@ -290,9 +290,8 @@ class Args(CamelCaseModel):
     tb_only: bool = False
     no_worker: bool = False
     no_downloads: bool = False
-    offline: bool = False
+    offline: bool = os.getenv("EXO_OFFLINE", "false").lower() == "true"
     fast_synch: bool | None = None  # None = auto, True = force on, False = force off
-    trust_remote_code: bool = False
 
     @classmethod
     def parse(cls) -> Self:
@@ -342,13 +341,9 @@ 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",
         )
-        parser.add_argument(
-            "--trust-remote-code",
-            action="store_true",
-            help="Allow models to execute custom code during tokenizer loading (security-sensitive, CLI-only)",
-        )
         fast_synch_group = parser.add_mutually_exclusive_group()
         fast_synch_group.add_argument(
             "--fast-synch",

src/exo/master/api.py

@@ -166,6 +166,13 @@ from exo.shared.types.openai_responses import (
     ResponsesRequest,
     ResponsesResponse,
 )
+from exo.shared.types.settings import (
+    ExoSettings,
+    load_settings,
+)
+from exo.shared.types.settings import (
+    save_settings as save_settings_to_file,
+)
 from exo.shared.types.state import State
 from exo.shared.types.worker.downloads import DownloadCompleted
 from exo.shared.types.worker.instances import Instance, InstanceId, InstanceMeta
@@ -349,6 +356,8 @@ class API:
         self.app.get("/v1/traces/{task_id}/raw")(self.get_trace_raw)
         self.app.get("/onboarding")(self.get_onboarding)
         self.app.post("/onboarding")(self.complete_onboarding)
+        self.app.get("/settings")(self.get_settings)
+        self.app.post("/settings")(self.save_settings)
 
     async def place_instance(self, payload: PlaceInstanceParams):
         command = PlaceInstance(
@@ -1825,3 +1834,13 @@ class API:
         ONBOARDING_COMPLETE_FILE.parent.mkdir(parents=True, exist_ok=True)
         ONBOARDING_COMPLETE_FILE.write_text("true")
         return JSONResponse({"completed": True})
+
+    async def get_settings(self) -> JSONResponse:
+        settings = load_settings()
+        return JSONResponse(settings.model_dump())
+
+    async def save_settings(self, request: Request) -> JSONResponse:
+        body = cast(object, await request.json())
+        settings = ExoSettings.model_validate(body)
+        save_settings_to_file(settings)
+        return JSONResponse(settings.model_dump())

src/exo/routing/router.py

@@ -174,28 +174,42 @@ class Router:
         logger.info(f"Unsubscribed from {topic}")
 
     async def _networking_recv(self):
-        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
+        try:
+            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)
+                router = self.topic_routers[topic]
+                await router.publish_bytes(data)
+        except Exception as exception:
+            logger.opt(exception=exception).error(
+                "Gossipsub receive loop terminated unexpectedly"
+            )
+            raise
 
     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}"
+        try:
+            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}"
+                )
+                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)
+        except Exception as exception:
+            logger.opt(exception=exception).error(
+                "Connection update receive loop terminated unexpectedly"
             )
-            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)
+            raise
 
     async def _networking_publish(self):
         with self.networking_receiver as networked_items:

src/exo/shared/constants.py

@@ -78,4 +78,6 @@ 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/types/memory.py

@@ -1,6 +1,11 @@
+import ctypes
+import sys
 from math import ceil
 from typing import Self, overload
 
+import psutil
+
+from exo.shared.logging import logger
 from exo.utils.pydantic_ext import FrozenModel
 
 
@@ -149,3 +154,67 @@ class Memory(FrozenModel):
             unit = "B"
 
         return f"{val:.2f} {unit}".rstrip("0").rstrip(".") + f" {unit}"
+
+
+def _load_memory_settings() -> tuple[float, "Memory"]:
+    """Load memory threshold and floor from settings (lazy import to avoid circular dep)."""
+    from exo.shared.types.settings import load_settings
+
+    s = load_settings()
+    return s.memory.memory_threshold, Memory.from_gb(s.memory.memory_floor_gb)
+
+_libc: ctypes.CDLL | None = None
+
+
+def _macos_memorystatus_level() -> int:
+    global _libc  # noqa: PLW0603
+    if _libc is None:
+        _libc = ctypes.CDLL("/usr/lib/libSystem.B.dylib")
+    level = ctypes.c_int(0)
+    size = ctypes.c_size_t(ctypes.sizeof(ctypes.c_int))
+    ret: int = _libc.sysctlbyname(  # pyright: ignore[reportAny]
+        b"kern.memorystatus_level",
+        ctypes.byref(level),
+        ctypes.byref(size),
+        None,
+        ctypes.c_size_t(0),
+    )
+    if ret != 0:
+        raise OSError("sysctlbyname kern.memorystatus_level failed")
+    return level.value
+
+
+def _get_macos_memory_pressure() -> float:
+    try:
+        return 1.0 - _macos_memorystatus_level() / 100.0
+    except (OSError, FileNotFoundError):
+        logger.warning("Using fallback memory pressure")
+        return _fallback_memory_pressure()
+
+
+def _fallback_memory_pressure() -> float:
+    vm = psutil.virtual_memory()
+    return 1.0 - vm.available / vm.total
+
+
+def get_memory_pressure() -> float:
+    if sys.platform == "darwin":
+        return _get_macos_memory_pressure()
+    return _fallback_memory_pressure()
+
+
+def get_memory_limit() -> Memory:
+    threshold, floor = _load_memory_settings()
+    total = psutil.virtual_memory().total
+    safety = min(int(total * (1 - threshold)), floor.in_bytes)
+    return Memory.from_bytes(total - safety)
+
+
+def get_memory_available_locally() -> Memory:
+    total = Memory.from_bytes(psutil.virtual_memory().total)
+    return get_memory_limit() - total * get_memory_pressure()
+
+
+def get_memory_pressure_threshold() -> float:
+    total = psutil.virtual_memory().total
+    return get_memory_limit().in_bytes / total

src/exo/shared/types/settings.py

@@ -0,0 +1,121 @@
+import os
+import tomllib
+from typing import Literal
+
+import psutil
+from pydantic import ConfigDict, Field, ValidationError
+
+from exo.shared.constants import EXO_CONFIG_FILE
+from exo.shared.logging import logger
+from exo.shared.types.memory import Memory
+from exo.utils.pydantic_ext import CamelCaseModel
+
+
+def _default_memory_threshold() -> float:
+    total_gb = Memory.from_bytes(psutil.virtual_memory().total).in_gb
+    if total_gb >= 128:
+        return 0.85
+    if total_gb >= 64:
+        return 0.80
+    if total_gb >= 32:
+        return 0.75
+    return 0.70
+
+
+class MemorySettings(CamelCaseModel):
+    model_config = ConfigDict(
+        alias_generator=None,
+        validate_by_name=True,
+        extra="forbid",
+        strict=False,
+    )
+
+    oom_prevention: bool = False
+    memory_threshold: float = Field(default_factory=_default_memory_threshold, ge=0.0, le=1.0)
+    memory_floor_gb: float = Field(default=5.0, ge=0.0)
+
+
+class GenerationSettings(CamelCaseModel):
+    model_config = ConfigDict(
+        alias_generator=None,
+        validate_by_name=True,
+        extra="forbid",
+        strict=False,
+    )
+
+    prefill_step_size: int = Field(default=4096, ge=1)
+    max_tokens: int = Field(default=32168, ge=1)
+    kv_cache_bits: Literal[4, 8] | None = None
+
+
+class ExoSettings(CamelCaseModel):
+    model_config = ConfigDict(
+        alias_generator=None,
+        validate_by_name=True,
+        extra="ignore",
+        strict=False,
+    )
+
+    memory: MemorySettings = Field(default_factory=MemorySettings)
+    generation: GenerationSettings = Field(default_factory=GenerationSettings)
+
+
+_cached_settings: ExoSettings | None = None
+_cached_mtime: float = 0.0
+
+
+def load_settings() -> ExoSettings:
+    global _cached_settings, _cached_mtime  # noqa: PLW0603
+
+    try:
+        mtime = EXO_CONFIG_FILE.stat().st_mtime
+        if _cached_settings is not None and mtime == _cached_mtime:
+            return _cached_settings
+        with open(EXO_CONFIG_FILE, "rb") as f:
+            data = tomllib.load(f)
+        settings = ExoSettings.model_validate(data)
+        _cached_mtime = mtime
+    except FileNotFoundError:
+        settings = ExoSettings()
+    except (tomllib.TOMLDecodeError, ValidationError) as e:
+        logger.warning(f"Invalid config file {EXO_CONFIG_FILE}: {e}")
+        settings = ExoSettings()
+
+    # Env vars override config file for backward compat.
+    env_threshold = os.environ.get("EXO_MEMORY_THRESHOLD")
+    if env_threshold is not None:
+        settings = settings.model_copy(
+            update={"memory": settings.memory.model_copy(update={"memory_threshold": float(env_threshold)})}
+        )
+    env_floor = os.environ.get("EXO_MEMORY_FLOOR")
+    if env_floor is not None:
+        settings = settings.model_copy(
+            update={"memory": settings.memory.model_copy(update={"memory_floor_gb": float(env_floor)})}
+        )
+
+    _cached_settings = settings
+    return settings
+
+
+def save_settings(settings: ExoSettings) -> None:
+    global _cached_settings, _cached_mtime  # noqa: PLW0603
+
+    EXO_CONFIG_FILE.parent.mkdir(parents=True, exist_ok=True)
+
+    lines = [
+        "[memory]",
+        f"oom_prevention = {'true' if settings.memory.oom_prevention else 'false'}",
+        f"memory_threshold = {settings.memory.memory_threshold}",
+        f"memory_floor_gb = {settings.memory.memory_floor_gb}",
+        "",
+        "[generation]",
+        f"prefill_step_size = {settings.generation.prefill_step_size}",
+        f"max_tokens = {settings.generation.max_tokens}",
+    ]
+    if settings.generation.kv_cache_bits is not None:
+        lines.append(f"kv_cache_bits = {settings.generation.kv_cache_bits}")
+
+    EXO_CONFIG_FILE.write_text("\n".join(lines) + "\n")
+
+    _cached_settings = settings
+    _cached_mtime = EXO_CONFIG_FILE.stat().st_mtime

src/exo/utils/info_gatherer/info_gatherer.py

@@ -12,7 +12,7 @@ from anyio import fail_after, open_process, to_thread
 from anyio.streams.buffered import BufferedByteReceiveStream
 from anyio.streams.text import TextReceiveStream
 from loguru import logger
-from pydantic import ValidationError
+from pydantic import ConfigDict, ValidationError
 
 from exo.shared.constants import EXO_CONFIG_FILE, EXO_MODELS_DIR
 from exo.shared.types.memory import Memory
@@ -295,6 +295,8 @@ class ThunderboltBridgeInfo(TaggedModel):
 class NodeConfig(TaggedModel):
     """Node configuration from EXO_CONFIG_FILE, reloaded from the file only at startup. Other changes should come in through the API and propagate from there"""
 
+    model_config = ConfigDict(extra="ignore")
+
     @classmethod
     async def gather(cls) -> Self | None:
         cfg_file = anyio.Path(EXO_CONFIG_FILE)

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

@@ -128,11 +128,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)
-            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)
+            mx.eval(x)
         return self.original_layer(x, *args, **kwargs)
 
 
@@ -158,6 +158,10 @@ 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:
             output = mx.distributed.send(
                 output, (self.r + 1) % self.s, group=self.group
@@ -167,15 +171,15 @@ class PipelineLastLayer(CustomMlxLayer):
                 # doesn't have .keys directly; access via first sub-cache.
                 _cache = cache[0] if hasattr(cache, "caches") else cache  # type: ignore
                 _cache.keys = mx.depends(_cache.keys, output)  # type: ignore
-            if self.is_prefill:
-                mx.eval(output)
-                if cache is not None:
-                    mx.eval(_cache.keys)  # type: ignore
+            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
 

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

@@ -1,8 +1,6 @@
-import os
 from copy import deepcopy
 
 import mlx.core as mx
-import psutil
 from mlx_lm.models.cache import (
     ArraysCache,
     CacheList,
@@ -12,31 +10,14 @@ from mlx_lm.models.cache import (
 )
 from mlx_lm.tokenizer_utils import TokenizerWrapper
 
-from exo.shared.types.memory import Memory
+from exo.shared.types.memory import Memory, get_memory_pressure
 from exo.shared.types.mlx import KVCacheType
+from exo.shared.types.settings import load_settings
 from exo.worker.engines.mlx import Model
-from exo.worker.engines.mlx.constants import CACHE_GROUP_SIZE, KV_CACHE_BITS
+from exo.worker.engines.mlx.constants import CACHE_GROUP_SIZE
 from exo.worker.runner.bootstrap import logger
 
 
-# Fraction of device memory above which LRU eviction kicks in.
-# Smaller machines need more aggressive eviction.
-def _default_memory_threshold() -> float:
-    total_gb = Memory.from_bytes(psutil.virtual_memory().total).in_gb
-    if total_gb >= 128:
-        return 0.85
-    if total_gb >= 64:
-        return 0.80
-    if total_gb >= 32:
-        return 0.75
-    return 0.70
-
-
-_MEMORY_THRESHOLD = float(
-    os.environ.get("EXO_MEMORY_THRESHOLD", _default_memory_threshold())
-)
-
-
 class CacheSnapshot:
     """Snapshot of states at a known token position."""
 
@@ -92,6 +73,15 @@ class KVPrefixCache:
         self._snapshots.clear()
         self._last_used.clear()
 
+    def force_evict_all(self) -> int:
+        count = len(self.caches)
+        self.clear()
+        if count > 0:
+            logger.info(
+                f"Force-evicted all {count} prefix cache entries due to memory pressure"
+            )
+        return count
+
     def add_kv_cache(
         self,
         prompt_tokens: mx.array,
@@ -217,7 +207,7 @@ class KVPrefixCache:
         # Evict LRU entries until below threshold
         while (
             len(self.caches) > 0
-            and self.get_memory_used_percentage() > _MEMORY_THRESHOLD
+            and self.get_memory_used_percentage() > load_settings().memory.memory_threshold
         ):
             lru_index = self._last_used.index(min(self._last_used))
             evicted_tokens = len(self.prompts[lru_index])
@@ -230,7 +220,7 @@ class KVPrefixCache:
             )
 
     def get_memory_used_percentage(self) -> float:
-        local_pressure: float = get_memory_used_percentage()
+        local_pressure: float = get_memory_pressure()
 
         if self._group is None:
             return local_pressure
@@ -299,15 +289,47 @@ def get_prefix_length(prompt: mx.array, cached_prompt: mx.array) -> int:
     return int(mx.sum(prefix_mask).item())
 
 
-def get_available_memory() -> Memory:
-    mem: int = psutil.virtual_memory().available
-    return Memory.from_bytes(mem)
+def _measure_single_cache_bytes(
+    entry: KVCache | RotatingKVCache | QuantizedKVCache | ArraysCache | CacheList,
+) -> int:
+    if isinstance(entry, CacheList):
+        return sum(
+            _measure_single_cache_bytes(c)  # pyright: ignore[reportArgumentType]
+            for c in entry.caches
+        )
+
+    total = 0
+    if isinstance(entry, ArraysCache):
+        state = entry.state  # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType]
+        for arr in state:  # pyright: ignore[reportUnknownVariableType]
+            if isinstance(arr, mx.array):
+                total += arr.nbytes
+        return total
+
+    total = 0
+    for attr_name in ("keys", "values"):
+        val: object = getattr(entry, attr_name, None)
+        if val is None:
+            continue
+        if isinstance(val, mx.array):
+            total += val.nbytes
+        elif isinstance(val, (tuple, list)):
+            for arr in val:  # pyright: ignore[reportUnknownVariableType]
+                if isinstance(arr, mx.array):
+                    total += arr.nbytes
+
+    return total
 
 
-def get_memory_used_percentage() -> float:
-    mem = psutil.virtual_memory()
-    # percent is 0-100
-    return float(mem.percent / 100)
+def measure_cache_bytes(cache: KVCacheType) -> int:
+    return sum(_measure_single_cache_bytes(c) for c in cache)
+
+
+def measure_kv_cache_bytes_per_token(cache: KVCacheType) -> Memory:
+    offset = cache_length(cache)
+    if offset == 0:
+        return Memory.from_bytes(0)
+    return Memory.from_bytes(measure_cache_bytes(cache) // offset)
 
 
 def make_kv_cache(
@@ -320,13 +342,14 @@ def make_kv_cache(
         return model.make_cache()  # type: ignore
 
     if max_kv_size is None:
-        if KV_CACHE_BITS is None:
+        kv_cache_bits = load_settings().generation.kv_cache_bits
+        if kv_cache_bits is None:
             logger.info("Using default KV cache")
             return [KVCache() for _ in model.layers]
         else:
             logger.info("Using quantized KV cache")
             return [
-                QuantizedKVCache(group_size=CACHE_GROUP_SIZE, bits=KV_CACHE_BITS)
+                QuantizedKVCache(group_size=CACHE_GROUP_SIZE, bits=kv_cache_bits)
                 for _ in model.layers
             ]
     else:

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

@@ -13,6 +13,5 @@ KV_CACHE_BITS: int | None = None
 
 DEFAULT_TOP_LOGPROBS: int = 5
 
-# True for built-in models with known model cards; custom models added via API default to False
-# and can be overridden with the --trust-remote-code CLI flag.
+# TODO: We should really make this opt-in, but Kimi requires trust_remote_code=True
 TRUST_REMOTE_CODE: bool = True

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

@@ -1,3 +1,4 @@
+import math
 import time
 from copy import deepcopy
 from typing import Callable, Generator, cast, get_args
@@ -17,8 +18,9 @@ from exo.shared.types.api import (
     Usage,
 )
 from exo.shared.types.common import ModelId
-from exo.shared.types.memory import Memory
+from exo.shared.types.memory import Memory, get_memory_available_locally
 from exo.shared.types.mlx import KVCacheType
+from exo.shared.types.settings import load_settings
 from exo.shared.types.text_generation import InputMessage, TextGenerationTaskParams
 from exo.shared.types.worker.runner_response import (
     GenerationResponse,
@@ -31,17 +33,18 @@ from exo.worker.engines.mlx.cache import (
     encode_prompt,
     has_non_kv_caches,
     make_kv_cache,
+    measure_kv_cache_bytes_per_token,
     snapshot_ssm_states,
 )
 from exo.worker.engines.mlx.constants import (
     DEFAULT_TOP_LOGPROBS,
     KV_BITS,
     KV_GROUP_SIZE,
-    MAX_TOKENS,
 )
 from exo.worker.engines.mlx.utils_mlx import (
     apply_chat_template,
     fix_unmatched_think_end_tokens,
+    mx_any,
     mx_barrier,
 )
 from exo.worker.runner.bootstrap import logger
@@ -109,7 +112,7 @@ def prefill(
             max_tokens=1,
             sampler=sampler,
             prompt_cache=cache,
-            prefill_step_size=4096,
+            prefill_step_size=load_settings().generation.prefill_step_size,
             kv_group_size=KV_GROUP_SIZE,
             kv_bits=KV_BITS,
             prompt_progress_callback=progress_callback,
@@ -147,7 +150,8 @@ def warmup_inference(
     model: Model,
     tokenizer: TokenizerWrapper,
     group: mx.distributed.Group | None,
-) -> int:
+) -> tuple[int, Memory]:
+    """Run warmup inference and tokens_generated and bytes_per_token"""
     content = "Prompt to warm up the inference engine. Repeat this."
 
     warmup_prompt = apply_chat_template(
@@ -186,9 +190,12 @@ def warmup_inference(
 
     logger.info("Generated ALL warmup tokens")
 
+    bytes_per_token = measure_kv_cache_bytes_per_token(cache)
+    logger.info(f"Measured KV cache cost: {bytes_per_token} per token")
+
     mx_barrier(group)
 
-    return tokens_generated
+    return tokens_generated, bytes_per_token
 
 
 def ban_token_ids(token_ids: list[int]) -> Callable[[mx.array, mx.array], mx.array]:
@@ -248,6 +255,9 @@ 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
@@ -263,6 +273,33 @@ def extract_top_logprobs(
     return selected_logprob, top_logprob_items
 
 
+def _check_memory_budget(
+    bytes_per_token: Memory,
+    total_sequence_tokens: int,
+    kv_prefix_cache: KVPrefixCache | None,
+    group: mx.distributed.Group | None,
+) -> str | None:
+    if bytes_per_token.in_bytes == 0:
+        return None
+
+    estimated = bytes_per_token * total_sequence_tokens
+    over_budget = estimated > get_memory_available_locally()
+
+    if not mx_any(over_budget, group):
+        return None
+
+    if kv_prefix_cache is not None and kv_prefix_cache.force_evict_all() > 0:
+        mx.clear_cache()
+        over_budget = estimated > get_memory_available_locally()
+        if not mx_any(over_budget, group):
+            return None
+
+    return (
+        "Not enough memory for this conversation. "
+        "Please start a new conversation or compact your messages."
+    )
+
+
 def mlx_generate(
     model: Model,
     tokenizer: TokenizerWrapper,
@@ -271,7 +308,10 @@ def mlx_generate(
     kv_prefix_cache: KVPrefixCache | None,
     group: mx.distributed.Group | None,
     on_prefill_progress: Callable[[int, int], None] | None = None,
+    bytes_per_token: Memory | None = None,
 ) -> Generator[GenerationResponse]:
+    if bytes_per_token is None:
+        bytes_per_token = Memory()
     # Ensure that generation stats only contains peak memory for this generation
     mx.reset_peak_memory()
     # TODO: Randomise task seed and set in taskparams, instead of hard coding as 42.
@@ -303,6 +343,23 @@ def mlx_generate(
                 f"KV cache hit: {prefix_hit_length}/{len(all_prompt_tokens)} tokens cached ({100 * prefix_hit_length / len(all_prompt_tokens):.1f}%)"
             )
 
+    if bytes_per_token.in_bytes > 0 and load_settings().memory.oom_prevention:
+        oom_error = _check_memory_budget(
+            bytes_per_token=bytes_per_token,
+            total_sequence_tokens=len(all_prompt_tokens),
+            kv_prefix_cache=kv_prefix_cache,
+            group=group,
+        )
+        if oom_error is not None:
+            logger.warning(f"OOM prevention (prefill): {oom_error}")
+            yield GenerationResponse(
+                text=oom_error,
+                token=0,
+                finish_reason="error",
+                usage=None,
+            )
+            return
+
     logits_processors: list[Callable[[mx.array, mx.array], mx.array]] = []
     if is_bench:
         # Only sample length eos tokens
@@ -338,7 +395,7 @@ def mlx_generate(
     # stream_generate starts from the last token
     last_token = prompt_tokens[-2:]
 
-    max_tokens = task.max_output_tokens or MAX_TOKENS
+    max_tokens = task.max_output_tokens or load_settings().generation.max_tokens
     accumulated_text = ""
     generated_text_parts: list[str] = []
     generation_start_time = time.perf_counter()

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

@@ -214,6 +214,8 @@ def load_mlx_items(
 
     set_wired_limit_for_model(get_weights_size(bound_instance.bound_shard))
 
+    mx.clear_cache()
+
     return cast(Model, model), tokenizer
 
 
@@ -291,14 +293,10 @@ def shard_and_load(
 
 def get_tokenizer(model_path: Path, shard_metadata: ShardMetadata) -> TokenizerWrapper:
     """Load tokenizer for a model shard. Delegates to load_tokenizer_for_model_id."""
-    trust_remote_code = (
-        shard_metadata.model_card.trust_remote_code
-        or os.environ.get("EXO_TRUST_REMOTE_CODE") == "1"
-    )
     return load_tokenizer_for_model_id(
         shard_metadata.model_card.model_id,
         model_path,
-        trust_remote_code=trust_remote_code,
+        trust_remote_code=shard_metadata.model_card.trust_remote_code,
     )
 
 

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

@@ -31,6 +31,12 @@ from exo.shared.types.events import (
     TaskAcknowledged,
     TaskStatusUpdated,
 )
+from exo.shared.types.memory import (
+    Memory,
+    get_memory_pressure,
+    get_memory_pressure_threshold,
+)
+from exo.shared.types.settings import load_settings
 from exo.shared.types.tasks import (
     ConnectToGroup,
     LoadModel,
@@ -97,6 +103,7 @@ def main(
         bound_instance.bound_runner_id,
         bound_instance.bound_shard,
     )
+    model_id = shard_metadata.model_card.model_id
     device_rank = shard_metadata.device_rank
     logger.info("hello from the runner")
     if getattr(shard_metadata, "immediate_exception", False):
@@ -113,6 +120,7 @@ def main(
     group = None
     kv_prefix_cache: KVPrefixCache | None = None
     check_for_cancel_every: int | None = None
+    bytes_per_token = Memory.from_bytes(0)
 
     current_status: RunnerStatus = RunnerIdle()
     logger.info("runner created")
@@ -224,12 +232,14 @@ def main(
                     assert tokenizer
 
                     t = time.monotonic()
-                    toks = warmup_inference(
+                    toks, bytes_per_token = warmup_inference(
                         model=cast(Model, inference_model),
                         tokenizer=tokenizer,
                         group=group,
                     )
-                    logger.info(f"warmed up by generating {toks} tokens")
+                    logger.info(
+                        f"warmed up by generating {toks} tokens, {bytes_per_token}/token for KV cache"
+                    )
                     check_for_cancel_every = min(
                         math.ceil(toks / min(time.monotonic() - t, 0.001)), 100
                     )
@@ -281,7 +291,7 @@ def main(
                                 ChunkGenerated(
                                     command_id=command_id,
                                     chunk=PrefillProgressChunk(
-                                        model=shard_metadata.model_card.model_id,
+                                        model=model_id,
                                         processed_tokens=processed,
                                         total_tokens=total,
                                     ),
@@ -309,6 +319,7 @@ def main(
                             kv_prefix_cache=kv_prefix_cache,
                             on_prefill_progress=on_prefill_progress,
                             group=group,
+                            bytes_per_token=bytes_per_token,
                         )
 
                         if tokenizer.has_thinking:
@@ -325,13 +336,17 @@ def main(
                         # Model-specific output parsing for tool calls.
                         if isinstance(inference_model, GptOssModel):
                             mlx_generator = parse_gpt_oss(mlx_generator)
-                        elif isinstance(inference_model, DeepseekV32Model):
+                        elif (
+                            isinstance(inference_model, DeepseekV32Model)
+                            and "deepseek" in model_id.normalize().lower()
+                        ):
                             mlx_generator = parse_deepseek_v32(mlx_generator)
                         elif tool_parser:
                             mlx_generator = parse_tool_calls(mlx_generator, tool_parser)
 
                         completion_tokens = 0
                         tokens_since_last_cancel_check = check_for_cancel_every
+                        oom_stopped = False
                         for response in mlx_generator:
                             tokens_since_last_cancel_check += 1
                             if tokens_since_last_cancel_check >= check_for_cancel_every:
@@ -340,7 +355,15 @@ def main(
                                 want_to_cancel = (task.task_id in cancelled_tasks) or (
                                     TaskId("CANCEL_CURRENT_TASK") in cancelled_tasks
                                 )
-                                if mx_any(want_to_cancel, group):
+                                oom_local = (
+                                    load_settings().memory.oom_prevention
+                                    and bytes_per_token.in_bytes > 0
+                                    and get_memory_pressure()
+                                    > get_memory_pressure_threshold()
+                                )
+                                if mx_any(want_to_cancel or oom_local, group):
+                                    if not want_to_cancel:
+                                        oom_stopped = True
                                     break
 
                             match response:
@@ -355,7 +378,7 @@ def main(
                                                 command_id=command_id,
                                                 chunk=ErrorChunk(
                                                     error_message=response.text,
-                                                    model=shard_metadata.model_card.model_id,
+                                                    model=model_id,
                                                 ),
                                             )
                                         )
@@ -370,7 +393,7 @@ def main(
                                             ChunkGenerated(
                                                 command_id=command_id,
                                                 chunk=TokenChunk(
-                                                    model=shard_metadata.model_card.model_id,
+                                                    model=model_id,
                                                     text=response.text,
                                                     token_id=response.token,
                                                     usage=response.usage,
@@ -389,13 +412,28 @@ def main(
                                                 command_id=command_id,
                                                 chunk=ToolCallChunk(
                                                     tool_calls=response.tool_calls,
-                                                    model=shard_metadata.model_card.model_id,
+                                                    model=model_id,
                                                     usage=response.usage,
                                                     stats=response.stats,
                                                 ),
                                             )
                                         )
 
+                        if oom_stopped and device_rank == 0:
+                            event_sender.send(
+                                ChunkGenerated(
+                                    command_id=command_id,
+                                    chunk=ErrorChunk(
+                                        model=model_id,
+                                        error_message=(
+                                            "Generation stopped: running out of memory. "
+                                            "Please start a new conversation or compact "
+                                            "your messages."
+                                        ),
+                                    ),
+                                )
+                            )
+
                     except PrefillCancelled:
                         logger.info(f"Prefill cancelled for task {task.task_id}")
                     # can we make this more explicit?
@@ -405,7 +443,7 @@ def main(
                                 ChunkGenerated(
                                     command_id=command_id,
                                     chunk=ErrorChunk(
-                                        model=shard_metadata.model_card.model_id,
+                                        model=model_id,
                                         finish_reason="error",
                                         error_message=str(e),
                                     ),
@@ -727,7 +765,7 @@ def parse_tool_calls(
     in_tool_call = False
     tool_call_text_parts: list[str] = []
     for response in responses:
-        if response.text.startswith(tool_parser.start_parsing):
+        if not in_tool_call and response.text.startswith(tool_parser.start_parsing):
             in_tool_call = True
 
         if in_tool_call:
@@ -765,10 +803,9 @@ def parse_tool_calls(
                 )
                 yield response
 
-            continue
-
-        # fallthrough
-        yield response
+        else:
+            # fallthrough
+            yield response
 
 
 EXO_RUNNER_MUST_FAIL = "EXO RUNNER MUST FAIL"

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

@@ -58,15 +58,16 @@ def _flatten(p: dict[str, Any]) -> dict[str, str]:
     }
 
 
-json_tool_parser = ToolParser(
-    start_parsing="<tool_call>",
-    end_parsing="</tool_call>",
-    parse_tool_calls=_parse_json_calls,
-)
+def make_json_parser() -> ToolParser:
+    return ToolParser(
+        start_parsing="<tool_call>",
+        end_parsing="</tool_call>",
+        parse_tool_calls=_parse_json_calls,
+    )
 
 
 def infer_tool_parser(chat_template: str) -> ToolParser | None:
     """Attempt to auto-infer a tool parser from the chat template."""
     if "<tool_call>" in chat_template and "tool_call.name" in chat_template:
-        return json_tool_parser
+        return make_json_parser()
     return None

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

@@ -15,6 +15,7 @@ from exo.shared.types.events import (
     TaskAcknowledged,
     TaskStatusUpdated,
 )
+from exo.shared.types.memory import Memory
 from exo.shared.types.tasks import (
     ConnectToGroup,
     LoadModel,
@@ -114,7 +115,9 @@ def patch_out_mlx(monkeypatch: pytest.MonkeyPatch):
     # initialize_mlx returns a mock group
     monkeypatch.setattr(mlx_runner, "initialize_mlx", make_nothin(MockGroup()))
     monkeypatch.setattr(mlx_runner, "load_mlx_items", make_nothin((1, MockTokenizer)))
-    monkeypatch.setattr(mlx_runner, "warmup_inference", make_nothin(1))
+    monkeypatch.setattr(
+        mlx_runner, "warmup_inference", make_nothin((1, Memory.from_bytes(0)))
+    )
     monkeypatch.setattr(mlx_runner, "_check_for_debug_prompts", nothin)
     monkeypatch.setattr(mlx_runner, "mx_any", make_nothin(False))
     # Mock apply_chat_template since we're using a fake tokenizer (integer 1).

tests/auto_bench.sh

@@ -17,6 +17,13 @@ git branch -r --contains "$commit" | grep -qE '^\s*origin/' || {
   exit 1
 }
 hosts=("$@")
+
+for host; do
+  ssh -T -o BatchMode=yes -o ServerAliveInterval=30 "$host@$host" \
+    "EXO_LIBP2P_NAMESPACE=$commit /nix/var/nix/profiles/default/bin/nix build github:exo-explore/exo/$commit" &
+done
+wait
+
 cleanup() {
   for host in "${hosts[@]}"; do
     ssh -T -o BatchMode=yes "$host@$host" "pkill -f bin/exo" &
@@ -26,11 +33,6 @@ cleanup() {
 }
 trap 'cleanup' EXIT INT TERM
 
-for host; do
-  ssh -T -o BatchMode=yes -o ServerAliveInterval=30 "$host@$host" \
-    "EXO_LIBP2P_NAMESPACE=$commit /nix/var/nix/profiles/default/bin/nix build github:exo-explore/exo/$commit" &
-done
-wait
 for host; do
   ssh -T -o BatchMode=yes -o ServerAliveInterval=30 "$host@$host" \
     "EXO_LIBP2P_NAMESPACE=$commit /nix/var/nix/profiles/default/bin/nix run github:exo-explore/exo/$commit" &>/dev/null &

uv.lock

@@ -363,7 +363,7 @@ wheels = [
 
 [[package]]
 name = "exo"
-version = "0.3.0"
+version = "0.3.68"
 source = { editable = "." }
 dependencies = [
     { name = "aiofiles", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
@@ -378,7 +378,7 @@ dependencies = [
     { name = "loguru", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
     { name = "mflux", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
     { name = "mlx", version = "0.30.6", source = { registry = "https://pypi.org/simple" }, extra = ["cpu"], marker = "sys_platform == 'linux'" },
-    { name = "mlx", version = "0.30.7.dev20260220+13998a05", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#13998a054715edcdc93618fb1496c79c7c25ff7c" }, marker = "sys_platform == 'darwin'" },
+    { name = "mlx", version = "0.30.7.dev20260225+257d5692", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#257d5692fc7af6bba3b8afaeb63c549b7d1e43d5" }, marker = "sys_platform == 'darwin'" },
     { name = "mlx-lm", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
     { name = "msgspec", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
     { name = "openai-harmony", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
@@ -1025,7 +1025,7 @@ dependencies = [
     { name = "huggingface-hub", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
     { name = "matplotlib", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
     { name = "mlx", version = "0.30.6", source = { registry = "https://pypi.org/simple" }, extra = ["cuda13"], marker = "sys_platform == 'linux'" },
-    { name = "mlx", version = "0.30.7.dev20260220+13998a05", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#13998a054715edcdc93618fb1496c79c7c25ff7c" }, marker = "sys_platform == 'darwin'" },
+    { name = "mlx", version = "0.30.7.dev20260225+257d5692", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#257d5692fc7af6bba3b8afaeb63c549b7d1e43d5" }, marker = "sys_platform == 'darwin'" },
     { name = "numpy", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
     { name = "opencv-python", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
     { name = "piexif", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
@@ -1072,8 +1072,8 @@ cuda13 = [
 
 [[package]]
 name = "mlx"
-version = "0.30.7.dev20260220+13998a05"
-source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#13998a054715edcdc93618fb1496c79c7c25ff7c" }
+version = "0.30.7.dev20260225+257d5692"
+source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#257d5692fc7af6bba3b8afaeb63c549b7d1e43d5" }
 resolution-markers = [
     "sys_platform == 'darwin'",
 ]
@@ -1108,7 +1108,7 @@ version = "0.30.7"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
     { name = "jinja2", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
-    { name = "mlx", version = "0.30.7.dev20260220+13998a05", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#13998a054715edcdc93618fb1496c79c7c25ff7c" }, marker = "sys_platform == 'darwin'" },
+    { name = "mlx", version = "0.30.7.dev20260225+257d5692", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#257d5692fc7af6bba3b8afaeb63c549b7d1e43d5" }, marker = "sys_platform == 'darwin'" },
     { name = "numpy", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
     { name = "protobuf", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
     { name = "pyyaml", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },