fix InstanceViewModel.swift

wasn't caught when we merged the API changes

.github/actions/typecheck/action.yml

@@ -1,12 +0,0 @@
-name: Type Check
-
-description: "Run type checker"
-
-runs:
-  using: "composite"
-  steps:
-    - name: Run type checker
-      run: |
-        nix --extra-experimental-features nix-command --extra-experimental-features flakes develop -c just sync
-        nix --extra-experimental-features nix-command --extra-experimental-features flakes develop -c just check
-      shell: bash

.github/workflows/pipeline.yml

@@ -26,73 +26,14 @@ jobs:
           name: exo
           authToken: "${{ secrets.CACHIX_AUTH_TOKEN }}"
 
-      - name: Configure git user
-        run: |
-          git config --local user.email "github-actions@users.noreply.github.com"
-          git config --local user.name  "github-actions bot"
-        shell: bash
+      - name: Load nix develop environment
+        run: nix run github:nicknovitski/nix-develop/v1
 
-      - name: Pull LFS files
-        run: |
-          echo "Pulling Git LFS files..."
-          git lfs pull
-        shell: bash
+      - name: Sync dependencies
+        run: uv sync --all-packages
 
-      - name: Setup Nix Environment
-        run: |
-          echo "Checking for nix installation..."
-          
-          # Check if nix binary exists directly
-          if [ -f /nix/var/nix/profiles/default/bin/nix ]; then
-            echo "Found nix binary at /nix/var/nix/profiles/default/bin/nix"
-            export PATH="/nix/var/nix/profiles/default/bin:$PATH"
-            echo "PATH=$PATH" >> $GITHUB_ENV
-            nix --version
-          elif [ -f /nix/var/nix/profiles/default/etc/profile.d/nix-daemon.sh ]; then
-            echo "Found nix profile script, sourcing..."
-            source /nix/var/nix/profiles/default/etc/profile.d/nix-daemon.sh
-            nix --version
-          elif command -v nix >/dev/null 2>&1; then
-            echo "Nix already in PATH"
-            nix --version
-          else
-            echo "Nix not found. Debugging info:"
-            echo "Contents of /nix/var/nix/profiles/default/:"
-            ls -la /nix/var/nix/profiles/default/ 2>/dev/null || echo "Directory not found"
-            echo "Contents of /nix/var/nix/profiles/default/bin/:"
-            ls -la /nix/var/nix/profiles/default/bin/ 2>/dev/null || echo "Directory not found"
-            exit 1
-          fi
-        shell: bash
-
-      - name: Configure basedpyright include for local MLX
-        run: |
-          RUNNER_LABELS='${{ toJSON(runner.labels) }}'
-          if echo "$RUNNER_LABELS" | grep -q "local_mlx"; then
-            if [ -d "/Users/Shared/mlx" ]; then
-              echo "Updating [tool.basedpyright].include to use /Users/Shared/mlx"
-              awk '
-                BEGIN { in=0 }
-                /^\[tool\.basedpyright\]/ { in=1; print; next }
-                in && /^\[/ { in=0 }  # next section
-                in && /^[ \t]*include[ \t]*=/ {
-                  print "include = [\"/Users/Shared/mlx\"]"
-                  next
-                }
-                { print }
-              ' pyproject.toml > pyproject.toml.tmp && mv pyproject.toml.tmp pyproject.toml
-
-              echo "New [tool.basedpyright] section:"
-              sed -n '/^\[tool\.basedpyright\]/,/^\[/p' pyproject.toml | sed '$d' || true
-            else
-              echo "local_mlx tag present but /Users/Shared/mlx not found; leaving pyproject unchanged."
-            fi
-          else
-            echo "Runner does not have 'local_mlx' tag; leaving pyproject unchanged."
-          fi
-        shell: bash
-
-      - uses: ./.github/actions/typecheck
+      - name: Run type checker
+        run: uv run basedpyright --project pyproject.toml
 
   nix:
     name: Build and check (${{ matrix.system }})
@@ -123,6 +64,63 @@ jobs:
           name: exo
           authToken: "${{ secrets.CACHIX_AUTH_TOKEN }}"
 
+      - name: Build Metal packages (macOS only)
+        if: runner.os == 'macOS'
+        run: |
+          # Try to build metal-toolchain first (may succeed via cachix cache hit)
+          if nix build .#metal-toolchain 2>/dev/null; then
+            echo "metal-toolchain built successfully (likely cache hit)"
+          else
+            echo "metal-toolchain build failed, extracting from Xcode..."
+
+            NAR_HASH="sha256-ayR5mXN4sZAddwKEG2OszGRF93k9ZFc7H0yi2xbylQw="
+            NAR_NAME="metal-toolchain-17C48.nar"
+
+            # Use RUNNER_TEMP to avoid /tmp symlink issues on macOS
+            WORK_DIR="${RUNNER_TEMP}/metal-work"
+            mkdir -p "$WORK_DIR"
+
+            # Download the Metal toolchain component
+            xcodebuild -downloadComponent MetalToolchain
+
+            # Find and mount the DMG
+            DMG_PATH=$(find /System/Library/AssetsV2/com_apple_MobileAsset_MetalToolchain -name '*.dmg' 2>/dev/null | head -1)
+            if [ -z "$DMG_PATH" ]; then
+              echo "Error: Could not find Metal toolchain DMG"
+              exit 1
+            fi
+
+            echo "Found DMG at: $DMG_PATH"
+            hdiutil attach "$DMG_PATH" -mountpoint "${WORK_DIR}/metal-dmg"
+
+            # Copy the toolchain
+            cp -R "${WORK_DIR}/metal-dmg/Metal.xctoolchain" "${WORK_DIR}/metal-export"
+            hdiutil detach "${WORK_DIR}/metal-dmg"
+
+            # Create NAR and add to store
+            nix nar pack "${WORK_DIR}/metal-export" > "${WORK_DIR}/${NAR_NAME}"
+            STORE_PATH=$(nix store add --mode flat "${WORK_DIR}/${NAR_NAME}")
+            echo "Added NAR to store: $STORE_PATH"
+
+            # Verify the hash matches
+            ACTUAL_HASH=$(nix hash file "${WORK_DIR}/${NAR_NAME}")
+            if [ "$ACTUAL_HASH" != "$NAR_HASH" ]; then
+              echo "Warning: NAR hash mismatch!"
+              echo "Expected: $NAR_HASH"
+              echo "Actual:   $ACTUAL_HASH"
+              echo "The metal-toolchain.nix may need updating"
+            fi
+
+            # Clean up
+            rm -rf "$WORK_DIR"
+
+            # Retry the build now that NAR is in store
+            nix build .#metal-toolchain
+          fi
+
+          # Build mlx (depends on metal-toolchain)
+          nix build .#mlx
+
       - name: Build all Nix outputs
         run: |
           nix flake show --json | jq -r '
@@ -134,3 +132,14 @@ jobs:
 
       - name: Run nix flake check
         run: nix flake check
+
+      - name: Run pytest (macOS only)
+        if: runner.os == 'macOS'
+        run: |
+          # Build the test environment (requires relaxed sandbox for uv2nix on macOS)
+          TEST_ENV=$(nix build '.#exo-test-env' --option sandbox relaxed --print-out-paths)
+
+          # Run pytest outside sandbox (needs GPU access for MLX)
+          export HOME="$RUNNER_TEMP"
+          export EXO_TESTS=1
+          $TEST_ENV/bin/python -m pytest src -m "not slow" --import-mode=importlib

.gitignore

@@ -28,3 +28,6 @@ target/
 dashboard/build/
 dashboard/node_modules/
 dashboard/.svelte-kit/
+
+# host config snapshots
+hosts_*.json

README.md

@@ -5,7 +5,7 @@
   <img alt="exo logo" src="/docs/imgs/exo-logo-transparent.png" width="50%" height="50%">
 </picture>
 
-exo: Run your own AI cluster at home with everyday devices. Maintained by [exo labs](https://x.com/exolabs).
+exo: Run frontier AI locally. Maintained by [exo labs](https://x.com/exolabs).
 
 <p align="center">
   <a href="https://discord.gg/TJ4P57arEm" target="_blank" rel="noopener noreferrer"><img src="https://img.shields.io/badge/Discord-Join%20Server-5865F2?logo=discord&logoColor=white" alt="Discord"></a>
@@ -107,6 +107,10 @@ uv run exo
 
 This starts the exo dashboard and API at http://localhost:52415/
 
+
+*Please view the section on RDMA to enable this feature on MacOS >=26.2!*
+
+
 ### Run from Source (Linux)
 
 **Prerequisites:**
@@ -230,7 +234,7 @@ This removes:
 
 RDMA is a new capability added to macOS 26.2. It works on any Mac with Thunderbolt 5 (M4 Pro Mac Mini, M4 Max Mac Studio, M4 Max MacBook Pro, M3 Ultra Mac Studio).
 
-Note that on Mac Studio, you cannot use the Thunderbolt 5 port next to the Ethernet port.
+Please refer to the caveats for immediate troubleshooting.
 
 To enable RDMA on macOS, follow these steps:
 
@@ -247,6 +251,14 @@ To enable RDMA on macOS, follow these steps:
 
 After that, RDMA will be enabled in macOS and exo will take care of the rest.
 
+**Important Caveats**
+
+1. Devices that wish to be part of an RDMA cluster must be connected to all other devices in the cluster.
+2. The cables must support TB5.
+3. On a Mac Studio, you cannot use the Thunderbolt 5 port next to the Ethernet port.
+4. If running from source, please use the script found at `tmp/set_rdma_network_config.sh`, which will disable Thunderbolt Bridge and set dhcp on each RDMA port.
+5. RDMA ports may be unable to discover each other on different versions of MacOS. Please ensure that OS versions match exactly (even beta version numbers) on all devices.
+
 ---
 
 ### Using the API

app/EXO/EXO.xcodeproj/project.pbxproj

@@ -342,6 +342,8 @@
 				SDKROOT = macosx;
 				SWIFT_ACTIVE_COMPILATION_CONDITIONS = "DEBUG $(inherited)";
 				SWIFT_OPTIMIZATION_LEVEL = "-Onone";
+				SWIFT_TREAT_WARNINGS_AS_ERRORS = YES;
+				GCC_TREAT_WARNINGS_AS_ERRORS = YES;
 			};
 			name = Debug;
 		};
@@ -397,6 +399,8 @@
 				MTL_FAST_MATH = YES;
 				SDKROOT = macosx;
 				SWIFT_COMPILATION_MODE = wholemodule;
+				SWIFT_TREAT_WARNINGS_AS_ERRORS = YES;
+				GCC_TREAT_WARNINGS_AS_ERRORS = YES;
 			};
 			name = Release;
 		};

app/EXO/EXO/ContentView.swift

@@ -14,6 +14,7 @@ struct ContentView: View {
     @EnvironmentObject private var networkStatusService: NetworkStatusService
     @EnvironmentObject private var localNetworkChecker: LocalNetworkChecker
     @EnvironmentObject private var updater: SparkleUpdater
+    @EnvironmentObject private var thunderboltBridgeService: ThunderboltBridgeService
     @State private var focusedNode: NodeViewModel?
     @State private var deletingInstanceIDs: Set<String> = []
     @State private var showAllNodes = false
@@ -24,6 +25,8 @@ struct ContentView: View {
     @State private var bugReportMessage: 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) {
@@ -303,6 +306,49 @@ struct ContentView: View {
                             .disabled(pendingNamespace == controller.customNamespace)
                         }
                     }
+                    VStack(alignment: .leading, spacing: 4) {
+                        Text("HuggingFace Token")
+                            .font(.caption2)
+                            .foregroundColor(.secondary)
+                        HStack {
+                            SecureField("optional", text: $pendingHFToken)
+                                .textFieldStyle(.roundedBorder)
+                                .font(.caption2)
+                                .onAppear {
+                                    pendingHFToken = controller.hfToken
+                                }
+                            Button("Save & Restart") {
+                                controller.hfToken = pendingHFToken
+                                if controller.status == .running || controller.status == .starting {
+                                    controller.restart()
+                                }
+                            }
+                            .font(.caption2)
+                            .disabled(pendingHFToken == controller.hfToken)
+                        }
+                    }
+                    Divider()
+                    HStack {
+                        Toggle(
+                            "Enable Image Models (experimental)", isOn: $pendingEnableImageModels
+                        )
+                        .toggleStyle(.switch)
+                        .font(.caption2)
+                        .onAppear {
+                            pendingEnableImageModels = controller.enableImageModels
+                        }
+
+                        Spacer()
+
+                        Button("Save & Restart") {
+                            controller.enableImageModels = pendingEnableImageModels
+                            if controller.status == .running || controller.status == .starting {
+                                controller.restart()
+                            }
+                        }
+                        .font(.caption2)
+                        .disabled(pendingEnableImageModels == controller.enableImageModels)
+                    }
                     HoverButton(title: "Check for Updates", small: true) {
                         updater.checkForUpdates()
                     }
@@ -423,6 +469,44 @@ struct ContentView: View {
         }
     }
 
+    /// 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) {
@@ -465,6 +549,7 @@ struct ContentView: View {
                     Text(thunderboltStatusText)
                         .font(.caption2)
                         .foregroundColor(thunderboltStatusColor)
+                    clusterThunderboltBridgeView
                     interfaceIpList
                     rdmaStatusView
                     sendBugReportButton

app/EXO/EXO/EXOApp.swift

@@ -21,6 +21,7 @@ struct EXOApp: App {
     @StateObject private var networkStatusService: NetworkStatusService
     @StateObject private var localNetworkChecker: LocalNetworkChecker
     @StateObject private var updater: SparkleUpdater
+    @StateObject private var thunderboltBridgeService: ThunderboltBridgeService
     private let terminationObserver: TerminationObserver
     private let ciContext = CIContext(options: nil)
 
@@ -41,10 +42,13 @@ struct EXOApp: App {
         let localNetwork = LocalNetworkChecker()
         _localNetworkChecker = StateObject(wrappedValue: localNetwork)
         _updater = StateObject(wrappedValue: updater)
+        let thunderboltBridge = ThunderboltBridgeService(clusterStateService: service)
+        _thunderboltBridgeService = StateObject(wrappedValue: thunderboltBridge)
         enableLaunchAtLoginIfNeeded()
-        NetworkSetupHelper.ensureLaunchDaemonInstalled()
-        // Check local network access BEFORE launching exo
-        localNetwork.check()
+        // Install LaunchDaemon to disable Thunderbolt Bridge on startup (prevents network loops)
+        NetworkSetupHelper.promptAndInstallIfNeeded()
+        // Check local network access periodically (warning disappears when user grants permission)
+        localNetwork.startPeriodicChecking(interval: 10)
         controller.scheduleLaunch(after: 15)
         service.startPolling()
         networkStatus.startPolling()
@@ -58,6 +62,7 @@ struct EXOApp: App {
                 .environmentObject(networkStatusService)
                 .environmentObject(localNetworkChecker)
                 .environmentObject(updater)
+                .environmentObject(thunderboltBridgeService)
         } label: {
             menuBarIcon
         }
@@ -130,6 +135,7 @@ struct EXOApp: App {
                 "Failed to register EXO for launch at login: \(error.localizedDescription)")
         }
     }
+
 }
 
 /// Helper for managing EXO's launch-at-login registration
@@ -219,7 +225,7 @@ private final class ExoUpdaterDelegate: NSObject, SPUUpdaterDelegate {
         }
     }
 
-    private func showNotification(title: String, body: String) {
+    nonisolated private func showNotification(title: String, body: String) {
         let center = UNUserNotificationCenter.current()
         let content = UNMutableNotificationContent()
         content.title = title

app/EXO/EXO/ExoProcessController.swift

@@ -3,6 +3,8 @@ import Combine
 import Foundation
 
 private let customNamespaceKey = "EXOCustomNamespace"
+private let hfTokenKey = "EXOHFToken"
+private let enableImageModelsKey = "EXOEnableImageModels"
 
 @MainActor
 final class ExoProcessController: ObservableObject {
@@ -37,6 +39,22 @@ final class ExoProcessController: ObservableObject {
             UserDefaults.standard.set(customNamespace, forKey: customNamespaceKey)
         }
     }
+    @Published var hfToken: String = {
+        return UserDefaults.standard.string(forKey: hfTokenKey) ?? ""
+    }()
+    {
+        didSet {
+            UserDefaults.standard.set(hfToken, forKey: hfTokenKey)
+        }
+    }
+    @Published var enableImageModels: Bool = {
+        return UserDefaults.standard.bool(forKey: enableImageModelsKey)
+    }()
+    {
+        didSet {
+            UserDefaults.standard.set(enableImageModels, forKey: enableImageModelsKey)
+        }
+    }
 
     private var process: Process?
     private var runtimeDirectoryURL: URL?
@@ -191,6 +209,12 @@ final class ExoProcessController: ObservableObject {
         var environment = ProcessInfo.processInfo.environment
         environment["EXO_RUNTIME_DIR"] = runtimeURL.path
         environment["EXO_LIBP2P_NAMESPACE"] = computeNamespace()
+        if !hfToken.isEmpty {
+            environment["HF_TOKEN"] = hfToken
+        }
+        if enableImageModels {
+            environment["EXO_ENABLE_IMAGE_MODELS"] = "true"
+        }
 
         var paths: [String] = []
         if let existing = environment["PATH"], !existing.isEmpty {

app/EXO/EXO/Models/ClusterState.swift

@@ -5,17 +5,43 @@ import Foundation
 struct ClusterState: Decodable {
     let instances: [String: ClusterInstance]
     let runners: [String: RunnerStatusSummary]
-    let nodeProfiles: [String: NodeProfile]
     let tasks: [String: ClusterTask]
     let topology: Topology?
     let downloads: [String: [NodeDownloadStatus]]
+    let thunderboltBridgeCycles: [[String]]
+
+    // Granular node state (split from the old nodeProfiles)
+    let nodeIdentities: [String: NodeIdentity]
+    let nodeMemory: [String: MemoryInfo]
+    let nodeSystem: [String: SystemInfo]
+    let nodeThunderboltBridge: [String: ThunderboltBridgeStatus]
+
+    /// Computed property for backwards compatibility - merges granular state into NodeProfile
+    var nodeProfiles: [String: NodeProfile] {
+        var profiles: [String: NodeProfile] = [:]
+        let allNodeIds = Set(nodeIdentities.keys)
+            .union(nodeMemory.keys)
+            .union(nodeSystem.keys)
+        for nodeId in allNodeIds {
+            let identity = nodeIdentities[nodeId]
+            let memory = nodeMemory[nodeId]
+            let system = nodeSystem[nodeId]
+            profiles[nodeId] = NodeProfile(
+                modelId: identity?.modelId,
+                chipId: identity?.chipId,
+                friendlyName: identity?.friendlyName,
+                memory: memory,
+                system: system
+            )
+        }
+        return profiles
+    }
 
     init(from decoder: Decoder) throws {
         let container = try decoder.container(keyedBy: CodingKeys.self)
         let rawInstances = try container.decode([String: TaggedInstance].self, forKey: .instances)
         self.instances = rawInstances.mapValues(\.instance)
         self.runners = try container.decode([String: RunnerStatusSummary].self, forKey: .runners)
-        self.nodeProfiles = try container.decode([String: NodeProfile].self, forKey: .nodeProfiles)
         let rawTasks =
             try container.decodeIfPresent([String: TaggedTask].self, forKey: .tasks) ?? [:]
         self.tasks = rawTasks.compactMapValues(\.task)
@@ -24,15 +50,34 @@ struct ClusterState: Decodable {
             try container.decodeIfPresent([String: [TaggedNodeDownload]].self, forKey: .downloads)
             ?? [:]
         self.downloads = rawDownloads.mapValues { $0.compactMap(\.status) }
+        self.thunderboltBridgeCycles =
+            try container.decodeIfPresent([[String]].self, forKey: .thunderboltBridgeCycles) ?? []
+
+        // Granular node state
+        self.nodeIdentities =
+            try container.decodeIfPresent([String: NodeIdentity].self, forKey: .nodeIdentities)
+            ?? [:]
+        self.nodeMemory =
+            try container.decodeIfPresent([String: MemoryInfo].self, forKey: .nodeMemory) ?? [:]
+        self.nodeSystem =
+            try container.decodeIfPresent([String: SystemInfo].self, forKey: .nodeSystem) ?? [:]
+        self.nodeThunderboltBridge =
+            try container.decodeIfPresent(
+                [String: ThunderboltBridgeStatus].self, forKey: .nodeThunderboltBridge
+            ) ?? [:]
     }
 
     private enum CodingKeys: String, CodingKey {
         case instances
         case runners
-        case nodeProfiles
         case topology
         case tasks
         case downloads
+        case thunderboltBridgeCycles
+        case nodeIdentities
+        case nodeMemory
+        case nodeSystem
+        case nodeThunderboltBridge
     }
 }
 
@@ -102,6 +147,18 @@ struct NodeProfile: Decodable {
     let system: SystemInfo?
 }
 
+struct NodeIdentity: Decodable {
+    let modelId: String?
+    let chipId: String?
+    let friendlyName: String?
+}
+
+struct ThunderboltBridgeStatus: Decodable {
+    let enabled: Bool
+    let exists: Bool
+    let serviceName: String?
+}
+
 struct MemoryInfo: Decodable {
     let ramTotal: MemoryValue?
     let ramAvailable: MemoryValue?
@@ -120,16 +177,51 @@ struct SystemInfo: Decodable {
 }
 
 struct Topology: Decodable {
-    let nodes: [TopologyNode]
-    let connections: [TopologyConnection]?
+    /// Node IDs in the topology
+    let nodes: [String]
+    /// Flattened list of connections (source -> sink pairs)
+    let connections: [TopologyConnection]
+
+    init(from decoder: Decoder) throws {
+        let container = try decoder.container(keyedBy: CodingKeys.self)
+        self.nodes = try container.decodeIfPresent([String].self, forKey: .nodes) ?? []
+
+        // Connections come as nested map: { source: { sink: [edges] } }
+        // We flatten to array of (source, sink) pairs
+        var flatConnections: [TopologyConnection] = []
+        if let nested = try container.decodeIfPresent(
+            [String: [String: [AnyCodable]]].self, forKey: .connections
+        ) {
+            for (source, sinks) in nested {
+                for sink in sinks.keys {
+                    flatConnections.append(
+                        TopologyConnection(localNodeId: source, sendBackNodeId: sink))
+                }
+            }
+        }
+        self.connections = flatConnections
+    }
+
+    private enum CodingKeys: String, CodingKey {
+        case nodes
+        case connections
+    }
 }
 
-struct TopologyNode: Decodable {
-    let nodeId: String
-    let nodeProfile: NodeProfile
+/// Placeholder for decoding arbitrary JSON values we don't need to inspect
+private struct AnyCodable: Decodable {
+    init(from decoder: Decoder) throws {
+        // Just consume the value without storing it
+        _ = try? decoder.singleValueContainer().decode(Bool.self)
+        _ = try? decoder.singleValueContainer().decode(Int.self)
+        _ = try? decoder.singleValueContainer().decode(Double.self)
+        _ = try? decoder.singleValueContainer().decode(String.self)
+        _ = try? decoder.singleValueContainer().decode([AnyCodable].self)
+        _ = try? decoder.singleValueContainer().decode([String: AnyCodable].self)
+    }
 }
 
-struct TopologyConnection: Decodable {
+struct TopologyConnection {
     let localNodeId: String
     let sendBackNodeId: String
 }
@@ -201,7 +293,7 @@ struct ClusterTask {
     let modelName: String?
     let promptPreview: String?
     let errorMessage: String?
-    let parameters: ChatCompletionTaskParameters?
+    let parameters: TextGenerationTaskParameters?
 
     var sortPriority: Int {
         switch status {
@@ -238,12 +330,12 @@ struct ClusterTaskPayload: Decodable {
     let taskStatus: TaskStatus?
     let instanceId: String?
     let commandId: String?
-    let taskParams: ChatCompletionTaskParameters?
+    let taskParams: TextGenerationTaskParameters?
     let errorType: String?
     let errorMessage: String?
 }
 
-struct ChatCompletionTaskParameters: Decodable, Equatable {
+struct TextGenerationTaskParameters: Decodable, Equatable {
     let model: String?
     let messages: [ChatCompletionMessage]?
     let maxTokens: Int?
@@ -282,7 +374,7 @@ extension ClusterTask {
         guard let id = payload.taskId else { return nil }
         let status = payload.taskStatus ?? .unknown
         switch kindKey {
-        case "ChatCompletion":
+        case "TextGeneration":
             self.init(
                 id: id,
                 status: status,

app/EXO/EXO/Services/BugReportService.swift

@@ -55,12 +55,16 @@ struct BugReportService {
         let stateData = try await stateResult
         let eventsData = try await eventsResult
 
+        // Extract cluster TB bridge status from exo state
+        let clusterTbBridgeStatus = extractClusterTbBridgeStatus(from: stateData)
+
         let reportJSON = makeReportJson(
             timestamp: timestamp,
             hostName: hostName,
             ifconfig: ifconfigText,
             debugInfo: debugInfo,
-            isManual: isManual
+            isManual: isManual,
+            clusterTbBridgeStatus: clusterTbBridgeStatus
         )
 
         let uploads: [(path: String, data: Data?)] = [
@@ -178,18 +182,19 @@ struct BugReportService {
     }
 
     private func readThunderboltBridgeDisabled() -> Bool? {
-        let result = runCommand([
-            "/usr/sbin/networksetup", "-getnetworkserviceenabled", "Thunderbolt Bridge",
-        ])
-        guard result.exitCode == 0 else { return nil }
-        let output = result.output.lowercased()
-        if output.contains("enabled") {
-            return false
+        // Dynamically find the Thunderbolt Bridge service (don't assume the name)
+        guard let serviceName = ThunderboltBridgeDetector.findThunderboltBridgeServiceName() else {
+            // No bridge containing Thunderbolt interfaces exists
+            return nil
         }
-        if output.contains("disabled") {
-            return true
+
+        guard let isEnabled = ThunderboltBridgeDetector.isServiceEnabled(serviceName: serviceName)
+        else {
+            return nil
         }
-        return nil
+
+        // Return true if disabled, false if enabled
+        return !isEnabled
     }
 
     private func readInterfaces() -> [DebugInfo.InterfaceStatus] {
@@ -268,11 +273,12 @@ struct BugReportService {
         hostName: String,
         ifconfig: String,
         debugInfo: DebugInfo,
-        isManual: Bool
+        isManual: Bool,
+        clusterTbBridgeStatus: [[String: Any]]?
     ) -> Data? {
         let system = readSystemMetadata()
         let exo = readExoMetadata()
-        let payload: [String: Any] = [
+        var payload: [String: Any] = [
             "timestamp": timestamp,
             "host": hostName,
             "ifconfig": ifconfig,
@@ -282,9 +288,38 @@ struct BugReportService {
             "exo_commit": exo.commit as Any,
             "report_type": isManual ? "manual" : "automated",
         ]
+        if let tbStatus = clusterTbBridgeStatus {
+            payload["cluster_thunderbolt_bridge"] = tbStatus
+        }
         return try? JSONSerialization.data(withJSONObject: payload, options: [.prettyPrinted])
     }
 
+    /// Extracts cluster-wide Thunderbolt Bridge status from exo state JSON
+    private func extractClusterTbBridgeStatus(from stateData: Data?) -> [[String: Any]]? {
+        guard let data = stateData,
+            let json = try? JSONSerialization.jsonObject(with: data) as? [String: Any],
+            let nodeThunderboltBridge = json["node_thunderbolt_bridge"] as? [String: [String: Any]]
+        else {
+            return nil
+        }
+
+        var result: [[String: Any]] = []
+        for (nodeId, status) in nodeThunderboltBridge {
+            var entry: [String: Any] = ["node_id": nodeId]
+            if let enabled = status["enabled"] as? Bool {
+                entry["enabled"] = enabled
+            }
+            if let exists = status["exists"] as? Bool {
+                entry["exists"] = exists
+            }
+            if let serviceName = status["service_name"] as? String {
+                entry["service_name"] = serviceName
+            }
+            result.append(entry)
+        }
+        return result.isEmpty ? nil : result
+    }
+
     private func readSystemMetadata() -> [String: Any] {
         let hostname = safeRunCommand(["/bin/hostname"])
         let computerName = safeRunCommand(["/usr/sbin/scutil", "--get", "ComputerName"])

app/EXO/EXO/Services/LocalNetworkChecker.swift

@@ -41,6 +41,7 @@ final class LocalNetworkChecker: ObservableObject {
 
     private var connection: NWConnection?
     private var checkTask: Task<Void, Never>?
+    private var periodicTask: Task<Void, Never>?
 
     /// Whether we've completed at least one check (stored in UserDefaults)
     private var hasCompletedInitialCheck: Bool {
@@ -48,10 +49,39 @@ final class LocalNetworkChecker: ObservableObject {
         set { UserDefaults.standard.set(newValue, forKey: Self.hasCompletedInitialCheckKey) }
     }
 
-    /// Checks if local network access is working.
+    /// Checks if local network access is working (one-time check).
     func check() {
+        performCheck()
+    }
+
+    /// Starts periodic checking of local network access.
+    /// Re-checks every `interval` seconds so the warning disappears when user grants permission.
+    func startPeriodicChecking(interval: TimeInterval = 10) {
+        stopPeriodicChecking()
+        // Do an immediate check first
+        performCheck()
+        // Then schedule periodic checks
+        periodicTask = Task { [weak self] in
+            while !Task.isCancelled {
+                try? await Task.sleep(nanoseconds: UInt64(interval * 1_000_000_000))
+                guard !Task.isCancelled else { break }
+                self?.performCheck()
+            }
+        }
+    }
+
+    /// Stops periodic checking.
+    func stopPeriodicChecking() {
+        periodicTask?.cancel()
+        periodicTask = nil
+    }
+
+    private func performCheck() {
         checkTask?.cancel()
-        status = .checking
+        // Only show "checking" status on first check to avoid UI flicker
+        if status == .unknown {
+            status = .checking
+        }
 
         // Use longer timeout on first launch to allow time for permission prompt
         let isFirstCheck = !hasCompletedInitialCheck
@@ -60,12 +90,15 @@ final class LocalNetworkChecker: ObservableObject {
         checkTask = Task { [weak self] in
             guard let self else { return }
 
-            Self.logger.info("Checking local network connectivity (first check: \(isFirstCheck))")
+            Self.logger.debug("Checking local network connectivity (first check: \(isFirstCheck))")
             let result = await self.checkConnectivity(timeout: timeout)
             self.status = result
             self.hasCompletedInitialCheck = true
 
-            Self.logger.info("Local network check complete: \(result.displayText)")
+            // Only log on state changes or first check to reduce noise
+            if isFirstCheck || result != self.status {
+                Self.logger.info("Local network check: \(result.displayText)")
+            }
         }
     }
 
@@ -141,6 +174,7 @@ final class LocalNetworkChecker: ObservableObject {
     }
 
     func stop() {
+        stopPeriodicChecking()
         checkTask?.cancel()
         checkTask = nil
         connection?.cancel()

app/EXO/EXO/Services/NetworkSetupHelper.swift

@@ -7,14 +7,20 @@ enum NetworkSetupHelper {
     private static let daemonLabel = "io.exo.networksetup"
     private static let scriptDestination =
         "/Library/Application Support/EXO/disable_bridge.sh"
+    // Legacy script path from older versions
+    private static let legacyScriptDestination =
+        "/Library/Application Support/EXO/disable_bridge_enable_dhcp.sh"
     private static let plistDestination = "/Library/LaunchDaemons/io.exo.networksetup.plist"
-    private static let requiredStartInterval: Int = 1791
+    private static let requiredStartInterval: Int = 1786
 
     private static let setupScript = """
         #!/usr/bin/env bash
 
         set -euo pipefail
 
+        # Wait for macOS to finish network setup after boot
+        sleep 20
+
         PREFS="/Library/Preferences/SystemConfiguration/preferences.plist"
 
         # Remove bridge0 interface
@@ -28,19 +34,69 @@ enum NetworkSetupHelper {
         # Remove Thunderbolt Bridge from VirtualNetworkInterfaces in preferences.plist
         /usr/libexec/PlistBuddy -c "Delete :VirtualNetworkInterfaces:Bridge:bridge0" "$PREFS" 2>/dev/null || true
 
+        networksetup -listlocations | grep -q exo || {
+          networksetup -createlocation exo
+        }
+
+        networksetup -switchtolocation exo
+        networksetup -listallhardwareports \\
+          | awk -F': ' '/Hardware Port: / {print $2}' \\
+          | while IFS=":" read -r name; do
+              case "$name" in
+                "Ethernet Adapter"*)
+                        ;;
+                "Thunderbolt Bridge")
+                        ;;
+                "Thunderbolt "*)
+                  networksetup -listallnetworkservices \\
+                    | grep -q "EXO $name" \\
+                      || networksetup -createnetworkservice "EXO $name" "$name" 2>/dev/null \\
+                      || continue
+                  networksetup -setdhcp "EXO $name"
+                        ;;
+                *)
+                  networksetup -listallnetworkservices \\
+                    | grep -q "$name" \\
+                      || networksetup -createnetworkservice "$name" "$name" 2>/dev/null \\
+                      || continue
+                        ;;
+              esac
+            done
+
         networksetup -listnetworkservices | grep -q "Thunderbolt Bridge" && {
           networksetup -setnetworkserviceenabled "Thunderbolt Bridge" off
         } || true
         """
 
-    static func ensureLaunchDaemonInstalled() {
+    /// Prompts user and installs the LaunchDaemon if not already installed.
+    /// Shows an alert explaining what will be installed before requesting admin privileges.
+    static func promptAndInstallIfNeeded() {
         // Use .utility priority to match NSAppleScript's internal QoS and avoid priority inversion
         Task.detached(priority: .utility) {
+            // If already correctly installed, skip
+            if daemonAlreadyInstalled() {
+                return
+            }
+
+            // Show alert on main thread
+            let shouldInstall = await MainActor.run {
+                let alert = NSAlert()
+                alert.messageText = "EXO Network Configuration"
+                alert.informativeText =
+                    "EXO needs to install a system service to configure local networking. This will disable Thunderbolt Bridge (preventing packet storms) and install a Network Location.\n\nYou will be prompted for your password."
+                alert.alertStyle = .informational
+                alert.addButton(withTitle: "Install")
+                alert.addButton(withTitle: "Not Now")
+                return alert.runModal() == .alertFirstButtonReturn
+            }
+
+            guard shouldInstall else {
+                logger.info("User deferred network setup daemon installation")
+                return
+            }
+
             do {
-                if daemonAlreadyInstalled() {
-                    return
-                }
-                try await installLaunchDaemon()
+                try installLaunchDaemon()
                 logger.info("Network setup launch daemon installed and started")
             } catch {
                 logger.error(
@@ -63,48 +119,9 @@ enum NetworkSetupHelper {
     static func hasInstalledComponents() -> Bool {
         let manager = FileManager.default
         let scriptExists = manager.fileExists(atPath: scriptDestination)
+        let legacyScriptExists = manager.fileExists(atPath: legacyScriptDestination)
         let plistExists = manager.fileExists(atPath: plistDestination)
-        return scriptExists || plistExists
-    }
-
-    private static func makeUninstallScript() -> String {
-        """
-        set -euo pipefail
-
-        LABEL="\(daemonLabel)"
-        SCRIPT_DEST="\(scriptDestination)"
-        PLIST_DEST="\(plistDestination)"
-        LOG_OUT="/var/log/\(daemonLabel).log"
-        LOG_ERR="/var/log/\(daemonLabel).err.log"
-
-        # Unload the LaunchDaemon if running
-        launchctl bootout system/"$LABEL" 2>/dev/null || true
-
-        # Remove LaunchDaemon plist
-        rm -f "$PLIST_DEST"
-
-        # Remove the script and parent directory if empty
-        rm -f "$SCRIPT_DEST"
-        rmdir "$(dirname "$SCRIPT_DEST")" 2>/dev/null || true
-
-        # Remove log files
-        rm -f "$LOG_OUT" "$LOG_ERR"
-
-        # Switch back to Automatic network location
-        networksetup -switchtolocation Automatic 2>/dev/null || true
-
-        # Delete the exo network location if it exists
-        networksetup -listlocations | grep -q '^exo$' && {
-          networksetup -deletelocation exo 2>/dev/null || true
-        } || true
-
-        # Re-enable Thunderbolt Bridge if it exists
-        networksetup -listnetworkservices | grep -q "Thunderbolt Bridge" && {
-          networksetup -setnetworkserviceenabled "Thunderbolt Bridge" on 2>/dev/null || true
-        } || true
-
-        echo "EXO network components removed successfully"
-        """
+        return scriptExists || legacyScriptExists || plistExists
     }
 
     private static func daemonAlreadyInstalled() -> Bool {
@@ -140,7 +157,7 @@ enum NetworkSetupHelper {
         return true
     }
 
-    private static func installLaunchDaemon() async throws {
+    private static func installLaunchDaemon() throws {
         let installerScript = makeInstallerScript()
         try runShellAsAdmin(installerScript)
     }
@@ -151,8 +168,19 @@ enum NetworkSetupHelper {
 
         LABEL="\(daemonLabel)"
         SCRIPT_DEST="\(scriptDestination)"
+        LEGACY_SCRIPT_DEST="\(legacyScriptDestination)"
         PLIST_DEST="\(plistDestination)"
+        LOG_OUT="/var/log/\(daemonLabel).log"
+        LOG_ERR="/var/log/\(daemonLabel).err.log"
 
+        # First, completely remove any existing installation
+        launchctl bootout system/"$LABEL" 2>/dev/null || true
+        rm -f "$PLIST_DEST"
+        rm -f "$SCRIPT_DEST"
+        rm -f "$LEGACY_SCRIPT_DEST"
+        rm -f "$LOG_OUT" "$LOG_ERR"
+
+        # Install fresh
         mkdir -p "$(dirname "$SCRIPT_DEST")"
 
         cat > "$SCRIPT_DEST" <<'EOF_SCRIPT'
@@ -184,13 +212,82 @@ enum NetworkSetupHelper {
         </plist>
         EOF_PLIST
 
-        launchctl bootout system/"$LABEL" >/dev/null 2>&1 || true
         launchctl bootstrap system "$PLIST_DEST"
         launchctl enable system/"$LABEL"
         launchctl kickstart -k system/"$LABEL"
         """
     }
 
+    private static func makeUninstallScript() -> String {
+        """
+        set -euo pipefail
+
+        LABEL="\(daemonLabel)"
+        SCRIPT_DEST="\(scriptDestination)"
+        LEGACY_SCRIPT_DEST="\(legacyScriptDestination)"
+        PLIST_DEST="\(plistDestination)"
+        LOG_OUT="/var/log/\(daemonLabel).log"
+        LOG_ERR="/var/log/\(daemonLabel).err.log"
+
+        # Unload the LaunchDaemon if running
+        launchctl bootout system/"$LABEL" 2>/dev/null || true
+
+        # Remove LaunchDaemon plist
+        rm -f "$PLIST_DEST"
+
+        # Remove the script (current and legacy paths) and parent directory if empty
+        rm -f "$SCRIPT_DEST"
+        rm -f "$LEGACY_SCRIPT_DEST"
+        rmdir "$(dirname "$SCRIPT_DEST")" 2>/dev/null || true
+
+        # Remove log files
+        rm -f "$LOG_OUT" "$LOG_ERR"
+
+        # Switch back to Automatic network location
+        networksetup -switchtolocation Automatic >/dev/null 2>&1 || true
+
+        # Delete the exo network location if it exists
+        networksetup -listlocations 2>/dev/null | grep -q '^exo$' && {
+          networksetup -deletelocation exo >/dev/null 2>&1 || true
+        } || true
+
+        # Re-enable any Thunderbolt Bridge service if it exists
+        # We find it dynamically by looking for bridges containing Thunderbolt interfaces
+        find_and_enable_thunderbolt_bridge() {
+          # Get Thunderbolt interface devices from hardware ports
+          tb_devices=$(networksetup -listallhardwareports 2>/dev/null | awk '
+            /^Hardware Port:/ { port = tolower(substr($0, 16)) }
+            /^Device:/ { if (port ~ /thunderbolt/) print substr($0, 9) }
+          ') || true
+          [ -z "$tb_devices" ] && return 0
+
+          # For each bridge device, check if it contains Thunderbolt interfaces
+          for bridge in bridge0 bridge1 bridge2; do
+            members=$(ifconfig "$bridge" 2>/dev/null | awk '/member:/ {print $2}') || true
+            [ -z "$members" ] && continue
+
+            for tb_dev in $tb_devices; do
+              if echo "$members" | grep -qx "$tb_dev"; then
+                # Find the service name for this bridge device
+                service_name=$(networksetup -listnetworkserviceorder 2>/dev/null | awk -v dev="$bridge" '
+                  /^\\([0-9*]/ { gsub(/^\\([0-9*]+\\) /, ""); svc = $0 }
+                  /Device:/ && $0 ~ dev { print svc; exit }
+                ') || true
+                if [ -n "$service_name" ]; then
+                  networksetup -setnetworkserviceenabled "$service_name" on 2>/dev/null || true
+                  return 0
+                fi
+              fi
+            done
+          done
+          return 0
+        }
+        find_and_enable_thunderbolt_bridge || true
+
+        echo "EXO network components removed successfully"
+        """
+    }
+
     private static func runShellAsAdmin(_ script: String) throws {
         let escapedScript =
             script

app/EXO/EXO/Services/NetworkStatusService.swift

@@ -153,22 +153,18 @@ private struct NetworkStatusFetcher {
     }
 
     private func readThunderboltBridgeState() -> ThunderboltState? {
-        let result = runCommand(["networksetup", "-getnetworkserviceenabled", "Thunderbolt Bridge"])
-        guard result.exitCode == 0 else {
-            let lower = result.output.lowercased() + result.error.lowercased()
-            if lower.contains("not a recognized network service") {
-                return .deleted
-            }
+        // Dynamically find the Thunderbolt Bridge service (don't assume the name)
+        guard let serviceName = ThunderboltBridgeDetector.findThunderboltBridgeServiceName() else {
+            // No bridge containing Thunderbolt interfaces exists
+            return .deleted
+        }
+
+        guard let isEnabled = ThunderboltBridgeDetector.isServiceEnabled(serviceName: serviceName)
+        else {
             return nil
         }
-        let output = result.output.lowercased()
-        if output.contains("enabled") {
-            return .enabled
-        }
-        if output.contains("disabled") {
-            return .disabled
-        }
-        return nil
+
+        return isEnabled ? .enabled : .disabled
     }
 
     private func readBridgeInactive() -> Bool? {

app/EXO/EXO/Services/ThunderboltBridgeDetector.swift

@@ -0,0 +1,194 @@
+import Foundation
+import os.log
+
+/// Utility for dynamically detecting Thunderbolt Bridge network services.
+/// This mirrors the Python logic in info_gatherer.py - we never assume the service
+/// is named "Thunderbolt Bridge", instead we find bridges containing Thunderbolt interfaces.
+enum ThunderboltBridgeDetector {
+    private static let logger = Logger(
+        subsystem: "io.exo.EXO", category: "ThunderboltBridgeDetector")
+
+    struct CommandResult {
+        let exitCode: Int32
+        let output: String
+        let error: String
+    }
+
+    /// Find the network service name of a bridge containing Thunderbolt interfaces.
+    /// Returns nil if no such bridge exists.
+    static func findThunderboltBridgeServiceName() -> String? {
+        // 1. Get all Thunderbolt interface devices (e.g., en2, en3)
+        guard let thunderboltDevices = getThunderboltDevices(), !thunderboltDevices.isEmpty else {
+            logger.debug("No Thunderbolt devices found")
+            return nil
+        }
+        logger.debug("Found Thunderbolt devices: \(thunderboltDevices.joined(separator: ", "))")
+
+        // 2. Get bridge services from network service order
+        guard let bridgeServices = getBridgeServices(), !bridgeServices.isEmpty else {
+            logger.debug("No bridge services found")
+            return nil
+        }
+        logger.debug("Found bridge services: \(bridgeServices.keys.joined(separator: ", "))")
+
+        // 3. Find a bridge that contains Thunderbolt interfaces
+        for (bridgeDevice, serviceName) in bridgeServices {
+            let members = getBridgeMembers(bridgeDevice: bridgeDevice)
+            logger.debug(
+                "Bridge \(bridgeDevice) (\(serviceName)) has members: \(members.joined(separator: ", "))"
+            )
+
+            // Check if any Thunderbolt device is a member of this bridge
+            if !members.isDisjoint(with: thunderboltDevices) {
+                logger.info(
+                    "Found Thunderbolt Bridge service: '\(serviceName)' (device: \(bridgeDevice))")
+                return serviceName
+            }
+        }
+
+        logger.debug("No bridge found containing Thunderbolt interfaces")
+        return nil
+    }
+
+    /// Get Thunderbolt interface device names (e.g., en2, en3) from hardware ports.
+    private static func getThunderboltDevices() -> Set<String>? {
+        let result = runCommand(["networksetup", "-listallhardwareports"])
+        guard result.exitCode == 0 else {
+            logger.warning("networksetup -listallhardwareports failed: \(result.error)")
+            return nil
+        }
+
+        var thunderboltDevices: Set<String> = []
+        var currentPort: String?
+
+        for line in result.output.components(separatedBy: .newlines) {
+            let trimmed = line.trimmingCharacters(in: .whitespaces)
+            if trimmed.hasPrefix("Hardware Port:") {
+                currentPort = String(trimmed.dropFirst("Hardware Port:".count)).trimmingCharacters(
+                    in: .whitespaces)
+            } else if trimmed.hasPrefix("Device:"), let port = currentPort {
+                let device = String(trimmed.dropFirst("Device:".count)).trimmingCharacters(
+                    in: .whitespaces)
+                if port.lowercased().contains("thunderbolt") {
+                    thunderboltDevices.insert(device)
+                }
+                currentPort = nil
+            }
+        }
+
+        return thunderboltDevices
+    }
+
+    /// Get mapping of bridge device -> service name from network service order.
+    private static func getBridgeServices() -> [String: String]? {
+        let result = runCommand(["networksetup", "-listnetworkserviceorder"])
+        guard result.exitCode == 0 else {
+            logger.warning("networksetup -listnetworkserviceorder failed: \(result.error)")
+            return nil
+        }
+
+        // Parse service order to find bridge devices and their service names
+        // Format: "(1) Service Name\n(Hardware Port: ..., Device: bridge0)\n"
+        var bridgeServices: [String: String] = [:]
+        var currentService: String?
+
+        for line in result.output.components(separatedBy: .newlines) {
+            let trimmed = line.trimmingCharacters(in: .whitespaces)
+
+            // Match "(N) Service Name" or "(*) Service Name" (disabled)
+            // but NOT "(Hardware Port: ...)" lines
+            if trimmed.hasPrefix("("), trimmed.contains(")"),
+                !trimmed.hasPrefix("(Hardware Port:")
+            {
+                if let parenEnd = trimmed.firstIndex(of: ")") {
+                    let afterParen = trimmed.index(after: parenEnd)
+                    if afterParen < trimmed.endIndex {
+                        currentService =
+                            String(trimmed[afterParen...])
+                            .trimmingCharacters(in: .whitespaces)
+                    }
+                }
+            }
+            // Match "(Hardware Port: ..., Device: bridgeX)"
+            else if let service = currentService, trimmed.contains("Device: bridge") {
+                // Extract device name from "..., Device: bridge0)"
+                if let deviceRange = trimmed.range(of: "Device: ") {
+                    let afterDevice = trimmed[deviceRange.upperBound...]
+                    if let parenIndex = afterDevice.firstIndex(of: ")") {
+                        let device = String(afterDevice[..<parenIndex])
+                        bridgeServices[device] = service
+                    }
+                }
+            }
+        }
+
+        return bridgeServices
+    }
+
+    /// Get member interfaces of a bridge device via ifconfig.
+    private static func getBridgeMembers(bridgeDevice: String) -> Set<String> {
+        let result = runCommand(["ifconfig", bridgeDevice])
+        guard result.exitCode == 0 else {
+            logger.debug("ifconfig \(bridgeDevice) failed")
+            return []
+        }
+
+        var members: Set<String> = []
+        for line in result.output.components(separatedBy: .newlines) {
+            let trimmed = line.trimmingCharacters(in: .whitespaces)
+            if trimmed.hasPrefix("member:") {
+                let parts = trimmed.split(separator: " ")
+                if parts.count > 1 {
+                    members.insert(String(parts[1]))
+                }
+            }
+        }
+
+        return members
+    }
+
+    /// Check if a network service is enabled.
+    static func isServiceEnabled(serviceName: String) -> Bool? {
+        let result = runCommand(["networksetup", "-getnetworkserviceenabled", serviceName])
+        guard result.exitCode == 0 else {
+            logger.warning("Failed to check if '\(serviceName)' is enabled: \(result.error)")
+            return nil
+        }
+
+        let output = result.output.lowercased().trimmingCharacters(in: .whitespacesAndNewlines)
+        if output.contains("enabled") {
+            return true
+        }
+        if output.contains("disabled") {
+            return false
+        }
+        return nil
+    }
+
+    private static func runCommand(_ arguments: [String]) -> CommandResult {
+        let process = Process()
+        process.launchPath = "/usr/bin/env"
+        process.arguments = arguments
+
+        let stdout = Pipe()
+        let stderr = Pipe()
+        process.standardOutput = stdout
+        process.standardError = stderr
+
+        do {
+            try process.run()
+        } catch {
+            return CommandResult(exitCode: -1, output: "", error: error.localizedDescription)
+        }
+        process.waitUntilExit()
+
+        let outputData = stdout.fileHandleForReading.readDataToEndOfFile()
+        let errorData = stderr.fileHandleForReading.readDataToEndOfFile()
+
+        return CommandResult(
+            exitCode: process.terminationStatus,
+            output: String(decoding: outputData, as: UTF8.self),
+            error: String(decoding: errorData, as: UTF8.self)
+        )
+    }
+}

app/EXO/EXO/Services/ThunderboltBridgeService.swift

@@ -0,0 +1,261 @@
+import AppKit
+import Combine
+import Foundation
+import Security
+import SystemConfiguration
+import os.log
+
+@MainActor
+final class ThunderboltBridgeService: ObservableObject {
+    private static let logger = Logger(subsystem: "io.exo.EXO", category: "ThunderboltBridge")
+
+    @Published private(set) var detectedCycle: [String]?
+    @Published private(set) var hasPromptedForCurrentCycle = false
+    @Published private(set) var lastError: String?
+
+    private weak var clusterStateService: ClusterStateService?
+    private var cancellables = Set<AnyCancellable>()
+    private var previousCycleSignature: String?
+
+    init(clusterStateService: ClusterStateService) {
+        self.clusterStateService = clusterStateService
+        setupObserver()
+    }
+
+    private func setupObserver() {
+        guard let service = clusterStateService else { return }
+
+        service.$latestSnapshot
+            .compactMap { $0 }
+            .sink { [weak self] snapshot in
+                self?.checkForCycles(snapshot: snapshot)
+            }
+            .store(in: &cancellables)
+    }
+
+    private func checkForCycles(snapshot: ClusterState) {
+        let cycles = snapshot.thunderboltBridgeCycles
+
+        // Only consider cycles with more than 2 nodes
+        guard let firstCycle = cycles.first, firstCycle.count > 2 else {
+            // No problematic cycles detected, reset state
+            if detectedCycle != nil {
+                detectedCycle = nil
+                previousCycleSignature = nil
+                hasPromptedForCurrentCycle = false
+            }
+            return
+        }
+
+        // Create a signature for this cycle to detect if it changed
+        let cycleSignature = firstCycle.sorted().joined(separator: ",")
+
+        // If this is a new/different cycle, reset the prompt state
+        if cycleSignature != previousCycleSignature {
+            previousCycleSignature = cycleSignature
+            hasPromptedForCurrentCycle = false
+        }
+
+        detectedCycle = firstCycle
+
+        // Only prompt once per cycle
+        if !hasPromptedForCurrentCycle {
+            showDisableBridgePrompt(nodeIds: firstCycle)
+        }
+    }
+
+    private func showDisableBridgePrompt(nodeIds: [String]) {
+        hasPromptedForCurrentCycle = true
+
+        // Get friendly names for the nodes if available
+        let nodeNames = nodeIds.map { nodeId -> String in
+            if let snapshot = clusterStateService?.latestSnapshot,
+                let profile = snapshot.nodeProfiles[nodeId],
+                let friendlyName = profile.friendlyName, !friendlyName.isEmpty
+            {
+                return friendlyName
+            }
+            return String(nodeId.prefix(8))  // Use first 8 chars of node ID as fallback
+        }
+        let machineNames = nodeNames.joined(separator: ", ")
+
+        let alert = NSAlert()
+        alert.messageText = "Thunderbolt Bridge Loop Detected"
+        alert.informativeText = """
+            A Thunderbolt Bridge loop has been detected between \(nodeNames.count) machines: \(machineNames).
+
+            This can cause network packet storms and connectivity issues. Would you like to disable Thunderbolt Bridge on this machine to break the loop?
+            """
+        alert.alertStyle = .warning
+        alert.addButton(withTitle: "Disable Bridge")
+        alert.addButton(withTitle: "Not Now")
+
+        let response = alert.runModal()
+
+        if response == .alertFirstButtonReturn {
+            Task {
+                await disableThunderboltBridge()
+            }
+        }
+    }
+
+    func disableThunderboltBridge() async {
+        Self.logger.info("Attempting to disable Thunderbolt Bridge via SCPreferences")
+        lastError = nil
+
+        do {
+            try await disableThunderboltBridgeWithSCPreferences()
+            Self.logger.info("Successfully disabled Thunderbolt Bridge")
+        } catch {
+            Self.logger.error(
+                "Failed to disable Thunderbolt Bridge: \(error.localizedDescription, privacy: .public)"
+            )
+            lastError = error.localizedDescription
+            showErrorAlert(message: error.localizedDescription)
+        }
+    }
+
+    private func disableThunderboltBridgeWithSCPreferences() async throws {
+        // 1. Create authorization reference
+        var authRef: AuthorizationRef?
+        var status = AuthorizationCreate(nil, nil, [], &authRef)
+        guard status == errAuthorizationSuccess, let authRef = authRef else {
+            throw ThunderboltBridgeError.authorizationFailed
+        }
+
+        defer { AuthorizationFree(authRef, [.destroyRights]) }
+
+        // 2. Request specific network configuration rights
+        let rightName = "system.services.systemconfiguration.network"
+        status = rightName.withCString { nameCString in
+            var item = AuthorizationItem(
+                name: nameCString,
+                valueLength: 0,
+                value: nil,
+                flags: 0
+            )
+            return withUnsafeMutablePointer(to: &item) { itemPointer in
+                var rights = AuthorizationRights(count: 1, items: itemPointer)
+                return AuthorizationCopyRights(
+                    authRef,
+                    &rights,
+                    nil,
+                    [.extendRights, .interactionAllowed],
+                    nil
+                )
+            }
+        }
+        guard status == errAuthorizationSuccess else {
+            if status == errAuthorizationCanceled {
+                throw ThunderboltBridgeError.authorizationCanceled
+            }
+            throw ThunderboltBridgeError.authorizationDenied
+        }
+
+        // 3. Create SCPreferences with authorization
+        guard
+            let prefs = SCPreferencesCreateWithAuthorization(
+                kCFAllocatorDefault,
+                "EXO" as CFString,
+                nil,
+                authRef
+            )
+        else {
+            throw ThunderboltBridgeError.preferencesCreationFailed
+        }
+
+        // 4. Lock, modify, commit
+        guard SCPreferencesLock(prefs, true) else {
+            throw ThunderboltBridgeError.lockFailed
+        }
+
+        defer {
+            SCPreferencesUnlock(prefs)
+        }
+
+        // 5. Find the Thunderbolt Bridge service dynamically (don't assume the name)
+        guard let targetServiceName = ThunderboltBridgeDetector.findThunderboltBridgeServiceName()
+        else {
+            throw ThunderboltBridgeError.serviceNotFound
+        }
+
+        guard let allServices = SCNetworkServiceCopyAll(prefs) as? [SCNetworkService] else {
+            throw ThunderboltBridgeError.servicesNotFound
+        }
+
+        var found = false
+        for service in allServices {
+            if let name = SCNetworkServiceGetName(service) as String?,
+                name == targetServiceName
+            {
+                guard SCNetworkServiceSetEnabled(service, false) else {
+                    throw ThunderboltBridgeError.disableFailed
+                }
+                found = true
+                Self.logger.info(
+                    "Found and disabled Thunderbolt Bridge service: '\(targetServiceName)'")
+                break
+            }
+        }
+
+        if !found {
+            throw ThunderboltBridgeError.serviceNotFound
+        }
+
+        // 6. Commit and apply
+        guard SCPreferencesCommitChanges(prefs) else {
+            throw ThunderboltBridgeError.commitFailed
+        }
+
+        guard SCPreferencesApplyChanges(prefs) else {
+            throw ThunderboltBridgeError.applyFailed
+        }
+    }
+
+    private func showErrorAlert(message: String) {
+        let alert = NSAlert()
+        alert.messageText = "Failed to Disable Thunderbolt Bridge"
+        alert.informativeText = message
+        alert.alertStyle = .critical
+        alert.addButton(withTitle: "OK")
+        alert.runModal()
+    }
+}
+
+enum ThunderboltBridgeError: LocalizedError {
+    case authorizationFailed
+    case authorizationCanceled
+    case authorizationDenied
+    case preferencesCreationFailed
+    case lockFailed
+    case servicesNotFound
+    case serviceNotFound
+    case disableFailed
+    case commitFailed
+    case applyFailed
+
+    var errorDescription: String? {
+        switch self {
+        case .authorizationFailed:
+            return "Failed to create authorization"
+        case .authorizationCanceled:
+            return "Authorization was canceled by user"
+        case .authorizationDenied:
+            return "Authorization was denied"
+        case .preferencesCreationFailed:
+            return "Failed to access network preferences"
+        case .lockFailed:
+            return "Failed to lock network preferences for modification"
+        case .servicesNotFound:
+            return "Could not retrieve network services"
+        case .serviceNotFound:
+            return "Thunderbolt Bridge service not found"
+        case .disableFailed:
+            return "Failed to disable Thunderbolt Bridge service"
+        case .commitFailed:
+            return "Failed to save network configuration changes"
+        case .applyFailed:
+            return "Failed to apply network configuration changes"
+        }
+    }
+}

app/EXO/EXO/ViewModels/InstanceViewModel.swift

@@ -216,7 +216,7 @@ struct InstanceTaskViewModel: Identifiable, Equatable {
     let promptPreview: String?
     let errorMessage: String?
     let subtitle: String?
-    let parameters: ChatCompletionTaskParameters?
+    let parameters: TextGenerationTaskParameters?
 
     var title: String {
         switch kind {

app/EXO/EXO/ViewModels/NodeViewModel.swift

@@ -86,7 +86,7 @@ struct TopologyViewModel {
 
 extension ClusterState {
     func topologyViewModel(localNodeId: String?) -> TopologyViewModel? {
-        let topologyNodeIds = Set(topology?.nodes.map(\.nodeId) ?? [])
+        let topologyNodeIds = Set(topology?.nodes ?? [])
         let allNodes = nodeViewModels().filter {
             topologyNodeIds.isEmpty || topologyNodeIds.contains($0.id)
         }
@@ -95,8 +95,8 @@ extension ClusterState {
         let nodesById = Dictionary(uniqueKeysWithValues: allNodes.map { ($0.id, $0) })
         var orderedNodes: [NodeViewModel] = []
         if let topologyNodes = topology?.nodes {
-            for topoNode in topologyNodes {
-                if let viewModel = nodesById[topoNode.nodeId] {
+            for nodeId in topologyNodes {
+                if let viewModel = nodesById[nodeId] {
                     orderedNodes.append(viewModel)
                 }
             }
@@ -116,7 +116,7 @@ extension ClusterState {
 
         let nodeIds = Set(orderedNodes.map(\.id))
         let edgesArray: [TopologyEdgeViewModel] =
-            topology?.connections?.compactMap { connection in
+            topology?.connections.compactMap { connection in
                 guard nodeIds.contains(connection.localNodeId),
                     nodeIds.contains(connection.sendBackNodeId)
                 else { return nil }

app/EXO/uninstall-exo.sh

@@ -29,21 +29,21 @@ YELLOW='\033[1;33m'
 NC='\033[0m' # No Color
 
 echo_info() {
-    echo -e "${GREEN}[INFO]${NC} $1"
+  echo -e "${GREEN}[INFO]${NC} $1"
 }
 
 echo_warn() {
-    echo -e "${YELLOW}[WARN]${NC} $1"
+  echo -e "${YELLOW}[WARN]${NC} $1"
 }
 
 echo_error() {
-    echo -e "${RED}[ERROR]${NC} $1"
+  echo -e "${RED}[ERROR]${NC} $1"
 }
 
 # Check if running as root
 if [[ $EUID -ne 0 ]]; then
-    echo_error "This script must be run as root (use sudo)"
-    exit 1
+  echo_error "This script must be run as root (use sudo)"
+  exit 1
 fi
 
 echo ""
@@ -55,64 +55,64 @@ echo ""
 # Unload the LaunchDaemon if running
 echo_info "Stopping network setup daemon..."
 if launchctl list | grep -q "$LABEL"; then
-    launchctl bootout system/"$LABEL" 2>/dev/null || true
-    echo_info "Daemon stopped"
+  launchctl bootout system/"$LABEL" 2>/dev/null || true
+  echo_info "Daemon stopped"
 else
-    echo_warn "Daemon was not running"
+  echo_warn "Daemon was not running"
 fi
 
 # Remove LaunchDaemon plist
-if [[ -f "$PLIST_DEST" ]]; then
-    rm -f "$PLIST_DEST"
-    echo_info "Removed LaunchDaemon plist"
+if [[ -f $PLIST_DEST ]]; then
+  rm -f "$PLIST_DEST"
+  echo_info "Removed LaunchDaemon plist"
 else
-    echo_warn "LaunchDaemon plist not found (already removed?)"
+  echo_warn "LaunchDaemon plist not found (already removed?)"
 fi
 
 # Remove the script and parent directory
-if [[ -f "$SCRIPT_DEST" ]]; then
-    rm -f "$SCRIPT_DEST"
-    echo_info "Removed network setup script"
+if [[ -f $SCRIPT_DEST ]]; then
+  rm -f "$SCRIPT_DEST"
+  echo_info "Removed network setup script"
 else
-    echo_warn "Network setup script not found (already removed?)"
+  echo_warn "Network setup script not found (already removed?)"
 fi
 
 # Remove EXO directory if empty
 if [[ -d "/Library/Application Support/EXO" ]]; then
-    rmdir "/Library/Application Support/EXO" 2>/dev/null && \
-        echo_info "Removed EXO support directory" || \
-        echo_warn "EXO support directory not empty, leaving in place"
+  rmdir "/Library/Application Support/EXO" 2>/dev/null &&
+    echo_info "Removed EXO support directory" ||
+    echo_warn "EXO support directory not empty, leaving in place"
 fi
 
 # Remove log files
-if [[ -f "$LOG_OUT" ]] || [[ -f "$LOG_ERR" ]]; then
-    rm -f "$LOG_OUT" "$LOG_ERR"
-    echo_info "Removed log files"
+if [[ -f $LOG_OUT ]] || [[ -f $LOG_ERR ]]; then
+  rm -f "$LOG_OUT" "$LOG_ERR"
+  echo_info "Removed log files"
 else
-    echo_warn "Log files not found (already removed?)"
+  echo_warn "Log files not found (already removed?)"
 fi
 
 # Switch back to Automatic network location
 echo_info "Restoring network configuration..."
 if networksetup -listlocations | grep -q "^Automatic$"; then
-    networksetup -switchtolocation Automatic 2>/dev/null || true
-    echo_info "Switched to Automatic network location"
+  networksetup -switchtolocation Automatic 2>/dev/null || true
+  echo_info "Switched to Automatic network location"
 else
-    echo_warn "Automatic network location not found"
+  echo_warn "Automatic network location not found"
 fi
 
 # Delete the exo network location if it exists
 if networksetup -listlocations | grep -q "^exo$"; then
-    networksetup -deletelocation exo 2>/dev/null || true
-    echo_info "Deleted 'exo' network location"
+  networksetup -deletelocation exo 2>/dev/null || true
+  echo_info "Deleted 'exo' network location"
 else
-    echo_warn "'exo' network location not found (already removed?)"
+  echo_warn "'exo' network location not found (already removed?)"
 fi
 
 # Re-enable Thunderbolt Bridge if it exists
 if networksetup -listnetworkservices 2>/dev/null | grep -q "Thunderbolt Bridge"; then
-    networksetup -setnetworkserviceenabled "Thunderbolt Bridge" on 2>/dev/null || true
-    echo_info "Re-enabled Thunderbolt Bridge"
+  networksetup -setnetworkserviceenabled "Thunderbolt Bridge" on 2>/dev/null || true
+  echo_info "Re-enabled Thunderbolt Bridge"
 fi
 
 # Note about launch at login registration
@@ -124,14 +124,14 @@ echo_warn "  System Settings → General → Login Items → Remove EXO"
 # Check if EXO.app exists in common locations
 APP_FOUND=false
 for app_path in "/Applications/EXO.app" "$HOME/Applications/EXO.app"; do
-    if [[ -d "$app_path" ]]; then
-        if [[ "$APP_FOUND" == false ]]; then
-            echo ""
-            APP_FOUND=true
-        fi
-        echo_warn "EXO.app found at: $app_path"
-        echo_warn "You may want to move it to Trash manually."
+  if [[ -d $app_path ]]; then
+    if [[ $APP_FOUND == false ]]; then
+      echo ""
+      APP_FOUND=true
     fi
+    echo_warn "EXO.app found at: $app_path"
+    echo_warn "You may want to move it to Trash manually."
+  fi
 done
 
 echo ""
@@ -151,4 +151,3 @@ echo ""
 echo "Manual step required:"
 echo "  Remove EXO from Login Items in System Settings → General → Login Items"
 echo ""
-

bench/exo_bench.py

@@ -5,10 +5,13 @@ from __future__ import annotations
 import argparse
 import contextlib
 import http.client
+import itertools
 import json
 import os
+import sys
 import time
 from collections.abc import Callable
+from pathlib import Path
 from statistics import mean
 from typing import Any
 from urllib.parse import urlencode
@@ -16,6 +19,84 @@ from urllib.parse import urlencode
 from loguru import logger
 from transformers import AutoTokenizer
 
+# Monkey-patch for transformers 5.x compatibility
+# Kimi's tokenization_kimi.py imports bytes_to_unicode from the old location
+# which was moved in transformers 5.0.0rc2
+try:
+    import transformers.models.gpt2.tokenization_gpt2 as gpt2_tokenization
+    from transformers.convert_slow_tokenizer import bytes_to_unicode
+
+    if not hasattr(gpt2_tokenization, "bytes_to_unicode"):
+        gpt2_tokenization.bytes_to_unicode = bytes_to_unicode  # type: ignore[attr-defined]
+except ImportError:
+    pass  # transformers < 5.0 or bytes_to_unicode not available
+
+
+def load_tokenizer_for_bench(model_id: str) -> Any:
+    """
+    Load tokenizer for benchmarking, with special handling for Kimi models.
+
+    Kimi uses a custom TikTokenTokenizer that transformers 5.x can't load via AutoTokenizer.
+    This function replicates the logic from utils_mlx.py for bench compatibility.
+    """
+    model_id_lower = model_id.lower()
+
+    if "kimi-k2" in model_id_lower:
+        import importlib.util
+        import types
+
+        from huggingface_hub import snapshot_download
+
+        # Download/get the model path
+        model_path = Path(
+            snapshot_download(
+                model_id,
+                allow_patterns=["*.json", "*.py", "*.tiktoken"],
+            )
+        )
+
+        sys.path.insert(0, str(model_path))
+
+        # Load tool_declaration_ts first (tokenization_kimi imports it with relative import)
+        tool_decl_path = model_path / "tool_declaration_ts.py"
+        if tool_decl_path.exists():
+            spec = importlib.util.spec_from_file_location(
+                "tool_declaration_ts", tool_decl_path
+            )
+            if spec and spec.loader:
+                tool_decl_module = importlib.util.module_from_spec(spec)
+                sys.modules["tool_declaration_ts"] = tool_decl_module
+                spec.loader.exec_module(tool_decl_module)
+
+        # Load tokenization_kimi with patched source (convert relative to absolute import)
+        tok_path = model_path / "tokenization_kimi.py"
+        source = tok_path.read_text()
+        source = source.replace("from .tool_declaration_ts", "from tool_declaration_ts")
+        spec = importlib.util.spec_from_file_location("tokenization_kimi", tok_path)
+        if spec:
+            tok_module = types.ModuleType("tokenization_kimi")
+            tok_module.__file__ = str(tok_path)
+            sys.modules["tokenization_kimi"] = tok_module
+            exec(compile(source, tok_path, "exec"), tok_module.__dict__)  # noqa: S102
+            TikTokenTokenizer = tok_module.TikTokenTokenizer  # noqa: N806
+        else:
+            from tokenization_kimi import TikTokenTokenizer  # type: ignore[import-not-found]  # noqa: I001
+
+        hf_tokenizer: Any = TikTokenTokenizer.from_pretrained(model_path)
+
+        # Patch encode to use internal tiktoken model directly
+        # transformers 5.x has a bug in the encode->pad path for slow tokenizers
+        def _patched_encode(text: str, **kwargs: object) -> list[int]:
+            # Pass allowed_special="all" to handle special tokens like <|im_user|>
+            return list(hf_tokenizer.model.encode(text, allowed_special="all"))
+
+        hf_tokenizer.encode = _patched_encode
+
+        return hf_tokenizer
+
+    # Default: use AutoTokenizer
+    return AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
+
 
 class ExoHttpError(RuntimeError):
     def __init__(self, status: int, reason: str, body_preview: str):
@@ -24,7 +105,7 @@ class ExoHttpError(RuntimeError):
 
 
 class ExoClient:
-    def __init__(self, host: str, port: int, timeout_s: float = 600.0):
+    def __init__(self, host: str, port: int, timeout_s: float = 7200.0):
         self.host = host
         self.port = port
         self.timeout_s = timeout_s
@@ -180,14 +261,7 @@ def parse_int_list(values: list[str]) -> list[int]:
             part = part.strip()
             if part:
                 items.append(int(part))
-
-    seen: set[int] = set()
-    out: list[int] = []
-    for x in items:
-        if x not in seen:
-            out.append(x)
-            seen.add(x)
-    return out
+    return items
 
 
 def resolve_model_short_id(client: ExoClient, model_arg: str) -> tuple[str, str]:
@@ -240,7 +314,11 @@ def run_one_completion(
 
     stats = out.get("generation_stats")
 
-    preview = (out.get("choices") or [{}])[0]["message"]["content"][:200]
+    # Extract preview, handling None content (common for thinking models)
+    choices = out.get("choices") or [{}]
+    message = choices[0].get("message", {}) if choices else {}
+    content = message.get("content") or ""
+    preview = content[:200] if content else ""
 
     return {
         "elapsed_s": elapsed,
@@ -277,12 +355,29 @@ class PromptSizer:
                 f"Target ({target}) is smaller than template overhead ({self.base_tokens})."
             )
 
-        content = ""
-        tok = self.count_fn(content)
+        # Estimate tokens per atom using a sample
+        sample_count = 100
+        sample_content = self.atom * sample_count
+        sample_tokens = self.count_fn(sample_content) - self.base_tokens
+        tokens_per_atom = sample_tokens / sample_count
 
-        while tok < target:
-            content += self.atom
-            tok = self.count_fn(content)
+        # Estimate starting point
+        needed_tokens = target - self.base_tokens
+        estimated_atoms = int(needed_tokens / tokens_per_atom)
+
+        # Binary search to find exact atom count
+        low, high = 0, estimated_atoms * 2 + 100
+        while low < high:
+            mid = (low + high) // 2
+            tok = self.count_fn(self.atom * mid)
+            if tok < target:
+                low = mid + 1
+            else:
+                high = mid
+
+        content = self.atom * low
+        tok = self.count_fn(content)
+        logger.info(f"{tok=}")
 
         if tok != target:
             raise RuntimeError(
@@ -348,7 +443,7 @@ def main() -> int:
         help="Warmup runs per placement (uses first pp/tg).",
     )
     ap.add_argument(
-        "--timeout", type=float, default=600.0, help="HTTP timeout (seconds)."
+        "--timeout", type=float, default=7200.0, help="HTTP timeout (seconds)."
     )
     ap.add_argument(
         "--json-out",
@@ -358,6 +453,11 @@ def main() -> int:
     ap.add_argument(
         "--dry-run", action="store_true", help="List selected placements and exit."
     )
+    ap.add_argument(
+        "--all-combinations",
+        action="store_true",
+        help="Force all pp×tg combinations (cartesian product) even when lists have equal length.",
+    )
     args = ap.parse_args()
 
     pp_list = parse_int_list(args.pp)
@@ -369,6 +469,15 @@ def main() -> int:
         logger.error("--repeat must be >= 1")
         return 2
 
+    # Log pairing mode
+    use_combinations = args.all_combinations or len(pp_list) != len(tg_list)
+    if use_combinations:
+        logger.info(
+            f"pp/tg mode: combinations (product) - {len(pp_list) * len(tg_list)} pairs"
+        )
+    else:
+        logger.info(f"pp/tg mode: tandem (zip) - {len(pp_list)} pairs")
+
     client = ExoClient(args.host, args.port, timeout_s=args.timeout)
     short_id, full_model_id = resolve_model_short_id(client, args.model)
 
@@ -377,10 +486,7 @@ def main() -> int:
     )
     previews = previews_resp.get("previews") or []
 
-    tokenizer = AutoTokenizer.from_pretrained(
-        full_model_id,
-        trust_remote_code=True,
-    )
+    tokenizer = load_tokenizer_for_bench(full_model_id)
     if tokenizer is None:
         raise RuntimeError("[exo-bench] tokenizer load failed")
 
@@ -486,60 +592,55 @@ def main() -> int:
                 )
                 logger.debug(f"  warmup {i + 1}/{args.warmup} done")
 
-            for pp in pp_list:
-                # if (
-                #     pp * n_nodes > 2048
-                #     and "ring" in instance_meta.lower()
-                #     and "tensor" in sharding.lower()
-                # ):
-                #     model_card = MODEL_CARDS[short_id]
-                #     if model_card.metadata.storage_size > Memory.from_gb(10):
-                #         logger.info(
-                #             f"Skipping tensor ring as this is too slow for model of size {model_card.metadata.storage_size} on {n_nodes=}"
-                #         )
-                #         continue
-                for tg in tg_list:
-                    runs: list[dict[str, Any]] = []
-                    for r in range(args.repeat):
-                        time.sleep(3)
-                        try:
-                            row, actual_pp_tokens = run_one_completion(
-                                client, full_model_id, pp, tg, prompt_sizer
-                            )
-                        except Exception as e:
-                            logger.error(e)
-                            continue
-                        row.update(
-                            {
-                                "model_short_id": short_id,
-                                "model_id": full_model_id,
-                                "placement_sharding": sharding,
-                                "placement_instance_meta": instance_meta,
-                                "placement_nodes": n_nodes,
-                                "instance_id": instance_id,
-                                "pp_tokens": actual_pp_tokens,
-                                "tg": tg,
-                                "repeat_index": r,
-                            }
-                        )
-                        runs.append(row)
-                        all_rows.append(row)
+            # If pp and tg lists have same length, run in tandem (zip)
+            # Otherwise (or if --all-combinations), run all combinations (cartesian product)
+            if use_combinations:
+                pp_tg_pairs = list(itertools.product(pp_list, tg_list))
+            else:
+                pp_tg_pairs = list(zip(pp_list, tg_list, strict=True))
 
-                    if runs:
-                        prompt_tps = mean(x["stats"]["prompt_tps"] for x in runs)
-                        gen_tps = mean(x["stats"]["generation_tps"] for x in runs)
-                        ptok = mean(x["stats"]["prompt_tokens"] for x in runs)
-                        gtok = mean(x["stats"]["generation_tokens"] for x in runs)
-                        peak = mean(
-                            x["stats"]["peak_memory_usage"]["inBytes"] for x in runs
+            for pp, tg in pp_tg_pairs:
+                runs: list[dict[str, Any]] = []
+                for r in range(args.repeat):
+                    time.sleep(3)
+                    try:
+                        row, actual_pp_tokens = run_one_completion(
+                            client, full_model_id, pp, tg, prompt_sizer
                         )
+                    except Exception as e:
+                        logger.error(e)
+                        continue
+                    row.update(
+                        {
+                            "model_short_id": short_id,
+                            "model_id": full_model_id,
+                            "placement_sharding": sharding,
+                            "placement_instance_meta": instance_meta,
+                            "placement_nodes": n_nodes,
+                            "instance_id": instance_id,
+                            "pp_tokens": actual_pp_tokens,
+                            "tg": tg,
+                            "repeat_index": r,
+                        }
+                    )
+                    runs.append(row)
+                    all_rows.append(row)
 
-                        logger.info(
-                            f"prompt_tps={prompt_tps:.2f} gen_tps={gen_tps:.2f}    "
-                            f"prompt_tokens={ptok} gen_tokens={gtok}    "
-                            f"peak_memory={format_peak_memory(peak)}\n"
-                        )
-                    time.sleep(2)
+                if runs:
+                    prompt_tps = mean(x["stats"]["prompt_tps"] for x in runs)
+                    gen_tps = mean(x["stats"]["generation_tps"] for x in runs)
+                    ptok = mean(x["stats"]["prompt_tokens"] for x in runs)
+                    gtok = mean(x["stats"]["generation_tokens"] for x in runs)
+                    peak = mean(
+                        x["stats"]["peak_memory_usage"]["inBytes"] for x in runs
+                    )
+
+                    logger.info(
+                        f"prompt_tps={prompt_tps:.2f} gen_tps={gen_tps:.2f}    "
+                        f"prompt_tokens={ptok} gen_tokens={gtok}    "
+                        f"peak_memory={format_peak_memory(peak)}\n"
+                    )
+                time.sleep(2)
         finally:
             try:
                 client.request_json("DELETE", f"/instance/{instance_id}")

dashboard/parts.nix

@@ -3,6 +3,61 @@
   perSystem =
     { pkgs, lib, ... }:
     let
+      # Filter source to ONLY include package.json and package-lock.json
+      # This ensures prettier-svelte only rebuilds when lockfiles change
+      dashboardLockfileSrc = lib.cleanSourceWith {
+        src = inputs.self;
+        filter =
+          path: type:
+          let
+            baseName = builtins.baseNameOf path;
+            isDashboardDir = baseName == "dashboard" && type == "directory";
+            isPackageFile =
+              (lib.hasInfix "/dashboard/" path || lib.hasSuffix "/dashboard" (builtins.dirOf path))
+              && (baseName == "package.json" || baseName == "package-lock.json");
+          in
+          isDashboardDir || isPackageFile;
+      };
+
+      # Stub source with lockfiles and minimal files for build to succeed
+      # This allows prettier-svelte to avoid rebuilding when dashboard source changes
+      dashboardStubSrc = pkgs.runCommand "dashboard-stub-src" { } ''
+        mkdir -p $out
+        cp ${dashboardLockfileSrc}/dashboard/package.json $out/
+        cp ${dashboardLockfileSrc}/dashboard/package-lock.json $out/
+        # Minimal files so vite build succeeds (produces empty output)
+        echo '<!DOCTYPE html><html><head></head><body></body></html>' > $out/index.html
+        mkdir -p $out/src
+        touch $out/src/app.html
+      '';
+
+      # Deps-only build using stub source (for prettier-svelte)
+      # Only rebuilds when package.json or package-lock.json change
+      dashboardDeps = inputs.dream2nix.lib.evalModules {
+        packageSets.nixpkgs = pkgs;
+        modules = [
+          ./dashboard.nix
+          {
+            paths.projectRoot = inputs.self;
+            paths.projectRootFile = "flake.nix";
+            paths.package = inputs.self + "/dashboard";
+          }
+          {
+            deps.dashboardSrc = lib.mkForce dashboardStubSrc;
+          }
+          # Override build phases to skip the actual build - just need node_modules
+          {
+            mkDerivation = {
+              buildPhase = lib.mkForce "true";
+              installPhase = lib.mkForce ''
+                runHook preInstall
+                runHook postInstall
+              '';
+            };
+          }
+        ];
+      };
+
       # Filter source to only include dashboard directory
       dashboardSrc = lib.cleanSourceWith {
         src = inputs.self;
@@ -42,11 +97,12 @@
       '';
 
       # Prettier with svelte plugin for treefmt
+      # Uses dashboardDeps instead of dashboardFull to avoid rebuilding on source changes
       packages.prettier-svelte = pkgs.writeShellScriptBin "prettier-svelte" ''
-        export NODE_PATH="${dashboardFull}/lib/node_modules/exo-dashboard/node_modules"
+        export NODE_PATH="${dashboardDeps}/lib/node_modules/exo-dashboard/node_modules"
         exec ${pkgs.nodejs}/bin/node \
-          ${dashboardFull}/lib/node_modules/exo-dashboard/node_modules/prettier/bin/prettier.cjs \
-          --plugin "${dashboardFull}/lib/node_modules/exo-dashboard/node_modules/prettier-plugin-svelte/plugin.js" \
+          ${dashboardDeps}/lib/node_modules/exo-dashboard/node_modules/prettier/bin/prettier.cjs \
+          --plugin "${dashboardDeps}/lib/node_modules/exo-dashboard/node_modules/prettier-plugin-svelte/plugin.js" \
           "$@"
       '';
     };

dashboard/src/lib/components/ChatForm.svelte

@@ -89,7 +89,10 @@
 
   const isImageModel = $derived(() => {
     if (!currentModel) return false;
-    return modelSupportsTextToImage(currentModel);
+    return (
+      modelSupportsTextToImage(currentModel) ||
+      modelSupportsImageEditing(currentModel)
+    );
   });
 
   const isEditOnlyWithoutImage = $derived(
@@ -646,6 +649,23 @@
             </svg>
             <span>EDIT</span>
           </span>
+        {:else if isEditOnlyWithoutImage}
+          <span class="inline-flex items-center gap-1.5">
+            <svg
+              class="w-3.5 h-3.5"
+              fill="none"
+              viewBox="0 0 24 24"
+              stroke="currentColor"
+              stroke-width="2"
+            >
+              <path
+                stroke-linecap="round"
+                stroke-linejoin="round"
+                d="M11 5H6a2 2 0 00-2 2v11a2 2 0 002 2h11a2 2 0 002-2v-5m-1.414-9.414a2 2 0 112.828 2.828L11.828 15H9v-2.828l8.586-8.586z"
+              />
+            </svg>
+            <span>EDIT</span>
+          </span>
         {:else if isImageModel()}
           <span class="inline-flex items-center gap-1.5">
             <svg

dashboard/src/lib/components/ImageParamsPanel.svelte

@@ -110,6 +110,36 @@
     setImageGenerationParams({ negativePrompt: value || null });
   }
 
+  function handleNumImagesChange(event: Event) {
+    const input = event.target as HTMLInputElement;
+    const value = input.value.trim();
+    if (value === "") {
+      setImageGenerationParams({ numImages: 1 });
+    } else {
+      const num = parseInt(value, 10);
+      if (!isNaN(num) && num >= 1) {
+        setImageGenerationParams({ numImages: num });
+      }
+    }
+  }
+
+  function handleStreamChange(enabled: boolean) {
+    setImageGenerationParams({ stream: enabled });
+  }
+
+  function handlePartialImagesChange(event: Event) {
+    const input = event.target as HTMLInputElement;
+    const value = input.value.trim();
+    if (value === "") {
+      setImageGenerationParams({ partialImages: 0 });
+    } else {
+      const num = parseInt(value, 10);
+      if (!isNaN(num) && num >= 0) {
+        setImageGenerationParams({ partialImages: num });
+      }
+    }
+  }
+
   function clearSteps() {
     setImageGenerationParams({ numInferenceSteps: null });
   }
@@ -134,90 +164,92 @@
 <div class="border-b border-exo-medium-gray/30 px-3 py-2">
   <!-- Basic params row -->
   <div class="flex items-center gap-3 flex-wrap">
-    <!-- Size -->
-    <div class="flex items-center gap-1.5">
-      <span class="text-xs text-exo-light-gray uppercase tracking-wider"
-        >SIZE:</span
-      >
-      <div class="relative">
-        <button
-          bind:this={sizeButtonRef}
-          type="button"
-          onclick={() => (isSizeDropdownOpen = !isSizeDropdownOpen)}
-          class="bg-exo-medium-gray/50 border border-exo-yellow/30 rounded pl-2 pr-6 py-1 text-xs font-mono text-exo-yellow cursor-pointer transition-all duration-200 hover:border-exo-yellow/50 focus:outline-none focus:border-exo-yellow/70 {isSizeDropdownOpen
-            ? 'border-exo-yellow/70'
-            : ''}"
+    <!-- Size (hidden in edit mode - output size comes from input image) -->
+    {#if !isEditMode}
+      <div class="flex items-center gap-1.5">
+        <span class="text-xs text-exo-light-gray uppercase tracking-wider"
+          >SIZE:</span
         >
-          {params.size}
-        </button>
-        <div
-          class="absolute right-1.5 top-1/2 -translate-y-1/2 pointer-events-none transition-transform duration-200 {isSizeDropdownOpen
-            ? 'rotate-180'
-            : ''}"
-        >
-          <svg
-            class="w-3 h-3 text-exo-yellow/60"
-            fill="none"
-            viewBox="0 0 24 24"
-            stroke="currentColor"
+        <div class="relative">
+          <button
+            bind:this={sizeButtonRef}
+            type="button"
+            onclick={() => (isSizeDropdownOpen = !isSizeDropdownOpen)}
+            class="bg-exo-medium-gray/50 border border-exo-yellow/30 rounded pl-2 pr-6 py-1 text-xs font-mono text-exo-yellow cursor-pointer transition-all duration-200 hover:border-exo-yellow/50 focus:outline-none focus:border-exo-yellow/70 {isSizeDropdownOpen
+              ? 'border-exo-yellow/70'
+              : ''}"
           >
-            <path
-              stroke-linecap="round"
-              stroke-linejoin="round"
-              stroke-width="2"
-              d="M19 9l-7 7-7-7"
-            />
-          </svg>
-        </div>
-      </div>
-
-      {#if isSizeDropdownOpen}
-        <!-- Backdrop to close dropdown -->
-        <button
-          type="button"
-          class="fixed inset-0 z-[9998] cursor-default"
-          onclick={() => (isSizeDropdownOpen = 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 min-w-max"
-          style="bottom: calc(100vh - {sizeDropdownPosition()
-            .top}px + 4px); left: {sizeDropdownPosition().left}px;"
-        >
-          <div class="py-1">
-            {#each sizeOptions as size}
-              <button
-                type="button"
-                onclick={() => selectSize(size)}
-                class="w-full px-3 py-1.5 text-left text-xs font-mono tracking-wide transition-colors duration-100 flex items-center gap-2 {params.size ===
-                size
-                  ? 'bg-transparent text-exo-yellow'
-                  : 'text-exo-light-gray hover:text-exo-yellow'}"
-              >
-                {#if params.size === size}
-                  <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}
-                <span>{size}</span>
-              </button>
-            {/each}
+            {params.size}
+          </button>
+          <div
+            class="absolute right-1.5 top-1/2 -translate-y-1/2 pointer-events-none transition-transform duration-200 {isSizeDropdownOpen
+              ? 'rotate-180'
+              : ''}"
+          >
+            <svg
+              class="w-3 h-3 text-exo-yellow/60"
+              fill="none"
+              viewBox="0 0 24 24"
+              stroke="currentColor"
+            >
+              <path
+                stroke-linecap="round"
+                stroke-linejoin="round"
+                stroke-width="2"
+                d="M19 9l-7 7-7-7"
+              />
+            </svg>
           </div>
         </div>
-      {/if}
-    </div>
+
+        {#if isSizeDropdownOpen}
+          <!-- Backdrop to close dropdown -->
+          <button
+            type="button"
+            class="fixed inset-0 z-[9998] cursor-default"
+            onclick={() => (isSizeDropdownOpen = 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 min-w-max"
+            style="bottom: calc(100vh - {sizeDropdownPosition()
+              .top}px + 4px); left: {sizeDropdownPosition().left}px;"
+          >
+            <div class="py-1">
+              {#each sizeOptions as size}
+                <button
+                  type="button"
+                  onclick={() => selectSize(size)}
+                  class="w-full px-3 py-1.5 text-left text-xs font-mono tracking-wide transition-colors duration-100 flex items-center gap-2 {params.size ===
+                  size
+                    ? 'bg-transparent text-exo-yellow'
+                    : 'text-exo-light-gray hover:text-exo-yellow'}"
+                >
+                  {#if params.size === size}
+                    <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}
+                  <span>{size}</span>
+                </button>
+              {/each}
+            </div>
+          </div>
+        {/if}
+      </div>
+    {/if}
 
     <!-- Quality -->
     <div class="flex items-center gap-1.5">
@@ -325,6 +357,59 @@
       </div>
     </div>
 
+    <!-- Number of Images (not in edit mode) -->
+    {#if !isEditMode}
+      <div class="flex items-center gap-1.5">
+        <span class="text-xs text-exo-light-gray uppercase tracking-wider"
+          >IMAGES:</span
+        >
+        <input
+          type="number"
+          min="1"
+          value={params.numImages}
+          oninput={handleNumImagesChange}
+          class="w-12 bg-exo-medium-gray/50 border border-exo-yellow/30 rounded px-2 py-1 text-xs font-mono text-exo-yellow text-center transition-all duration-200 hover:border-exo-yellow/50 focus:outline-none focus:border-exo-yellow/70"
+        />
+      </div>
+    {/if}
+
+    <!-- Stream toggle -->
+    <div class="flex items-center gap-1.5">
+      <span class="text-xs text-exo-light-gray uppercase tracking-wider"
+        >STREAM:</span
+      >
+      <button
+        type="button"
+        onclick={() => handleStreamChange(!params.stream)}
+        class="w-8 h-4 rounded-full transition-all duration-200 cursor-pointer relative {params.stream
+          ? 'bg-exo-yellow'
+          : 'bg-exo-medium-gray/50 border border-exo-yellow/30'}"
+        title={params.stream ? "Streaming enabled" : "Streaming disabled"}
+      >
+        <div
+          class="absolute top-0.5 w-3 h-3 rounded-full transition-all duration-200 {params.stream
+            ? 'right-0.5 bg-exo-black'
+            : 'left-0.5 bg-exo-light-gray'}"
+        ></div>
+      </button>
+    </div>
+
+    <!-- Partial Images (only when streaming) -->
+    {#if params.stream}
+      <div class="flex items-center gap-1.5">
+        <span class="text-xs text-exo-light-gray uppercase tracking-wider"
+          >PARTIALS:</span
+        >
+        <input
+          type="number"
+          min="0"
+          value={params.partialImages}
+          oninput={handlePartialImagesChange}
+          class="w-12 bg-exo-medium-gray/50 border border-exo-yellow/30 rounded px-2 py-1 text-xs font-mono text-exo-yellow text-center transition-all duration-200 hover:border-exo-yellow/50 focus:outline-none focus:border-exo-yellow/70"
+        />
+      </div>
+    {/if}
+
     <!-- Input Fidelity (edit mode only) -->
     {#if isEditMode}
       <div class="flex items-center gap-1.5">

dashboard/src/lib/components/TopologyGraph.svelte

@@ -5,22 +5,32 @@
     topologyData,
     isTopologyMinimized,
     debugMode,
+    nodeThunderboltBridge,
     type NodeInfo,
   } from "$lib/stores/app.svelte";
 
   interface Props {
     class?: string;
     highlightedNodes?: Set<string>;
+    filteredNodes?: Set<string>;
+    onNodeClick?: (nodeId: string) => void;
   }
 
-  let { class: className = "", highlightedNodes = new Set() }: Props = $props();
+  let {
+    class: className = "",
+    highlightedNodes = new Set(),
+    filteredNodes = new Set(),
+    onNodeClick,
+  }: Props = $props();
 
   let svgContainer: SVGSVGElement | undefined = $state();
   let resizeObserver: ResizeObserver | undefined;
+  let hoveredNodeId = $state<string | null>(null);
 
   const isMinimized = $derived(isTopologyMinimized());
   const data = $derived(topologyData());
   const debugEnabled = $derived(debugMode());
+  const tbBridgeData = $derived(nodeThunderboltBridge());
 
   function getNodeLabel(nodeId: string): string {
     const node = data?.nodes?.[nodeId];
@@ -522,10 +532,72 @@
         }
       }
 
+      let iconBaseWidth = nodeRadius * 1.2;
+      let iconBaseHeight = nodeRadius * 1.0;
+      const clipPathId = `clip-${nodeInfo.id.replace(/[^a-zA-Z0-9]/g, "-")}`;
+
+      const modelLower = modelId.toLowerCase();
+
+      // Check node states for styling
+      const isHighlighted = highlightedNodes.has(nodeInfo.id);
+      const isInFilter =
+        filteredNodes.size > 0 && filteredNodes.has(nodeInfo.id);
+      const isFilteredOut =
+        filteredNodes.size > 0 && !filteredNodes.has(nodeInfo.id);
+      const isHovered = hoveredNodeId === nodeInfo.id && !isInFilter;
+
+      // Holographic wireframe colors - bright yellow for filter, subtle yellow for hover, grey for filtered out
+      const wireColor = isInFilter
+        ? "rgba(255,215,0,1)" // Bright yellow for filter selection
+        : isHovered
+          ? "rgba(255,215,0,0.7)" // Subtle yellow for hover
+          : isHighlighted
+            ? "rgba(255,215,0,0.9)" // Yellow for instance highlight
+            : isFilteredOut
+              ? "rgba(140,140,140,0.6)" // Grey for filtered out
+              : "rgba(179,179,179,0.8)"; // Default
+      const wireColorBright = "rgba(255,255,255,0.9)";
+      const fillColor = isInFilter
+        ? "rgba(255,215,0,0.25)"
+        : isHovered
+          ? "rgba(255,215,0,0.12)"
+          : isHighlighted
+            ? "rgba(255,215,0,0.15)"
+            : "rgba(255,215,0,0.08)";
+      const strokeWidth = isInFilter
+        ? 3
+        : isHovered
+          ? 2
+          : isHighlighted
+            ? 2.5
+            : 1.5;
+      const screenFill = "rgba(0,20,40,0.9)";
+      const glowColor = "rgba(255,215,0,0.3)";
+
       const nodeG = nodesGroup
         .append("g")
         .attr("class", "graph-node")
-        .style("cursor", "pointer");
+        .style("cursor", onNodeClick ? "pointer" : "default")
+        .style("opacity", isFilteredOut ? 0.5 : 1);
+
+      // Add click and hover handlers - hover just updates state, styling is applied during render
+      nodeG
+        .on("click", (event: MouseEvent) => {
+          if (onNodeClick) {
+            event.stopPropagation();
+            onNodeClick(nodeInfo.id);
+          }
+        })
+        .on("mouseenter", () => {
+          if (onNodeClick) {
+            hoveredNodeId = nodeInfo.id;
+          }
+        })
+        .on("mouseleave", () => {
+          if (hoveredNodeId === nodeInfo.id) {
+            hoveredNodeId = null;
+          }
+        });
 
       // Add tooltip
       nodeG
@@ -534,27 +606,6 @@
           `${friendlyName}\nID: ${nodeInfo.id.slice(-8)}\nMemory: ${formatBytes(ramUsed)}/${formatBytes(ramTotal)}`,
         );
 
-      let iconBaseWidth = nodeRadius * 1.2;
-      let iconBaseHeight = nodeRadius * 1.0;
-      const clipPathId = `clip-${nodeInfo.id.replace(/[^a-zA-Z0-9]/g, "-")}`;
-
-      const modelLower = modelId.toLowerCase();
-
-      // Check if this node should be highlighted (from hovered instance)
-      const isHighlighted = highlightedNodes.has(nodeInfo.id);
-
-      // Holographic wireframe colors - yellow border when highlighted
-      const wireColor = isHighlighted
-        ? "rgba(255,215,0,0.9)"
-        : "rgba(179,179,179,0.8)";
-      const wireColorBright = "rgba(255,255,255,0.9)";
-      const fillColor = isHighlighted
-        ? "rgba(255,215,0,0.15)"
-        : "rgba(255,215,0,0.08)";
-      const strokeWidth = isHighlighted ? 2.5 : 1.5;
-      const screenFill = "rgba(0,20,40,0.9)";
-      const glowColor = "rgba(255,215,0,0.3)";
-
       if (modelLower === "mac studio") {
         // Mac Studio - classic cube with memory fill
         iconBaseWidth = nodeRadius * 1.25;
@@ -579,6 +630,7 @@
         // Main body (uniform color)
         nodeG
           .append("rect")
+          .attr("class", "node-outline")
           .attr("x", x)
           .attr("y", y)
           .attr("width", iconBaseWidth)
@@ -661,6 +713,7 @@
         // Main body (uniform color)
         nodeG
           .append("rect")
+          .attr("class", "node-outline")
           .attr("x", x)
           .attr("y", y)
           .attr("width", iconBaseWidth)
@@ -738,6 +791,7 @@
         // Screen outer frame
         nodeG
           .append("rect")
+          .attr("class", "node-outline")
           .attr("x", screenX)
           .attr("y", y)
           .attr("width", screenWidth)
@@ -846,6 +900,7 @@
         // Main shape
         nodeG
           .append("polygon")
+          .attr("class", "node-outline")
           .attr("points", hexPoints)
           .attr("fill", fillColor)
           .attr("stroke", wireColor)
@@ -1064,11 +1119,41 @@
           .attr("fill", "rgba(179,179,179,0.7)")
           .text(` (${ramUsagePercent.toFixed(0)}%)`);
       }
+
+      // Debug mode: Show TB bridge status
+      if (debugEnabled) {
+        const tbStatus = tbBridgeData[nodeInfo.id];
+        if (tbStatus) {
+          const tbY =
+            nodeInfo.y +
+            iconBaseHeight / 2 +
+            (showFullLabels ? 32 : showCompactLabels ? 26 : 22);
+          const tbFontSize = showFullLabels ? 9 : 7;
+          const tbColor = tbStatus.enabled
+            ? "rgba(234,179,8,0.9)"
+            : "rgba(100,100,100,0.7)";
+          const tbText = tbStatus.enabled ? "TB:ON" : "TB:OFF";
+          nodeG
+            .append("text")
+            .attr("x", nodeInfo.x)
+            .attr("y", tbY)
+            .attr("text-anchor", "middle")
+            .attr("fill", tbColor)
+            .attr("font-size", tbFontSize)
+            .attr("font-family", "SF Mono, Monaco, monospace")
+            .text(tbText);
+        }
+      }
     });
   }
 
   $effect(() => {
-    if (data) {
+    // Track all reactive dependencies that affect rendering
+    const _data = data;
+    const _hoveredNodeId = hoveredNodeId;
+    const _filteredNodes = filteredNodes;
+    const _highlightedNodes = highlightedNodes;
+    if (_data) {
       renderGraph();
     }
   });
@@ -1091,12 +1176,8 @@
 
 <style>
   :global(.graph-node) {
-    transition:
-      transform 0.2s ease,
-      opacity 0.2s ease;
-  }
-  :global(.graph-node:hover) {
-    filter: brightness(1.1);
+    /* Only transition opacity for filtered-out nodes, no transition on hover stroke changes */
+    transition: opacity 0.2s ease;
   }
   :global(.graph-link) {
     stroke: var(--exo-light-gray, #b3b3b3);

dashboard/src/routes/+page.svelte

@@ -19,6 +19,9 @@
     selectedPreviewModelId,
     isLoadingPreviews,
     selectPreviewModel,
+    togglePreviewNodeFilter,
+    clearPreviewNodeFilter,
+    previewNodeFilter,
     createConversation,
     setSelectedChatModel,
     selectedChatModel,
@@ -28,6 +31,8 @@
     toggleTopologyOnlyMode,
     chatSidebarVisible,
     toggleChatSidebarVisible,
+    thunderboltBridgeCycles,
+    nodeThunderboltBridge,
     type DownloadProgress,
     type PlacementPreview,
   } from "$lib/stores/app.svelte";
@@ -49,6 +54,41 @@
   const debugEnabled = $derived(debugMode());
   const topologyOnlyEnabled = $derived(topologyOnlyMode());
   const sidebarVisible = $derived(chatSidebarVisible());
+  const tbBridgeCycles = $derived(thunderboltBridgeCycles());
+  const tbBridgeData = $derived(nodeThunderboltBridge());
+  const nodeFilter = $derived(previewNodeFilter());
+
+  // Helper to get friendly node name from node ID
+  function getNodeName(nodeId: string): string {
+    const node = data?.nodes?.[nodeId];
+    return node?.friendly_name || nodeId.slice(0, 8) + "...";
+  }
+
+  // Helper to get the thunderbolt bridge service name from a cycle
+  function getTbBridgeServiceName(cycle: string[]): string {
+    // Try to find service name from any node in the cycle
+    for (const nodeId of cycle) {
+      const nodeData = tbBridgeData?.[nodeId];
+      if (nodeData?.serviceName) {
+        return nodeData.serviceName;
+      }
+    }
+    return "Thunderbolt Bridge"; // Fallback if no service name found
+  }
+
+  // Copy to clipboard state and function
+  let copiedCommand = $state(false);
+  async function copyToClipboard(text: string) {
+    try {
+      await navigator.clipboard.writeText(text);
+      copiedCommand = true;
+      setTimeout(() => {
+        copiedCommand = false;
+      }, 2000);
+    } catch (err) {
+      console.error("Failed to copy:", err);
+    }
+  }
 
   let mounted = $state(false);
 
@@ -90,6 +130,15 @@
       model.tasks.includes("ImageToImage")
     );
   }
+
+  // Helper to check if a model supports image editing
+  function modelSupportsImageEditing(modelId: string): boolean {
+    const model = models.find(
+      (m) => m.id === modelId || m.hugging_face_id === modelId,
+    );
+    if (!model?.tasks) return false;
+    return model.tasks.includes("ImageToImage");
+  }
   let selectedSharding = $state<"Pipeline" | "Tensor">("Pipeline");
   type InstanceMeta = "MlxRing" | "MlxIbv" | "MlxJaccl";
 
@@ -181,6 +230,9 @@
   // Preview card hover state for highlighting nodes in topology
   let hoveredPreviewNodes = $state<Set<string>>(new Set());
 
+  // Computed: Check if filter is active (from store)
+  const isFilterActive = $derived(() => nodeFilter.size > 0);
+
   // Helper to unwrap tagged instance for hover highlighting
   function unwrapInstanceNodes(instanceWrapped: unknown): Set<string> {
     if (!instanceWrapped || typeof instanceWrapped !== "object")
@@ -732,6 +784,8 @@
     instanceWrapped: unknown,
   ): {
     isDownloading: boolean;
+    isFailed: boolean;
+    errorMessage: string | null;
     progress: DownloadProgress | null;
     statusText: string;
     perNode: Array<{
@@ -743,6 +797,8 @@
     if (!downloadsData || Object.keys(downloadsData).length === 0) {
       return {
         isDownloading: false,
+        isFailed: false,
+        errorMessage: null,
         progress: null,
         statusText: "RUNNING",
         perNode: [],
@@ -754,6 +810,8 @@
     if (!instance || typeof instance !== "object") {
       return {
         isDownloading: false,
+        isFailed: false,
+        errorMessage: null,
         progress: null,
         statusText: "PREPARING",
         perNode: [],
@@ -809,6 +867,26 @@
           downloadKind
         ] as Record<string, unknown>;
 
+        // Handle DownloadFailed - return immediately with error info
+        if (downloadKind === "DownloadFailed") {
+          const downloadModelId = extractModelIdFromDownload(downloadPayload);
+          if (
+            instanceModelId &&
+            downloadModelId &&
+            downloadModelId === instanceModelId
+          ) {
+            return {
+              isDownloading: false,
+              isFailed: true,
+              errorMessage:
+                (downloadPayload.errorMessage as string) || "Download failed",
+              progress: null,
+              statusText: "FAILED",
+              perNode: [],
+            };
+          }
+        }
+
         if (downloadKind !== "DownloadOngoing") continue;
         if (!downloadPayload) continue;
 
@@ -844,6 +922,8 @@
       const statusInfo = deriveInstanceStatus(instanceWrapped);
       return {
         isDownloading: false,
+        isFailed: statusInfo.statusText === "FAILED",
+        errorMessage: null,
         progress: null,
         statusText: statusInfo.statusText,
         perNode: [],
@@ -856,6 +936,8 @@
 
     return {
       isDownloading: true,
+      isFailed: false,
+      errorMessage: null,
       progress: {
         totalBytes,
         downloadedBytes,
@@ -1451,7 +1533,7 @@
 
   // Get ALL filtered previews based on current settings (matching minimum nodes)
   // Note: previewsData already contains previews for the selected model (fetched via API)
-  // We filter by sharding/instance type and min nodes, returning ALL eligible previews
+  // Backend handles node_ids filtering, we filter by sharding/instance type and min nodes
   const filteredPreviews = $derived(() => {
     if (!selectedModelId || previewsData.length === 0) return [];
 
@@ -1584,7 +1666,86 @@
           <TopologyGraph
             class="w-full h-full"
             highlightedNodes={highlightedNodes()}
+            filteredNodes={nodeFilter}
+            onNodeClick={togglePreviewNodeFilter}
           />
+
+          <!-- Thunderbolt Bridge Cycle Warning -->
+          {#if tbBridgeCycles.length > 0}
+            {@const cycle = tbBridgeCycles[0]}
+            {@const serviceName = getTbBridgeServiceName(cycle)}
+            {@const disableCmd = `sudo networksetup -setnetworkserviceenabled "${serviceName}" off`}
+            <div class="absolute top-4 left-4 group" role="alert">
+              <div
+                class="flex items-center gap-2 px-3 py-2 rounded border border-yellow-500/50 bg-yellow-500/10 backdrop-blur-sm cursor-help"
+              >
+                <svg
+                  class="w-5 h-5 text-yellow-400 flex-shrink-0"
+                  fill="none"
+                  viewBox="0 0 24 24"
+                  stroke="currentColor"
+                  stroke-width="2"
+                >
+                  <path
+                    stroke-linecap="round"
+                    stroke-linejoin="round"
+                    d="M12 9v2m0 4h.01m-6.938 4h13.856c1.54 0 2.502-1.667 1.732-3L13.732 4c-.77-1.333-2.694-1.333-3.464 0L3.34 16c-.77 1.333.192 3 1.732 3z"
+                  />
+                </svg>
+                <span class="text-sm font-mono text-yellow-200">
+                  THUNDERBOLT BRIDGE CYCLE DETECTED
+                </span>
+              </div>
+
+              <!-- Tooltip on hover -->
+              <div
+                class="absolute top-full left-0 mt-2 w-80 p-3 rounded border border-yellow-500/30 bg-exo-dark-gray/95 backdrop-blur-sm opacity-0 invisible group-hover:opacity-100 group-hover:visible transition-all duration-200 z-50 shadow-lg"
+              >
+                <p class="text-xs text-white/80 mb-2">
+                  A network routing cycle was detected between nodes connected
+                  via Thunderbolt Bridge. This can cause connectivity issues.
+                </p>
+                <p class="text-xs text-white/60 mb-2">
+                  <span class="text-yellow-300">Affected nodes:</span>
+                  {cycle.map(getNodeName).join(" → ")}
+                </p>
+                <p class="text-xs text-white/60 mb-1">
+                  <span class="text-yellow-300">To fix:</span> Disable the Thunderbolt
+                  Bridge on one of the affected nodes:
+                </p>
+                <button
+                  type="button"
+                  onclick={() => copyToClipboard(disableCmd)}
+                  class="w-full flex items-center gap-2 text-[10px] font-mono bg-exo-black/60 px-2 py-1.5 rounded text-exo-yellow break-all text-left hover:bg-exo-black/80 transition-colors cursor-pointer group/copy"
+                  title="Click to copy"
+                >
+                  <span class="flex-1">{disableCmd}</span>
+                  <svg
+                    class="w-3.5 h-3.5 flex-shrink-0 text-white/40 group-hover/copy:text-exo-yellow transition-colors"
+                    fill="none"
+                    viewBox="0 0 24 24"
+                    stroke="currentColor"
+                    stroke-width="2"
+                  >
+                    {#if copiedCommand}
+                      <path
+                        stroke-linecap="round"
+                        stroke-linejoin="round"
+                        d="M5 13l4 4L19 7"
+                      />
+                    {:else}
+                      <path
+                        stroke-linecap="round"
+                        stroke-linejoin="round"
+                        d="M8 16H6a2 2 0 01-2-2V6a2 2 0 012-2h8a2 2 0 012 2v2m-6 12h8a2 2 0 002-2v-8a2 2 0 00-2-2h-8a2 2 0 00-2 2v8a2 2 0 002 2z"
+                      />
+                    {/if}
+                  </svg>
+                </button>
+              </div>
+            </div>
+          {/if}
+
           <!-- Exit topology-only mode button -->
           <button
             type="button"
@@ -1624,7 +1785,111 @@
             <TopologyGraph
               class="w-full h-full"
               highlightedNodes={highlightedNodes()}
+              filteredNodes={nodeFilter}
+              onNodeClick={togglePreviewNodeFilter}
             />
+
+            <!-- Thunderbolt Bridge Cycle Warning -->
+            {#if tbBridgeCycles.length > 0}
+              {@const cycle = tbBridgeCycles[0]}
+              {@const serviceName = getTbBridgeServiceName(cycle)}
+              {@const disableCmd = `sudo networksetup -setnetworkserviceenabled "${serviceName}" off`}
+              <div class="absolute top-4 left-4 group" role="alert">
+                <div
+                  class="flex items-center gap-2 px-3 py-2 rounded border border-yellow-500/50 bg-yellow-500/10 backdrop-blur-sm cursor-help"
+                >
+                  <svg
+                    class="w-5 h-5 text-yellow-400 flex-shrink-0"
+                    fill="none"
+                    viewBox="0 0 24 24"
+                    stroke="currentColor"
+                    stroke-width="2"
+                  >
+                    <path
+                      stroke-linecap="round"
+                      stroke-linejoin="round"
+                      d="M12 9v2m0 4h.01m-6.938 4h13.856c1.54 0 2.502-1.667 1.732-3L13.732 4c-.77-1.333-2.694-1.333-3.464 0L3.34 16c-.77 1.333.192 3 1.732 3z"
+                    />
+                  </svg>
+                  <span class="text-sm font-mono text-yellow-200">
+                    THUNDERBOLT BRIDGE CYCLE DETECTED
+                  </span>
+                </div>
+
+                <!-- Tooltip on hover -->
+                <div
+                  class="absolute top-full left-0 mt-2 w-80 p-3 rounded border border-yellow-500/30 bg-exo-dark-gray/95 backdrop-blur-sm opacity-0 invisible group-hover:opacity-100 group-hover:visible transition-all duration-200 z-50 shadow-lg"
+                >
+                  <p class="text-xs text-white/80 mb-2">
+                    A network routing cycle was detected between nodes connected
+                    via Thunderbolt Bridge. This can cause connectivity issues.
+                  </p>
+                  <p class="text-xs text-white/60 mb-2">
+                    <span class="text-yellow-300">Affected nodes:</span>
+                    {cycle.map(getNodeName).join(" → ")}
+                  </p>
+                  <p class="text-xs text-white/60 mb-1">
+                    <span class="text-yellow-300">To fix:</span> Disable the Thunderbolt
+                    Bridge on one of the affected nodes:
+                  </p>
+                  <button
+                    type="button"
+                    onclick={() => copyToClipboard(disableCmd)}
+                    class="w-full flex items-center gap-2 text-[10px] font-mono bg-exo-black/60 px-2 py-1.5 rounded text-exo-yellow break-all text-left hover:bg-exo-black/80 transition-colors cursor-pointer group/copy"
+                    title="Click to copy"
+                  >
+                    <span class="flex-1">{disableCmd}</span>
+                    <svg
+                      class="w-3.5 h-3.5 flex-shrink-0 text-white/40 group-hover/copy:text-exo-yellow transition-colors"
+                      fill="none"
+                      viewBox="0 0 24 24"
+                      stroke="currentColor"
+                      stroke-width="2"
+                    >
+                      {#if copiedCommand}
+                        <path
+                          stroke-linecap="round"
+                          stroke-linejoin="round"
+                          d="M5 13l4 4L19 7"
+                        />
+                      {:else}
+                        <path
+                          stroke-linecap="round"
+                          stroke-linejoin="round"
+                          d="M8 16H6a2 2 0 01-2-2V6a2 2 0 012-2h8a2 2 0 012 2v2m-6 12h8a2 2 0 002-2v-8a2 2 0 00-2-2h-8a2 2 0 00-2 2v8a2 2 0 002 2z"
+                        />
+                      {/if}
+                    </svg>
+                  </button>
+                </div>
+              </div>
+            {/if}
+
+            <!-- Node Filter Indicator (top-right corner) -->
+            {#if isFilterActive()}
+              <button
+                onclick={clearPreviewNodeFilter}
+                class="absolute top-2 right-2 flex items-center gap-1.5 px-2 py-1 bg-exo-dark-gray/80 border border-exo-yellow/40 rounded text-exo-yellow hover:border-exo-yellow/60 transition-colors cursor-pointer backdrop-blur-sm"
+                title="Clear filter"
+              >
+                <span class="text-[10px] font-mono tracking-wider">
+                  FILTER: {nodeFilter.size}
+                </span>
+                <svg
+                  class="w-3 h-3"
+                  fill="none"
+                  viewBox="0 0 24 24"
+                  stroke="currentColor"
+                  stroke-width="2"
+                >
+                  <path
+                    stroke-linecap="round"
+                    stroke-linejoin="round"
+                    d="M6 18L18 6M6 6l12 12"
+                  />
+                </svg>
+              </button>
+            {/if}
           </div>
 
           <!-- Chat Input - Below topology -->
@@ -2061,6 +2326,13 @@
                           >
                             {downloadInfo.statusText}
                           </div>
+                          {#if downloadInfo.isFailed && downloadInfo.errorMessage}
+                            <div
+                              class="text-xs text-red-400/80 font-mono mt-1 break-words"
+                            >
+                              {downloadInfo.errorMessage}
+                            </div>
+                          {/if}
                         {/if}
                       </div>
                     </div>
@@ -2106,6 +2378,9 @@
                     {@const isImageModel = modelSupportsImageGeneration(
                       foundModel.id,
                     )}
+                    {@const isImageEditModel = modelSupportsImageEditing(
+                      foundModel.id,
+                    )}
                     <span
                       class="flex items-center justify-between gap-2 w-full pr-4"
                     >
@@ -2132,6 +2407,22 @@
                             <polyline points="21 15 16 10 5 21" />
                           </svg>
                         {/if}
+                        {#if isImageEditModel}
+                          <svg
+                            class="w-4 h-4 flex-shrink-0 text-exo-yellow"
+                            fill="none"
+                            viewBox="0 0 24 24"
+                            stroke="currentColor"
+                            stroke-width="2"
+                          >
+                            <path
+                              d="M11 4H4a2 2 0 0 0-2 2v14a2 2 0 0 0 2 2h14a2 2 0 0 0 2-2v-7"
+                            />
+                            <path
+                              d="M18.5 2.5a2.121 2.121 0 0 1 3 3L12 15l-4 1 1-4 9.5-9.5z"
+                            />
+                          </svg>
+                        {/if}
                         <span class="truncate"
                           >{foundModel.name || foundModel.id}</span
                         >
@@ -2204,6 +2495,9 @@
                       {@const isImageModel = modelSupportsImageGeneration(
                         model.id,
                       )}
+                      {@const isImageEditModel = modelSupportsImageEditing(
+                        model.id,
+                      )}
                       <button
                         type="button"
                         onclick={() => {
@@ -2244,6 +2538,23 @@
                               <polyline points="21 15 16 10 5 21" />
                             </svg>
                           {/if}
+                          {#if isImageEditModel}
+                            <svg
+                              class="w-4 h-4 flex-shrink-0 text-exo-yellow"
+                              fill="none"
+                              viewBox="0 0 24 24"
+                              stroke="currentColor"
+                              stroke-width="2"
+                              aria-label="Image editing model"
+                            >
+                              <path
+                                d="M11 4H4a2 2 0 0 0-2 2v14a2 2 0 0 0 2 2h14a2 2 0 0 0 2-2v-7"
+                              />
+                              <path
+                                d="M18.5 2.5a2.121 2.121 0 0 1 3 3L12 15l-4 1 1-4 9.5-9.5z"
+                              />
+                            </svg>
+                          {/if}
                           <span class="truncate">{model.name || model.id}</span>
                         </span>
                         <span
@@ -2564,7 +2875,36 @@
               <div
                 class="relative aspect-square bg-exo-dark-gray rounded-lg overflow-hidden"
               >
-                <TopologyGraph highlightedNodes={highlightedNodes()} />
+                <TopologyGraph
+                  highlightedNodes={highlightedNodes()}
+                  filteredNodes={nodeFilter}
+                  onNodeClick={togglePreviewNodeFilter}
+                />
+
+                <!-- Thunderbolt Bridge Cycle Warning (compact) -->
+                {#if tbBridgeCycles.length > 0}
+                  <div
+                    class="absolute top-2 left-2 flex items-center gap-1.5 px-2 py-1 rounded border border-yellow-500/50 bg-yellow-500/10 backdrop-blur-sm"
+                    title="Thunderbolt Bridge cycle detected - click to view details"
+                  >
+                    <svg
+                      class="w-3.5 h-3.5 text-yellow-400"
+                      fill="none"
+                      viewBox="0 0 24 24"
+                      stroke="currentColor"
+                      stroke-width="2"
+                    >
+                      <path
+                        stroke-linecap="round"
+                        stroke-linejoin="round"
+                        d="M12 9v2m0 4h.01m-6.938 4h13.856c1.54 0 2.502-1.667 1.732-3L13.732 4c-.77-1.333-2.694-1.333-3.464 0L3.34 16c-.77 1.333.192 3 1.732 3z"
+                      />
+                    </svg>
+                    <span class="text-[10px] font-mono text-yellow-200"
+                      >TB CYCLE</span
+                    >
+                  </div>
+                {/if}
               </div>
             </button>
 
@@ -2993,6 +3333,13 @@
                             >
                               {downloadInfo.statusText}
                             </div>
+                            {#if downloadInfo.isFailed && downloadInfo.errorMessage}
+                              <div
+                                class="text-xs text-red-400/80 font-mono mt-1 break-words"
+                              >
+                                {downloadInfo.errorMessage}
+                              </div>
+                            {/if}
                           {/if}
                         </div>
                       </div>

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

@@ -6,6 +6,8 @@
     type DownloadProgress,
     refreshState,
     lastUpdate as lastUpdateStore,
+    startDownload,
+    deleteDownload,
   } from "$lib/stores/app.svelte";
   import HeaderNav from "$lib/components/HeaderNav.svelte";
 
@@ -28,6 +30,7 @@
     etaMs: number;
     status: "completed" | "downloading";
     files: FileProgress[];
+    shardMetadata?: Record<string, unknown>;
   };
 
   type NodeEntry = {
@@ -172,33 +175,6 @@
   }
 
   let downloadOverview = $state<NodeEntry[]>([]);
-  let models = $state<Array<{ id: string; storage_size_megabytes?: number }>>(
-    [],
-  );
-
-  async function fetchModels() {
-    try {
-      const response = await fetch("/models");
-      if (response.ok) {
-        const data = await response.json();
-        models = data.data || [];
-      }
-    } catch (error) {
-      console.error("Failed to fetch models:", error);
-    }
-  }
-
-  function getModelTotalBytes(
-    modelId: string,
-    downloadTotalBytes: number,
-  ): number {
-    if (downloadTotalBytes > 0) return downloadTotalBytes;
-    const model = models.find((m) => m.id === modelId);
-    if (model?.storage_size_megabytes) {
-      return model.storage_size_megabytes * 1024 * 1024;
-    }
-    return 0;
-  }
 
   $effect(() => {
     try {
@@ -296,6 +272,12 @@
             }
           }
 
+          // Extract shard_metadata for use with download actions
+          const shardMetadata = (downloadPayload.shard_metadata ??
+            downloadPayload.shardMetadata) as
+            | Record<string, unknown>
+            | undefined;
+
           const entry: ModelEntry = {
             modelId,
             prettyName,
@@ -312,6 +294,7 @@
                 ? "completed"
                 : "downloading",
             files,
+            shardMetadata,
           };
 
           const existing = modelMap.get(modelId);
@@ -373,7 +356,6 @@
   onMount(() => {
     // Ensure we fetch at least once when visiting downloads directly
     refreshState();
-    fetchModels();
   });
 </script>
 
@@ -482,7 +464,7 @@
                     {#if model.status !== "completed"}
                       <div class="text-[11px] text-exo-light-gray font-mono">
                         {formatBytes(model.downloadedBytes)} / {formatBytes(
-                          getModelTotalBytes(model.modelId, model.totalBytes),
+                          model.totalBytes,
                         )}
                       </div>
                     {/if}
@@ -497,6 +479,52 @@
                     >
                       {pct.toFixed(1)}%
                     </span>
+                    {#if model.status !== "completed" && model.shardMetadata}
+                      <button
+                        type="button"
+                        class="text-exo-light-gray hover:text-exo-yellow transition-colors"
+                        onclick={() =>
+                          startDownload(node.nodeId, model.shardMetadata!)}
+                        title="Start download"
+                      >
+                        <svg
+                          class="w-4 h-4"
+                          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>
+                    {/if}
+                    {#if model.status === "completed"}
+                      <button
+                        type="button"
+                        class="text-exo-light-gray hover:text-red-400 transition-colors"
+                        onclick={() =>
+                          deleteDownload(node.nodeId, model.modelId)}
+                        title="Delete download"
+                      >
+                        <svg
+                          class="w-4 h-4"
+                          viewBox="0 0 20 20"
+                          fill="none"
+                          stroke="currentColor"
+                          stroke-width="2"
+                        >
+                          <path
+                            d="M4 6h12M8 6V4h4v2m1 0v10a1 1 0 01-1 1H8a1 1 0 01-1-1V6h6"
+                            stroke-linecap="round"
+                            stroke-linejoin="round"
+                          ></path>
+                        </svg>
+                      </button>
+                    {/if}
                     <button
                       type="button"
                       class="text-exo-light-gray hover:text-exo-yellow transition-colors"

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

@@ -0,0 +1,190 @@
+<script lang="ts">
+  import { onMount } from "svelte";
+  import {
+    listTraces,
+    getTraceRawUrl,
+    type TraceListItem,
+  } from "$lib/stores/app.svelte";
+  import HeaderNav from "$lib/components/HeaderNav.svelte";
+
+  let traces = $state<TraceListItem[]>([]);
+  let loading = $state(true);
+  let error = $state<string | null>(null);
+
+  function formatBytes(bytes: number): string {
+    if (!bytes || bytes <= 0) return "0B";
+    const units = ["B", "KB", "MB", "GB"];
+    const i = Math.min(
+      Math.floor(Math.log(bytes) / Math.log(1024)),
+      units.length - 1,
+    );
+    const val = bytes / Math.pow(1024, i);
+    return `${val.toFixed(val >= 10 ? 0 : 1)}${units[i]}`;
+  }
+
+  function formatDate(isoString: string): string {
+    const date = new Date(isoString);
+    return date.toLocaleString();
+  }
+
+  async function downloadTrace(taskId: string) {
+    const response = await fetch(getTraceRawUrl(taskId));
+    const blob = await response.blob();
+    const url = URL.createObjectURL(blob);
+    const a = document.createElement("a");
+    a.href = url;
+    a.download = `trace_${taskId}.json`;
+    a.click();
+    URL.revokeObjectURL(url);
+  }
+
+  async function openInPerfetto(taskId: string) {
+    // Fetch trace data from our local API
+    const response = await fetch(getTraceRawUrl(taskId));
+    const traceData = await response.arrayBuffer();
+
+    // Open Perfetto UI
+    const perfettoWindow = window.open("https://ui.perfetto.dev");
+    if (!perfettoWindow) {
+      alert("Failed to open Perfetto. Please allow popups.");
+      return;
+    }
+
+    // Wait for Perfetto to be ready, then send trace via postMessage
+    const onMessage = (e: MessageEvent) => {
+      if (e.data === "PONG") {
+        window.removeEventListener("message", onMessage);
+        perfettoWindow.postMessage(
+          {
+            perfetto: {
+              buffer: traceData,
+              title: `Trace ${taskId}`,
+            },
+          },
+          "https://ui.perfetto.dev",
+        );
+      }
+    };
+    window.addEventListener("message", onMessage);
+
+    // Ping Perfetto until it responds
+    const pingInterval = setInterval(() => {
+      perfettoWindow.postMessage("PING", "https://ui.perfetto.dev");
+    }, 50);
+
+    // Clean up after 10 seconds
+    setTimeout(() => {
+      clearInterval(pingInterval);
+      window.removeEventListener("message", onMessage);
+    }, 10000);
+  }
+
+  async function refresh() {
+    loading = true;
+    error = null;
+    try {
+      const response = await listTraces();
+      traces = response.traces;
+    } catch (e) {
+      error = e instanceof Error ? e.message : "Failed to load traces";
+    } finally {
+      loading = false;
+    }
+  }
+
+  onMount(() => {
+    refresh();
+  });
+</script>
+
+<div class="min-h-screen bg-exo-dark-gray text-white">
+  <HeaderNav showHome={true} />
+  <div class="max-w-7xl mx-auto px-4 lg:px-8 py-6 space-y-6">
+    <div class="flex items-center justify-between gap-4 flex-wrap">
+      <div>
+        <h1
+          class="text-2xl font-mono tracking-[0.2em] uppercase text-exo-yellow"
+        >
+          Traces
+        </h1>
+      </div>
+      <div class="flex items-center gap-3">
+        <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"
+          onclick={refresh}
+          disabled={loading}
+        >
+          Refresh
+        </button>
+      </div>
+    </div>
+
+    {#if loading}
+      <div
+        class="rounded border border-exo-medium-gray/30 bg-exo-black/30 p-6 text-center text-exo-light-gray"
+      >
+        <div class="text-sm">Loading traces...</div>
+      </div>
+    {:else if error}
+      <div
+        class="rounded border border-red-500/30 bg-red-500/10 p-6 text-center text-red-400"
+      >
+        <div class="text-sm">{error}</div>
+      </div>
+    {:else if traces.length === 0}
+      <div
+        class="rounded border border-exo-medium-gray/30 bg-exo-black/30 p-6 text-center text-exo-light-gray space-y-2"
+      >
+        <div class="text-sm">No traces found.</div>
+        <div class="text-xs text-exo-light-gray/70">
+          Run exo with EXO_TRACING_ENABLED=1 to collect traces.
+        </div>
+      </div>
+    {:else}
+      <div class="space-y-3">
+        {#each traces as trace}
+          <div
+            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 text-white hover:text-exo-yellow transition-colors truncate block"
+              >
+                {trace.taskId}
+              </a>
+              <div class="text-xs text-exo-light-gray font-mono mt-1">
+                {formatDate(trace.createdAt)} &bull; {formatBytes(
+                  trace.fileSize,
+                )}
+              </div>
+            </div>
+            <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"
+              >
+                View Stats
+              </a>
+              <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"
+                onclick={() => downloadTrace(trace.taskId)}
+              >
+                Download
+              </button>
+              <button
+                type="button"
+                class="text-xs font-mono text-exo-dark-gray bg-exo-yellow hover:bg-exo-yellow/90 transition-colors uppercase px-2 py-1 rounded font-semibold"
+                onclick={() => openInPerfetto(trace.taskId)}
+              >
+                View Trace
+              </button>
+            </div>
+          </div>
+        {/each}
+      </div>
+    {/if}
+  </div>
+</div>

dashboard/src/routes/traces/[taskId]/+page.svelte

@@ -0,0 +1,367 @@
+<script lang="ts">
+  import { page } from "$app/stores";
+  import { onMount } from "svelte";
+  import {
+    fetchTraceStats,
+    getTraceRawUrl,
+    type TraceStatsResponse,
+    type TraceCategoryStats,
+  } from "$lib/stores/app.svelte";
+  import HeaderNav from "$lib/components/HeaderNav.svelte";
+
+  const taskId = $derived($page.params.taskId);
+
+  let stats = $state<TraceStatsResponse | null>(null);
+  let loading = $state(true);
+  let error = $state<string | null>(null);
+
+  function formatDuration(us: number): string {
+    if (us < 1000) return `${us.toFixed(0)}us`;
+    if (us < 1_000_000) return `${(us / 1000).toFixed(2)}ms`;
+    return `${(us / 1_000_000).toFixed(2)}s`;
+  }
+
+  function formatPercentage(part: number, total: number): string {
+    if (total === 0) return "0.0%";
+    return `${((part / total) * 100).toFixed(1)}%`;
+  }
+
+  // Parse hierarchical categories like "sync/compute" into phases
+  type PhaseData = {
+    name: string;
+    subcategories: { name: string; stats: TraceCategoryStats }[];
+    totalUs: number; // From outer span (e.g., "sync" category)
+    stepCount: number; // Count of outer span events
+  };
+
+  function parsePhases(
+    byCategory: Record<string, TraceCategoryStats>,
+  ): PhaseData[] {
+    const phases = new Map<
+      string,
+      {
+        subcats: Map<string, TraceCategoryStats>;
+        outerStats: TraceCategoryStats | null;
+      }
+    >();
+
+    for (const [category, catStats] of Object.entries(byCategory)) {
+      if (category.includes("/")) {
+        const [phase, subcat] = category.split("/", 2);
+        if (!phases.has(phase)) {
+          phases.set(phase, { subcats: new Map(), outerStats: null });
+        }
+        phases.get(phase)!.subcats.set(subcat, catStats);
+      } else {
+        // Outer span - this IS the phase total
+        if (!phases.has(category)) {
+          phases.set(category, { subcats: new Map(), outerStats: null });
+        }
+        phases.get(category)!.outerStats = catStats;
+      }
+    }
+
+    return Array.from(phases.entries())
+      .filter(([_, data]) => data.outerStats !== null) // Only phases with outer spans
+      .map(([name, data]) => ({
+        name,
+        subcategories: Array.from(data.subcats.entries())
+          .map(([subName, subStats]) => ({ name: subName, stats: subStats }))
+          .sort((a, b) => b.stats.totalUs - a.stats.totalUs),
+        totalUs: data.outerStats!.totalUs, // Outer span total
+        stepCount: data.outerStats!.count, // Number of steps
+      }))
+      .sort((a, b) => b.totalUs - a.totalUs);
+  }
+
+  async function downloadTrace() {
+    if (!taskId) return;
+    const response = await fetch(getTraceRawUrl(taskId));
+    const blob = await response.blob();
+    const url = URL.createObjectURL(blob);
+    const a = document.createElement("a");
+    a.href = url;
+    a.download = `trace_${taskId}.json`;
+    a.click();
+    URL.revokeObjectURL(url);
+  }
+
+  async function openInPerfetto() {
+    if (!taskId) return;
+
+    // Fetch trace data from our local API
+    const response = await fetch(getTraceRawUrl(taskId));
+    const traceData = await response.arrayBuffer();
+
+    // Open Perfetto UI
+    const perfettoWindow = window.open("https://ui.perfetto.dev");
+    if (!perfettoWindow) {
+      alert("Failed to open Perfetto. Please allow popups.");
+      return;
+    }
+
+    // Wait for Perfetto to be ready, then send trace via postMessage
+    const onMessage = (e: MessageEvent) => {
+      if (e.data === "PONG") {
+        window.removeEventListener("message", onMessage);
+        perfettoWindow.postMessage(
+          {
+            perfetto: {
+              buffer: traceData,
+              title: `Trace ${taskId}`,
+            },
+          },
+          "https://ui.perfetto.dev",
+        );
+      }
+    };
+    window.addEventListener("message", onMessage);
+
+    // Ping Perfetto until it responds
+    const pingInterval = setInterval(() => {
+      perfettoWindow.postMessage("PING", "https://ui.perfetto.dev");
+    }, 50);
+
+    // Clean up after 10 seconds
+    setTimeout(() => {
+      clearInterval(pingInterval);
+      window.removeEventListener("message", onMessage);
+    }, 10000);
+  }
+
+  onMount(async () => {
+    if (!taskId) {
+      error = "No task ID provided";
+      loading = false;
+      return;
+    }
+
+    try {
+      stats = await fetchTraceStats(taskId);
+    } catch (e) {
+      error = e instanceof Error ? e.message : "Failed to load trace";
+    } finally {
+      loading = false;
+    }
+  });
+
+  const phases = $derived(stats ? parsePhases(stats.byCategory) : []);
+  const sortedRanks = $derived(
+    stats
+      ? Object.keys(stats.byRank)
+          .map(Number)
+          .sort((a, b) => a - b)
+      : [],
+  );
+  const nodeCount = $derived(sortedRanks.length || 1);
+</script>
+
+<div class="min-h-screen bg-exo-dark-gray text-white">
+  <HeaderNav showHome={true} />
+  <div class="max-w-7xl mx-auto px-4 lg:px-8 py-6 space-y-6">
+    <div class="flex items-center justify-between gap-4 flex-wrap">
+      <div>
+        <h1
+          class="text-2xl font-mono tracking-[0.2em] uppercase text-exo-yellow"
+        >
+          Trace
+        </h1>
+        <p class="text-sm text-exo-light-gray font-mono truncate max-w-lg">
+          {taskId}
+        </p>
+      </div>
+      <div class="flex items-center gap-3">
+        <a
+          href="#/traces"
+          class="text-xs font-mono text-exo-light-gray hover:text-exo-yellow transition-colors uppercase border border-exo-medium-gray/40 px-3 py-1.5 rounded"
+        >
+          All Traces
+        </a>
+        <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-3 py-1.5 rounded"
+          onclick={downloadTrace}
+          disabled={loading || !!error}
+        >
+          Download
+        </button>
+        <button
+          type="button"
+          class="text-xs font-mono text-exo-dark-gray bg-exo-yellow hover:bg-exo-yellow/90 transition-colors uppercase px-3 py-1.5 rounded font-semibold"
+          onclick={openInPerfetto}
+          disabled={loading || !!error}
+        >
+          View Trace
+        </button>
+      </div>
+    </div>
+
+    {#if loading}
+      <div
+        class="rounded border border-exo-medium-gray/30 bg-exo-black/30 p-6 text-center text-exo-light-gray"
+      >
+        <div class="text-sm">Loading trace data...</div>
+      </div>
+    {:else if error}
+      <div
+        class="rounded border border-red-500/30 bg-red-500/10 p-6 text-center text-red-400"
+      >
+        <div class="text-sm">{error}</div>
+      </div>
+    {:else if stats}
+      <!-- Wall Time Summary -->
+      <div
+        class="rounded border border-exo-medium-gray/30 bg-exo-black/30 p-4 space-y-2"
+      >
+        <h2
+          class="text-sm font-mono uppercase tracking-wider text-exo-light-gray"
+        >
+          Summary
+        </h2>
+        <div class="text-3xl font-mono text-exo-yellow">
+          {formatDuration(stats.totalWallTimeUs)}
+        </div>
+        <div class="text-xs text-exo-light-gray">Total wall time</div>
+      </div>
+
+      <!-- By Phase -->
+      {#if phases.length > 0}
+        <div
+          class="rounded border border-exo-medium-gray/30 bg-exo-black/30 p-4 space-y-4"
+        >
+          <h2
+            class="text-sm font-mono uppercase tracking-wider text-exo-light-gray"
+          >
+            By Phase <span class="text-exo-light-gray/50">(avg per node)</span>
+          </h2>
+          <div class="space-y-4">
+            {#each phases as phase}
+              {@const normalizedTotal = phase.totalUs / nodeCount}
+              {@const normalizedStepCount = phase.stepCount / nodeCount}
+              <div class="space-y-2">
+                <div class="flex items-center justify-between">
+                  <span class="text-sm font-mono text-white">{phase.name}</span>
+                  <span class="text-sm font-mono">
+                    <span class="text-exo-yellow"
+                      >{formatDuration(normalizedTotal)}</span
+                    >
+                    <span class="text-exo-light-gray ml-2">
+                      ({normalizedStepCount} steps, {formatDuration(
+                        normalizedTotal / normalizedStepCount,
+                      )}/step)
+                    </span>
+                  </span>
+                </div>
+                {#if phase.subcategories.length > 0}
+                  <div class="pl-4 space-y-1.5">
+                    {#each phase.subcategories as subcat}
+                      {@const normalizedSubcat =
+                        subcat.stats.totalUs / nodeCount}
+                      {@const pct = formatPercentage(
+                        normalizedSubcat,
+                        normalizedTotal,
+                      )}
+                      {@const perStep = normalizedSubcat / normalizedStepCount}
+                      <div
+                        class="flex items-center justify-between text-xs font-mono"
+                      >
+                        <span class="text-exo-light-gray">{subcat.name}</span>
+                        <span class="text-white">
+                          {formatDuration(normalizedSubcat)}
+                          <span class="text-exo-light-gray ml-2">({pct})</span>
+                          <span class="text-exo-light-gray/60 ml-2"
+                            >{formatDuration(perStep)}/step</span
+                          >
+                        </span>
+                      </div>
+                      <!-- Progress bar -->
+                      <div
+                        class="relative h-1.5 bg-exo-black/60 rounded-sm overflow-hidden"
+                      >
+                        <div
+                          class="absolute inset-y-0 left-0 bg-gradient-to-r from-exo-yellow to-exo-yellow/70 transition-all duration-300"
+                          style="width: {pct}"
+                        ></div>
+                      </div>
+                    {/each}
+                  </div>
+                {/if}
+              </div>
+            {/each}
+          </div>
+        </div>
+      {/if}
+
+      <!-- By Rank -->
+      {#if sortedRanks.length > 0}
+        <div
+          class="rounded border border-exo-medium-gray/30 bg-exo-black/30 p-4 space-y-4"
+        >
+          <h2
+            class="text-sm font-mono uppercase tracking-wider text-exo-light-gray"
+          >
+            By Rank
+          </h2>
+          <div class="grid grid-cols-1 md:grid-cols-2 lg:grid-cols-3 gap-4">
+            {#each sortedRanks as rank}
+              {@const rankStats = stats.byRank[rank]}
+              {@const rankPhases = parsePhases(rankStats.byCategory)}
+              <div
+                class="rounded border border-exo-medium-gray/20 bg-exo-dark-gray/60 p-3 space-y-3"
+              >
+                <div class="text-sm font-mono text-exo-yellow">
+                  Rank {rank}
+                </div>
+                <div class="space-y-2">
+                  {#each rankPhases as phase}
+                    <div class="space-y-1">
+                      <div class="flex items-center justify-between text-xs">
+                        <span class="font-mono text-exo-light-gray"
+                          >{phase.name}</span
+                        >
+                        <span class="font-mono text-white">
+                          {formatDuration(phase.totalUs)}
+                          <span class="text-exo-light-gray/50 ml-1">
+                            ({phase.stepCount}x)
+                          </span>
+                        </span>
+                      </div>
+                      {#if phase.subcategories.length > 0}
+                        <div class="pl-2 space-y-0.5">
+                          {#each phase.subcategories as subcat}
+                            {@const pct = formatPercentage(
+                              subcat.stats.totalUs,
+                              phase.totalUs,
+                            )}
+                            {@const perStep =
+                              subcat.stats.totalUs / phase.stepCount}
+                            <div
+                              class="flex items-center justify-between text-[10px] font-mono"
+                            >
+                              <span class="text-exo-light-gray/70"
+                                >{subcat.name}</span
+                              >
+                              <span class="text-exo-light-gray">
+                                {formatDuration(subcat.stats.totalUs)}
+                                <span class="text-exo-light-gray/50"
+                                  >({pct})</span
+                                >
+                                <span class="text-exo-light-gray/30 ml-1"
+                                  >{formatDuration(perStep)}/step</span
+                                >
+                              </span>
+                            </div>
+                          {/each}
+                        </div>
+                      {/if}
+                    </div>
+                  {/each}
+                </div>
+              </div>
+            {/each}
+          </div>
+        </div>
+      {/if}
+    {/if}
+  </div>
+</div>

flake.lock

@@ -21,7 +21,9 @@
           "nixpkgs"
         ],
         "purescript-overlay": "purescript-overlay",
-        "pyproject-nix": "pyproject-nix"
+        "pyproject-nix": [
+          "pyproject-nix"
+        ]
       },
       "locked": {
         "lastModified": 1765953015,
@@ -149,19 +151,44 @@
         "type": "github"
       }
     },
+    "pyproject-build-systems": {
+      "inputs": {
+        "nixpkgs": [
+          "nixpkgs"
+        ],
+        "pyproject-nix": [
+          "pyproject-nix"
+        ],
+        "uv2nix": [
+          "uv2nix"
+        ]
+      },
+      "locked": {
+        "lastModified": 1763662255,
+        "narHash": "sha256-4bocaOyLa3AfiS8KrWjZQYu+IAta05u3gYZzZ6zXbT0=",
+        "owner": "pyproject-nix",
+        "repo": "build-system-pkgs",
+        "rev": "042904167604c681a090c07eb6967b4dd4dae88c",
+        "type": "github"
+      },
+      "original": {
+        "owner": "pyproject-nix",
+        "repo": "build-system-pkgs",
+        "type": "github"
+      }
+    },
     "pyproject-nix": {
       "inputs": {
         "nixpkgs": [
-          "dream2nix",
           "nixpkgs"
         ]
       },
       "locked": {
-        "lastModified": 1763017646,
-        "narHash": "sha256-Z+R2lveIp6Skn1VPH3taQIuMhABg1IizJd8oVdmdHsQ=",
+        "lastModified": 1764134915,
+        "narHash": "sha256-xaKvtPx6YAnA3HQVp5LwyYG1MaN4LLehpQI8xEdBvBY=",
         "owner": "pyproject-nix",
         "repo": "pyproject.nix",
-        "rev": "47bd6f296502842643078d66128f7b5e5370790c",
+        "rev": "2c8df1383b32e5443c921f61224b198a2282a657",
         "type": "github"
       },
       "original": {
@@ -178,7 +205,10 @@
         "flake-parts": "flake-parts",
         "nixpkgs": "nixpkgs",
         "nixpkgs-swift": "nixpkgs-swift",
-        "treefmt-nix": "treefmt-nix"
+        "pyproject-build-systems": "pyproject-build-systems",
+        "pyproject-nix": "pyproject-nix",
+        "treefmt-nix": "treefmt-nix",
+        "uv2nix": "uv2nix"
       }
     },
     "rust-analyzer-src": {
@@ -239,6 +269,29 @@
         "repo": "treefmt-nix",
         "type": "github"
       }
+    },
+    "uv2nix": {
+      "inputs": {
+        "nixpkgs": [
+          "nixpkgs"
+        ],
+        "pyproject-nix": [
+          "pyproject-nix"
+        ]
+      },
+      "locked": {
+        "lastModified": 1767701098,
+        "narHash": "sha256-CJhKZnWb3gumR9oTRjFvCg/6lYTGbZRU7xtvcyWIRwU=",
+        "owner": "pyproject-nix",
+        "repo": "uv2nix",
+        "rev": "9d357f0d2ce6f5f35ec7959d7e704452352eb4da",
+        "type": "github"
+      },
+      "original": {
+        "owner": "pyproject-nix",
+        "repo": "uv2nix",
+        "type": "github"
+      }
     }
   },
   "root": "root",

flake.nix

@@ -24,6 +24,26 @@
     dream2nix = {
       url = "github:nix-community/dream2nix";
       inputs.nixpkgs.follows = "nixpkgs";
+      inputs.pyproject-nix.follows = "pyproject-nix";
+    };
+
+    # Python packaging with uv2nix
+    pyproject-nix = {
+      url = "github:pyproject-nix/pyproject.nix";
+      inputs.nixpkgs.follows = "nixpkgs";
+    };
+
+    uv2nix = {
+      url = "github:pyproject-nix/uv2nix";
+      inputs.pyproject-nix.follows = "pyproject-nix";
+      inputs.nixpkgs.follows = "nixpkgs";
+    };
+
+    pyproject-build-systems = {
+      url = "github:pyproject-nix/build-system-pkgs";
+      inputs.pyproject-nix.follows = "pyproject-nix";
+      inputs.uv2nix.follows = "uv2nix";
+      inputs.nixpkgs.follows = "nixpkgs";
     };
 
     # Pinned nixpkgs for swift-format (swift is broken on x86_64-linux in newer nixpkgs)
@@ -48,6 +68,7 @@
         inputs.treefmt-nix.flakeModule
         ./dashboard/parts.nix
         ./rust/parts.nix
+        ./python/parts.nix
       ];
 
       perSystem =
@@ -58,6 +79,11 @@
           pkgsSwift = import inputs.nixpkgs-swift { inherit system; };
         in
         {
+          # Allow unfree for metal-toolchain (needed for Darwin Metal packages)
+          _module.args.pkgs = import inputs.nixpkgs {
+            inherit system;
+            config.allowUnfreePredicate = pkg: (pkg.pname or "") == "metal-toolchain";
+          };
           treefmt = {
             projectRootFile = "flake.nix";
             programs = {
@@ -79,14 +105,24 @@
                 enable = true;
                 package = pkgsSwift.swiftPackages.swift-format;
               };
+              shfmt.enable = true;
             };
           };
 
-          checks.lint = pkgs.runCommand "lint-check" { } ''
-            export RUFF_CACHE_DIR="$TMPDIR/ruff-cache"
-            ${pkgs.ruff}/bin/ruff check ${inputs.self}/
-            touch $out
-          '';
+          packages = lib.optionalAttrs pkgs.stdenv.hostPlatform.isDarwin (
+            let
+              uvLock = builtins.fromTOML (builtins.readFile ./uv.lock);
+              mlxPackage = builtins.head (builtins.filter (p: p.name == "mlx") uvLock.package);
+              uvLockMlxVersion = mlxPackage.version;
+            in
+            {
+              metal-toolchain = pkgs.callPackage ./nix/metal-toolchain.nix { };
+              mlx = pkgs.callPackage ./nix/mlx.nix {
+                metal-toolchain = self'.packages.metal-toolchain;
+                inherit uvLockMlxVersion;
+              };
+            }
+          );
 
           devShells.default = with pkgs; pkgs.mkShell {
             inputsFrom = [ self'.checks.cargo-build ];

justfile

@@ -1,7 +1,7 @@
 export NIX_CONFIG := "extra-experimental-features = nix-command flakes"
 
 fmt:
-    nix fmt
+    treefmt || nix fmt
 
 lint:
     uv run ruff check --fix

nix/darwin-build-fixes.patch

@@ -0,0 +1,79 @@
+diff --git a/CMakeLists.txt b/CMakeLists.txt
+index 0ed30932..d8528132 100644
+--- a/CMakeLists.txt
++++ b/CMakeLists.txt
+@@ -177,11 +177,7 @@ if(MLX_BUILD_METAL)
+     add_compile_definitions(MLX_METAL_DEBUG)
+   endif()
+
+-  # Throw an error if xcrun not found
+-  execute_process(
+-    COMMAND zsh "-c" "/usr/bin/xcrun -sdk macosx --show-sdk-version"
+-    OUTPUT_VARIABLE MACOS_SDK_VERSION
+-    OUTPUT_STRIP_TRAILING_WHITESPACE COMMAND_ERROR_IS_FATAL ANY)
++  set(MACOS_SDK_VERSION @sdkVersion@)
+
+   if(${MACOS_SDK_VERSION} LESS 14.0)
+     message(
+@@ -199,11 +195,8 @@ if(MLX_BUILD_METAL)
+     endif()
+     set(XCRUN_FLAGS "-mmacosx-version-min=${CMAKE_OSX_DEPLOYMENT_TARGET}")
+   endif()
+-  execute_process(
+-    COMMAND
+-      zsh "-c"
+-      "echo \"__METAL_VERSION__\" | xcrun -sdk macosx metal ${XCRUN_FLAGS} -E -x metal -P - | tail -1 | tr -d '\n'"
+-    OUTPUT_VARIABLE MLX_METAL_VERSION COMMAND_ERROR_IS_FATAL ANY)
++  set(
++    MLX_METAL_VERSION @metalVersion@)
+   FetchContent_Declare(metal_cpp URL ${METAL_CPP_URL})
+   FetchContent_MakeAvailable(metal_cpp)
+   target_include_directories(
+diff --git a/cmake/extension.cmake b/cmake/extension.cmake
+index 13db804a..5b385132 100644
+--- a/cmake/extension.cmake
++++ b/cmake/extension.cmake
+@@ -36,7 +36,7 @@ macro(mlx_build_metallib)
+   add_custom_command(
+     OUTPUT ${MTLLIB_BUILD_TARGET}
+     COMMAND
+-      xcrun -sdk macosx metal
++      metal -fmodules-cache-path=${CMAKE_BINARY_DIR}/metal-cache
+       "$<LIST:TRANSFORM,${MTLLIB_INCLUDE_DIRS},PREPEND,-I>"
+       ${MTLLIB_COMPILE_OPTIONS} ${MTLLIB_SOURCES} -o ${MTLLIB_BUILD_TARGET}
+     DEPENDS ${MTLLIB_DEPS} ${MTLLIB_SOURCES}
+diff --git a/mlx/backend/metal/kernels/CMakeLists.txt b/mlx/backend/metal/kernels/CMakeLists.txt
+index 262b0495..5c7446ad 100644
+--- a/mlx/backend/metal/kernels/CMakeLists.txt
++++ b/mlx/backend/metal/kernels/CMakeLists.txt
+@@ -29,7 +29,7 @@ function(build_kernel_base TARGET SRCFILE DEPS)
+                     "-mmacosx-version-min=${CMAKE_OSX_DEPLOYMENT_TARGET}")
+   endif()
+   add_custom_command(
+-    COMMAND xcrun -sdk macosx metal ${METAL_FLAGS} -c ${SRCFILE}
++    COMMAND metal -fmodules-cache-path=${CMAKE_BINARY_DIR}/metal-cache ${METAL_FLAGS} -c ${SRCFILE}
+             -I${PROJECT_SOURCE_DIR} -o ${TARGET}.air
+     DEPENDS ${SRCFILE} ${DEPS} ${BASE_HEADERS}
+     OUTPUT ${TARGET}.air
+@@ -170,7 +170,7 @@ endif()
+
+ add_custom_command(
+   OUTPUT ${MLX_METAL_PATH}/mlx.metallib
+-  COMMAND xcrun -sdk macosx metallib ${KERNEL_AIR} -o
++  COMMAND metallib ${KERNEL_AIR} -o
+           ${MLX_METAL_PATH}/mlx.metallib
+   DEPENDS ${KERNEL_AIR}
+   COMMENT "Building mlx.metallib"
+diff --git a/mlx/backend/metal/make_compiled_preamble.sh b/mlx/backend/metal/make_compiled_preamble.sh
+index bb55ed3a..94ea7dd7 100644
+--- a/mlx/backend/metal/make_compiled_preamble.sh
++++ b/mlx/backend/metal/make_compiled_preamble.sh
+@@ -31,7 +31,7 @@ OUTPUT_FILE=${OUTPUT_DIR}/${SRC_NAME}.cpp
+ mkdir -p "$OUTPUT_DIR"
+
+ # Use the metal compiler to get a list of headers (with depth)
+-CCC="xcrun -sdk macosx metal -x metal"
++CCC="metal -x metal -fmodules-cache-path=${OUTPUT_DIR}/metal-cache"
+ HDRS=$( $CCC -I"$SRC_DIR" -I"$JIT_INCLUDES" -DMLX_METAL_JIT -E -P -CC -C -H "$INPUT_FILE" $CFLAGS -w 2>&1 1>/dev/null )
+
+ # Remove any included system frameworks (for MetalPerformancePrimitive headers)

nix/metal-toolchain.nix

@@ -0,0 +1,56 @@
+{ lib, stdenvNoCC, requireFile, nix }:
+
+let
+  narFile = requireFile {
+    name = "metal-toolchain-17C48.nar";
+    message = ''
+      The Metal Toolchain NAR must be available.
+
+      If you have cachix configured for exo.cachix.org, this should be automatic.
+
+      Otherwise:
+        1. Install Xcode 26+ from the App Store
+        2. Run: xcodebuild -downloadComponent MetalToolchain
+        3. Export the toolchain:
+           hdiutil attach "$(find /System/Library/AssetsV2/com_apple_MobileAsset_MetalToolchain -name '*.dmg' | head -1)" -mountpoint /tmp/metal-dmg
+           cp -R /tmp/metal-dmg/Metal.xctoolchain /tmp/metal-export
+           hdiutil detach /tmp/metal-dmg
+        4. Create NAR and add to store:
+           nix nar pack /tmp/metal-export > /tmp/metal-toolchain-17C48.nar
+           nix store add --mode flat /tmp/metal-toolchain-17C48.nar
+    '';
+    hash = "sha256-ayR5mXN4sZAddwKEG2OszGRF93k9ZFc7H0yi2xbylQw=";
+  };
+in
+stdenvNoCC.mkDerivation {
+  pname = "metal-toolchain";
+  version = "17C48";
+
+  dontUnpack = true;
+  dontBuild = true;
+  dontFixup = true;
+
+  nativeBuildInputs = [ nix ];
+
+  installPhase = ''
+    runHook preInstall
+
+    nix-store --restore $out < ${narFile}
+
+    # Create bin directory with symlinks for PATH
+    mkdir -p $out/bin
+    ln -s $out/usr/bin/metal $out/bin/metal
+    ln -s $out/usr/bin/metallib $out/bin/metallib
+
+    runHook postInstall
+  '';
+
+  # Metal language version for CMake (from: echo __METAL_VERSION__ | metal -E -x metal -P -)
+  passthru.metalVersion = "400";
+
+  meta = {
+    description = "Apple Metal compiler toolchain";
+    platforms = [ "aarch64-darwin" ];
+    license = lib.licenses.unfree;
+  };
+}

nix/mlx.nix

@@ -0,0 +1,158 @@
+{ stdenv
+, lib
+, fetchFromGitHub
+, replaceVars
+, fetchzip
+, cmake
+, nlohmann_json
+, apple-sdk_26
+, metal-toolchain
+, runCommand
+, fmt
+, python313Packages
+, uvLockMlxVersion
+}:
+
+assert stdenv.isDarwin;
+
+let
+  python = python313Packages.python;
+
+  # Static dependencies included directly during compilation
+  gguf-tools = fetchFromGitHub {
+    owner = "antirez";
+    repo = "gguf-tools";
+    rev = "8fa6eb65236618e28fd7710a0fba565f7faa1848";
+    hash = "sha256-15FvyPOFqTOr5vdWQoPnZz+mYH919++EtghjozDlnSA=";
+  };
+
+  metal_cpp = fetchzip {
+    url = "https://developer.apple.com/metal/cpp/files/metal-cpp_26.zip";
+    hash = "sha256-7n2eI2lw/S+Us6l7YPAATKwcIbRRpaQ8VmES7S8ZjY8=";
+  };
+
+  nanobind = fetchFromGitHub {
+    owner = "wjakob";
+    repo = "nanobind";
+    rev = "v2.10.2";
+    hash = "sha256-io44YhN+VpfHFWyvvLWSanRgbzA0whK8WlDNRi3hahU=";
+    fetchSubmodules = true;
+  };
+
+  mlx = stdenv.mkDerivation rec {
+    pname = "mlx";
+    version = let v = "0.30.4"; 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;
+
+    src = fetchFromGitHub {
+      owner = "ml-explore";
+      repo = "mlx";
+      tag = "v${version}";
+      hash = "sha256-OJk6jPlbaSlsUdk3ADz3tWcRzTWXRof3/q8Soe1AO6w=";
+    };
+
+    patches = [
+      (replaceVars ./darwin-build-fixes.patch {
+        sdkVersion = apple-sdk_26.version;
+        metalVersion = metal-toolchain.metalVersion;
+      })
+    ];
+
+    postPatch = ''
+      substituteInPlace mlx/backend/cpu/jit_compiler.cpp \
+        --replace-fail "g++" "$CXX"
+    '';
+
+    dontUseCmakeConfigure = true;
+
+    enableParallelBuilding = true;
+
+    # Allows multiple cores to be used in Python builds.
+    postUnpack = ''
+      export MAKEFLAGS+="''${enableParallelBuilding:+-j$NIX_BUILD_CORES}"
+    '';
+
+    # Updates the wrong fetcher rev attribute
+    passthru.skipBulkUpdate = true;
+
+    env = {
+      DEV_RELEASE = 1;
+      CMAKE_ARGS = toString [
+        (lib.cmakeBool "USE_SYSTEM_FMT" true)
+        (lib.cmakeOptionType "filepath" "FETCHCONTENT_SOURCE_DIR_GGUFLIB" "${gguf-tools}")
+        (lib.cmakeOptionType "filepath" "FETCHCONTENT_SOURCE_DIR_JSON" "${nlohmann_json.src}")
+        (lib.cmakeOptionType "filepath" "FETCHCONTENT_SOURCE_DIR_NANOBIND" "${nanobind}")
+        (lib.cmakeBool "FETCHCONTENT_FULLY_DISCONNECTED" true)
+        (lib.cmakeBool "MLX_BUILD_METAL" true)
+        (lib.cmakeOptionType "filepath" "FETCHCONTENT_SOURCE_DIR_METAL_CPP" "${metal_cpp}")
+        (lib.cmakeOptionType "string" "CMAKE_OSX_DEPLOYMENT_TARGET" "${apple-sdk_26.version}")
+        (lib.cmakeOptionType "filepath" "CMAKE_OSX_SYSROOT" "${apple-sdk_26.passthru.sdkroot}")
+      ];
+      SDKROOT = apple-sdk_26.passthru.sdkroot;
+      MACOSX_DEPLOYMENT_TARGET = apple-sdk_26.version;
+    };
+
+    build-system = [
+      python313Packages.setuptools
+    ];
+
+    nativeBuildInputs = [
+      cmake
+      metal-toolchain
+      python313Packages.pypaBuildHook
+      python313Packages.pypaInstallHook
+      python313Packages.setuptools
+      python313Packages.typing-extensions
+      python313Packages.wheel
+      python313Packages.cmake
+      python313Packages.ninja
+    ];
+
+    buildInputs = [
+      fmt
+      gguf-tools
+      python313Packages.nanobind
+      python313Packages.pybind11
+      apple-sdk_26
+    ];
+
+    # Tests require Metal GPU access which isn't available in the Nix sandbox.
+    # To run tests, build with: nix build --option sandbox false .#mlx.passthru.tests.mlxTest
+    doCheck = false;
+
+    pythonImportsCheck = [ "mlx" ];
+
+    passthru.tests = {
+      # Runs example scripts to verify MLX works. Requires --option sandbox false
+      # since Metal GPU access is needed.
+      mlxTest =
+        runCommand "run-mlx-examples"
+          {
+            buildInputs = [ mlx ];
+            nativeBuildInputs = [ python ];
+          }
+          ''
+            cp ${src}/examples/python/logistic_regression.py .
+            ${python.interpreter} logistic_regression.py
+            rm logistic_regression.py
+
+            cp ${src}/examples/python/linear_regression.py .
+            ${python.interpreter} linear_regression.py
+            rm linear_regression.py
+
+            touch $out
+          '';
+    };
+
+    meta = {
+      homepage = "https://github.com/ml-explore/mlx";
+      description = "Array framework for Apple silicon";
+      changelog = "https://github.com/ml-explore/mlx/releases/tag/${src.tag}";
+      license = lib.licenses.mit;
+      platforms = [ "aarch64-darwin" ];
+    };
+  };
+in
+mlx

pyproject.toml

@@ -17,16 +17,16 @@ dependencies = [
     "loguru>=0.7.3",
     "exo_pyo3_bindings", # rust bindings
     "anyio==4.11.0",
-    "mlx==0.30.3; sys_platform == 'darwin'",
-    "mlx[cpu]==0.30.3; sys_platform == 'linux'",
-    "mlx-lm @ git+https://github.com/AlexCheema/mlx-lm.git@fix-transformers-5.0.0rc2",
+    "mlx==0.30.4; sys_platform == 'darwin'",
+    "mlx[cpu]==0.30.4; 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.14.2",
+    "mflux==0.15.4",
     "python-multipart>=0.0.21",
 ]
 
@@ -63,6 +63,7 @@ members = [
 
 [tool.uv.sources]
 exo_pyo3_bindings = { workspace = true }
+mlx-lm = { git = "https://github.com/ml-explore/mlx-lm", branch = "main" }
 # 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 }

python/parts.nix

@@ -0,0 +1,94 @@
+{ inputs, ... }:
+{
+  perSystem =
+    { config, self', pkgs, lib, system, ... }:
+    let
+      # Load workspace from uv.lock
+      workspace = inputs.uv2nix.lib.workspace.loadWorkspace {
+        workspaceRoot = inputs.self;
+      };
+
+      # Create overlay from workspace
+      # Use wheels from PyPI for most packages; we override mlx with our pure Nix Metal build
+      overlay = workspace.mkPyprojectOverlay { sourcePreference = "wheel"; };
+
+      # Override overlay to inject Nix-built components
+      exoOverlay = final: prev: {
+        # Replace workspace exo_pyo3_bindings with Nix-built wheel
+        exo-pyo3-bindings = pkgs.stdenv.mkDerivation {
+          pname = "exo-pyo3-bindings";
+          version = "0.1.0";
+          src = self'.packages.exo_pyo3_bindings;
+          # Install from pre-built wheel
+          nativeBuildInputs = [ final.pyprojectWheelHook ];
+          dontStrip = true;
+        };
+      };
+
+      python = pkgs.python313;
+
+      # Overlay to provide build systems and custom packages
+      buildSystemsOverlay = final: prev: {
+        # Use our pure Nix-built MLX with Metal support
+        mlx = self'.packages.mlx;
+
+        # mlx-lm is a git dependency that needs setuptools
+        mlx-lm = prev.mlx-lm.overrideAttrs (old: {
+          nativeBuildInputs = (old.nativeBuildInputs or [ ]) ++ [
+            final.setuptools
+          ];
+        });
+      };
+
+      pythonSet = (pkgs.callPackage inputs.pyproject-nix.build.packages {
+        inherit python;
+      }).overrideScope (
+        lib.composeManyExtensions [
+          inputs.pyproject-build-systems.overlays.default
+          overlay
+          exoOverlay
+          buildSystemsOverlay
+        ]
+      );
+      exoVenv = pythonSet.mkVirtualEnv "exo-env" workspace.deps.default;
+
+      # Virtual environment with dev dependencies for testing
+      testVenv = pythonSet.mkVirtualEnv "exo-test-env" (
+        workspace.deps.default // {
+          exo = [ "dev" ]; # Include pytest, pytest-asyncio, pytest-env
+        }
+      );
+
+      exoPackage = pkgs.runCommand "exo"
+        {
+          nativeBuildInputs = [ pkgs.makeWrapper ];
+        }
+        ''
+          mkdir -p $out/bin
+
+          # Create wrapper scripts
+          for script in exo exo-master exo-worker; do
+            makeWrapper ${exoVenv}/bin/$script $out/bin/$script \
+              --set DASHBOARD_DIR ${self'.packages.dashboard} \
+              ${lib.optionalString pkgs.stdenv.isDarwin "--prefix PATH : ${pkgs.macmon}/bin"}
+          done
+        '';
+    in
+    {
+      # Python package only available on macOS (requires MLX/Metal)
+      packages = lib.optionalAttrs pkgs.stdenv.hostPlatform.isDarwin {
+        exo = exoPackage;
+        # Test environment for running pytest outside of Nix sandbox (needs GPU access)
+        exo-test-env = testVenv;
+      };
+
+      checks = {
+        # Ruff linting (works on all platforms)
+        lint = pkgs.runCommand "ruff-lint" { } ''
+          export RUFF_CACHE_DIR="$TMPDIR/ruff-cache"
+          ${pkgs.ruff}/bin/ruff check ${inputs.self}/
+          touch $out
+        '';
+      };
+    };
+}

src/exo/download/coordinator.py

@@ -0,0 +1,284 @@
+import asyncio
+from dataclasses import dataclass, field
+from typing import Iterator
+
+import anyio
+from anyio import current_time
+from anyio.abc import TaskGroup
+from loguru import logger
+
+from exo.download.download_utils import (
+    RepoDownloadProgress,
+    delete_model,
+    map_repo_download_progress_to_download_progress_data,
+)
+from exo.download.shard_downloader import ShardDownloader
+from exo.shared.models.model_cards import ModelId
+from exo.shared.types.commands import (
+    DeleteDownload,
+    ForwarderDownloadCommand,
+    StartDownload,
+)
+from exo.shared.types.common import NodeId, SessionId
+from exo.shared.types.events import (
+    Event,
+    ForwarderEvent,
+    NodeDownloadProgress,
+)
+from exo.shared.types.worker.downloads import (
+    DownloadCompleted,
+    DownloadFailed,
+    DownloadOngoing,
+    DownloadPending,
+    DownloadProgress,
+)
+from exo.shared.types.worker.shards import ShardMetadata
+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]
+    local_event_sender: Sender[ForwarderEvent]
+    event_index_counter: Iterator[int]
+
+    # Local state
+    download_status: dict[ModelId, DownloadProgress] = field(default_factory=dict)
+    active_downloads: dict[ModelId, asyncio.Task[None]] = field(default_factory=dict)
+
+    # 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)
+
+    def __post_init__(self) -> None:
+        self.event_sender, self.event_receiver = channel[Event]()
+        self._tg = anyio.create_task_group()
+
+    async def run(self) -> None:
+        logger.info("Starting DownloadCoordinator")
+        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)
+
+    def shutdown(self) -> None:
+        self._tg.cancel_scope.cancel()
+
+    async def _command_processor(self) -> None:
+        with self.download_command_receiver as commands:
+            async for cmd in commands:
+                # Only process commands targeting this node
+                if cmd.command.target_node_id != self.node_id:
+                    continue
+
+                match cmd.command:
+                    case StartDownload(shard_metadata=shard):
+                        await self._start_download(shard)
+                    case DeleteDownload(model_id=model_id):
+                        await self._delete_download(model_id)
+
+    async def _start_download(self, shard: ShardMetadata) -> None:
+        model_id = shard.model_card.model_id
+
+        # Check if already downloading or complete
+        if model_id in self.download_status:
+            status = self.download_status[model_id]
+            if isinstance(status, (DownloadOngoing, DownloadCompleted)):
+                logger.debug(
+                    f"Download for {model_id} already in progress or complete, skipping"
+                )
+                return
+
+        # Emit pending status
+        progress = DownloadPending(shard_metadata=shard, node_id=self.node_id)
+        self.download_status[model_id] = progress
+        await self.event_sender.send(NodeDownloadProgress(download_progress=progress))
+
+        # Check initial status from downloader
+        initial_progress = (
+            await self.shard_downloader.get_shard_download_status_for_shard(shard)
+        )
+
+        if initial_progress.status == "complete":
+            completed = DownloadCompleted(
+                shard_metadata=shard,
+                node_id=self.node_id,
+                total_bytes=initial_progress.total_bytes,
+            )
+            self.download_status[model_id] = completed
+            await self.event_sender.send(
+                NodeDownloadProgress(download_progress=completed)
+            )
+            return
+
+        # Start actual download
+        self._start_download_task(shard, initial_progress)
+
+    def _start_download_task(
+        self, shard: ShardMetadata, initial_progress: RepoDownloadProgress
+    ) -> None:
+        model_id = shard.model_card.model_id
+
+        # Emit ongoing status
+        status = DownloadOngoing(
+            node_id=self.node_id,
+            shard_metadata=shard,
+            download_progress=map_repo_download_progress_to_download_progress_data(
+                initial_progress
+            ),
+        )
+        self.download_status[model_id] = status
+        self.event_sender.send_nowait(NodeDownloadProgress(download_progress=status))
+
+        last_progress_time = 0.0
+        throttle_interval_secs = 1.0
+
+        async def download_progress_callback(
+            callback_shard: ShardMetadata, progress: RepoDownloadProgress
+        ) -> None:
+            nonlocal last_progress_time
+
+            if progress.status == "complete":
+                completed = DownloadCompleted(
+                    shard_metadata=callback_shard,
+                    node_id=self.node_id,
+                    total_bytes=progress.total_bytes,
+                )
+                self.download_status[callback_shard.model_card.model_id] = completed
+                await self.event_sender.send(
+                    NodeDownloadProgress(download_progress=completed)
+                )
+                # Clean up active download tracking
+                if callback_shard.model_card.model_id in self.active_downloads:
+                    del self.active_downloads[callback_shard.model_card.model_id]
+            elif (
+                progress.status == "in_progress"
+                and current_time() - last_progress_time > throttle_interval_secs
+            ):
+                ongoing = DownloadOngoing(
+                    node_id=self.node_id,
+                    shard_metadata=callback_shard,
+                    download_progress=map_repo_download_progress_to_download_progress_data(
+                        progress
+                    ),
+                )
+                self.download_status[callback_shard.model_card.model_id] = ongoing
+                await self.event_sender.send(
+                    NodeDownloadProgress(download_progress=ongoing)
+                )
+                last_progress_time = current_time()
+
+        self.shard_downloader.on_progress(download_progress_callback)
+
+        async def download_wrapper() -> None:
+            try:
+                await self.shard_downloader.ensure_shard(shard)
+            except Exception as e:
+                logger.error(f"Download failed for {model_id}: {e}")
+                failed = DownloadFailed(
+                    shard_metadata=shard,
+                    node_id=self.node_id,
+                    error_message=str(e),
+                )
+                self.download_status[model_id] = failed
+                await self.event_sender.send(
+                    NodeDownloadProgress(download_progress=failed)
+                )
+            finally:
+                if model_id in self.active_downloads:
+                    del self.active_downloads[model_id]
+
+        task = asyncio.create_task(download_wrapper())
+        self.active_downloads[model_id] = task
+
+    async def _delete_download(self, model_id: ModelId) -> None:
+        # Cancel if active
+        if model_id in self.active_downloads:
+            logger.info(f"Cancelling active download for {model_id} before deletion")
+            self.active_downloads[model_id].cancel()
+            del self.active_downloads[model_id]
+
+        # Delete from disk
+        logger.info(f"Deleting model files for {model_id}")
+        deleted = await delete_model(model_id)
+
+        if deleted:
+            logger.info(f"Successfully deleted model {model_id}")
+        else:
+            logger.warning(f"Model {model_id} was not found on disk")
+
+        # Emit pending status to reset UI state, then remove from local tracking
+        if model_id in self.download_status:
+            current_status = self.download_status[model_id]
+            pending = DownloadPending(
+                shard_metadata=current_status.shard_metadata,
+                node_id=self.node_id,
+            )
+            await self.event_sender.send(
+                NodeDownloadProgress(download_progress=pending)
+            )
+            del self.download_status[model_id]
+
+    async def _forward_events(self) -> None:
+        with self.event_receiver as events:
+            async for event in events:
+                idx = next(self.event_index_counter)
+                fe = ForwarderEvent(
+                    origin_idx=idx,
+                    origin=self.node_id,
+                    session=self.session_id,
+                    event=event,
+                )
+                logger.debug(
+                    f"DownloadCoordinator published event {idx}: {str(event)[:100]}"
+                )
+                await self.local_event_sender.send(fe)
+
+    async def _emit_existing_download_progress(self) -> None:
+        try:
+            while True:
+                logger.info(
+                    "DownloadCoordinator: Fetching and emitting existing download progress..."
+                )
+                async for (
+                    _,
+                    progress,
+                ) in self.shard_downloader.get_shard_download_status():
+                    if progress.status == "complete":
+                        status: DownloadProgress = DownloadCompleted(
+                            node_id=self.node_id,
+                            shard_metadata=progress.shard,
+                            total_bytes=progress.total_bytes,
+                        )
+                    elif progress.status in ["in_progress", "not_started"]:
+                        if progress.downloaded_bytes_this_session.in_bytes == 0:
+                            status = DownloadPending(
+                                node_id=self.node_id, shard_metadata=progress.shard
+                            )
+                        else:
+                            status = DownloadOngoing(
+                                node_id=self.node_id,
+                                shard_metadata=progress.shard,
+                                download_progress=map_repo_download_progress_to_download_progress_data(
+                                    progress
+                                ),
+                            )
+                    else:
+                        continue
+
+                    self.download_status[progress.shard.model_card.model_id] = status
+                    await self.event_sender.send(
+                        NodeDownloadProgress(download_progress=status)
+                    )
+                logger.info(
+                    "DownloadCoordinator: Done emitting existing download progress."
+                )
+                await anyio.sleep(5 * 60)  # 5 minutes
+        except Exception as e:
+            logger.error(
+                f"DownloadCoordinator: Error emitting existing download progress: {e}"
+            )

src/exo/download/download_utils.py

@@ -24,7 +24,15 @@ from pydantic import (
     TypeAdapter,
 )
 
+from exo.download.huggingface_utils import (
+    filter_repo_objects,
+    get_allow_patterns,
+    get_auth_headers,
+    get_hf_endpoint,
+    get_hf_token,
+)
 from exo.shared.constants import EXO_MODELS_DIR
+from exo.shared.models.model_cards import ModelTask
 from exo.shared.types.common import ModelId
 from exo.shared.types.memory import Memory
 from exo.shared.types.worker.downloads import (
@@ -35,12 +43,27 @@ from exo.shared.types.worker.downloads import (
     RepoFileDownloadProgress,
 )
 from exo.shared.types.worker.shards import ShardMetadata
-from exo.worker.download.huggingface_utils import (
-    filter_repo_objects,
-    get_allow_patterns,
-    get_auth_headers,
-    get_hf_endpoint,
-)
+
+
+class HuggingFaceAuthenticationError(Exception):
+    """Raised when HuggingFace returns 401/403 for a model download."""
+
+
+async def _build_auth_error_message(status_code: int, model_id: ModelId) -> str:
+    token = await get_hf_token()
+    if status_code == 401 and token is None:
+        return (
+            f"Model '{model_id}' requires authentication. "
+            f"Set HF_TOKEN in the app's Advanced settings, set the HF_TOKEN environment variable, or run `hf auth login`. "
+            f"Get a token at https://huggingface.co/settings/tokens"
+        )
+    elif status_code == 403:
+        return (
+            f"Access denied to '{model_id}'. "
+            f"Please accept the model terms at https://huggingface.co/{model_id}"
+        )
+    else:
+        return f"Authentication failed for '{model_id}' (HTTP {status_code})"
 
 
 def trim_etag(etag: str) -> str:
@@ -98,11 +121,20 @@ async def ensure_models_dir() -> Path:
 
 
 async def delete_model(model_id: ModelId) -> bool:
-    model_dir = await ensure_models_dir() / model_id.normalize()
-    if not await aios.path.exists(model_dir):
-        return False
-    await asyncio.to_thread(shutil.rmtree, model_dir, ignore_errors=False)
-    return True
+    models_dir = await ensure_models_dir()
+    model_dir = models_dir / model_id.normalize()
+    cache_dir = models_dir / "caches" / model_id.normalize()
+
+    deleted = False
+    if await aios.path.exists(model_dir):
+        await asyncio.to_thread(shutil.rmtree, model_dir, ignore_errors=False)
+        deleted = True
+
+    # Also clear cache
+    if await aios.path.exists(cache_dir):
+        await asyncio.to_thread(shutil.rmtree, cache_dir, ignore_errors=False)
+
+    return deleted
 
 
 async def seed_models(seed_dir: str | Path):
@@ -128,16 +160,28 @@ async def fetch_file_list_with_cache(
     target_dir = (await ensure_models_dir()) / "caches" / model_id.normalize()
     await aios.makedirs(target_dir, exist_ok=True)
     cache_file = target_dir / f"{model_id.normalize()}--{revision}--file_list.json"
-    if await aios.path.exists(cache_file):
-        async with aiofiles.open(cache_file, "r") as f:
-            return TypeAdapter(list[FileListEntry]).validate_json(await f.read())
-    file_list = await fetch_file_list_with_retry(
-        model_id, revision, recursive=recursive
-    )
-    await aios.makedirs(cache_file.parent, exist_ok=True)
-    async with aiofiles.open(cache_file, "w") as f:
-        await f.write(TypeAdapter(list[FileListEntry]).dump_json(file_list).decode())
-    return file_list
+
+    # Always try fresh first
+    try:
+        file_list = await fetch_file_list_with_retry(
+            model_id, revision, recursive=recursive
+        )
+        # Update cache with fresh data
+        async with aiofiles.open(cache_file, "w") as f:
+            await f.write(
+                TypeAdapter(list[FileListEntry]).dump_json(file_list).decode()
+            )
+        return file_list
+    except Exception as e:
+        # Fetch failed - try cache fallback
+        if await aios.path.exists(cache_file):
+            logger.warning(
+                f"Failed to fetch file list for {model_id}, using cached data: {e}"
+            )
+            async with aiofiles.open(cache_file, "r") as f:
+                return TypeAdapter(list[FileListEntry]).validate_json(await f.read())
+        # No cache available, propagate the error
+        raise
 
 
 async def fetch_file_list_with_retry(
@@ -147,6 +191,8 @@ async def fetch_file_list_with_retry(
     for attempt in range(n_attempts):
         try:
             return await _fetch_file_list(model_id, revision, path, recursive)
+        except HuggingFaceAuthenticationError:
+            raise
         except Exception as e:
             if attempt == n_attempts - 1:
                 raise e
@@ -167,6 +213,9 @@ async def _fetch_file_list(
         create_http_session(timeout_profile="short") as session,
         session.get(url, headers=headers) as response,
     ):
+        if response.status in [401, 403]:
+            msg = await _build_auth_error_message(response.status, model_id)
+            raise HuggingFaceAuthenticationError(msg)
         if response.status == 200:
             data_json = await response.text()
             data = TypeAdapter(list[FileListEntry]).validate_json(data_json)
@@ -256,6 +305,9 @@ async def file_meta(
             # Otherwise, follow the redirect to get authoritative size/hash
             redirected_location = r.headers.get("location")
             return await file_meta(model_id, revision, path, redirected_location)
+        if r.status in [401, 403]:
+            msg = await _build_auth_error_message(r.status, model_id)
+            raise HuggingFaceAuthenticationError(msg)
         content_length = int(
             r.headers.get("x-linked-size") or r.headers.get("content-length") or 0
         )
@@ -279,6 +331,8 @@ async def download_file_with_retry(
             return await _download_file(
                 model_id, revision, path, target_dir, on_progress
             )
+        except HuggingFaceAuthenticationError:
+            raise
         except Exception as e:
             if isinstance(e, FileNotFoundError) or attempt == n_attempts - 1:
                 raise e
@@ -299,8 +353,28 @@ async def _download_file(
     target_dir: Path,
     on_progress: Callable[[int, int, bool], None] = lambda _, __, ___: None,
 ) -> Path:
-    if await aios.path.exists(target_dir / path):
-        return target_dir / path
+    target_path = target_dir / path
+
+    if await aios.path.exists(target_path):
+        local_size = (await aios.stat(target_path)).st_size
+
+        # Try to verify against remote, but allow offline operation
+        try:
+            remote_size, _ = await file_meta(model_id, revision, path)
+            if local_size != remote_size:
+                logger.info(
+                    f"File {path} size mismatch (local={local_size}, remote={remote_size}), re-downloading"
+                )
+                await aios.remove(target_path)
+            else:
+                return target_path
+        except Exception as e:
+            # Offline or network error - trust local file
+            logger.debug(
+                f"Could not verify {path} against remote (offline?): {e}, using local file"
+            )
+            return target_path
+
     await aios.makedirs((target_dir / path).parent, exist_ok=True)
     length, etag = await file_meta(model_id, revision, path)
     remote_hash = etag[:-5] if etag.endswith("-gzip") else etag
@@ -322,6 +396,9 @@ async def _download_file(
         ):
             if r.status == 404:
                 raise FileNotFoundError(f"File not found: {url}")
+            if r.status in [401, 403]:
+                msg = await _build_auth_error_message(r.status, model_id)
+                raise HuggingFaceAuthenticationError(msg)
             assert r.status in [200, 206], (
                 f"Failed to download {path} from {url}: {r.status}"
             )
@@ -446,6 +523,11 @@ async def resolve_allow_patterns(shard: ShardMetadata) -> list[str]:
         return ["*"]
 
 
+def is_image_model(shard: ShardMetadata) -> bool:
+    tasks = shard.model_card.tasks
+    return ModelTask.TextToImage in tasks or ModelTask.ImageToImage in tasks
+
+
 async def get_downloaded_size(path: Path) -> int:
     partial_path = path.with_suffix(path.suffix + ".partial")
     if await aios.path.exists(path):
@@ -463,7 +545,7 @@ async def download_shard(
     allow_patterns: list[str] | None = None,
 ) -> tuple[Path, RepoDownloadProgress]:
     if not skip_download:
-        logger.info(f"Downloading {shard.model_card.model_id=}")
+        logger.debug(f"Downloading {shard.model_card.model_id=}")
 
     revision = "main"
     target_dir = await ensure_models_dir() / str(shard.model_card.model_id).replace(
@@ -476,7 +558,7 @@ async def download_shard(
         allow_patterns = await resolve_allow_patterns(shard)
 
     if not skip_download:
-        logger.info(f"Downloading {shard.model_card.model_id=} with {allow_patterns=}")
+        logger.debug(f"Downloading {shard.model_card.model_id=} with {allow_patterns=}")
 
     all_start_time = time.time()
     file_list = await fetch_file_list_with_cache(
@@ -487,22 +569,40 @@ async def download_shard(
             file_list, allow_patterns=allow_patterns, key=lambda x: x.path
         )
     )
+
+    # For image models, skip root-level safetensors files since weights
+    # are stored in component subdirectories (e.g., transformer/, vae/)
+    if is_image_model(shard):
+        filtered_file_list = [
+            f
+            for f in filtered_file_list
+            if "/" in f.path or not f.path.endswith(".safetensors")
+        ]
     file_progress: dict[str, RepoFileDownloadProgress] = {}
 
     async def on_progress_wrapper(
         file: FileListEntry, curr_bytes: int, total_bytes: int, is_renamed: bool
     ) -> None:
-        start_time = (
-            file_progress[file.path].start_time
-            if file.path in file_progress
-            else time.time()
-        )
-        downloaded_this_session = (
-            file_progress[file.path].downloaded_this_session.in_bytes
-            + (curr_bytes - file_progress[file.path].downloaded.in_bytes)
-            if file.path in file_progress
-            else curr_bytes
+        previous_progress = file_progress.get(file.path)
+
+        # Detect re-download: curr_bytes < previous downloaded means file was deleted and restarted
+        is_redownload = (
+            previous_progress is not None
+            and curr_bytes < previous_progress.downloaded.in_bytes
         )
+
+        if is_redownload or previous_progress is None:
+            # Fresh download or re-download: reset tracking
+            start_time = time.time()
+            downloaded_this_session = curr_bytes
+        else:
+            # Continuing download: accumulate
+            start_time = previous_progress.start_time
+            downloaded_this_session = (
+                previous_progress.downloaded_this_session.in_bytes
+                + (curr_bytes - previous_progress.downloaded.in_bytes)
+            )
+
         speed = (
             downloaded_this_session / (time.time() - start_time)
             if time.time() - start_time > 0

src/exo/download/huggingface_utils.py

@@ -68,7 +68,11 @@ def get_hf_home() -> Path:
 
 
 async def get_hf_token() -> str | None:
-    """Retrieve the Hugging Face token from the user's HF_HOME directory."""
+    """Retrieve the Hugging Face token from HF_TOKEN env var or HF_HOME directory."""
+    # Check environment variable first
+    if token := os.environ.get("HF_TOKEN"):
+        return token
+    # Fall back to file-based token
     token_path = get_hf_home() / "token"
     if await aios.path.exists(token_path):
         async with aiofiles.open(token_path, "r") as f:

src/exo/download/impl_shard_downloader.py

@@ -3,13 +3,15 @@ from collections.abc import Awaitable
 from pathlib import Path
 from typing import AsyncIterator, Callable
 
+from loguru import logger
+
+from exo.download.download_utils import RepoDownloadProgress, download_shard
+from exo.download.shard_downloader import ShardDownloader
 from exo.shared.models.model_cards import MODEL_CARDS, ModelCard, ModelId
 from exo.shared.types.worker.shards import (
     PipelineShardMetadata,
     ShardMetadata,
 )
-from exo.worker.download.download_utils import RepoDownloadProgress, download_shard
-from exo.worker.download.shard_downloader import ShardDownloader
 
 
 def exo_shard_downloader(max_parallel_downloads: int = 8) -> ShardDownloader:
@@ -164,9 +166,8 @@ class ResumableShardDownloader(ShardDownloader):
         for task in asyncio.as_completed(tasks):
             try:
                 yield await task
-            # TODO: except Exception
             except Exception as e:
-                print("Error downloading shard:", e)
+                logger.warning(f"Error downloading shard: {type(e).__name__}")
 
     async def get_shard_download_status_for_shard(
         self, shard: ShardMetadata

src/exo/download/shard_downloader.py

@@ -5,13 +5,13 @@ from datetime import timedelta
 from pathlib import Path
 from typing import AsyncIterator, Callable
 
+from exo.download.download_utils import RepoDownloadProgress
 from exo.shared.models.model_cards import ModelCard, ModelId, ModelTask
 from exo.shared.types.memory import Memory
 from exo.shared.types.worker.shards import (
     PipelineShardMetadata,
     ShardMetadata,
 )
-from exo.worker.download.download_utils import RepoDownloadProgress
 
 
 # TODO: the PipelineShardMetadata getting reinstantiated is a bit messy. Should this be a classmethod?

src/exo/download/tests/test_download_verification.py

@@ -0,0 +1,451 @@
+"""Tests for download verification and cache behavior."""
+
+import time
+from collections.abc import AsyncIterator
+from datetime import timedelta
+from pathlib import Path
+from unittest.mock import AsyncMock, MagicMock, patch
+
+import aiofiles
+import aiofiles.os as aios
+import pytest
+from pydantic import TypeAdapter
+
+from exo.download.download_utils import (
+    delete_model,
+    fetch_file_list_with_cache,
+)
+from exo.shared.types.common import ModelId
+from exo.shared.types.memory import Memory
+from exo.shared.types.worker.downloads import FileListEntry, RepoFileDownloadProgress
+
+
+@pytest.fixture
+def model_id() -> ModelId:
+    return ModelId("test-org/test-model")
+
+
+@pytest.fixture
+async def temp_models_dir(tmp_path: Path) -> AsyncIterator[Path]:
+    """Set up a temporary models directory for testing."""
+    models_dir = tmp_path / "models"
+    await aios.makedirs(models_dir, exist_ok=True)
+    with patch("exo.download.download_utils.EXO_MODELS_DIR", models_dir):
+        yield models_dir
+
+
+class TestFileVerification:
+    """Tests for file size verification in _download_file."""
+
+    async def test_redownload_when_file_size_changes_upstream(
+        self, model_id: ModelId, tmp_path: Path
+    ) -> None:
+        """Test that files with mismatched sizes are re-downloaded."""
+        # Import inside test to allow patching
+        from exo.download.download_utils import (
+            _download_file,  # pyright: ignore[reportPrivateUsage]
+        )
+
+        target_dir = tmp_path / "downloads"
+        await aios.makedirs(target_dir, exist_ok=True)
+
+        # Create a local file with wrong size
+        local_file = target_dir / "test.safetensors"
+        async with aiofiles.open(local_file, "wb") as f:
+            await f.write(b"local content")  # 13 bytes
+
+        remote_size = 1000  # Different from local
+        remote_hash = "abc123"
+
+        with (
+            patch(
+                "exo.download.download_utils.file_meta",
+                new_callable=AsyncMock,
+                return_value=(remote_size, remote_hash),
+            ) as mock_file_meta,
+            patch(
+                "exo.download.download_utils.create_http_session"
+            ) as mock_session_factory,
+        ):
+            # Set up mock HTTP response for re-download
+            mock_response = MagicMock()
+            mock_response.status = 200
+            mock_response.content.read = AsyncMock(  # pyright: ignore[reportAny]
+                side_effect=[b"x" * remote_size, b""]
+            )
+
+            mock_session = MagicMock()
+            mock_session.get.return_value.__aenter__ = AsyncMock(  # pyright: ignore[reportAny]
+                return_value=mock_response
+            )
+            mock_session.get.return_value.__aexit__ = AsyncMock(  # pyright: ignore[reportAny]
+                return_value=None
+            )
+            mock_session_factory.return_value.__aenter__ = AsyncMock(  # pyright: ignore[reportAny]
+                return_value=mock_session
+            )
+            mock_session_factory.return_value.__aexit__ = AsyncMock(  # pyright: ignore[reportAny]
+                return_value=None
+            )
+
+            # Mock calc_hash to return the expected hash
+            with patch(
+                "exo.download.download_utils.calc_hash",
+                new_callable=AsyncMock,
+                return_value=remote_hash,
+            ):
+                await _download_file(model_id, "main", "test.safetensors", target_dir)
+
+            # file_meta should be called twice: once for verification, once for download
+            assert mock_file_meta.call_count == 2
+
+    async def test_skip_download_when_file_size_matches(
+        self, model_id: ModelId, tmp_path: Path
+    ) -> None:
+        """Test that files with matching sizes are not re-downloaded."""
+        from exo.download.download_utils import (
+            _download_file,  # pyright: ignore[reportPrivateUsage]
+        )
+
+        target_dir = tmp_path / "downloads"
+        await aios.makedirs(target_dir, exist_ok=True)
+
+        # Create a local file
+        local_file = target_dir / "test.safetensors"
+        local_content = b"local content"
+        async with aiofiles.open(local_file, "wb") as f:
+            await f.write(local_content)
+
+        remote_size = len(local_content)  # Same as local
+        remote_hash = "abc123"
+
+        with (
+            patch(
+                "exo.download.download_utils.file_meta",
+                new_callable=AsyncMock,
+                return_value=(remote_size, remote_hash),
+            ) as mock_file_meta,
+            patch(
+                "exo.download.download_utils.create_http_session"
+            ) as mock_session_factory,
+        ):
+            result = await _download_file(
+                model_id, "main", "test.safetensors", target_dir
+            )
+
+            # Should return immediately without downloading
+            assert result == local_file
+            mock_file_meta.assert_called_once()
+            mock_session_factory.assert_not_called()
+
+    async def test_offline_fallback_uses_local_file(
+        self, model_id: ModelId, tmp_path: Path
+    ) -> None:
+        """Test that local files are used when network is unavailable."""
+        from exo.download.download_utils import (
+            _download_file,  # pyright: ignore[reportPrivateUsage]
+        )
+
+        target_dir = tmp_path / "downloads"
+        await aios.makedirs(target_dir, exist_ok=True)
+
+        # Create a local file
+        local_file = target_dir / "test.safetensors"
+        async with aiofiles.open(local_file, "wb") as f:
+            await f.write(b"local content")
+
+        with (
+            patch(
+                "exo.download.download_utils.file_meta",
+                new_callable=AsyncMock,
+                side_effect=Exception("Network error"),
+            ),
+            patch(
+                "exo.download.download_utils.create_http_session"
+            ) as mock_session_factory,
+        ):
+            result = await _download_file(
+                model_id, "main", "test.safetensors", target_dir
+            )
+
+            # Should return local file without attempting download
+            assert result == local_file
+            mock_session_factory.assert_not_called()
+
+
+class TestFileListCache:
+    """Tests for file list caching behavior."""
+
+    async def test_fetch_fresh_and_update_cache(
+        self, model_id: ModelId, tmp_path: Path
+    ) -> None:
+        """Test that fresh data is fetched and cache is updated."""
+        models_dir = tmp_path / "models"
+
+        file_list = [
+            FileListEntry(type="file", path="model.safetensors", size=1000),
+            FileListEntry(type="file", path="config.json", size=100),
+        ]
+
+        with (
+            patch("exo.download.download_utils.EXO_MODELS_DIR", models_dir),
+            patch(
+                "exo.download.download_utils.fetch_file_list_with_retry",
+                new_callable=AsyncMock,
+                return_value=file_list,
+            ) as mock_fetch,
+        ):
+            result = await fetch_file_list_with_cache(model_id, "main")
+
+            assert result == file_list
+            mock_fetch.assert_called_once()
+
+            # Verify cache was written
+            cache_file = (
+                models_dir
+                / "caches"
+                / model_id.normalize()
+                / f"{model_id.normalize()}--main--file_list.json"
+            )
+            assert await aios.path.exists(cache_file)
+
+            async with aiofiles.open(cache_file, "r") as f:
+                cached_data = TypeAdapter(list[FileListEntry]).validate_json(
+                    await f.read()
+                )
+            assert cached_data == file_list
+
+    async def test_fallback_to_cache_when_fetch_fails(
+        self, model_id: ModelId, tmp_path: Path
+    ) -> None:
+        """Test that cached data is used when fetch fails."""
+        models_dir = tmp_path / "models"
+        cache_dir = models_dir / "caches" / model_id.normalize()
+        await aios.makedirs(cache_dir, exist_ok=True)
+
+        # Create cache file
+        cached_file_list = [
+            FileListEntry(type="file", path="model.safetensors", size=1000),
+        ]
+        cache_file = cache_dir / f"{model_id.normalize()}--main--file_list.json"
+        async with aiofiles.open(cache_file, "w") as f:
+            await f.write(
+                TypeAdapter(list[FileListEntry]).dump_json(cached_file_list).decode()
+            )
+
+        with (
+            patch("exo.download.download_utils.EXO_MODELS_DIR", models_dir),
+            patch(
+                "exo.download.download_utils.fetch_file_list_with_retry",
+                new_callable=AsyncMock,
+                side_effect=Exception("Network error"),
+            ),
+        ):
+            result = await fetch_file_list_with_cache(model_id, "main")
+
+            assert result == cached_file_list
+
+    async def test_error_propagates_when_no_cache(
+        self, model_id: ModelId, tmp_path: Path
+    ) -> None:
+        """Test that errors propagate when fetch fails and no cache exists."""
+        models_dir = tmp_path / "models"
+
+        with (
+            patch("exo.download.download_utils.EXO_MODELS_DIR", models_dir),
+            patch(
+                "exo.download.download_utils.fetch_file_list_with_retry",
+                new_callable=AsyncMock,
+                side_effect=Exception("Network error"),
+            ),
+            pytest.raises(Exception, match="Network error"),
+        ):
+            await fetch_file_list_with_cache(model_id, "main")
+
+
+class TestModelDeletion:
+    """Tests for model deletion including cache cleanup."""
+
+    async def test_delete_model_clears_cache(
+        self, model_id: ModelId, tmp_path: Path
+    ) -> None:
+        """Test that deleting a model also deletes its cache."""
+        models_dir = tmp_path / "models"
+        model_dir = models_dir / model_id.normalize()
+        cache_dir = models_dir / "caches" / model_id.normalize()
+
+        # Create model and cache directories
+        await aios.makedirs(model_dir, exist_ok=True)
+        await aios.makedirs(cache_dir, exist_ok=True)
+
+        # Add some files
+        async with aiofiles.open(model_dir / "model.safetensors", "w") as f:
+            await f.write("model data")
+        async with aiofiles.open(cache_dir / "file_list.json", "w") as f:
+            await f.write("[]")
+
+        with patch("exo.download.download_utils.EXO_MODELS_DIR", models_dir):
+            result = await delete_model(model_id)
+
+            assert result is True
+            assert not await aios.path.exists(model_dir)
+            assert not await aios.path.exists(cache_dir)
+
+    async def test_delete_model_only_cache_exists(
+        self, model_id: ModelId, tmp_path: Path
+    ) -> None:
+        """Test deleting when only cache exists (model already deleted)."""
+        models_dir = tmp_path / "models"
+        cache_dir = models_dir / "caches" / model_id.normalize()
+
+        # Only create cache directory
+        await aios.makedirs(cache_dir, exist_ok=True)
+        async with aiofiles.open(cache_dir / "file_list.json", "w") as f:
+            await f.write("[]")
+
+        with patch("exo.download.download_utils.EXO_MODELS_DIR", models_dir):
+            result = await delete_model(model_id)
+
+            # Returns False because model dir didn't exist
+            assert result is False
+            # But cache should still be cleaned up
+            assert not await aios.path.exists(cache_dir)
+
+    async def test_delete_nonexistent_model(
+        self, model_id: ModelId, tmp_path: Path
+    ) -> None:
+        """Test deleting a model that doesn't exist."""
+        models_dir = tmp_path / "models"
+        await aios.makedirs(models_dir, exist_ok=True)
+
+        with patch("exo.download.download_utils.EXO_MODELS_DIR", models_dir):
+            result = await delete_model(model_id)
+
+            assert result is False
+
+
+class TestProgressResetOnRedownload:
+    """Tests for progress tracking when files are re-downloaded."""
+
+    async def test_progress_resets_correctly_on_redownload(
+        self, model_id: ModelId
+    ) -> None:
+        """Test that progress tracking resets when a file is re-downloaded.
+
+        When a file is deleted and re-downloaded (due to size mismatch),
+        the progress tracking should reset rather than calculating negative
+        downloaded_this_session values.
+        """
+        # Simulate file_progress dict as it exists in download_shard
+        file_progress: dict[str, RepoFileDownloadProgress] = {}
+
+        # Initialize with old file progress (simulating existing large file)
+        old_file_size = 1_500_000_000  # 1.5 GB
+        file_progress["model.safetensors"] = RepoFileDownloadProgress(
+            repo_id=model_id,
+            repo_revision="main",
+            file_path="model.safetensors",
+            downloaded=Memory.from_bytes(old_file_size),
+            downloaded_this_session=Memory.from_bytes(0),
+            total=Memory.from_bytes(old_file_size),
+            speed=0,
+            eta=timedelta(0),
+            status="not_started",
+            start_time=time.time() - 10,  # Started 10 seconds ago
+        )
+
+        # Simulate the logic from on_progress_wrapper after re-download starts
+        # This is the exact logic from the fixed on_progress_wrapper
+        curr_bytes = 100_000  # 100 KB - new download just started
+        previous_progress = file_progress.get("model.safetensors")
+
+        # Detect re-download: curr_bytes < previous downloaded
+        is_redownload = (
+            previous_progress is not None
+            and curr_bytes < previous_progress.downloaded.in_bytes
+        )
+
+        if is_redownload or previous_progress is None:
+            # Fresh download or re-download: reset tracking
+            start_time = time.time()
+            downloaded_this_session = curr_bytes
+        else:
+            # Continuing download: accumulate
+            start_time = previous_progress.start_time
+            downloaded_this_session = (
+                previous_progress.downloaded_this_session.in_bytes
+                + (curr_bytes - previous_progress.downloaded.in_bytes)
+            )
+
+        # Key assertions
+        assert is_redownload is True, "Should detect re-download scenario"
+        assert downloaded_this_session == curr_bytes, (
+            "downloaded_this_session should equal curr_bytes on re-download"
+        )
+        assert downloaded_this_session > 0, (
+            "downloaded_this_session should be positive, not negative"
+        )
+
+        # Calculate speed (should be positive)
+        elapsed = time.time() - start_time
+        speed = downloaded_this_session / elapsed if elapsed > 0 else 0
+        assert speed >= 0, "Speed should be non-negative"
+
+    async def test_progress_accumulates_on_continuing_download(
+        self, model_id: ModelId
+    ) -> None:
+        """Test that progress accumulates correctly for continuing downloads.
+
+        When a download continues from where it left off (resume),
+        the progress should accumulate correctly.
+        """
+        file_progress: dict[str, RepoFileDownloadProgress] = {}
+
+        # Initialize with partial download progress
+        initial_downloaded = 500_000  # 500 KB already downloaded
+        start_time = time.time() - 5  # Started 5 seconds ago
+        file_progress["model.safetensors"] = RepoFileDownloadProgress(
+            repo_id=model_id,
+            repo_revision="main",
+            file_path="model.safetensors",
+            downloaded=Memory.from_bytes(initial_downloaded),
+            downloaded_this_session=Memory.from_bytes(initial_downloaded),
+            total=Memory.from_bytes(1_000_000),
+            speed=100_000,
+            eta=timedelta(seconds=5),
+            status="in_progress",
+            start_time=start_time,
+        )
+
+        # Progress callback with more bytes downloaded
+        curr_bytes = 600_000  # 600 KB - continuing download
+        previous_progress = file_progress.get("model.safetensors")
+
+        # This is NOT a re-download (curr_bytes > previous downloaded)
+        is_redownload = (
+            previous_progress is not None
+            and curr_bytes < previous_progress.downloaded.in_bytes
+        )
+
+        if is_redownload or previous_progress is None:
+            downloaded_this_session = curr_bytes
+            used_start_time = time.time()
+        else:
+            used_start_time = previous_progress.start_time
+            downloaded_this_session = (
+                previous_progress.downloaded_this_session.in_bytes
+                + (curr_bytes - previous_progress.downloaded.in_bytes)
+            )
+
+        # Key assertions
+        assert is_redownload is False, (
+            "Should NOT detect re-download for continuing download"
+        )
+        assert used_start_time == start_time, "Should preserve original start_time"
+        expected_session = initial_downloaded + (curr_bytes - initial_downloaded)
+        assert downloaded_this_session == expected_session, (
+            f"Should accumulate: {downloaded_this_session} == {expected_session}"
+        )
+        assert downloaded_this_session == 600_000, (
+            "downloaded_this_session should equal total downloaded so far"
+        )

src/exo/main.py

@@ -1,10 +1,11 @@
 import argparse
+import itertools
 import multiprocessing as mp
 import os
 import resource
 import signal
 from dataclasses import dataclass, field
-from typing import Self
+from typing import Iterator, Self
 
 import anyio
 from anyio.abc import TaskGroup
@@ -12,6 +13,8 @@ from loguru import logger
 from pydantic import PositiveInt
 
 import exo.routing.topics as topics
+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.router import Router, get_node_id_keypair
@@ -21,7 +24,6 @@ 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.worker.download.impl_shard_downloader import exo_shard_downloader
 from exo.worker.main import Worker
 
 
@@ -29,6 +31,7 @@ from exo.worker.main import Worker
 @dataclass
 class Node:
     router: Router
+    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.
     election_result_receiver: Receiver[ElectionResult]
@@ -36,6 +39,7 @@ class Node:
     api: API | None
 
     node_id: NodeId
+    event_index_counter: Iterator[int]
     _tg: TaskGroup = field(init=False, default_factory=anyio.create_task_group)
 
     @classmethod
@@ -49,8 +53,26 @@ class Node:
         await router.register_topic(topics.COMMANDS)
         await router.register_topic(topics.ELECTION_MESSAGES)
         await router.register_topic(topics.CONNECTION_MESSAGES)
+        await router.register_topic(topics.DOWNLOAD_COMMANDS)
 
         logger.info(f"Starting node {node_id}")
+
+        # Create shared event index counter for Worker and DownloadCoordinator
+        event_index_counter = itertools.count()
+
+        # Create DownloadCoordinator (unless --no-downloads)
+        if not args.no_downloads:
+            download_coordinator = DownloadCoordinator(
+                node_id,
+                session_id,
+                exo_shard_downloader(),
+                download_command_receiver=router.receiver(topics.DOWNLOAD_COMMANDS),
+                local_event_sender=router.sender(topics.LOCAL_EVENTS),
+                event_index_counter=event_index_counter,
+            )
+        else:
+            download_coordinator = None
+
         if args.spawn_api:
             api = API(
                 node_id,
@@ -58,6 +80,7 @@ class Node:
                 port=args.api_port,
                 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),
             )
         else:
@@ -67,11 +90,11 @@ class Node:
             worker = Worker(
                 node_id,
                 session_id,
-                exo_shard_downloader(),
-                connection_message_receiver=router.receiver(topics.CONNECTION_MESSAGES),
                 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),
+                event_index_counter=event_index_counter,
             )
         else:
             worker = None
@@ -99,13 +122,25 @@ class Node:
             election_result_sender=er_send,
         )
 
-        return cls(router, worker, election, er_recv, master, api, node_id)
+        return cls(
+            router,
+            download_coordinator,
+            worker,
+            election,
+            er_recv,
+            master,
+            api,
+            node_id,
+            event_index_counter,
+        )
 
     async def run(self):
         async with self._tg as tg:
             signal.signal(signal.SIGINT, lambda _, __: self.shutdown())
             tg.start_soon(self.router.run)
             tg.start_soon(self.election.run)
+            if self.download_coordinator:
+                tg.start_soon(self.download_coordinator.run)
             if self.worker:
                 tg.start_soon(self.worker.run)
             if self.master:
@@ -170,21 +205,36 @@ class Node:
                     )
                 if result.is_new_master:
                     await anyio.sleep(0)
+                    # Fresh counter for new session (buffer expects indices from 0)
+                    self.event_index_counter = itertools.count()
+                    if self.download_coordinator:
+                        self.download_coordinator.shutdown()
+                        self.download_coordinator = DownloadCoordinator(
+                            self.node_id,
+                            result.session_id,
+                            exo_shard_downloader(),
+                            download_command_receiver=self.router.receiver(
+                                topics.DOWNLOAD_COMMANDS
+                            ),
+                            local_event_sender=self.router.sender(topics.LOCAL_EVENTS),
+                            event_index_counter=self.event_index_counter,
+                        )
+                        self._tg.start_soon(self.download_coordinator.run)
                     if self.worker:
                         self.worker.shutdown()
                         # TODO: add profiling etc to resource monitor
                         self.worker = Worker(
                             self.node_id,
                             result.session_id,
-                            exo_shard_downloader(),
-                            connection_message_receiver=self.router.receiver(
-                                topics.CONNECTION_MESSAGES
-                            ),
                             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
+                            ),
+                            event_index_counter=self.event_index_counter,
                         )
                         self._tg.start_soon(self.worker.run)
                     if self.api:
@@ -226,6 +276,7 @@ class Args(CamelCaseModel):
     api_port: PositiveInt = 52415
     tb_only: bool = False
     no_worker: bool = False
+    no_downloads: bool = False
     fast_synch: bool | None = None  # None = auto, True = force on, False = force off
 
     @classmethod
@@ -268,6 +319,11 @@ class Args(CamelCaseModel):
             "--no-worker",
             action="store_true",
         )
+        parser.add_argument(
+            "--no-downloads",
+            action="store_true",
+            help="Disable the download coordinator (node won't download models)",
+        )
         fast_synch_group = parser.add_mutually_exclusive_group()
         fast_synch_group.add_argument(
             "--fast-synch",

src/exo/master/adapters/__init__.py

@@ -0,0 +1 @@
+"""API adapters for different API formats (Claude, OpenAI Responses, etc.)."""

src/exo/master/adapters/chat_completions.py

@@ -0,0 +1,212 @@
+"""OpenAI Chat Completions API adapter for converting requests/responses."""
+
+import time
+from collections.abc import AsyncGenerator
+from typing import Any
+from uuid import uuid4
+
+from exo.shared.types.api import (
+    ChatCompletionChoice,
+    ChatCompletionMessage,
+    ChatCompletionMessageText,
+    ChatCompletionRequest,
+    ChatCompletionResponse,
+    ErrorInfo,
+    ErrorResponse,
+    FinishReason,
+    StreamingChoiceResponse,
+    ToolCall,
+)
+from exo.shared.types.chunks import ErrorChunk, TokenChunk, ToolCallChunk
+from exo.shared.types.common import CommandId
+from exo.shared.types.text_generation import InputMessage, TextGenerationTaskParams
+
+
+def chat_request_to_text_generation(
+    request: ChatCompletionRequest,
+) -> TextGenerationTaskParams:
+    instructions: str | None = None
+    input_messages: list[InputMessage] = []
+    chat_template_messages: list[dict[str, Any]] = []
+
+    for msg in request.messages:
+        # Normalize content to string
+        content: str
+        if msg.content is None:
+            content = ""
+        elif isinstance(msg.content, str):
+            content = msg.content
+        elif isinstance(msg.content, ChatCompletionMessageText):
+            content = msg.content.text
+        else:
+            # List of ChatCompletionMessageText
+            content = "\n".join(item.text for item in msg.content)
+
+        # Extract system message as instructions
+        if msg.role == "system":
+            if instructions is None:
+                instructions = content
+            else:
+                # Append additional system messages
+                instructions = f"{instructions}\n{content}"
+            chat_template_messages.append({"role": "system", "content": content})
+        else:
+            # Skip messages with no meaningful content
+            if msg.content is None and msg.thinking is None and msg.tool_calls is None:
+                continue
+
+            if msg.role in ("user", "assistant", "developer"):
+                input_messages.append(InputMessage(role=msg.role, content=content))
+
+            # Build full message dict for chat template (preserves tool_calls etc.)
+            # Normalize content for model_dump
+            msg_copy = msg.model_copy(update={"content": content})
+            dumped: dict[str, Any] = msg_copy.model_dump(exclude_none=True)
+            chat_template_messages.append(dumped)
+
+    return TextGenerationTaskParams(
+        model=request.model,
+        input=input_messages if input_messages else "",
+        instructions=instructions,
+        max_output_tokens=request.max_tokens,
+        temperature=request.temperature,
+        top_p=request.top_p,
+        top_k=request.top_k,
+        stop=request.stop,
+        seed=request.seed,
+        stream=request.stream,
+        tools=request.tools,
+        chat_template_messages=chat_template_messages
+        if chat_template_messages
+        else None,
+    )
+
+
+def chunk_to_response(
+    chunk: TokenChunk, command_id: CommandId
+) -> ChatCompletionResponse:
+    """Convert a TokenChunk to a streaming ChatCompletionResponse."""
+    return ChatCompletionResponse(
+        id=command_id,
+        created=int(time.time()),
+        model=chunk.model,
+        choices=[
+            StreamingChoiceResponse(
+                index=0,
+                delta=ChatCompletionMessage(role="assistant", content=chunk.text),
+                finish_reason=chunk.finish_reason,
+            )
+        ],
+    )
+
+
+async def generate_chat_stream(
+    command_id: CommandId,
+    chunk_stream: AsyncGenerator[ErrorChunk | ToolCallChunk | TokenChunk, None],
+) -> AsyncGenerator[str, None]:
+    """Generate Chat Completions API streaming events from chunks."""
+    async for chunk in chunk_stream:
+        if isinstance(chunk, ErrorChunk):
+            error_response = ErrorResponse(
+                error=ErrorInfo(
+                    message=chunk.error_message or "Internal server error",
+                    type="InternalServerError",
+                    code=500,
+                )
+            )
+            yield f"data: {error_response.model_dump_json()}\n\n"
+            yield "data: [DONE]\n\n"
+            return
+
+        if isinstance(chunk, ToolCallChunk):
+            tool_call_deltas = [
+                ToolCall(
+                    id=str(uuid4()),
+                    index=i,
+                    function=tool,
+                )
+                for i, tool in enumerate(chunk.tool_calls)
+            ]
+            tool_response = ChatCompletionResponse(
+                id=command_id,
+                created=int(time.time()),
+                model=chunk.model,
+                choices=[
+                    StreamingChoiceResponse(
+                        index=0,
+                        delta=ChatCompletionMessage(
+                            role="assistant",
+                            tool_calls=tool_call_deltas,
+                        ),
+                        finish_reason="tool_calls",
+                    )
+                ],
+            )
+            yield f"data: {tool_response.model_dump_json()}\n\n"
+            yield "data: [DONE]\n\n"
+            return
+
+        chunk_response = chunk_to_response(chunk, command_id)
+        yield f"data: {chunk_response.model_dump_json()}\n\n"
+
+        if chunk.finish_reason is not None:
+            yield "data: [DONE]\n\n"
+
+
+async def collect_chat_response(
+    command_id: CommandId,
+    chunk_stream: AsyncGenerator[ErrorChunk | ToolCallChunk | TokenChunk, None],
+) -> ChatCompletionResponse:
+    """Collect all token chunks and return a single ChatCompletionResponse."""
+    text_parts: list[str] = []
+    tool_calls: list[ToolCall] = []
+    model: str | None = None
+    finish_reason: FinishReason | None = None
+    error_message: str | None = None
+
+    async for chunk in chunk_stream:
+        if isinstance(chunk, ErrorChunk):
+            error_message = chunk.error_message or "Internal server error"
+            break
+
+        if model is None:
+            model = chunk.model
+
+        if isinstance(chunk, TokenChunk):
+            text_parts.append(chunk.text)
+
+        if isinstance(chunk, ToolCallChunk):
+            tool_calls.extend(
+                ToolCall(
+                    id=str(uuid4()),
+                    index=i,
+                    function=tool,
+                )
+                for i, tool in enumerate(chunk.tool_calls)
+            )
+
+        if chunk.finish_reason is not None:
+            finish_reason = chunk.finish_reason
+
+    if error_message is not None:
+        raise ValueError(error_message)
+
+    combined_text = "".join(text_parts)
+    assert model is not None
+
+    return ChatCompletionResponse(
+        id=command_id,
+        created=int(time.time()),
+        model=model,
+        choices=[
+            ChatCompletionChoice(
+                index=0,
+                message=ChatCompletionMessage(
+                    role="assistant",
+                    content=combined_text,
+                    tool_calls=tool_calls if tool_calls else None,
+                ),
+                finish_reason=finish_reason,
+            )
+        ],
+    )

src/exo/master/adapters/claude.py

@@ -0,0 +1,321 @@
+"""Claude Messages API adapter for converting requests/responses."""
+
+import json
+from collections.abc import AsyncGenerator
+from typing import Any
+from uuid import uuid4
+
+from exo.shared.types.api import FinishReason
+from exo.shared.types.chunks import ErrorChunk, TokenChunk, ToolCallChunk
+from exo.shared.types.claude_api import (
+    ClaudeContentBlock,
+    ClaudeContentBlockDeltaEvent,
+    ClaudeContentBlockStartEvent,
+    ClaudeContentBlockStopEvent,
+    ClaudeInputJsonDelta,
+    ClaudeMessageDelta,
+    ClaudeMessageDeltaEvent,
+    ClaudeMessageDeltaUsage,
+    ClaudeMessagesRequest,
+    ClaudeMessagesResponse,
+    ClaudeMessageStart,
+    ClaudeMessageStartEvent,
+    ClaudeMessageStopEvent,
+    ClaudeStopReason,
+    ClaudeTextBlock,
+    ClaudeTextDelta,
+    ClaudeToolResultBlock,
+    ClaudeToolUseBlock,
+    ClaudeUsage,
+)
+from exo.shared.types.common import CommandId
+from exo.shared.types.text_generation import InputMessage, TextGenerationTaskParams
+
+
+def finish_reason_to_claude_stop_reason(
+    finish_reason: FinishReason | None,
+) -> ClaudeStopReason | None:
+    """Map OpenAI finish_reason to Claude stop_reason."""
+    if finish_reason is None:
+        return None
+    mapping: dict[FinishReason, ClaudeStopReason] = {
+        "stop": "end_turn",
+        "length": "max_tokens",
+        "tool_calls": "tool_use",
+        "content_filter": "end_turn",
+        "function_call": "tool_use",
+    }
+    return mapping.get(finish_reason, "end_turn")
+
+
+def _extract_tool_result_text(block: ClaudeToolResultBlock) -> str:
+    """Extract plain text from a tool_result content field."""
+    if block.content is None:
+        return ""
+    if isinstance(block.content, str):
+        return block.content
+    return "".join(sub_block.text for sub_block in block.content)
+
+
+def claude_request_to_text_generation(
+    request: ClaudeMessagesRequest,
+) -> TextGenerationTaskParams:
+    # Handle system message
+    instructions: str | None = None
+    chat_template_messages: list[dict[str, Any]] = []
+
+    if request.system:
+        if isinstance(request.system, str):
+            instructions = request.system
+        else:
+            instructions = "".join(block.text for block in request.system)
+        chat_template_messages.append({"role": "system", "content": instructions})
+
+    # Convert messages to input
+    input_messages: list[InputMessage] = []
+    for msg in request.messages:
+        if isinstance(msg.content, str):
+            input_messages.append(InputMessage(role=msg.role, content=msg.content))
+            chat_template_messages.append({"role": msg.role, "content": msg.content})
+            continue
+
+        # Process structured content blocks
+        text_parts: list[str] = []
+        tool_calls: list[dict[str, Any]] = []
+        tool_results: list[ClaudeToolResultBlock] = []
+
+        for block in msg.content:
+            if isinstance(block, ClaudeTextBlock):
+                text_parts.append(block.text)
+            elif isinstance(block, ClaudeToolUseBlock):
+                tool_calls.append(
+                    {
+                        "id": block.id,
+                        "type": "function",
+                        "function": {
+                            "name": block.name,
+                            "arguments": json.dumps(block.input),
+                        },
+                    }
+                )
+            elif isinstance(block, ClaudeToolResultBlock):
+                tool_results.append(block)
+
+        content = "".join(text_parts)
+
+        # Build InputMessage from text content
+        if msg.role in ("user", "assistant"):
+            input_messages.append(InputMessage(role=msg.role, content=content))
+
+        # Build chat_template_messages preserving tool structure
+        if tool_calls:
+            chat_template_messages.append(
+                {"role": "assistant", "content": content, "tool_calls": tool_calls}
+            )
+        elif tool_results:
+            for tr in tool_results:
+                chat_template_messages.append(
+                    {
+                        "role": "tool",
+                        "tool_call_id": tr.tool_use_id,
+                        "content": _extract_tool_result_text(tr),
+                    }
+                )
+        else:
+            chat_template_messages.append({"role": msg.role, "content": content})
+
+    # Convert Claude tool definitions to OpenAI-style function tools
+    tools: list[dict[str, Any]] | None = None
+    if request.tools:
+        tools = [
+            {
+                "type": "function",
+                "function": {
+                    "name": tool.name,
+                    "description": tool.description or "",
+                    "parameters": tool.input_schema,
+                },
+            }
+            for tool in request.tools
+        ]
+
+    return TextGenerationTaskParams(
+        model=request.model,
+        input=input_messages if input_messages else "",
+        instructions=instructions,
+        max_output_tokens=request.max_tokens,
+        temperature=request.temperature,
+        top_p=request.top_p,
+        top_k=request.top_k,
+        stop=request.stop_sequences,
+        stream=request.stream,
+        tools=tools,
+        chat_template_messages=chat_template_messages
+        if chat_template_messages
+        else None,
+    )
+
+
+async def collect_claude_response(
+    command_id: CommandId,
+    model: str,
+    chunk_stream: AsyncGenerator[ErrorChunk | ToolCallChunk | TokenChunk, None],
+) -> ClaudeMessagesResponse:
+    """Collect all token chunks and return a single ClaudeMessagesResponse."""
+    text_parts: list[str] = []
+    tool_use_blocks: list[ClaudeToolUseBlock] = []
+    stop_reason: ClaudeStopReason | None = None
+    last_stats = None
+    error_message: str | None = None
+
+    async for chunk in chunk_stream:
+        if isinstance(chunk, ErrorChunk):
+            error_message = chunk.error_message or "Internal server error"
+            break
+
+        if isinstance(chunk, ToolCallChunk):
+            for tool in chunk.tool_calls:
+                tool_use_blocks.append(
+                    ClaudeToolUseBlock(
+                        id=f"toolu_{uuid4().hex[:24]}",
+                        name=tool.name,
+                        input=json.loads(tool.arguments),  # pyright: ignore[reportAny]
+                    )
+                )
+            last_stats = chunk.stats or last_stats
+            stop_reason = "tool_use"
+            continue
+
+        text_parts.append(chunk.text)
+        last_stats = chunk.stats or last_stats
+
+        if chunk.finish_reason is not None:
+            stop_reason = finish_reason_to_claude_stop_reason(chunk.finish_reason)
+
+    if error_message is not None:
+        raise ValueError(error_message)
+
+    combined_text = "".join(text_parts)
+
+    # Build content blocks
+    content: list[ClaudeContentBlock] = []
+    if combined_text:
+        content.append(ClaudeTextBlock(text=combined_text))
+    content.extend(tool_use_blocks)
+
+    # If no content at all, include empty text block
+    if not content:
+        content.append(ClaudeTextBlock(text=""))
+
+    # Use actual usage data from stats if available
+    input_tokens = last_stats.prompt_tokens if last_stats else 0
+    output_tokens = last_stats.generation_tokens if last_stats else 0
+
+    return ClaudeMessagesResponse(
+        id=f"msg_{command_id}",
+        model=model,
+        content=content,
+        stop_reason=stop_reason,
+        usage=ClaudeUsage(
+            input_tokens=input_tokens,
+            output_tokens=output_tokens,
+        ),
+    )
+
+
+async def generate_claude_stream(
+    command_id: CommandId,
+    model: str,
+    chunk_stream: AsyncGenerator[ErrorChunk | ToolCallChunk | TokenChunk, None],
+) -> AsyncGenerator[str, None]:
+    """Generate Claude Messages API streaming events from TokenChunks."""
+    # Initial message_start event
+    initial_message = ClaudeMessageStart(
+        id=f"msg_{command_id}",
+        model=model,
+        content=[],
+        stop_reason=None,
+        usage=ClaudeUsage(input_tokens=0, output_tokens=0),
+    )
+    start_event = ClaudeMessageStartEvent(message=initial_message)
+    yield f"event: message_start\ndata: {start_event.model_dump_json()}\n\n"
+
+    # content_block_start for text block at index 0
+    block_start = ClaudeContentBlockStartEvent(
+        index=0, content_block=ClaudeTextBlock(text="")
+    )
+    yield f"event: content_block_start\ndata: {block_start.model_dump_json()}\n\n"
+
+    output_tokens = 0
+    stop_reason: ClaudeStopReason | None = None
+    last_stats = None
+    next_block_index = 1  # text block is 0, tool blocks start at 1
+
+    async for chunk in chunk_stream:
+        if isinstance(chunk, ErrorChunk):
+            # Close text block and bail
+            break
+
+        if isinstance(chunk, ToolCallChunk):
+            last_stats = chunk.stats or last_stats
+            stop_reason = "tool_use"
+
+            # Emit tool_use content blocks
+            for tool in chunk.tool_calls:
+                tool_id = f"toolu_{uuid4().hex[:24]}"
+                tool_input_json = tool.arguments
+
+                # content_block_start for tool_use
+                tool_block_start = ClaudeContentBlockStartEvent(
+                    index=next_block_index,
+                    content_block=ClaudeToolUseBlock(
+                        id=tool_id, name=tool.name, input={}
+                    ),
+                )
+                yield f"event: content_block_start\ndata: {tool_block_start.model_dump_json()}\n\n"
+
+                # content_block_delta with input_json_delta
+                tool_delta_event = ClaudeContentBlockDeltaEvent(
+                    index=next_block_index,
+                    delta=ClaudeInputJsonDelta(partial_json=tool_input_json),
+                )
+                yield f"event: content_block_delta\ndata: {tool_delta_event.model_dump_json()}\n\n"
+
+                # content_block_stop
+                tool_block_stop = ClaudeContentBlockStopEvent(index=next_block_index)
+                yield f"event: content_block_stop\ndata: {tool_block_stop.model_dump_json()}\n\n"
+
+                next_block_index += 1
+            continue
+
+        output_tokens += 1  # Count each chunk as one token
+        last_stats = chunk.stats or last_stats
+
+        # content_block_delta
+        delta_event = ClaudeContentBlockDeltaEvent(
+            index=0,
+            delta=ClaudeTextDelta(text=chunk.text),
+        )
+        yield f"event: content_block_delta\ndata: {delta_event.model_dump_json()}\n\n"
+
+        if chunk.finish_reason is not None:
+            stop_reason = finish_reason_to_claude_stop_reason(chunk.finish_reason)
+
+    # Use actual token count from stats if available
+    if last_stats is not None:
+        output_tokens = last_stats.generation_tokens
+
+    # content_block_stop for text block
+    block_stop = ClaudeContentBlockStopEvent(index=0)
+    yield f"event: content_block_stop\ndata: {block_stop.model_dump_json()}\n\n"
+
+    # message_delta
+    message_delta = ClaudeMessageDeltaEvent(
+        delta=ClaudeMessageDelta(stop_reason=stop_reason),
+        usage=ClaudeMessageDeltaUsage(output_tokens=output_tokens),
+    )
+    yield f"event: message_delta\ndata: {message_delta.model_dump_json()}\n\n"
+
+    # message_stop
+    message_stop = ClaudeMessageStopEvent()
+    yield f"event: message_stop\ndata: {message_stop.model_dump_json()}\n\n"

src/exo/master/adapters/responses.py

@@ -0,0 +1,369 @@
+"""OpenAI Responses API adapter for converting requests/responses."""
+
+from collections.abc import AsyncGenerator
+from itertools import count
+from typing import Any
+from uuid import uuid4
+
+from exo.shared.types.chunks import ErrorChunk, TokenChunk, ToolCallChunk
+from exo.shared.types.common import CommandId
+from exo.shared.types.openai_responses import (
+    FunctionCallInputItem,
+    ResponseCompletedEvent,
+    ResponseContentPart,
+    ResponseContentPartAddedEvent,
+    ResponseContentPartDoneEvent,
+    ResponseCreatedEvent,
+    ResponseFunctionCallArgumentsDeltaEvent,
+    ResponseFunctionCallArgumentsDoneEvent,
+    ResponseFunctionCallItem,
+    ResponseInProgressEvent,
+    ResponseInputMessage,
+    ResponseItem,
+    ResponseMessageItem,
+    ResponseOutputItemAddedEvent,
+    ResponseOutputItemDoneEvent,
+    ResponseOutputText,
+    ResponsesRequest,
+    ResponsesResponse,
+    ResponseTextDeltaEvent,
+    ResponseTextDoneEvent,
+    ResponseUsage,
+)
+from exo.shared.types.text_generation import InputMessage, TextGenerationTaskParams
+
+
+def _extract_content(content: str | list[ResponseContentPart]) -> str:
+    """Extract plain text from a content field that may be a string or list of parts."""
+    if isinstance(content, str):
+        return content
+    return "".join(part.text for part in content)
+
+
+def responses_request_to_text_generation(
+    request: ResponsesRequest,
+) -> TextGenerationTaskParams:
+    input_value: str | list[InputMessage]
+    built_chat_template: list[dict[str, Any]] | None = None
+    if isinstance(request.input, str):
+        input_value = request.input
+    else:
+        input_messages: list[InputMessage] = []
+        chat_template_messages: list[dict[str, Any]] = []
+
+        if request.instructions is not None:
+            chat_template_messages.append(
+                {"role": "system", "content": request.instructions}
+            )
+
+        for item in request.input:
+            if isinstance(item, ResponseInputMessage):
+                content = _extract_content(item.content)
+                if item.role in ("user", "assistant", "developer"):
+                    input_messages.append(InputMessage(role=item.role, content=content))
+                if item.role == "system":
+                    chat_template_messages.append(
+                        {"role": "system", "content": content}
+                    )
+                else:
+                    chat_template_messages.append(
+                        {"role": item.role, "content": content}
+                    )
+            elif isinstance(item, FunctionCallInputItem):
+                chat_template_messages.append(
+                    {
+                        "role": "assistant",
+                        "content": "",
+                        "tool_calls": [
+                            {
+                                "id": item.call_id,
+                                "type": "function",
+                                "function": {
+                                    "name": item.name,
+                                    "arguments": item.arguments,
+                                },
+                            }
+                        ],
+                    }
+                )
+            else:
+                chat_template_messages.append(
+                    {
+                        "role": "tool",
+                        "tool_call_id": item.call_id,
+                        "content": item.output,
+                    }
+                )
+
+        input_value = input_messages if input_messages else ""
+        built_chat_template = chat_template_messages if chat_template_messages else None
+
+    return TextGenerationTaskParams(
+        model=request.model,
+        input=input_value,
+        instructions=request.instructions,
+        max_output_tokens=request.max_output_tokens,
+        temperature=request.temperature,
+        top_p=request.top_p,
+        stream=request.stream,
+        tools=request.tools,
+        top_k=request.top_k,
+        stop=request.stop,
+        seed=request.seed,
+        chat_template_messages=built_chat_template or request.chat_template_messages,
+    )
+
+
+async def collect_responses_response(
+    command_id: CommandId,
+    model: str,
+    chunk_stream: AsyncGenerator[ErrorChunk | ToolCallChunk | TokenChunk, None],
+) -> ResponsesResponse:
+    """Collect all token chunks and return a single ResponsesResponse."""
+    response_id = f"resp_{command_id}"
+    item_id = f"item_{command_id}"
+    accumulated_text = ""
+    function_call_items: list[ResponseFunctionCallItem] = []
+    last_stats = None
+    error_message: str | None = None
+
+    async for chunk in chunk_stream:
+        if isinstance(chunk, ErrorChunk):
+            error_message = chunk.error_message or "Internal server error"
+            break
+
+        if isinstance(chunk, ToolCallChunk):
+            for tool in chunk.tool_calls:
+                function_call_items.append(
+                    ResponseFunctionCallItem(
+                        id=f"fc_{uuid4().hex[:24]}",
+                        call_id=f"call_{uuid4().hex[:24]}",
+                        name=tool.name,
+                        arguments=tool.arguments,
+                    )
+                )
+            last_stats = chunk.stats or last_stats
+            continue
+
+        accumulated_text += chunk.text
+        last_stats = chunk.stats or last_stats
+
+    if error_message is not None:
+        raise ValueError(error_message)
+
+    # Create usage from stats if available
+    usage = None
+    if last_stats is not None:
+        usage = ResponseUsage(
+            input_tokens=last_stats.prompt_tokens,
+            output_tokens=last_stats.generation_tokens,
+            total_tokens=last_stats.prompt_tokens + last_stats.generation_tokens,
+        )
+
+    output: list[ResponseItem] = [
+        ResponseMessageItem(
+            id=item_id,
+            content=[ResponseOutputText(text=accumulated_text)],
+            status="completed",
+        )
+    ]
+    output.extend(function_call_items)
+
+    return ResponsesResponse(
+        id=response_id,
+        model=model,
+        status="completed",
+        output=output,
+        output_text=accumulated_text,
+        usage=usage,
+    )
+
+
+async def generate_responses_stream(
+    command_id: CommandId,
+    model: str,
+    chunk_stream: AsyncGenerator[ErrorChunk | ToolCallChunk | TokenChunk, None],
+) -> AsyncGenerator[str, None]:
+    """Generate OpenAI Responses API streaming events from TokenChunks."""
+    response_id = f"resp_{command_id}"
+    item_id = f"item_{command_id}"
+    seq = count(1)
+
+    # response.created
+    initial_response = ResponsesResponse(
+        id=response_id,
+        model=model,
+        status="in_progress",
+        output=[],
+        output_text="",
+    )
+    created_event = ResponseCreatedEvent(
+        sequence_number=next(seq), response=initial_response
+    )
+    yield f"event: response.created\ndata: {created_event.model_dump_json()}\n\n"
+
+    # response.in_progress
+    in_progress_event = ResponseInProgressEvent(
+        sequence_number=next(seq), response=initial_response
+    )
+    yield f"event: response.in_progress\ndata: {in_progress_event.model_dump_json()}\n\n"
+
+    # response.output_item.added
+    initial_item = ResponseMessageItem(
+        id=item_id,
+        content=[ResponseOutputText(text="")],
+        status="in_progress",
+    )
+    item_added = ResponseOutputItemAddedEvent(
+        sequence_number=next(seq), output_index=0, item=initial_item
+    )
+    yield f"event: response.output_item.added\ndata: {item_added.model_dump_json()}\n\n"
+
+    # response.content_part.added
+    initial_part = ResponseOutputText(text="")
+    part_added = ResponseContentPartAddedEvent(
+        sequence_number=next(seq),
+        item_id=item_id,
+        output_index=0,
+        content_index=0,
+        part=initial_part,
+    )
+    yield f"event: response.content_part.added\ndata: {part_added.model_dump_json()}\n\n"
+
+    accumulated_text = ""
+    function_call_items: list[ResponseFunctionCallItem] = []
+    last_stats = None
+    next_output_index = 1  # message item is at 0
+
+    async for chunk in chunk_stream:
+        if isinstance(chunk, ErrorChunk):
+            break
+
+        if isinstance(chunk, ToolCallChunk):
+            last_stats = chunk.stats or last_stats
+            for tool in chunk.tool_calls:
+                fc_id = f"fc_{uuid4().hex[:24]}"
+                call_id = f"call_{uuid4().hex[:24]}"
+
+                # response.output_item.added for function_call
+                fc_item = ResponseFunctionCallItem(
+                    id=fc_id,
+                    call_id=call_id,
+                    name=tool.name,
+                    arguments="",
+                    status="in_progress",
+                )
+                fc_added = ResponseOutputItemAddedEvent(
+                    sequence_number=next(seq),
+                    output_index=next_output_index,
+                    item=fc_item,
+                )
+                yield f"event: response.output_item.added\ndata: {fc_added.model_dump_json()}\n\n"
+
+                # response.function_call_arguments.delta
+                args_delta = ResponseFunctionCallArgumentsDeltaEvent(
+                    sequence_number=next(seq),
+                    item_id=fc_id,
+                    output_index=next_output_index,
+                    delta=tool.arguments,
+                )
+                yield f"event: response.function_call_arguments.delta\ndata: {args_delta.model_dump_json()}\n\n"
+
+                # response.function_call_arguments.done
+                args_done = ResponseFunctionCallArgumentsDoneEvent(
+                    sequence_number=next(seq),
+                    item_id=fc_id,
+                    output_index=next_output_index,
+                    name=tool.name,
+                    arguments=tool.arguments,
+                )
+                yield f"event: response.function_call_arguments.done\ndata: {args_done.model_dump_json()}\n\n"
+
+                # response.output_item.done
+                fc_done_item = ResponseFunctionCallItem(
+                    id=fc_id,
+                    call_id=call_id,
+                    name=tool.name,
+                    arguments=tool.arguments,
+                    status="completed",
+                )
+                fc_item_done = ResponseOutputItemDoneEvent(
+                    sequence_number=next(seq),
+                    output_index=next_output_index,
+                    item=fc_done_item,
+                )
+                yield f"event: response.output_item.done\ndata: {fc_item_done.model_dump_json()}\n\n"
+
+                function_call_items.append(fc_done_item)
+                next_output_index += 1
+            continue
+
+        accumulated_text += chunk.text
+        last_stats = chunk.stats or last_stats
+
+        # response.output_text.delta
+        delta_event = ResponseTextDeltaEvent(
+            sequence_number=next(seq),
+            item_id=item_id,
+            output_index=0,
+            content_index=0,
+            delta=chunk.text,
+        )
+        yield f"event: response.output_text.delta\ndata: {delta_event.model_dump_json()}\n\n"
+
+    # response.output_text.done
+    text_done = ResponseTextDoneEvent(
+        sequence_number=next(seq),
+        item_id=item_id,
+        output_index=0,
+        content_index=0,
+        text=accumulated_text,
+    )
+    yield f"event: response.output_text.done\ndata: {text_done.model_dump_json()}\n\n"
+
+    # response.content_part.done
+    final_part = ResponseOutputText(text=accumulated_text)
+    part_done = ResponseContentPartDoneEvent(
+        sequence_number=next(seq),
+        item_id=item_id,
+        output_index=0,
+        content_index=0,
+        part=final_part,
+    )
+    yield f"event: response.content_part.done\ndata: {part_done.model_dump_json()}\n\n"
+
+    # response.output_item.done
+    final_message_item = ResponseMessageItem(
+        id=item_id,
+        content=[ResponseOutputText(text=accumulated_text)],
+        status="completed",
+    )
+    item_done = ResponseOutputItemDoneEvent(
+        sequence_number=next(seq), output_index=0, item=final_message_item
+    )
+    yield f"event: response.output_item.done\ndata: {item_done.model_dump_json()}\n\n"
+
+    # Create usage from stats if available
+    usage = None
+    if last_stats is not None:
+        usage = ResponseUsage(
+            input_tokens=last_stats.prompt_tokens,
+            output_tokens=last_stats.generation_tokens,
+            total_tokens=last_stats.prompt_tokens + last_stats.generation_tokens,
+        )
+
+    # response.completed
+    output: list[ResponseItem] = [final_message_item]
+    output.extend(function_call_items)
+    final_response = ResponsesResponse(
+        id=response_id,
+        model=model,
+        status="completed",
+        output=output,
+        output_text=accumulated_text,
+        usage=usage,
+    )
+    completed_event = ResponseCompletedEvent(
+        sequence_number=next(seq), response=final_response
+    )
+    yield f"event: response.completed\ndata: {completed_event.model_dump_json()}\n\n"

src/exo/master/api.py

@@ -1,14 +1,18 @@
 import base64
+import contextlib
 import json
 import time
-from collections.abc import AsyncGenerator
+from collections.abc import AsyncGenerator, Awaitable, Callable
+from datetime import datetime, timezone
 from http import HTTPStatus
-from typing import Literal, cast
+from pathlib import Path
+from typing import Annotated, Literal, cast
+from uuid import uuid4
 
 import anyio
 from anyio import BrokenResourceError, create_task_group
 from anyio.abc import TaskGroup
-from fastapi import FastAPI, File, Form, HTTPException, Request, UploadFile
+from fastapi import FastAPI, File, Form, HTTPException, Query, Request, UploadFile
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.responses import FileResponse, JSONResponse, StreamingResponse
 from fastapi.staticfiles import StaticFiles
@@ -17,10 +21,29 @@ from hypercorn.config import Config
 from hypercorn.typing import ASGIFramework
 from loguru import logger
 
+from exo.master.adapters.chat_completions import (
+    chat_request_to_text_generation,
+    collect_chat_response,
+    generate_chat_stream,
+)
+from exo.master.adapters.claude import (
+    claude_request_to_text_generation,
+    collect_claude_response,
+    generate_claude_stream,
+)
+from exo.master.adapters.responses import (
+    collect_responses_response,
+    generate_responses_stream,
+    responses_request_to_text_generation,
+)
 from exo.master.image_store import ImageStore
 from exo.master.placement import place_instance as get_instance_placements
 from exo.shared.apply import apply
-from exo.shared.constants import EXO_IMAGE_CACHE_DIR, EXO_MAX_CHUNK_SIZE
+from exo.shared.constants import (
+    EXO_IMAGE_CACHE_DIR,
+    EXO_MAX_CHUNK_SIZE,
+    EXO_TRACING_CACHE_DIR,
+)
 from exo.shared.election import ElectionMessage
 from exo.shared.logging import InterceptLogger
 from exo.shared.models.model_cards import (
@@ -28,23 +51,27 @@ from exo.shared.models.model_cards import (
     ModelCard,
     ModelId,
 )
+from exo.shared.tracing import TraceEvent, compute_stats, export_trace, load_trace_file
 from exo.shared.types.api import (
+    AdvancedImageParams,
+    BenchChatCompletionRequest,
     BenchChatCompletionResponse,
-    BenchChatCompletionTaskParams,
     BenchImageGenerationResponse,
     BenchImageGenerationTaskParams,
     ChatCompletionChoice,
     ChatCompletionMessage,
+    ChatCompletionRequest,
     ChatCompletionResponse,
     CreateInstanceParams,
     CreateInstanceResponse,
+    DeleteDownloadResponse,
     DeleteInstanceResponse,
     ErrorInfo,
     ErrorResponse,
     FinishReason,
     GenerationStats,
     ImageData,
-    ImageEditsInternalParams,
+    ImageEditsTaskParams,
     ImageGenerationResponse,
     ImageGenerationStats,
     ImageGenerationTaskParams,
@@ -55,20 +82,43 @@ from exo.shared.types.api import (
     PlaceInstanceParams,
     PlacementPreview,
     PlacementPreviewResponse,
-    StreamingChoiceResponse,
+    StartDownloadParams,
+    StartDownloadResponse,
+    ToolCall,
+    TraceCategoryStats,
+    TraceEventResponse,
+    TraceListItem,
+    TraceListResponse,
+    TraceRankStats,
+    TraceResponse,
+    TraceStatsResponse,
+)
+from exo.shared.types.chunks import (
+    ErrorChunk,
+    ImageChunk,
+    InputImageChunk,
+    TokenChunk,
+    ToolCallChunk,
+)
+from exo.shared.types.claude_api import (
+    ClaudeMessagesRequest,
+    ClaudeMessagesResponse,
 )
-from exo.shared.types.chunks import ImageChunk, InputImageChunk, TokenChunk
 from exo.shared.types.commands import (
-    ChatCompletion,
     Command,
     CreateInstance,
+    DeleteDownload,
     DeleteInstance,
+    DownloadCommand,
     ForwarderCommand,
+    ForwarderDownloadCommand,
     ImageEdits,
     ImageGeneration,
     PlaceInstance,
     SendInputChunk,
+    StartDownload,
     TaskFinished,
+    TextGeneration,
 )
 from exo.shared.types.common import CommandId, Id, NodeId, SessionId
 from exo.shared.types.events import (
@@ -76,10 +126,14 @@ from exo.shared.types.events import (
     Event,
     ForwarderEvent,
     IndexedEvent,
+    TracesMerged,
 )
 from exo.shared.types.memory import Memory
+from exo.shared.types.openai_responses import (
+    ResponsesRequest,
+    ResponsesResponse,
+)
 from exo.shared.types.state import State
-from exo.shared.types.tasks import ChatCompletionTaskParams
 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
@@ -92,23 +146,6 @@ def _format_to_content_type(image_format: Literal["png", "jpeg", "webp"] | None)
     return f"image/{image_format or 'png'}"
 
 
-def chunk_to_response(
-    chunk: TokenChunk, command_id: CommandId
-) -> ChatCompletionResponse:
-    return ChatCompletionResponse(
-        id=command_id,
-        created=int(time.time()),
-        model=chunk.model,
-        choices=[
-            StreamingChoiceResponse(
-                index=0,
-                delta=ChatCompletionMessage(role="assistant", content=chunk.text),
-                finish_reason=chunk.finish_reason,
-            )
-        ],
-    )
-
-
 async def resolve_model_card(model_id: ModelId) -> ModelCard:
     if model_id in MODEL_CARDS:
         model_card = MODEL_CARDS[model_id]
@@ -131,12 +168,14 @@ class API:
         # Ideally this would be a MasterForwarderEvent but type system says no :(
         global_event_receiver: Receiver[ForwarderEvent],
         command_sender: Sender[ForwarderCommand],
+        download_command_sender: Sender[ForwarderDownloadCommand],
         # This lets us pause the API if an election is running
         election_receiver: Receiver[ElectionMessage],
     ) -> None:
         self.state = State()
         self._event_log: list[Event] = []
         self.command_sender = command_sender
+        self.download_command_sender = download_command_sender
         self.global_event_receiver = global_event_receiver
         self.election_receiver = election_receiver
         self.event_buffer: OrderedBuffer[Event] = OrderedBuffer[Event]()
@@ -149,6 +188,15 @@ class API:
         self.paused_ev: anyio.Event = anyio.Event()
 
         self.app = FastAPI()
+
+        @self.app.middleware("http")
+        async def _log_requests(  # pyright: ignore[reportUnusedFunction]
+            request: Request,
+            call_next: Callable[[Request], Awaitable[StreamingResponse]],
+        ) -> StreamingResponse:
+            logger.debug(f"API request: {request.method} {request.url.path}")
+            return await call_next(request)
+
         self._setup_exception_handlers()
         self._setup_cors()
         self._setup_routes()
@@ -162,8 +210,12 @@ class API:
             name="dashboard",
         )
 
-        self._chat_completion_queues: dict[CommandId, Sender[TokenChunk]] = {}
-        self._image_generation_queues: dict[CommandId, Sender[ImageChunk]] = {}
+        self._text_generation_queues: dict[
+            CommandId, Sender[TokenChunk | ErrorChunk | ToolCallChunk]
+        ] = {}
+        self._image_generation_queues: dict[
+            CommandId, Sender[ImageChunk | ErrorChunk]
+        ] = {}
         self._image_store = ImageStore(EXO_IMAGE_CACHE_DIR)
         self._tg: TaskGroup | None = None
 
@@ -172,7 +224,7 @@ class API:
         self.state = State()
         self.session_id = new_session_id
         self.event_buffer = OrderedBuffer[Event]()
-        self._chat_completion_queues = {}
+        self._text_generation_queues = {}
         self._image_generation_queues = {}
         self.unpause(result_clock)
 
@@ -229,12 +281,20 @@ class API:
         self.app.post("/bench/images/edits")(self.bench_image_edits)
         self.app.get("/images")(self.list_images)
         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.get("/state")(lambda: self.state)
         self.app.get("/events")(lambda: self._event_log)
+        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.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)
 
     async def place_instance(self, payload: PlaceInstanceParams):
         command = PlaceInstance(
-            model_card=await resolve_model_card(payload.model_id),
+            model_card=await ModelCard.load(payload.model_id),
             sharding=payload.sharding,
             instance_meta=payload.instance_meta,
             min_nodes=payload.min_nodes,
@@ -251,7 +311,7 @@ class API:
         self, payload: CreateInstanceParams
     ) -> CreateInstanceResponse:
         instance = payload.instance
-        model_card = await resolve_model_card(instance.shard_assignments.model_id)
+        model_card = await ModelCard.load(instance.shard_assignments.model_id)
         required_memory = model_card.storage_size
         available_memory = self._calculate_total_available_memory()
 
@@ -279,7 +339,7 @@ class API:
         instance_meta: InstanceMeta = InstanceMeta.MlxRing,
         min_nodes: int = 1,
     ) -> Instance:
-        model_card = await resolve_model_card(model_id)
+        model_card = await ModelCard.load(model_id)
 
         try:
             placements = get_instance_placements(
@@ -310,10 +370,14 @@ class API:
         return placements[new_ids[0]]
 
     async def get_placement_previews(
-        self, model_id: ModelId
+        self,
+        model_id: ModelId,
+        node_ids: Annotated[list[NodeId] | None, Query()] = None,
     ) -> PlacementPreviewResponse:
         seen: set[tuple[ModelId, Sharding, InstanceMeta, int]] = set()
         previews: list[PlacementPreview] = []
+        required_nodes = set(node_ids) if node_ids else None
+
         if len(list(self.state.topology.list_nodes())) == 0:
             return PlacementPreviewResponse(previews=[])
 
@@ -349,6 +413,7 @@ class API:
                         node_network=self.state.node_network,
                         topology=self.state.topology,
                         current_instances=self.state.instances,
+                        required_nodes=required_nodes,
                     )
                 except ValueError as exc:
                     if (model_card.model_id, sharding, instance_meta, 0) not in seen:
@@ -387,14 +452,16 @@ class API:
 
                 instance = new_instances[0]
                 shard_assignments = instance.shard_assignments
-                node_ids = list(shard_assignments.node_to_runner.keys())
+                placement_node_ids = list(shard_assignments.node_to_runner.keys())
 
                 memory_delta_by_node: dict[str, int] = {}
-                if node_ids:
+                if placement_node_ids:
                     total_bytes = model_card.storage_size.in_bytes
-                    per_node = total_bytes // len(node_ids)
-                    remainder = total_bytes % len(node_ids)
-                    for index, node_id in enumerate(sorted(node_ids, key=str)):
+                    per_node = total_bytes // len(placement_node_ids)
+                    remainder = total_bytes % len(placement_node_ids)
+                    for index, node_id in enumerate(
+                        sorted(placement_node_ids, key=str)
+                    ):
                         extra = 1 if index < remainder else 0
                         memory_delta_by_node[str(node_id)] = per_node + extra
 
@@ -402,7 +469,7 @@ class API:
                     model_card.model_id,
                     sharding,
                     instance_meta,
-                    len(node_ids),
+                    len(placement_node_ids),
                 ) not in seen:
                     previews.append(
                         PlacementPreview(
@@ -414,7 +481,14 @@ class API:
                             error=None,
                         )
                     )
-                seen.add((model_card.model_id, sharding, instance_meta, len(node_ids)))
+                seen.add(
+                    (
+                        model_card.model_id,
+                        sharding,
+                        instance_meta,
+                        len(placement_node_ids),
+                    )
+                )
 
         return PlacementPreviewResponse(previews=previews)
 
@@ -437,13 +511,17 @@ class API:
             instance_id=instance_id,
         )
 
-    async def _chat_chunk_stream(
+    async def _token_chunk_stream(
         self, command_id: CommandId
-    ) -> AsyncGenerator[TokenChunk, None]:
-        """Yield `TokenChunk`s for a given command until completion."""
+    ) -> AsyncGenerator[ErrorChunk | ToolCallChunk | TokenChunk, None]:
+        """Yield chunks for a given command until completion.
 
+        This is the internal low-level stream used by all API adapters.
+        """
         try:
-            self._chat_completion_queues[command_id], recv = channel[TokenChunk]()
+            self._text_generation_queues[command_id], recv = channel[
+                ErrorChunk | ToolCallChunk | TokenChunk
+            ]()
 
             with recv as token_chunks:
                 async for chunk in token_chunks:
@@ -462,89 +540,20 @@ class API:
         finally:
             command = TaskFinished(finished_command_id=command_id)
             await self._send(command)
-            del self._chat_completion_queues[command_id]
+            if command_id in self._text_generation_queues:
+                del self._text_generation_queues[command_id]
 
-    async def _generate_chat_stream(
-        self, command_id: CommandId
-    ) -> AsyncGenerator[str, None]:
-        """Generate chat completion stream as JSON strings."""
-
-        async for chunk in self._chat_chunk_stream(command_id):
-            if chunk.finish_reason == "error":
-                error_response = ErrorResponse(
-                    error=ErrorInfo(
-                        message=chunk.error_message or "Internal server error",
-                        type="InternalServerError",
-                        code=500,
-                    )
-                )
-                yield f"data: {error_response.model_dump_json()}\n\n"
-                yield "data: [DONE]\n\n"
-                return
-
-            chunk_response: ChatCompletionResponse = chunk_to_response(
-                chunk, command_id
-            )
-            logger.debug(f"chunk_response: {chunk_response}")
-
-            yield f"data: {chunk_response.model_dump_json()}\n\n"
-
-            if chunk.finish_reason is not None:
-                yield "data: [DONE]\n\n"
-
-    async def _collect_chat_completion(
-        self, command_id: CommandId
-    ) -> ChatCompletionResponse:
-        """Collect all token chunks for a chat completion and return a single response."""
-
-        text_parts: list[str] = []
-        model: str | None = None
-        finish_reason: FinishReason | None = None
-
-        async for chunk in self._chat_chunk_stream(command_id):
-            if chunk.finish_reason == "error":
-                raise HTTPException(
-                    status_code=500,
-                    detail=chunk.error_message or "Internal server error",
-                )
-
-            if model is None:
-                model = chunk.model
-
-            text_parts.append(chunk.text)
-
-            if chunk.finish_reason is not None:
-                finish_reason = chunk.finish_reason
-
-        combined_text = "".join(text_parts)
-        assert model is not None
-
-        return ChatCompletionResponse(
-            id=command_id,
-            created=int(time.time()),
-            model=model,
-            choices=[
-                ChatCompletionChoice(
-                    index=0,
-                    message=ChatCompletionMessage(
-                        role="assistant",
-                        content=combined_text,
-                    ),
-                    finish_reason=finish_reason,
-                )
-            ],
-        )
-
-    async def _collect_chat_completion_with_stats(
+    async def _collect_text_generation_with_stats(
         self, command_id: CommandId
     ) -> BenchChatCompletionResponse:
         text_parts: list[str] = []
-        model: str | None = None
+        tool_calls: list[ToolCall] = []
+        model: ModelId | None = None
         finish_reason: FinishReason | None = None
 
         stats: GenerationStats | None = None
 
-        async for chunk in self._chat_chunk_stream(command_id):
+        async for chunk in self._token_chunk_stream(command_id):
             if chunk.finish_reason == "error":
                 raise HTTPException(
                     status_code=500,
@@ -554,7 +563,19 @@ class API:
             if model is None:
                 model = chunk.model
 
-            text_parts.append(chunk.text)
+            if isinstance(chunk, TokenChunk):
+                text_parts.append(chunk.text)
+
+            if isinstance(chunk, ToolCallChunk):
+                tool_calls.extend(
+                    ToolCall(
+                        id=str(uuid4()),
+                        index=i,
+                        function=tool,
+                    )
+                    for i, tool in enumerate(chunk.tool_calls)
+                )
+
             stats = chunk.stats or stats
 
             if chunk.finish_reason is not None:
@@ -571,7 +592,9 @@ class API:
                 ChatCompletionChoice(
                     index=0,
                     message=ChatCompletionMessage(
-                        role="assistant", content=combined_text
+                        role="assistant",
+                        content=combined_text,
+                        tool_calls=tool_calls if tool_calls else None,
                     ),
                     finish_reason=finish_reason,
                 )
@@ -580,68 +603,79 @@ class API:
         )
         return resp
 
-    async def _trigger_notify_user_to_download_model(self, model_id: str) -> None:
+    async def _trigger_notify_user_to_download_model(self, model_id: ModelId) -> None:
         logger.warning(
             "TODO: we should send a notification to the user to download the model"
         )
 
     async def chat_completions(
-        self, payload: ChatCompletionTaskParams
+        self, payload: ChatCompletionRequest
     ) -> ChatCompletionResponse | StreamingResponse:
-        """Handle chat completions, supporting both streaming and non-streaming responses."""
-        model_card = await resolve_model_card(ModelId(payload.model))
-        payload.model = model_card.model_id
-
-        if not any(
-            instance.shard_assignments.model_id == payload.model
-            for instance in self.state.instances.values()
-        ):
-            await self._trigger_notify_user_to_download_model(payload.model)
-            raise HTTPException(
-                status_code=404, detail=f"No instance found for model {payload.model}"
-            )
-
-        command = ChatCompletion(
-            request_params=payload,
+        """OpenAI Chat Completions API - adapter."""
+        task_params = chat_request_to_text_generation(payload)
+        resolved_model = await self._resolve_and_validate_text_model(
+            ModelId(task_params.model)
         )
+        task_params = task_params.model_copy(update={"model": resolved_model})
+
+        command = TextGeneration(task_params=task_params)
         await self._send(command)
+
         if payload.stream:
             return StreamingResponse(
-                self._generate_chat_stream(command.command_id),
+                generate_chat_stream(
+                    command.command_id,
+                    self._token_chunk_stream(command.command_id),
+                ),
                 media_type="text/event-stream",
             )
 
-        return await self._collect_chat_completion(command.command_id)
+        return await collect_chat_response(
+            command.command_id,
+            self._token_chunk_stream(command.command_id),
+        )
 
     async def bench_chat_completions(
-        self, payload: BenchChatCompletionTaskParams
+        self, payload: BenchChatCompletionRequest
     ) -> BenchChatCompletionResponse:
-        model_card = await resolve_model_card(ModelId(payload.model))
-        payload.model = model_card.model_id
+        task_params = chat_request_to_text_generation(payload)
+        resolved_model = await self._resolve_and_validate_text_model(
+            ModelId(task_params.model)
+        )
+        task_params = task_params.model_copy(update={"model": resolved_model})
 
-        if not any(
-            instance.shard_assignments.model_id == payload.model
-            for instance in self.state.instances.values()
-        ):
-            await self._trigger_notify_user_to_download_model(payload.model)
-            raise HTTPException(
-                status_code=404, detail=f"No instance found for model {payload.model}"
-            )
+        task_params = task_params.model_copy(update={"stream": False, "bench": True})
 
-        payload.stream = False
-
-        command = ChatCompletion(request_params=payload)
+        command = TextGeneration(task_params=task_params)
         await self._send(command)
 
-        response = await self._collect_chat_completion_with_stats(command.command_id)
+        response = await self._collect_text_generation_with_stats(command.command_id)
         return response
 
-    async def _validate_image_model(self, model: str) -> ModelId:
+    async def _resolve_and_validate_text_model(self, model: ModelId) -> ModelId:
+        """Validate a text model exists and return the resolved model ID.
+
+        Raises HTTPException 404 if no instance is found for the model.
+        """
+        model_card = await resolve_model_card(model)
+        resolved = model_card.model_id
+        if not any(
+            instance.shard_assignments.model_id == resolved
+            for instance in self.state.instances.values()
+        ):
+            await self._trigger_notify_user_to_download_model(resolved)
+            raise HTTPException(
+                status_code=404,
+                detail=f"No instance found for model {resolved}",
+            )
+        return resolved
+
+    async def _validate_image_model(self, model: ModelId) -> ModelId:
         """Validate model exists and return resolved model ID.
 
         Raises HTTPException 404 if no instance is found for the model.
         """
-        model_card = await resolve_model_card(ModelId(model))
+        model_card = await ModelCard.load(model)
         resolved_model = model_card.model_id
         if not any(
             instance.shard_assignments.model_id == resolved_model
@@ -687,10 +721,10 @@ class API:
         When stream=True and partial_images > 0, returns a StreamingResponse
         with SSE-formatted events for partial and final images.
         """
-        payload.model = await self._validate_image_model(payload.model)
+        payload.model = await self._validate_image_model(ModelId(payload.model))
 
         command = ImageGeneration(
-            request_params=payload,
+            task_params=payload,
         )
         await self._send(command)
 
@@ -729,7 +763,9 @@ class API:
         images_complete = 0
 
         try:
-            self._image_generation_queues[command_id], recv = channel[ImageChunk]()
+            self._image_generation_queues[command_id], recv = channel[
+                ImageChunk | ErrorChunk
+            ]()
 
             with recv as chunks:
                 async for chunk in chunks:
@@ -769,6 +805,7 @@ class API:
                             # Yield partial image event (always use b64_json for partials)
                             event_data = {
                                 "type": "partial",
+                                "image_index": chunk.image_index,
                                 "partial_index": partial_idx,
                                 "total_partials": total_partials,
                                 "format": str(chunk.format),
@@ -838,7 +875,9 @@ class API:
         stats: ImageGenerationStats | None = None
 
         try:
-            self._image_generation_queues[command_id], recv = channel[ImageChunk]()
+            self._image_generation_queues[command_id], recv = channel[
+                ImageChunk | ErrorChunk
+            ]()
 
             while images_complete < num_images:
                 with recv as chunks:
@@ -927,13 +966,13 @@ class API:
     async def bench_image_generations(
         self, request: Request, payload: BenchImageGenerationTaskParams
     ) -> BenchImageGenerationResponse:
-        payload.model = await self._validate_image_model(payload.model)
+        payload.model = await self._validate_image_model(ModelId(payload.model))
 
         payload.stream = False
         payload.partial_images = 0
 
         command = ImageGeneration(
-            request_params=payload,
+            task_params=payload,
         )
         await self._send(command)
 
@@ -948,7 +987,7 @@ class API:
         self,
         image: UploadFile,
         prompt: str,
-        model: str,
+        model: ModelId,
         n: int,
         size: str,
         response_format: Literal["url", "b64_json"],
@@ -956,6 +995,9 @@ class API:
         stream: bool,
         partial_images: int,
         bench: bool,
+        quality: Literal["high", "medium", "low"],
+        output_format: Literal["png", "jpeg", "webp"],
+        advanced_params: AdvancedImageParams | None,
     ) -> ImageEdits:
         """Prepare and send an image edits command with chunked image upload."""
         resolved_model = await self._validate_image_model(model)
@@ -972,7 +1014,7 @@ class API:
         total_chunks = len(data_chunks)
 
         command = ImageEdits(
-            request_params=ImageEditsInternalParams(
+            task_params=ImageEditsTaskParams(
                 image_data="",
                 total_input_chunks=total_chunks,
                 prompt=prompt,
@@ -984,6 +1026,9 @@ class API:
                 stream=stream,
                 partial_images=partial_images,
                 bench=bench,
+                quality=quality,
+                output_format=output_format,
+                advanced_params=advanced_params,
             ),
         )
 
@@ -994,7 +1039,6 @@ class API:
             await self._send(
                 SendInputChunk(
                     chunk=InputImageChunk(
-                        idx=chunk_index,
                         model=resolved_model,
                         command_id=command.command_id,
                         data=chunk_data,
@@ -1019,16 +1063,26 @@ class API:
         input_fidelity: Literal["low", "high"] = Form("low"),
         stream: str = Form("false"),
         partial_images: str = Form("0"),
+        quality: Literal["high", "medium", "low"] = Form("medium"),
+        output_format: Literal["png", "jpeg", "webp"] = Form("png"),
+        advanced_params: str | None = Form(None),
     ) -> ImageGenerationResponse | StreamingResponse:
         """Handle image editing requests (img2img)."""
         # Parse string form values to proper types
         stream_bool = stream.lower() in ("true", "1", "yes")
         partial_images_int = int(partial_images) if partial_images.isdigit() else 0
 
+        parsed_advanced_params: AdvancedImageParams | None = None
+        if advanced_params:
+            with contextlib.suppress(Exception):
+                parsed_advanced_params = AdvancedImageParams.model_validate_json(
+                    advanced_params
+                )
+
         command = await self._send_image_edits_command(
             image=image,
             prompt=prompt,
-            model=model,
+            model=ModelId(model),
             n=n,
             size=size,
             response_format=response_format,
@@ -1036,6 +1090,9 @@ class API:
             stream=stream_bool,
             partial_images=partial_images_int,
             bench=False,
+            quality=quality,
+            output_format=output_format,
+            advanced_params=parsed_advanced_params,
         )
 
         if stream_bool and partial_images_int > 0:
@@ -1066,12 +1123,22 @@ class API:
         size: str = Form("1024x1024"),
         response_format: Literal["url", "b64_json"] = Form("b64_json"),
         input_fidelity: Literal["low", "high"] = Form("low"),
+        quality: Literal["high", "medium", "low"] = Form("medium"),
+        output_format: Literal["png", "jpeg", "webp"] = Form("png"),
+        advanced_params: str | None = Form(None),
     ) -> BenchImageGenerationResponse:
         """Handle benchmark image editing requests with generation stats."""
+        parsed_advanced_params: AdvancedImageParams | None = None
+        if advanced_params:
+            with contextlib.suppress(Exception):
+                parsed_advanced_params = AdvancedImageParams.model_validate_json(
+                    advanced_params
+                )
+
         command = await self._send_image_edits_command(
             image=image,
             prompt=prompt,
-            model=model,
+            model=ModelId(model),
             n=n,
             size=size,
             response_format=response_format,
@@ -1079,6 +1146,9 @@ class API:
             stream=False,
             partial_images=0,
             bench=True,
+            quality=quality,
+            output_format=output_format,
+            advanced_params=parsed_advanced_params,
         )
 
         return await self._collect_image_generation_with_stats(
@@ -1088,6 +1158,62 @@ class API:
             response_format=response_format,
         )
 
+    async def claude_messages(
+        self, payload: ClaudeMessagesRequest
+    ) -> ClaudeMessagesResponse | StreamingResponse:
+        """Claude Messages API - adapter."""
+        task_params = claude_request_to_text_generation(payload)
+        resolved_model = await self._resolve_and_validate_text_model(
+            ModelId(task_params.model)
+        )
+        task_params = task_params.model_copy(update={"model": resolved_model})
+
+        command = TextGeneration(task_params=task_params)
+        await self._send(command)
+
+        if payload.stream:
+            return StreamingResponse(
+                generate_claude_stream(
+                    command.command_id,
+                    payload.model,
+                    self._token_chunk_stream(command.command_id),
+                ),
+                media_type="text/event-stream",
+            )
+
+        return await collect_claude_response(
+            command.command_id,
+            payload.model,
+            self._token_chunk_stream(command.command_id),
+        )
+
+    async def openai_responses(
+        self, payload: ResponsesRequest
+    ) -> ResponsesResponse | StreamingResponse:
+        """OpenAI Responses API."""
+        task_params = responses_request_to_text_generation(payload)
+        resolved_model = await self._resolve_and_validate_text_model(task_params.model)
+        task_params = task_params.model_copy(update={"model": resolved_model})
+
+        command = TextGeneration(task_params=task_params)
+        await self._send(command)
+
+        if payload.stream:
+            return StreamingResponse(
+                generate_responses_stream(
+                    command.command_id,
+                    payload.model,
+                    self._token_chunk_stream(command.command_id),
+                ),
+                media_type="text/event-stream",
+            )
+
+        return await collect_responses_response(
+            command.command_id,
+            payload.model,
+            self._token_chunk_stream(command.command_id),
+        )
+
     def _calculate_total_available_memory(self) -> Memory:
         """Calculate total available memory across all nodes in bytes."""
         total_available = Memory()
@@ -1148,27 +1274,44 @@ class API:
                 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 event.command_id in self._chat_completion_queues:
-                            assert isinstance(event.chunk, TokenChunk)
-                            queue = self._chat_completion_queues.get(event.command_id)
-                            if queue is not None:
-                                try:
-                                    await queue.send(event.chunk)
-                                except BrokenResourceError:
-                                    self._chat_completion_queues.pop(
-                                        event.command_id, None
-                                    )
-                        elif event.command_id in self._image_generation_queues:
+                        if queue := self._image_generation_queues.get(
+                            event.command_id, None
+                        ):
                             assert isinstance(event.chunk, ImageChunk)
-                            queue = self._image_generation_queues.get(event.command_id)
-                            if queue is not None:
-                                try:
-                                    await queue.send(event.chunk)
-                                except BrokenResourceError:
-                                    self._image_generation_queues.pop(
-                                        event.command_id, None
-                                    )
+                            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 = [
+            TraceEvent(
+                name=t.name,
+                start_us=t.start_us,
+                duration_us=t.duration_us,
+                rank=t.rank,
+                category=t.category,
+            )
+            for t in event.traces
+        ]
+        output_path = EXO_TRACING_CACHE_DIR / f"trace_{event.task_id}.json"
+        export_trace(traces, output_path)
+        logger.debug(f"Saved merged trace to {output_path}")
 
     async def _pause_on_new_election(self):
         with self.election_receiver as ems:
@@ -1191,3 +1334,128 @@ class API:
         await self.command_sender.send(
             ForwarderCommand(origin=self.node_id, command=command)
         )
+
+    async def _send_download(self, command: DownloadCommand):
+        await self.download_command_sender.send(
+            ForwarderDownloadCommand(origin=self.node_id, command=command)
+        )
+
+    async def start_download(
+        self, payload: StartDownloadParams
+    ) -> StartDownloadResponse:
+        command = StartDownload(
+            target_node_id=payload.target_node_id,
+            shard_metadata=payload.shard_metadata,
+        )
+        await self._send_download(command)
+        return StartDownloadResponse(command_id=command.command_id)
+
+    async def delete_download(
+        self, node_id: NodeId, model_id: ModelId
+    ) -> DeleteDownloadResponse:
+        command = DeleteDownload(
+            target_node_id=node_id,
+            model_id=ModelId(model_id),
+        )
+        await self._send_download(command)
+        return DeleteDownloadResponse(command_id=command.command_id)
+
+    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] = []
+
+        for trace_file in sorted(
+            EXO_TRACING_CACHE_DIR.glob("trace_*.json"),
+            key=lambda p: p.stat().st_mtime,
+            reverse=True,
+        ):
+            # Extract task_id from filename (trace_{task_id}.json)
+            task_id = trace_file.stem.removeprefix("trace_")
+            stat = trace_file.stat()
+            created_at = datetime.fromtimestamp(
+                stat.st_mtime, tz=timezone.utc
+            ).isoformat()
+            traces.append(
+                TraceListItem(
+                    task_id=task_id,
+                    created_at=created_at,
+                    file_size=stat.st_size,
+                )
+            )
+
+        return TraceListResponse(traces=traces)
+
+    async def get_trace(self, task_id: str) -> TraceResponse:
+        trace_path = self._get_trace_path(task_id)
+
+        if not trace_path.exists():
+            raise HTTPException(status_code=404, detail=f"Trace not found: {task_id}")
+
+        trace_events = load_trace_file(trace_path)
+
+        return TraceResponse(
+            task_id=task_id,
+            traces=[
+                TraceEventResponse(
+                    name=event.name,
+                    start_us=event.start_us,
+                    duration_us=event.duration_us,
+                    rank=event.rank,
+                    category=event.category,
+                )
+                for event in trace_events
+            ],
+        )
+
+    async def get_trace_stats(self, task_id: str) -> TraceStatsResponse:
+        trace_path = self._get_trace_path(task_id)
+
+        if not trace_path.exists():
+            raise HTTPException(status_code=404, detail=f"Trace not found: {task_id}")
+
+        trace_events = load_trace_file(trace_path)
+        stats = compute_stats(trace_events)
+
+        return TraceStatsResponse(
+            task_id=task_id,
+            total_wall_time_us=stats.total_wall_time_us,
+            by_category={
+                category: TraceCategoryStats(
+                    total_us=cat_stats.total_us,
+                    count=cat_stats.count,
+                    min_us=cat_stats.min_us,
+                    max_us=cat_stats.max_us,
+                    avg_us=cat_stats.avg_us,
+                )
+                for category, cat_stats in stats.by_category.items()
+            },
+            by_rank={
+                rank: TraceRankStats(
+                    by_category={
+                        category: TraceCategoryStats(
+                            total_us=cat_stats.total_us,
+                            count=cat_stats.count,
+                            min_us=cat_stats.min_us,
+                            max_us=cat_stats.max_us,
+                            avg_us=cat_stats.avg_us,
+                        )
+                        for category, cat_stats in rank_stats.items()
+                    }
+                )
+                for rank, rank_stats in stats.by_rank.items()
+            },
+        )
+
+    async def get_trace_raw(self, task_id: str) -> FileResponse:
+        trace_path = self._get_trace_path(task_id)
+
+        if not trace_path.exists():
+            raise HTTPException(status_code=404, detail=f"Trace not found: {task_id}")
+
+        return FileResponse(
+            path=trace_path,
+            media_type="application/json",
+            filename=f"trace_{task_id}.json",
+        )

src/exo/master/main.py

@@ -11,8 +11,8 @@ from exo.master.placement import (
     place_instance,
 )
 from exo.shared.apply import apply
+from exo.shared.constants import EXO_TRACING_ENABLED
 from exo.shared.types.commands import (
-    ChatCompletion,
     CreateInstance,
     DeleteInstance,
     ForwarderCommand,
@@ -23,6 +23,7 @@ from exo.shared.types.commands import (
     SendInputChunk,
     TaskFinished,
     TestCommand,
+    TextGeneration,
 )
 from exo.shared.types.common import CommandId, NodeId, SessionId
 from exo.shared.types.events import (
@@ -35,11 +36,11 @@ from exo.shared.types.events import (
     NodeTimedOut,
     TaskCreated,
     TaskDeleted,
+    TraceEventData,
+    TracesCollected,
+    TracesMerged,
 )
 from exo.shared.types.state import State
-from exo.shared.types.tasks import (
-    ChatCompletion as ChatCompletionTask,
-)
 from exo.shared.types.tasks import (
     ImageEdits as ImageEditsTask,
 )
@@ -50,6 +51,9 @@ from exo.shared.types.tasks import (
     TaskId,
     TaskStatus,
 )
+from exo.shared.types.tasks import (
+    TextGeneration as TextGenerationTask,
+)
 from exo.shared.types.worker.instances import InstanceId
 from exo.utils.channels import Receiver, Sender, channel
 from exo.utils.event_buffer import MultiSourceBuffer
@@ -86,6 +90,8 @@ class Master:
         self._multi_buffer = MultiSourceBuffer[NodeId, Event]()
         # TODO: not have this
         self._event_log: list[Event] = []
+        self._pending_traces: dict[TaskId, dict[int, list[TraceEventData]]] = {}
+        self._expected_ranks: dict[TaskId, set[int]] = {}
 
     async def run(self):
         logger.info("Starting Master")
@@ -117,11 +123,11 @@ class Master:
                     match command:
                         case TestCommand():
                             pass
-                        case ChatCompletion():
+                        case TextGeneration():
                             for instance in self.state.instances.values():
                                 if (
                                     instance.shard_assignments.model_id
-                                    == command.request_params.model
+                                    == command.task_params.model
                                 ):
                                     task_count = sum(
                                         1
@@ -134,7 +140,7 @@ class Master:
 
                             if not instance_task_counts:
                                 raise ValueError(
-                                    f"No instance found for model {command.request_params.model}"
+                                    f"No instance found for model {command.task_params.model}"
                                 )
 
                             available_instance_ids = sorted(
@@ -148,12 +154,12 @@ class Master:
                             generated_events.append(
                                 TaskCreated(
                                     task_id=task_id,
-                                    task=ChatCompletionTask(
+                                    task=TextGenerationTask(
                                         task_id=task_id,
                                         command_id=command.command_id,
                                         instance_id=available_instance_ids[0],
                                         task_status=TaskStatus.Pending,
-                                        task_params=command.request_params,
+                                        task_params=command.task_params,
                                     ),
                                 )
                             )
@@ -163,7 +169,7 @@ class Master:
                             for instance in self.state.instances.values():
                                 if (
                                     instance.shard_assignments.model_id
-                                    == command.request_params.model
+                                    == command.task_params.model
                                 ):
                                     task_count = sum(
                                         1
@@ -176,7 +182,7 @@ class Master:
 
                             if not instance_task_counts:
                                 raise ValueError(
-                                    f"No instance found for model {command.request_params.model}"
+                                    f"No instance found for model {command.task_params.model}"
                                 )
 
                             available_instance_ids = sorted(
@@ -187,25 +193,37 @@ class Master:
                             )
 
                             task_id = TaskId()
+                            selected_instance_id = available_instance_ids[0]
                             generated_events.append(
                                 TaskCreated(
                                     task_id=task_id,
                                     task=ImageGenerationTask(
                                         task_id=task_id,
                                         command_id=command.command_id,
-                                        instance_id=available_instance_ids[0],
+                                        instance_id=selected_instance_id,
                                         task_status=TaskStatus.Pending,
-                                        task_params=command.request_params,
+                                        task_params=command.task_params,
                                     ),
                                 )
                             )
 
                             self.command_task_mapping[command.command_id] = task_id
+
+                            if EXO_TRACING_ENABLED:
+                                selected_instance = self.state.instances.get(
+                                    selected_instance_id
+                                )
+                                if selected_instance:
+                                    ranks = set(
+                                        shard.device_rank
+                                        for shard in selected_instance.shard_assignments.runner_to_shard.values()
+                                    )
+                                    self._expected_ranks[task_id] = ranks
                         case ImageEdits():
                             for instance in self.state.instances.values():
                                 if (
                                     instance.shard_assignments.model_id
-                                    == command.request_params.model
+                                    == command.task_params.model
                                 ):
                                     task_count = sum(
                                         1
@@ -218,7 +236,7 @@ class Master:
 
                             if not instance_task_counts:
                                 raise ValueError(
-                                    f"No instance found for model {command.request_params.model}"
+                                    f"No instance found for model {command.task_params.model}"
                                 )
 
                             available_instance_ids = sorted(
@@ -229,20 +247,32 @@ class Master:
                             )
 
                             task_id = TaskId()
+                            selected_instance_id = available_instance_ids[0]
                             generated_events.append(
                                 TaskCreated(
                                     task_id=task_id,
                                     task=ImageEditsTask(
                                         task_id=task_id,
                                         command_id=command.command_id,
-                                        instance_id=available_instance_ids[0],
+                                        instance_id=selected_instance_id,
                                         task_status=TaskStatus.Pending,
-                                        task_params=command.request_params,
+                                        task_params=command.task_params,
                                     ),
                                 )
                             )
 
                             self.command_task_mapping[command.command_id] = task_id
+
+                            if EXO_TRACING_ENABLED:
+                                selected_instance = self.state.instances.get(
+                                    selected_instance_id
+                                )
+                                if selected_instance:
+                                    ranks = set(
+                                        shard.device_rank
+                                        for shard in selected_instance.shard_assignments.runner_to_shard.values()
+                                    )
+                                    self._expected_ranks[task_id] = ranks
                         case DeleteInstance():
                             placement = delete_instance(command, self.state.instances)
                             transition_events = get_transition_events(
@@ -335,6 +365,10 @@ class Master:
                     local_event.origin,
                 )
                 for event in self._multi_buffer.drain():
+                    if isinstance(event, TracesCollected):
+                        await self._handle_traces_collected(event)
+                        continue
+
                     logger.debug(f"Master indexing event: {str(event)[:100]}")
                     indexed = IndexedEvent(event=event, idx=len(self._event_log))
                     self.state = apply(self.state, indexed)
@@ -373,3 +407,29 @@ class Master:
                 event=event.event,
             )
         )
+
+    async def _handle_traces_collected(self, event: TracesCollected) -> None:
+        task_id = event.task_id
+        if task_id not in self._pending_traces:
+            self._pending_traces[task_id] = {}
+        self._pending_traces[task_id][event.rank] = event.traces
+
+        if (
+            task_id in self._expected_ranks
+            and set(self._pending_traces[task_id].keys())
+            >= self._expected_ranks[task_id]
+        ):
+            await self._merge_and_save_traces(task_id)
+
+    async def _merge_and_save_traces(self, task_id: TaskId) -> None:
+        all_trace_data: list[TraceEventData] = []
+        for trace_data in self._pending_traces[task_id].values():
+            all_trace_data.extend(trace_data)
+
+        await self.event_sender.send(
+            TracesMerged(task_id=task_id, traces=all_trace_data)
+        )
+
+        del self._pending_traces[task_id]
+        if task_id in self._expected_ranks:
+            del self._expected_ranks[task_id]

src/exo/master/placement.py

@@ -35,7 +35,7 @@ from exo.shared.types.worker.shards import Sharding
 
 def random_ephemeral_port() -> int:
     port = random.randint(49153, 65535)
-    return port - 1 if port <= 52415 else 52414
+    return port - 1 if port <= 52415 else port
 
 
 def add_instance_to_placements(
@@ -54,9 +54,18 @@ def place_instance(
     current_instances: Mapping[InstanceId, Instance],
     node_memory: Mapping[NodeId, MemoryUsage],
     node_network: Mapping[NodeId, NodeNetworkInfo],
+    required_nodes: set[NodeId] | None = None,
 ) -> dict[InstanceId, Instance]:
     cycles = topology.get_cycles()
     candidate_cycles = list(filter(lambda it: len(it) >= command.min_nodes, cycles))
+
+    # Filter to cycles containing all required nodes (subset matching)
+    if required_nodes:
+        candidate_cycles = [
+            cycle
+            for cycle in candidate_cycles
+            if required_nodes.issubset(cycle.node_ids)
+        ]
     cycles_with_sufficient_memory = filter_cycles_by_memory(
         candidate_cycles, node_memory, command.model_card.storage_size
     )
@@ -87,12 +96,12 @@ def place_instance(
 
     smallest_cycles = get_smallest_cycles(cycles_with_sufficient_memory)
 
-    smallest_tb_cycles = [
-        cycle for cycle in smallest_cycles if topology.is_thunderbolt_cycle(cycle)
+    smallest_rdma_cycles = [
+        cycle for cycle in smallest_cycles if topology.is_rdma_cycle(cycle)
     ]
 
-    if smallest_tb_cycles != []:
-        smallest_cycles = smallest_tb_cycles
+    if command.instance_meta == InstanceMeta.MlxJaccl and smallest_rdma_cycles != []:
+        smallest_cycles = smallest_rdma_cycles
 
     cycles_with_leaf_nodes: list[Cycle] = [
         cycle

src/exo/master/placement_utils.py

@@ -197,49 +197,6 @@ def get_shard_assignments(
             )
 
 
-def get_hosts_from_subgraph(cycle_digraph: Topology) -> list[Host]:
-    cycles = cycle_digraph.get_cycles()
-    expected_length = len(list(cycle_digraph.list_nodes()))
-    cycles = [cycle for cycle in cycles if len(cycle) == expected_length]
-    if not cycles:
-        if expected_length > 1:
-            logger.warning(
-                f"No cycles of length {expected_length} found even though chosen subgraph contained {expected_length} nodes"
-            )
-        return []
-
-    cycle = cycles[0]
-
-    get_thunderbolt = False
-    if cycle_digraph.is_thunderbolt_cycle(cycle):
-        get_thunderbolt = True
-
-    logger.debug(f"Using thunderbolt cycle: {get_thunderbolt}")
-
-    hosts: list[Host] = []
-    for i in range(len(cycle)):
-        current_node = cycle.node_ids[i]
-        next_node = cycle.node_ids[(i + 1) % len(cycle)]
-
-        for connection in cycle_digraph.get_all_connections_between(
-            source=current_node, sink=next_node
-        ):
-            if not isinstance(connection, SocketConnection):
-                continue
-
-            if get_thunderbolt and not connection.is_thunderbolt():
-                continue
-
-            host = Host(
-                ip=connection.sink_multiaddr.ip_address,
-                port=connection.sink_multiaddr.port,
-            )
-            hosts.append(host)
-            break
-
-    return hosts
-
-
 def get_mlx_jaccl_devices_matrix(
     selected_cycle: list[NodeId],
     cycle_digraph: Topology,
@@ -265,9 +222,6 @@ def get_mlx_jaccl_devices_matrix(
                     matrix[i][j] = conn.source_rdma_iface
                     break
             else:
-                logger.warning(
-                    f"Failed to find interface name between {node_i} and {node_j}"
-                )
                 raise ValueError(
                     "Current jaccl backend requires all-to-all RDMA connections"
                 )
@@ -279,22 +233,11 @@ def _find_connection_ip(
     node_i: NodeId,
     node_j: NodeId,
     cycle_digraph: Topology,
-) -> Generator[tuple[str, bool]]:
+) -> Generator[str, None, None]:
     """Find all IP addresses that connect node i to node j."""
     for connection in cycle_digraph.get_all_connections_between(node_i, node_j):
         if isinstance(connection, SocketConnection):
-            yield connection.sink_multiaddr.ip_address, connection.is_thunderbolt()
-
-
-def _find_interface_name_for_ip(
-    ip_address: str, node_network: NodeNetworkInfo
-) -> str | None:
-    """Find the interface name for an IP address on a node (any interface)."""
-    for interface in node_network.interfaces:
-        if interface.ip_address == ip_address:
-            return interface.name
-
-    return None
+            yield connection.sink_multiaddr.ip_address
 
 
 def _find_ip_prioritised(
@@ -303,43 +246,25 @@ def _find_ip_prioritised(
     cycle_digraph: Topology,
     node_network: Mapping[NodeId, NodeNetworkInfo],
 ) -> str | None:
-    # TODO: Actually prioritize in the correct Ethernet > Wifi > Non-TB > TB order.
     """Find an IP address between nodes with prioritization.
 
-    Priority order:
-    1. en0 (Ethernet on Mac Studio, WiFi on MacBook)
-    2. en1 (WiFi on Mac Studio, Ethernet on MacBook)
-    3. Non-Thunderbolt connections
-    4. Any other IP address
+    Priority: ethernet > wifi > unknown > thunderbolt
     """
     ips = list(_find_connection_ip(node_id, other_node_id, cycle_digraph))
-    # We expect a unique iface -> ip mapping
-    iface_map = {
-        _find_interface_name_for_ip(
-            ip, node_network.get(other_node_id, NodeNetworkInfo())
-        ): ip
-        for ip, _ in ips
+    if not ips:
+        return None
+    other_network = node_network.get(other_node_id, NodeNetworkInfo())
+    ip_to_type = {
+        iface.ip_address: iface.interface_type for iface in other_network.interfaces
     }
-
-    en0_ip = iface_map.get("en0")
-    if en0_ip:
-        return en0_ip
-
-    en1_ip = iface_map.get("en1")
-    if en1_ip:
-        return en1_ip
-
-    non_thunderbolt_ip = next(
-        (ip for (ip, is_thunderbolt) in ips if not is_thunderbolt), None
-    )
-
-    if non_thunderbolt_ip:
-        return non_thunderbolt_ip
-
-    if ips:
-        return ips[0][0]
-
-    return None
+    priority = {
+        "ethernet": 0,
+        "wifi": 1,
+        "unknown": 2,
+        "maybe_ethernet": 3,
+        "thunderbolt": 4,
+    }
+    return min(ips, key=lambda ip: priority.get(ip_to_type.get(ip, "unknown"), 2))
 
 
 def get_mlx_ring_hosts_by_node(
@@ -381,9 +306,6 @@ def get_mlx_ring_hosts_by_node(
                 node_id, other_node_id, cycle_digraph, node_network
             )
             if connection_ip is None:
-                logger.warning(
-                    f"Failed to find prioritised connection IP between {node_id} and {other_node_id}"
-                )
                 raise ValueError(
                     "MLX ring backend requires connectivity between neighbouring nodes"
                 )
@@ -416,9 +338,6 @@ def get_mlx_jaccl_coordinators(
         if ip is not None:
             return ip
 
-        logger.warning(
-            f"Failed to find directly connected ip between {n} and {coordinator}"
-        )
         raise ValueError(
             "Current jaccl backend requires all participating devices to be able to communicate"
         )

src/exo/master/tests/test_api_error_handling.py

@@ -1,13 +1,9 @@
 # pyright: reportUnusedFunction=false, reportAny=false
-from typing import Any, get_args
+from typing import Any
 
 from fastapi import FastAPI, HTTPException
 from fastapi.testclient import TestClient
 
-from exo.shared.types.api import ErrorInfo, ErrorResponse, FinishReason
-from exo.shared.types.chunks import ImageChunk, TokenChunk
-from exo.worker.tests.constants import MODEL_A_ID
-
 
 def test_http_exception_handler_formats_openai_style() -> None:
     """Test that HTTPException is converted to OpenAI-style error format."""
@@ -48,95 +44,3 @@ def test_http_exception_handler_formats_openai_style() -> None:
     assert data["error"]["message"] == "Resource not found"
     assert data["error"]["type"] == "Not Found"
     assert data["error"]["code"] == 404
-
-
-def test_finish_reason_includes_error() -> None:
-    valid_reasons = get_args(FinishReason)
-    assert "error" in valid_reasons
-
-
-def test_token_chunk_with_error_fields() -> None:
-    chunk = TokenChunk(
-        idx=0,
-        model=MODEL_A_ID,
-        text="",
-        token_id=0,
-        finish_reason="error",
-        error_message="Something went wrong",
-    )
-
-    assert chunk.finish_reason == "error"
-    assert chunk.error_message == "Something went wrong"
-
-
-def test_token_chunk_without_error() -> None:
-    chunk = TokenChunk(
-        idx=1,
-        model=MODEL_A_ID,
-        text="Hello",
-        token_id=42,
-        finish_reason=None,
-    )
-
-    assert chunk.finish_reason is None
-    assert chunk.error_message is None
-
-
-def test_error_response_construction() -> None:
-    error_response = ErrorResponse(
-        error=ErrorInfo(
-            message="Generation failed",
-            type="InternalServerError",
-            code=500,
-        )
-    )
-
-    assert error_response.error.message == "Generation failed"
-    assert error_response.error.code == 500
-
-
-def test_normal_finish_reasons_still_work() -> None:
-    for reason in ["stop", "length", "tool_calls", "content_filter", "function_call"]:
-        chunk = TokenChunk(
-            idx=0,
-            model=MODEL_A_ID,
-            text="done",
-            token_id=100,
-            finish_reason=reason,  # type: ignore[arg-type]
-        )
-        assert chunk.finish_reason == reason
-
-
-def test_image_chunk_with_error_fields() -> None:
-    chunk = ImageChunk(
-        idx=0,
-        model=MODEL_A_ID,
-        data="",
-        chunk_index=0,
-        total_chunks=1,
-        image_index=0,
-        finish_reason="error",
-        error_message="Image generation failed",
-    )
-
-    assert chunk.finish_reason == "error"
-    assert chunk.error_message == "Image generation failed"
-    assert chunk.data == ""
-    assert chunk.chunk_index == 0
-    assert chunk.total_chunks == 1
-    assert chunk.image_index == 0
-
-
-def test_image_chunk_without_error() -> None:
-    chunk = ImageChunk(
-        idx=0,
-        model=MODEL_A_ID,
-        data="base64encodeddata",
-        chunk_index=0,
-        total_chunks=1,
-        image_index=0,
-    )
-
-    assert chunk.finish_reason is None
-    assert chunk.error_message is None
-    assert chunk.data == "base64encodeddata"

src/exo/master/tests/test_claude_api.py

@@ -0,0 +1,182 @@
+"""Tests for Claude Messages API conversion functions and types."""
+
+import pydantic
+import pytest
+
+from exo.master.adapters.claude import (
+    claude_request_to_text_generation,
+    finish_reason_to_claude_stop_reason,
+)
+from exo.shared.types.claude_api import (
+    ClaudeMessage,
+    ClaudeMessagesRequest,
+    ClaudeTextBlock,
+)
+from exo.shared.types.common import ModelId
+
+
+class TestFinishReasonToClaudeStopReason:
+    """Tests for finish_reason to Claude stop_reason mapping."""
+
+    def test_stop_maps_to_end_turn(self):
+        assert finish_reason_to_claude_stop_reason("stop") == "end_turn"
+
+    def test_length_maps_to_max_tokens(self):
+        assert finish_reason_to_claude_stop_reason("length") == "max_tokens"
+
+    def test_tool_calls_maps_to_tool_use(self):
+        assert finish_reason_to_claude_stop_reason("tool_calls") == "tool_use"
+
+    def test_function_call_maps_to_tool_use(self):
+        assert finish_reason_to_claude_stop_reason("function_call") == "tool_use"
+
+    def test_content_filter_maps_to_end_turn(self):
+        assert finish_reason_to_claude_stop_reason("content_filter") == "end_turn"
+
+    def test_none_returns_none(self):
+        assert finish_reason_to_claude_stop_reason(None) is None
+
+
+class TestClaudeRequestToInternal:
+    """Tests for converting Claude Messages API requests to TextGenerationTaskParams."""
+
+    def test_basic_request_conversion(self):
+        request = ClaudeMessagesRequest(
+            model=ModelId("claude-3-opus"),
+            max_tokens=100,
+            messages=[
+                ClaudeMessage(role="user", content="Hello"),
+            ],
+        )
+        params = claude_request_to_text_generation(request)
+
+        assert params.model == "claude-3-opus"
+        assert params.max_output_tokens == 100
+        assert isinstance(params.input, list)
+        assert len(params.input) == 1
+        assert params.input[0].role == "user"
+        assert params.input[0].content == "Hello"
+        assert params.instructions is None
+
+    def test_request_with_system_string(self):
+        request = ClaudeMessagesRequest(
+            model=ModelId("claude-3-opus"),
+            max_tokens=100,
+            system="You are a helpful assistant.",
+            messages=[
+                ClaudeMessage(role="user", content="Hello"),
+            ],
+        )
+        params = claude_request_to_text_generation(request)
+
+        assert params.instructions == "You are a helpful assistant."
+        assert isinstance(params.input, list)
+        assert len(params.input) == 1
+        assert params.input[0].role == "user"
+        assert params.input[0].content == "Hello"
+
+    def test_request_with_system_text_blocks(self):
+        request = ClaudeMessagesRequest(
+            model=ModelId("claude-3-opus"),
+            max_tokens=100,
+            system=[
+                ClaudeTextBlock(text="You are helpful. "),
+                ClaudeTextBlock(text="Be concise."),
+            ],
+            messages=[
+                ClaudeMessage(role="user", content="Hello"),
+            ],
+        )
+        params = claude_request_to_text_generation(request)
+
+        assert params.instructions == "You are helpful. Be concise."
+        assert isinstance(params.input, list)
+        assert len(params.input) == 1
+
+    def test_request_with_content_blocks(self):
+        request = ClaudeMessagesRequest(
+            model=ModelId("claude-3-opus"),
+            max_tokens=100,
+            messages=[
+                ClaudeMessage(
+                    role="user",
+                    content=[
+                        ClaudeTextBlock(text="First part. "),
+                        ClaudeTextBlock(text="Second part."),
+                    ],
+                ),
+            ],
+        )
+        params = claude_request_to_text_generation(request)
+
+        assert isinstance(params.input, list)
+        assert len(params.input) == 1
+        assert params.input[0].content == "First part. Second part."
+
+    def test_request_with_multi_turn_conversation(self):
+        request = ClaudeMessagesRequest(
+            model=ModelId("claude-3-opus"),
+            max_tokens=100,
+            messages=[
+                ClaudeMessage(role="user", content="Hello"),
+                ClaudeMessage(role="assistant", content="Hi there!"),
+                ClaudeMessage(role="user", content="How are you?"),
+            ],
+        )
+        params = claude_request_to_text_generation(request)
+
+        assert isinstance(params.input, list)
+        assert len(params.input) == 3
+        assert params.input[0].role == "user"
+        assert params.input[1].role == "assistant"
+        assert params.input[2].role == "user"
+
+    def test_request_with_optional_parameters(self):
+        request = ClaudeMessagesRequest(
+            model=ModelId("claude-3-opus"),
+            max_tokens=100,
+            messages=[ClaudeMessage(role="user", content="Hello")],
+            temperature=0.7,
+            top_p=0.9,
+            top_k=40,
+            stop_sequences=["STOP", "END"],
+            stream=True,
+        )
+        params = claude_request_to_text_generation(request)
+
+        assert params.temperature == 0.7
+        assert params.top_p == 0.9
+        assert params.top_k == 40
+        assert params.stop == ["STOP", "END"]
+        assert params.stream is True
+
+
+class TestClaudeMessagesRequestValidation:
+    """Tests for Claude Messages API request validation."""
+
+    def test_request_requires_model(self):
+        with pytest.raises(pydantic.ValidationError):
+            ClaudeMessagesRequest.model_validate(
+                {
+                    "max_tokens": 100,
+                    "messages": [{"role": "user", "content": "Hello"}],
+                }
+            )
+
+    def test_request_requires_max_tokens(self):
+        with pytest.raises(pydantic.ValidationError):
+            ClaudeMessagesRequest.model_validate(
+                {
+                    "model": "claude-3-opus",
+                    "messages": [{"role": "user", "content": "Hello"}],
+                }
+            )
+
+    def test_request_requires_messages(self):
+        with pytest.raises(pydantic.ValidationError):
+            ClaudeMessagesRequest.model_validate(
+                {
+                    "model": "claude-3-opus",
+                    "max_tokens": 100,
+                }
+            )

src/exo/master/tests/test_claude_tool_use.py

@@ -0,0 +1,265 @@
+"""Tests for Claude Messages API tool_use support in the adapter."""
+
+import json
+from collections.abc import AsyncGenerator
+from typing import Any, cast
+
+from exo.master.adapters.claude import collect_claude_response, generate_claude_stream
+from exo.shared.types.api import ToolCallItem
+from exo.shared.types.chunks import ErrorChunk, TokenChunk, ToolCallChunk
+from exo.shared.types.common import CommandId, ModelId
+
+
+async def _chunks_to_stream(
+    chunks: list[ErrorChunk | ToolCallChunk | TokenChunk],
+) -> AsyncGenerator[ErrorChunk | ToolCallChunk | TokenChunk, None]:
+    for chunk in chunks:
+        yield chunk
+
+
+MODEL = ModelId("test-model")
+COMMAND_ID = CommandId("cmd_test123")
+
+
+def _parse_sse_events(events: list[str]) -> list[dict[str, Any]]:
+    """Parse SSE event strings into JSON dicts."""
+    parsed: list[dict[str, Any]] = []
+    for event_str in events:
+        for line in event_str.strip().split("\n"):
+            if line.startswith("data: "):
+                parsed.append(cast(dict[str, Any], json.loads(line[6:])))
+    return parsed
+
+
+class TestCollectClaudeResponseToolUse:
+    """Tests for non-streaming tool_use response collection."""
+
+    async def test_tool_call_chunk_produces_tool_use_blocks(self):
+        chunks: list[ErrorChunk | ToolCallChunk | TokenChunk] = [
+            ToolCallChunk(
+                model=MODEL,
+                usage=None,
+                tool_calls=[
+                    ToolCallItem(
+                        name="get_weather",
+                        arguments='{"location": "San Francisco"}',
+                    )
+                ],
+            ),
+        ]
+        response = await collect_claude_response(
+            COMMAND_ID, "test-model", _chunks_to_stream(chunks)
+        )
+
+        assert response.stop_reason == "tool_use"
+        tool_blocks = [b for b in response.content if b.type == "tool_use"]
+        assert len(tool_blocks) == 1
+        block = tool_blocks[0]
+        assert block.type == "tool_use"
+        assert block.name == "get_weather"
+        assert block.input == {"location": "San Francisco"}
+        assert block.id.startswith("toolu_")
+
+    async def test_multiple_tool_calls(self):
+        chunks: list[ErrorChunk | ToolCallChunk | TokenChunk] = [
+            ToolCallChunk(
+                model=MODEL,
+                usage=None,
+                tool_calls=[
+                    ToolCallItem(
+                        name="get_weather",
+                        arguments='{"location": "SF"}',
+                    ),
+                    ToolCallItem(
+                        name="get_time",
+                        arguments='{"timezone": "PST"}',
+                    ),
+                ],
+            ),
+        ]
+        response = await collect_claude_response(
+            COMMAND_ID, "test-model", _chunks_to_stream(chunks)
+        )
+
+        assert response.stop_reason == "tool_use"
+        tool_blocks = [b for b in response.content if b.type == "tool_use"]
+        assert len(tool_blocks) == 2
+        assert tool_blocks[0].name == "get_weather"
+        assert tool_blocks[1].name == "get_time"
+
+    async def test_mixed_text_and_tool_use(self):
+        chunks: list[ErrorChunk | ToolCallChunk | TokenChunk] = [
+            TokenChunk(model=MODEL, text="Let me check ", token_id=1, usage=None),
+            TokenChunk(model=MODEL, text="the weather.", token_id=2, usage=None),
+            ToolCallChunk(
+                model=MODEL,
+                usage=None,
+                tool_calls=[
+                    ToolCallItem(
+                        name="get_weather",
+                        arguments='{"location": "NYC"}',
+                    )
+                ],
+            ),
+        ]
+        response = await collect_claude_response(
+            COMMAND_ID, "test-model", _chunks_to_stream(chunks)
+        )
+
+        assert response.stop_reason == "tool_use"
+        text_blocks = [b for b in response.content if b.type == "text"]
+        tool_blocks = [b for b in response.content if b.type == "tool_use"]
+        assert len(text_blocks) == 1
+        assert text_blocks[0].text == "Let me check the weather."
+        assert len(tool_blocks) == 1
+        assert tool_blocks[0].name == "get_weather"
+
+    async def test_no_content_produces_empty_text_block(self):
+        chunks: list[ErrorChunk | ToolCallChunk | TokenChunk] = []
+        response = await collect_claude_response(
+            COMMAND_ID, "test-model", _chunks_to_stream(chunks)
+        )
+        assert len(response.content) == 1
+        assert response.content[0].type == "text"
+
+
+class TestGenerateClaudeStreamToolUse:
+    """Tests for streaming tool_use event generation."""
+
+    async def test_tool_call_emits_tool_use_events(self):
+        chunks: list[ErrorChunk | ToolCallChunk | TokenChunk] = [
+            ToolCallChunk(
+                model=MODEL,
+                usage=None,
+                tool_calls=[
+                    ToolCallItem(
+                        name="get_weather",
+                        arguments='{"location": "SF"}',
+                    )
+                ],
+            ),
+        ]
+        events: list[str] = []
+        async for event in generate_claude_stream(
+            COMMAND_ID, "test-model", _chunks_to_stream(chunks)
+        ):
+            events.append(event)
+
+        parsed = _parse_sse_events(events)
+
+        # Find tool_use content_block_start
+        tool_starts = [
+            e
+            for e in parsed
+            if e.get("type") == "content_block_start"
+            and cast(dict[str, Any], e.get("content_block", {})).get("type")
+            == "tool_use"
+        ]
+        assert len(tool_starts) == 1
+        content_block = cast(dict[str, Any], tool_starts[0]["content_block"])
+        assert content_block["name"] == "get_weather"
+        assert content_block["input"] == {}
+        assert cast(str, content_block["id"]).startswith("toolu_")
+
+        # Find input_json_delta
+        json_deltas = [
+            e
+            for e in parsed
+            if e.get("type") == "content_block_delta"
+            and cast(dict[str, Any], e.get("delta", {})).get("type")
+            == "input_json_delta"
+        ]
+        assert len(json_deltas) == 1
+        delta = cast(dict[str, Any], json_deltas[0]["delta"])
+        assert json.loads(cast(str, delta["partial_json"])) == {"location": "SF"}
+
+        # Find message_delta with tool_use stop reason
+        msg_deltas = [e for e in parsed if e.get("type") == "message_delta"]
+        assert len(msg_deltas) == 1
+        assert cast(dict[str, Any], msg_deltas[0]["delta"])["stop_reason"] == "tool_use"
+
+    async def test_streaming_mixed_text_and_tool_use(self):
+        chunks: list[ErrorChunk | ToolCallChunk | TokenChunk] = [
+            TokenChunk(model=MODEL, text="Hello ", token_id=1, usage=None),
+            ToolCallChunk(
+                model=MODEL,
+                usage=None,
+                tool_calls=[
+                    ToolCallItem(
+                        name="search",
+                        arguments='{"query": "test"}',
+                    )
+                ],
+            ),
+        ]
+        events: list[str] = []
+        async for event in generate_claude_stream(
+            COMMAND_ID, "test-model", _chunks_to_stream(chunks)
+        ):
+            events.append(event)
+
+        parsed = _parse_sse_events(events)
+
+        # Should have text delta at index 0
+        text_deltas = [
+            e
+            for e in parsed
+            if e.get("type") == "content_block_delta"
+            and cast(dict[str, Any], e.get("delta", {})).get("type") == "text_delta"
+        ]
+        assert len(text_deltas) == 1
+        assert text_deltas[0]["index"] == 0
+        assert cast(dict[str, Any], text_deltas[0]["delta"])["text"] == "Hello "
+
+        # Tool block at index 1
+        tool_starts = [
+            e
+            for e in parsed
+            if e.get("type") == "content_block_start"
+            and cast(dict[str, Any], e.get("content_block", {})).get("type")
+            == "tool_use"
+        ]
+        assert len(tool_starts) == 1
+        assert tool_starts[0]["index"] == 1
+
+        # Stop reason should be tool_use
+        msg_deltas = [e for e in parsed if e.get("type") == "message_delta"]
+        assert cast(dict[str, Any], msg_deltas[0]["delta"])["stop_reason"] == "tool_use"
+
+    async def test_streaming_tool_block_stop_events(self):
+        chunks: list[ErrorChunk | ToolCallChunk | TokenChunk] = [
+            ToolCallChunk(
+                model=MODEL,
+                usage=None,
+                tool_calls=[
+                    ToolCallItem(name="fn1", arguments="{}"),
+                    ToolCallItem(name="fn2", arguments='{"a": 1}'),
+                ],
+            ),
+        ]
+        events: list[str] = []
+        async for event in generate_claude_stream(
+            COMMAND_ID, "test-model", _chunks_to_stream(chunks)
+        ):
+            events.append(event)
+
+        parsed = _parse_sse_events(events)
+
+        # Two tool block starts (at indices 1 and 2)
+        tool_starts = [
+            e
+            for e in parsed
+            if e.get("type") == "content_block_start"
+            and cast(dict[str, Any], e.get("content_block", {})).get("type")
+            == "tool_use"
+        ]
+        assert len(tool_starts) == 2
+        assert tool_starts[0]["index"] == 1
+        assert tool_starts[1]["index"] == 2
+
+        # Two tool block stops (at indices 1 and 2), plus text block stop at 0
+        block_stops = [e for e in parsed if e.get("type") == "content_block_stop"]
+        stop_indices = [e["index"] for e in block_stops]
+        assert 0 in stop_indices
+        assert 1 in stop_indices
+        assert 2 in stop_indices

src/exo/master/tests/test_master.py

@@ -7,15 +7,14 @@ from loguru import logger
 
 from exo.master.main import Master
 from exo.routing.router import get_node_id_keypair
-from exo.shared.models.model_cards import ModelCard, ModelId, ModelTask
-from exo.shared.types.api import ChatCompletionMessage, ChatCompletionTaskParams
+from exo.shared.models.model_cards import ModelCard, ModelTask
 from exo.shared.types.commands import (
-    ChatCompletion,
     CommandId,
     ForwarderCommand,
     PlaceInstance,
+    TextGeneration,
 )
-from exo.shared.types.common import NodeId, SessionId
+from exo.shared.types.common import ModelId, NodeId, SessionId
 from exo.shared.types.events import (
     ForwarderEvent,
     IndexedEvent,
@@ -27,8 +26,9 @@ from exo.shared.types.memory import Memory
 from exo.shared.types.profiling import (
     MemoryUsage,
 )
-from exo.shared.types.tasks import ChatCompletion as ChatCompletionTask
 from exo.shared.types.tasks import TaskStatus
+from exo.shared.types.tasks import TextGeneration as TextGenerationTask
+from exo.shared.types.text_generation import TextGenerationTaskParams
 from exo.shared.types.worker.instances import (
     InstanceMeta,
     MlxRingInstance,
@@ -127,20 +127,16 @@ async def test_master():
         logger.info("wait for an instance")
         while len(master.state.instances.keys()) == 0:
             await anyio.sleep(0.001)
-        logger.info("inject a ChatCompletion Command")
+        logger.info("inject a TextGeneration Command")
         await command_sender.send(
             ForwarderCommand(
                 origin=node_id,
                 command=(
-                    ChatCompletion(
+                    TextGeneration(
                         command_id=CommandId(),
-                        request_params=ChatCompletionTaskParams(
-                            model="llama-3.2-1b",
-                            messages=[
-                                ChatCompletionMessage(
-                                    role="user", content="Hello, how are you?"
-                                )
-                            ],
+                        task_params=TextGenerationTaskParams(
+                            model=ModelId("llama-3.2-1b"),
+                            input="Hello, how are you?",
                         ),
                     )
                 ),
@@ -190,12 +186,10 @@ async def test_master():
         assert created_instance.ephemeral_port > 0
         assert isinstance(events[2].event, TaskCreated)
         assert events[2].event.task.task_status == TaskStatus.Pending
-        assert isinstance(events[2].event.task, ChatCompletionTask)
-        assert events[2].event.task.task_params == ChatCompletionTaskParams(
-            model="llama-3.2-1b",
-            messages=[
-                ChatCompletionMessage(role="user", content="Hello, how are you?")
-            ],
+        assert isinstance(events[2].event.task, TextGenerationTask)
+        assert events[2].event.task.task_params == TextGenerationTaskParams(
+            model=ModelId("llama-3.2-1b"),
+            input="Hello, how are you?",
         )
 
         await master.shutdown()

src/exo/master/tests/test_openai_responses_api.py

@@ -0,0 +1,48 @@
+"""Tests for OpenAI Responses API wire types.
+
+ResponsesRequest is the API wire type for the Responses endpoint.
+The responses adapter converts it to TextGenerationTaskParams for the pipeline.
+"""
+
+import pydantic
+import pytest
+
+from exo.shared.types.common import ModelId
+from exo.shared.types.openai_responses import (
+    ResponseInputMessage,
+    ResponsesRequest,
+)
+
+
+class TestResponsesRequestValidation:
+    """Tests for OpenAI Responses API request validation."""
+
+    def test_request_requires_model(self):
+        with pytest.raises(pydantic.ValidationError):
+            ResponsesRequest.model_validate(
+                {
+                    "input": "Hello",
+                }
+            )
+
+    def test_request_requires_input(self):
+        with pytest.raises(pydantic.ValidationError):
+            ResponsesRequest.model_validate(
+                {
+                    "model": "gpt-4o",
+                }
+            )
+
+    def test_request_accepts_string_input(self):
+        request = ResponsesRequest(
+            model=ModelId("gpt-4o"),
+            input="Hello",
+        )
+        assert request.input == "Hello"
+
+    def test_request_accepts_message_array_input(self):
+        request = ResponsesRequest(
+            model=ModelId("gpt-4o"),
+            input=[ResponseInputMessage(role="user", content="Hello")],
+        )
+        assert len(request.input) == 1

src/exo/master/tests/test_placement_utils.py

@@ -3,7 +3,6 @@ import pytest
 from exo.master.placement_utils import (
     allocate_layers_proportionally,
     filter_cycles_by_memory,
-    get_hosts_from_subgraph,
     get_mlx_jaccl_coordinators,
     get_shard_assignments,
     get_smallest_cycles,
@@ -14,7 +13,7 @@ from exo.master.tests.conftest import (
 )
 from exo.shared.models.model_cards import ModelCard, ModelId, ModelTask
 from exo.shared.topology import Topology
-from exo.shared.types.common import Host, NodeId
+from exo.shared.types.common import NodeId
 from exo.shared.types.memory import Memory
 from exo.shared.types.profiling import (
     NetworkInterfaceInfo,
@@ -273,45 +272,6 @@ def test_get_shard_assignments(
     )
 
 
-def test_get_hosts_from_subgraph():
-    # arrange
-    node_a_id = NodeId()
-    node_b_id = NodeId()
-    node_c_id = NodeId()
-    topology = Topology()
-
-    topology.add_node(node_a_id)
-    topology.add_node(node_b_id)
-    topology.add_node(node_c_id)
-
-    connection1 = Connection(
-        source=node_a_id, sink=node_b_id, edge=create_socket_connection(1)
-    )
-    connection2 = Connection(
-        source=node_b_id, sink=node_c_id, edge=create_socket_connection(2)
-    )
-    connection3 = Connection(
-        source=node_c_id, sink=node_a_id, edge=create_socket_connection(3)
-    )
-
-    topology.add_connection(connection1)
-    topology.add_connection(connection2)
-    topology.add_connection(connection3)
-
-    # act
-    hosts = get_hosts_from_subgraph(topology)
-
-    # assert
-    assert len(hosts) == 3
-    expected_hosts = [
-        Host(ip="169.254.0.1", port=1234),
-        Host(ip="169.254.0.2", port=1234),
-        Host(ip="169.254.0.3", port=1234),
-    ]
-    for expected_host in expected_hosts:
-        assert expected_host in hosts
-
-
 def test_get_mlx_jaccl_coordinators():
     # arrange
     node_a_id = NodeId()

src/exo/routing/router.py

@@ -216,6 +216,8 @@ def get_node_id_keypair(
     Obtains the :class:`Keypair` associated with this node-ID.
     Obtain the :class:`PeerId` by from it.
     """
+    # TODO(evan): bring back node id persistence once we figure out how to deal with duplicates
+    return Keypair.generate_ed25519()
 
     def lock_path(path: str | bytes | PathLike[str] | PathLike[bytes]) -> Path:
         return Path(str(path) + ".lock")

src/exo/routing/topics.py

@@ -3,7 +3,7 @@ from enum import Enum
 
 from exo.routing.connection_message import ConnectionMessage
 from exo.shared.election import ElectionMessage
-from exo.shared.types.commands import ForwarderCommand
+from exo.shared.types.commands import ForwarderCommand, ForwarderDownloadCommand
 from exo.shared.types.events import (
     ForwarderEvent,
 )
@@ -45,3 +45,6 @@ ELECTION_MESSAGES = TypedTopic(
 CONNECTION_MESSAGES = TypedTopic(
     "connection_messages", PublishPolicy.Never, ConnectionMessage
 )
+DOWNLOAD_COMMANDS = TypedTopic(
+    "download_commands", PublishPolicy.Always, ForwarderDownloadCommand
+)

src/exo/shared/apply.py

@@ -25,11 +25,14 @@ from exo.shared.types.events import (
     TestEvent,
     TopologyEdgeCreated,
     TopologyEdgeDeleted,
+    TracesCollected,
+    TracesMerged,
 )
 from exo.shared.types.profiling import (
     NodeIdentity,
     NodeNetworkInfo,
     NodeThunderboltInfo,
+    ThunderboltBridgeStatus,
 )
 from exo.shared.types.state import State
 from exo.shared.types.tasks import Task, TaskId, TaskStatus
@@ -46,6 +49,7 @@ from exo.utils.info_gatherer.info_gatherer import (
     NodeConfig,
     NodeNetworkInterfaces,
     StaticNodeInformation,
+    ThunderboltBridgeInfo,
 )
 
 
@@ -53,7 +57,12 @@ def event_apply(event: Event, state: State) -> State:
     """Apply an event to state."""
     match event:
         case (
-            TestEvent() | ChunkGenerated() | TaskAcknowledged() | InputChunkReceived()
+            TestEvent()
+            | ChunkGenerated()
+            | TaskAcknowledged()
+            | InputChunkReceived()
+            | TracesCollected()
+            | TracesMerged()
         ):  # Pass-through events that don't modify state
             return state
         case InstanceCreated():
@@ -225,6 +234,21 @@ def apply_node_timed_out(event: NodeTimedOut, state: State) -> State:
         for key, value in state.node_thunderbolt.items()
         if key != event.node_id
     }
+    node_thunderbolt_bridge = {
+        key: value
+        for key, value in state.node_thunderbolt_bridge.items()
+        if key != event.node_id
+    }
+    # Only recompute cycles if the leaving node had TB bridge enabled
+    leaving_node_status = state.node_thunderbolt_bridge.get(event.node_id)
+    leaving_node_had_tb_enabled = (
+        leaving_node_status is not None and leaving_node_status.enabled
+    )
+    thunderbolt_bridge_cycles = (
+        topology.get_thunderbolt_bridge_cycles(node_thunderbolt_bridge, node_network)
+        if leaving_node_had_tb_enabled
+        else [list(cycle) for cycle in state.thunderbolt_bridge_cycles]
+    )
     return state.model_copy(
         update={
             "downloads": downloads,
@@ -235,6 +259,8 @@ def apply_node_timed_out(event: NodeTimedOut, state: State) -> State:
             "node_system": node_system,
             "node_network": node_network,
             "node_thunderbolt": node_thunderbolt,
+            "node_thunderbolt_bridge": node_thunderbolt_bridge,
+            "thunderbolt_bridge_cycles": thunderbolt_bridge_cycles,
         }
     )
 
@@ -312,6 +338,22 @@ def apply_node_gathered_info(event: NodeGatheredInfo, state: State) -> State:
                 if tb_conn.sink_uuid in conn_map
             ]
             topology.replace_all_out_rdma_connections(event.node_id, as_rdma_conns)
+        case ThunderboltBridgeInfo():
+            new_tb_bridge: dict[NodeId, ThunderboltBridgeStatus] = {
+                **state.node_thunderbolt_bridge,
+                event.node_id: info.status,
+            }
+            update["node_thunderbolt_bridge"] = new_tb_bridge
+            # Only recompute cycles if the enabled status changed
+            old_status = state.node_thunderbolt_bridge.get(event.node_id)
+            old_enabled = old_status.enabled if old_status else False
+            new_enabled = info.status.enabled
+            if old_enabled != new_enabled:
+                update["thunderbolt_bridge_cycles"] = (
+                    topology.get_thunderbolt_bridge_cycles(
+                        new_tb_bridge, state.node_network
+                    )
+                )
 
     return state.model_copy(update=update)
 

src/exo/shared/constants.py

@@ -49,3 +49,10 @@ LIBP2P_COMMANDS_TOPIC = "commands"
 EXO_MAX_CHUNK_SIZE = 512 * 1024
 
 EXO_IMAGE_CACHE_DIR = EXO_CACHE_HOME / "images"
+EXO_TRACING_CACHE_DIR = EXO_CACHE_HOME / "traces"
+
+EXO_ENABLE_IMAGE_MODELS = (
+    os.getenv("EXO_ENABLE_IMAGE_MODELS", "false").lower() == "true"
+)
+
+EXO_TRACING_ENABLED = os.getenv("EXO_TRACING_ENABLED", "false").lower() == "true"

src/exo/shared/models/model_cards.py

@@ -1,5 +1,5 @@
 from enum import Enum
-from typing import Annotated
+from typing import Annotated, Any
 
 import aiofiles
 import aiofiles.os as aios
@@ -7,8 +7,16 @@ import tomlkit
 from anyio import Path, open_file
 from huggingface_hub import model_info
 from loguru import logger
-from pydantic import BaseModel, Field, PositiveInt, field_validator
+from pydantic import (
+    AliasChoices,
+    BaseModel,
+    Field,
+    PositiveInt,
+    field_validator,
+    model_validator,
+)
 
+from exo.shared.constants import EXO_ENABLE_IMAGE_MODELS
 from exo.shared.types.common import ModelId
 from exo.shared.types.memory import Memory
 from exo.utils.pydantic_ext import CamelCaseModel
@@ -120,6 +128,14 @@ MODEL_CARDS: dict[str, ModelCard] = {
         supports_tensor=True,
         tasks=[ModelTask.TextGeneration],
     ),
+    "kimi-k2.5": ModelCard(
+        model_id=ModelId("mlx-community/Kimi-K2.5"),
+        storage_size=Memory.from_gb(617),
+        n_layers=61,
+        hidden_size=7168,
+        supports_tensor=True,
+        tasks=[ModelTask.TextGeneration],
+    ),
     # llama-3.1
     "llama-3.1-8b": ModelCard(
         model_id=ModelId("mlx-community/Meta-Llama-3.1-8B-Instruct-4bit"),
@@ -254,7 +270,7 @@ MODEL_CARDS: dict[str, ModelCard] = {
     ),
     "qwen3-80b-a3B-thinking-4bit": ModelCard(
         model_id=ModelId("mlx-community/Qwen3-Next-80B-A3B-Thinking-4bit"),
-        storage_size=Memory.from_mb(84700),
+        storage_size=Memory.from_mb(44900),
         n_layers=48,
         hidden_size=2048,
         supports_tensor=True,
@@ -410,174 +426,310 @@ MODEL_CARDS: dict[str, ModelCard] = {
         supports_tensor=True,
         tasks=[ModelTask.TextGeneration],
     ),
-    # Image models commented out - feature not stable (see https://github.com/exo-explore/exo/issues/1242)
-    # "flux1-schnell": ModelCard(
-    #     model_id=ModelId("black-forest-labs/FLUX.1-schnell"),
-    #     storage_size=Memory.from_bytes(23782357120 + 9524621312),
-    #     n_layers=57,
-    #     hidden_size=1,
-    #     supports_tensor=False,
-    #     tasks=[ModelTask.TextToImage],
-    #     components=[
-    #         ComponentInfo(
-    #             component_name="text_encoder",
-    #             component_path="text_encoder/",
-    #             storage_size=Memory.from_kb(0),
-    #             n_layers=12,
-    #             can_shard=False,
-    #             safetensors_index_filename=None,  # Single file
-    #         ),
-    #         ComponentInfo(
-    #             component_name="text_encoder_2",
-    #             component_path="text_encoder_2/",
-    #             storage_size=Memory.from_bytes(9524621312),
-    #             n_layers=24,
-    #             can_shard=False,
-    #             safetensors_index_filename="model.safetensors.index.json",
-    #         ),
-    #         ComponentInfo(
-    #             component_name="transformer",
-    #             component_path="transformer/",
-    #             storage_size=Memory.from_bytes(23782357120),
-    #             n_layers=57,  # 19 transformer_blocks + 38 single_transformer_blocks
-    #             can_shard=True,
-    #             safetensors_index_filename="diffusion_pytorch_model.safetensors.index.json",
-    #         ),
-    #         ComponentInfo(
-    #             component_name="vae",
-    #             component_path="vae/",
-    #             storage_size=Memory.from_kb(0),
-    #             n_layers=None,
-    #             can_shard=False,
-    #             safetensors_index_filename=None,
-    #         ),
-    #     ],
-    # ),
-    # "flux1-dev": ModelCard(
-    #     model_id=ModelId("black-forest-labs/FLUX.1-dev"),
-    #     storage_size=Memory.from_bytes(23782357120 + 9524621312),
-    #     n_layers=57,
-    #     hidden_size=1,
-    #     supports_tensor=False,
-    #     tasks=[ModelTask.TextToImage, ModelTask.ImageToImage],
-    #     components=[
-    #         ComponentInfo(
-    #             component_name="text_encoder",
-    #             component_path="text_encoder/",
-    #             storage_size=Memory.from_kb(0),
-    #             n_layers=12,
-    #             can_shard=False,
-    #             safetensors_index_filename=None,  # Single file
-    #         ),
-    #         ComponentInfo(
-    #             component_name="text_encoder_2",
-    #             component_path="text_encoder_2/",
-    #             storage_size=Memory.from_bytes(9524621312),
-    #             n_layers=24,
-    #             can_shard=False,
-    #             safetensors_index_filename="model.safetensors.index.json",
-    #         ),
-    #         ComponentInfo(
-    #             component_name="transformer",
-    #             component_path="transformer/",
-    #             storage_size=Memory.from_bytes(23802816640),
-    #             n_layers=57,  # 19 transformer_blocks + 38 single_transformer_blocks
-    #             can_shard=True,
-    #             safetensors_index_filename="diffusion_pytorch_model.safetensors.index.json",
-    #         ),
-    #         ComponentInfo(
-    #             component_name="vae",
-    #             component_path="vae/",
-    #             storage_size=Memory.from_kb(0),
-    #             n_layers=None,
-    #             can_shard=False,
-    #             safetensors_index_filename=None,
-    #         ),
-    #     ],
-    # ),
-    # "qwen-image": ModelCard(
-    #     model_id=ModelId("Qwen/Qwen-Image"),
-    #     storage_size=Memory.from_bytes(16584333312 + 40860802176),
-    #     n_layers=60,  # Qwen has 60 transformer blocks (all joint-style)
-    #     hidden_size=1,
-    #     supports_tensor=False,
-    #     tasks=[ModelTask.TextToImage, ModelTask.ImageToImage],
-    #     components=[
-    #         ComponentInfo(
-    #             component_name="text_encoder",
-    #             component_path="text_encoder/",
-    #             storage_size=Memory.from_kb(16584333312),
-    #             n_layers=12,
-    #             can_shard=False,
-    #             safetensors_index_filename=None,  # Single file
-    #         ),
-    #         ComponentInfo(
-    #             component_name="transformer",
-    #             component_path="transformer/",
-    #             storage_size=Memory.from_bytes(40860802176),
-    #             n_layers=60,
-    #             can_shard=True,
-    #             safetensors_index_filename="diffusion_pytorch_model.safetensors.index.json",
-    #         ),
-    #         ComponentInfo(
-    #             component_name="vae",
-    #             component_path="vae/",
-    #             storage_size=Memory.from_kb(0),
-    #             n_layers=None,
-    #             can_shard=False,
-    #             safetensors_index_filename=None,
-    #         ),
-    #     ],
-    # ),
-    # "qwen-image-edit-2509": ModelCard(
-    #     model_id=ModelId("Qwen/Qwen-Image-Edit-2509"),
-    #     storage_size=Memory.from_bytes(16584333312 + 40860802176),
-    #     n_layers=60,  # Qwen has 60 transformer blocks (all joint-style)
-    #     hidden_size=1,
-    #     supports_tensor=False,
-    #     tasks=[ModelTask.ImageToImage],
-    #     components=[
-    #         ComponentInfo(
-    #             component_name="text_encoder",
-    #             component_path="text_encoder/",
-    #             storage_size=Memory.from_kb(16584333312),
-    #             n_layers=12,
-    #             can_shard=False,
-    #             safetensors_index_filename=None,  # Single file
-    #         ),
-    #         ComponentInfo(
-    #             component_name="transformer",
-    #             component_path="transformer/",
-    #             storage_size=Memory.from_bytes(40860802176),
-    #             n_layers=60,
-    #             can_shard=True,
-    #             safetensors_index_filename="diffusion_pytorch_model.safetensors.index.json",
-    #         ),
-    #         ComponentInfo(
-    #             component_name="vae",
-    #             component_path="vae/",
-    #             storage_size=Memory.from_kb(0),
-    #             n_layers=None,
-    #             can_shard=False,
-    #             safetensors_index_filename=None,
-    #         ),
-    #     ],
-    # ),
 }
 
+_IMAGE_BASE_MODEL_CARDS: dict[str, ModelCard] = {
+    "flux1-schnell": ModelCard(
+        model_id=ModelId("exolabs/FLUX.1-schnell"),
+        storage_size=Memory.from_bytes(23782357120 + 9524621312),
+        n_layers=57,
+        hidden_size=1,
+        supports_tensor=False,
+        tasks=[ModelTask.TextToImage],
+        components=[
+            ComponentInfo(
+                component_name="text_encoder",
+                component_path="text_encoder/",
+                storage_size=Memory.from_kb(0),
+                n_layers=12,
+                can_shard=False,
+                safetensors_index_filename=None,
+            ),
+            ComponentInfo(
+                component_name="text_encoder_2",
+                component_path="text_encoder_2/",
+                storage_size=Memory.from_bytes(9524621312),
+                n_layers=24,
+                can_shard=False,
+                safetensors_index_filename="model.safetensors.index.json",
+            ),
+            ComponentInfo(
+                component_name="transformer",
+                component_path="transformer/",
+                storage_size=Memory.from_bytes(23782357120),
+                n_layers=57,
+                can_shard=True,
+                safetensors_index_filename="diffusion_pytorch_model.safetensors.index.json",
+            ),
+            ComponentInfo(
+                component_name="vae",
+                component_path="vae/",
+                storage_size=Memory.from_kb(0),
+                n_layers=None,
+                can_shard=False,
+                safetensors_index_filename=None,
+            ),
+        ],
+    ),
+    "flux1-dev": ModelCard(
+        model_id=ModelId("exolabs/FLUX.1-dev"),
+        storage_size=Memory.from_bytes(23782357120 + 9524621312),
+        n_layers=57,
+        hidden_size=1,
+        supports_tensor=False,
+        tasks=[ModelTask.TextToImage],
+        components=[
+            ComponentInfo(
+                component_name="text_encoder",
+                component_path="text_encoder/",
+                storage_size=Memory.from_kb(0),
+                n_layers=12,
+                can_shard=False,
+                safetensors_index_filename=None,
+            ),
+            ComponentInfo(
+                component_name="text_encoder_2",
+                component_path="text_encoder_2/",
+                storage_size=Memory.from_bytes(9524621312),
+                n_layers=24,
+                can_shard=False,
+                safetensors_index_filename="model.safetensors.index.json",
+            ),
+            ComponentInfo(
+                component_name="transformer",
+                component_path="transformer/",
+                storage_size=Memory.from_bytes(23802816640),
+                n_layers=57,
+                can_shard=True,
+                safetensors_index_filename="diffusion_pytorch_model.safetensors.index.json",
+            ),
+            ComponentInfo(
+                component_name="vae",
+                component_path="vae/",
+                storage_size=Memory.from_kb(0),
+                n_layers=None,
+                can_shard=False,
+                safetensors_index_filename=None,
+            ),
+        ],
+    ),
+    "flux1-krea-dev": ModelCard(
+        model_id=ModelId("exolabs/FLUX.1-Krea-dev"),
+        storage_size=Memory.from_bytes(23802816640 + 9524621312),  # Same as dev
+        n_layers=57,
+        hidden_size=1,
+        supports_tensor=False,
+        tasks=[ModelTask.TextToImage],
+        components=[
+            ComponentInfo(
+                component_name="text_encoder",
+                component_path="text_encoder/",
+                storage_size=Memory.from_kb(0),
+                n_layers=12,
+                can_shard=False,
+                safetensors_index_filename=None,
+            ),
+            ComponentInfo(
+                component_name="text_encoder_2",
+                component_path="text_encoder_2/",
+                storage_size=Memory.from_bytes(9524621312),
+                n_layers=24,
+                can_shard=False,
+                safetensors_index_filename="model.safetensors.index.json",
+            ),
+            ComponentInfo(
+                component_name="transformer",
+                component_path="transformer/",
+                storage_size=Memory.from_bytes(23802816640),
+                n_layers=57,
+                can_shard=True,
+                safetensors_index_filename="diffusion_pytorch_model.safetensors.index.json",
+            ),
+            ComponentInfo(
+                component_name="vae",
+                component_path="vae/",
+                storage_size=Memory.from_kb(0),
+                n_layers=None,
+                can_shard=False,
+                safetensors_index_filename=None,
+            ),
+        ],
+    ),
+    "qwen-image": ModelCard(
+        model_id=ModelId("exolabs/Qwen-Image"),
+        storage_size=Memory.from_bytes(16584333312 + 40860802176),
+        n_layers=60,
+        hidden_size=1,
+        supports_tensor=False,
+        tasks=[ModelTask.TextToImage],
+        components=[
+            ComponentInfo(
+                component_name="text_encoder",
+                component_path="text_encoder/",
+                storage_size=Memory.from_bytes(16584333312),
+                n_layers=12,
+                can_shard=False,
+                safetensors_index_filename=None,
+            ),
+            ComponentInfo(
+                component_name="transformer",
+                component_path="transformer/",
+                storage_size=Memory.from_bytes(40860802176),
+                n_layers=60,
+                can_shard=True,
+                safetensors_index_filename="diffusion_pytorch_model.safetensors.index.json",
+            ),
+            ComponentInfo(
+                component_name="vae",
+                component_path="vae/",
+                storage_size=Memory.from_kb(0),
+                n_layers=None,
+                can_shard=False,
+                safetensors_index_filename=None,
+            ),
+        ],
+    ),
+    "qwen-image-edit-2509": ModelCard(
+        model_id=ModelId("exolabs/Qwen-Image-Edit-2509"),
+        storage_size=Memory.from_bytes(16584333312 + 40860802176),
+        n_layers=60,
+        hidden_size=1,
+        supports_tensor=False,
+        tasks=[ModelTask.ImageToImage],
+        components=[
+            ComponentInfo(
+                component_name="text_encoder",
+                component_path="text_encoder/",
+                storage_size=Memory.from_bytes(16584333312),
+                n_layers=12,
+                can_shard=False,
+                safetensors_index_filename=None,
+            ),
+            ComponentInfo(
+                component_name="transformer",
+                component_path="transformer/",
+                storage_size=Memory.from_bytes(40860802176),
+                n_layers=60,
+                can_shard=True,
+                safetensors_index_filename="diffusion_pytorch_model.safetensors.index.json",
+            ),
+            ComponentInfo(
+                component_name="vae",
+                component_path="vae/",
+                storage_size=Memory.from_kb(0),
+                n_layers=None,
+                can_shard=False,
+                safetensors_index_filename=None,
+            ),
+        ],
+    ),
+}
+
+
+def _generate_image_model_quant_variants(
+    base_name: str,
+    base_card: ModelCard,
+) -> dict[str, ModelCard]:
+    """Create quantized variants of an image model card.
+
+    Only the transformer component is quantized; text encoders stay at bf16.
+    Sizes are calculated exactly from the base card's component sizes.
+    """
+    if base_card.components is None:
+        raise ValueError(f"Image model {base_name} must have components defined")
+
+    # quantizations = [8, 6, 5, 4, 3]
+    quantizations = [8, 4]
+
+    num_transformer_bytes = next(
+        c.storage_size.in_bytes
+        for c in base_card.components
+        if c.component_name == "transformer"
+    )
+
+    transformer_bytes = Memory.from_bytes(num_transformer_bytes)
+
+    remaining_bytes = Memory.from_bytes(
+        sum(
+            c.storage_size.in_bytes
+            for c in base_card.components
+            if c.component_name != "transformer"
+        )
+    )
+
+    def with_transformer_size(new_size: Memory) -> list[ComponentInfo]:
+        assert base_card.components is not None
+        return [
+            ComponentInfo(
+                component_name=c.component_name,
+                component_path=c.component_path,
+                storage_size=new_size
+                if c.component_name == "transformer"
+                else c.storage_size,
+                n_layers=c.n_layers,
+                can_shard=c.can_shard,
+                safetensors_index_filename=c.safetensors_index_filename,
+            )
+            for c in base_card.components
+        ]
+
+    variants = {
+        base_name: ModelCard(
+            model_id=base_card.model_id,
+            storage_size=transformer_bytes + remaining_bytes,
+            n_layers=base_card.n_layers,
+            hidden_size=base_card.hidden_size,
+            supports_tensor=base_card.supports_tensor,
+            tasks=base_card.tasks,
+            components=with_transformer_size(transformer_bytes),
+        )
+    }
+
+    for quant in quantizations:
+        quant_transformer_bytes = Memory.from_bytes(
+            (num_transformer_bytes * quant) // 16
+        )
+        total_bytes = remaining_bytes + quant_transformer_bytes
+
+        model_id = ModelId(base_card.model_id + f"-{quant}bit")
+
+        variants[f"{base_name}-{quant}bit"] = ModelCard(
+            model_id=model_id,
+            storage_size=total_bytes,
+            n_layers=base_card.n_layers,
+            hidden_size=base_card.hidden_size,
+            supports_tensor=base_card.supports_tensor,
+            tasks=base_card.tasks,
+            components=with_transformer_size(quant_transformer_bytes),
+        )
+
+    return variants
+
+
+_image_model_cards: dict[str, ModelCard] = {}
+for _base_name, _base_card in _IMAGE_BASE_MODEL_CARDS.items():
+    _image_model_cards |= _generate_image_model_quant_variants(_base_name, _base_card)
+_IMAGE_MODEL_CARDS = _image_model_cards
+
+if EXO_ENABLE_IMAGE_MODELS:
+    MODEL_CARDS.update(_IMAGE_MODEL_CARDS)
+
 
 class ConfigData(BaseModel):
     model_config = {"extra": "ignore"}  # Allow unknown fields
 
-    # Common field names for number of layers across different architectures
-    num_hidden_layers: Annotated[int, Field(ge=0)] | None = None
-    num_layers: Annotated[int, Field(ge=0)] | None = None
-    n_layer: Annotated[int, Field(ge=0)] | None = None
-    n_layers: Annotated[int, Field(ge=0)] | None = None  # Sometimes used
-    num_decoder_layers: Annotated[int, Field(ge=0)] | None = None  # Transformer models
-    decoder_layers: Annotated[int, Field(ge=0)] | None = None  # Some architectures
-    hidden_size: Annotated[int, Field(ge=0)] | None = None
     architectures: list[str] | None = None
+    hidden_size: Annotated[int, Field(ge=0)] | None = None
+    layer_count: int = Field(
+        validation_alias=AliasChoices(
+            "num_hidden_layers",
+            "num_layers",
+            "n_layer",
+            "n_layers",
+            "num_decoder_layers",
+            "decoder_layers",
+        )
+    )
 
     @property
     def supports_tensor(self) -> bool:
@@ -592,30 +744,32 @@ class ConfigData(BaseModel):
             ["GptOssForCausalLM"],
         ]
 
-    @property
-    def layer_count(self) -> int:
-        # Check common field names for layer count
-        layer_fields = [
-            self.num_hidden_layers,
-            self.num_layers,
-            self.n_layer,
-            self.n_layers,
-            self.num_decoder_layers,
-            self.decoder_layers,
-        ]
+    @model_validator(mode="before")
+    @classmethod
+    def defer_to_text_config(cls, data: dict[str, Any]):
+        text_config = data.get("text_config")
+        if text_config is None:
+            return data
 
-        for layer_count in layer_fields:
-            if layer_count is not None:
-                return layer_count
+        for field in [
+            "architectures",
+            "hidden_size",
+            "num_hidden_layers",
+            "num_layers",
+            "n_layer",
+            "n_layers",
+            "num_decoder_layers",
+            "decoder_layers",
+        ]:
+            if (val := text_config.get(field)) is not None:  # pyright: ignore[reportAny]
+                data[field] = val
 
-        raise ValueError(
-            f"No layer count found in config.json: {self.model_dump_json()}"
-        )
+        return data
 
 
 async def get_config_data(model_id: ModelId) -> ConfigData:
     """Downloads and parses config.json for a model."""
-    from exo.worker.download.download_utils import (
+    from exo.download.download_utils import (
         download_file_with_retry,
         ensure_models_dir,
     )
@@ -627,7 +781,7 @@ async def get_config_data(model_id: ModelId) -> ConfigData:
         "main",
         "config.json",
         target_dir,
-        lambda curr_bytes, total_bytes, is_renamed: logger.info(
+        lambda curr_bytes, total_bytes, is_renamed: logger.debug(
             f"Downloading config.json for {model_id}: {curr_bytes}/{total_bytes} ({is_renamed=})"
         ),
     )
@@ -637,11 +791,11 @@ async def get_config_data(model_id: ModelId) -> ConfigData:
 
 async def get_safetensors_size(model_id: ModelId) -> Memory:
     """Gets model size from safetensors index or falls back to HF API."""
-    from exo.shared.types.worker.downloads import ModelSafetensorsIndex
-    from exo.worker.download.download_utils import (
+    from exo.download.download_utils import (
         download_file_with_retry,
         ensure_models_dir,
     )
+    from exo.shared.types.worker.downloads import ModelSafetensorsIndex
 
     target_dir = (await ensure_models_dir()) / model_id.normalize()
     await aios.makedirs(target_dir, exist_ok=True)
@@ -650,7 +804,7 @@ async def get_safetensors_size(model_id: ModelId) -> Memory:
         "main",
         "model.safetensors.index.json",
         target_dir,
-        lambda curr_bytes, total_bytes, is_renamed: logger.info(
+        lambda curr_bytes, total_bytes, is_renamed: logger.debug(
             f"Downloading model.safetensors.index.json for {model_id}: {curr_bytes}/{total_bytes} ({is_renamed=})"
         ),
     )

src/exo/shared/tests/test_node_id_persistence.py

@@ -8,7 +8,7 @@ from multiprocessing.synchronize import Event as EventT
 from multiprocessing.synchronize import Semaphore as SemaphoreT
 
 from loguru import logger
-from pytest import LogCaptureFixture
+from pytest import LogCaptureFixture, mark
 
 from exo.routing.router import get_node_id_keypair
 from exo.shared.constants import EXO_NODE_ID_KEYPAIR
@@ -74,6 +74,7 @@ def _delete_if_exists(p: str | bytes | os.PathLike[str] | os.PathLike[bytes]):
         os.remove(p)
 
 
+@mark.skip(reason="this functionality is currently disabled but may return in future")
 def test_node_id_fetching(caplog: LogCaptureFixture):
     reps = 10
 

src/exo/shared/topology.py

@@ -7,6 +7,11 @@ import rustworkx as rx
 from pydantic import BaseModel, ConfigDict
 
 from exo.shared.types.common import NodeId
+from exo.shared.types.profiling import (
+    InterfaceType,
+    NodeNetworkInfo,
+    ThunderboltBridgeStatus,
+)
 from exo.shared.types.topology import (
     Connection,
     Cycle,
@@ -188,24 +193,25 @@ class Topology:
             cycles.append(Cycle(node_ids=[node_id]))
         return cycles
 
-    def get_cycles_tb(self) -> list[Cycle]:
-        tb_edges = [
+    def get_rdma_cycles(self) -> list[Cycle]:
+        rdma_edges = [
             (u, v, conn)
             for u, v, conn in self._graph.weighted_edge_list()
-            if conn.is_thunderbolt()
+            if isinstance(conn, RDMAConnection)
         ]
 
-        tb_graph: rx.PyDiGraph[NodeId, SocketConnection] = rx.PyDiGraph()
-        tb_graph.add_nodes_from(self._graph.nodes())
+        rdma_graph: rx.PyDiGraph[NodeId, SocketConnection | RDMAConnection] = (
+            rx.PyDiGraph()
+        )
+        rdma_graph.add_nodes_from(self._graph.nodes())
 
-        for u, v, conn in tb_edges:
-            if isinstance(conn, SocketConnection):
-                tb_graph.add_edge(u, v, conn)
+        for u, v, conn in rdma_edges:
+            rdma_graph.add_edge(u, v, conn)
 
-        cycle_idxs = rx.simple_cycles(tb_graph)
+        cycle_idxs = rx.simple_cycles(rdma_graph)
         cycles: list[Cycle] = []
         for cycle_idx in cycle_idxs:
-            cycle = Cycle(node_ids=[tb_graph[idx] for idx in cycle_idx])
+            cycle = Cycle(node_ids=[rdma_graph[idx] for idx in cycle_idx])
             cycles.append(cycle)
 
         return cycles
@@ -219,18 +225,83 @@ class Topology:
                 topology.add_connection(connection)
         return topology
 
-    def is_thunderbolt_cycle(self, cycle: Cycle) -> bool:
+    def is_rdma_cycle(self, cycle: Cycle) -> bool:
         node_idxs = [node for node in cycle]
         rx_idxs = [self._vertex_indices[idx] for idx in node_idxs]
         for rid in rx_idxs:
             for neighbor_rid in self._graph.neighbors(rid):
                 if neighbor_rid not in rx_idxs:
                     continue
-                has_tb = False
+                has_rdma = False
                 for edge in self._graph.get_all_edge_data(rid, neighbor_rid):
-                    if edge.is_thunderbolt():
-                        has_tb = True
+                    if isinstance(edge, RDMAConnection):
+                        has_rdma = True
                         break
-                if not has_tb:
+                if not has_rdma:
                     return False
         return True
+
+    def get_thunderbolt_bridge_cycles(
+        self,
+        node_tb_bridge_status: Mapping[NodeId, ThunderboltBridgeStatus],
+        node_network: Mapping[NodeId, NodeNetworkInfo],
+    ) -> list[list[NodeId]]:
+        """
+        Find cycles in the Thunderbolt topology where all nodes have TB bridge enabled.
+        Only returns cycles with >=2 nodes (2+ machines in a loop), as
+        1 node doesn't cause the broadcast storm problem.
+        """
+        enabled_nodes = {
+            node_id
+            for node_id, status in node_tb_bridge_status.items()
+            if status.enabled
+        }
+
+        if len(enabled_nodes) < 2:
+            return []
+
+        thunderbolt_ips = _get_ips_with_interface_type(
+            enabled_nodes, node_network, "thunderbolt"
+        )
+
+        # Build subgraph with only TB bridge enabled nodes and thunderbolt connections
+        graph: rx.PyDiGraph[NodeId, SocketConnection | RDMAConnection] = rx.PyDiGraph()
+        node_to_idx: dict[NodeId, int] = {}
+
+        for node_id in enabled_nodes:
+            if node_id in self._vertex_indices:
+                node_to_idx[node_id] = graph.add_node(node_id)
+
+        for u, v, conn in self._graph.weighted_edge_list():
+            source_id, sink_id = self._graph[u], self._graph[v]
+            if source_id not in node_to_idx or sink_id not in node_to_idx:
+                continue
+            # Include connection if it's over a thunderbolt interface
+            if (
+                isinstance(conn, SocketConnection)
+                and conn.sink_multiaddr.ip_address in thunderbolt_ips
+            ):
+                graph.add_edge(node_to_idx[source_id], node_to_idx[sink_id], conn)
+            if isinstance(conn, RDMAConnection):
+                graph.add_edge(node_to_idx[source_id], node_to_idx[sink_id], conn)
+
+        return [
+            [graph[idx] for idx in cycle]
+            for cycle in rx.simple_cycles(graph)
+            if len(cycle) >= 2
+        ]
+
+
+def _get_ips_with_interface_type(
+    node_ids: set[NodeId],
+    node_network: Mapping[NodeId, NodeNetworkInfo],
+    interface_type: InterfaceType,
+) -> set[str]:
+    """Get all IP addresses on interfaces of the specified type for the given nodes."""
+    ips: set[str] = set()
+    for node_id in node_ids:
+        network_info = node_network.get(node_id, NodeNetworkInfo())
+        for iface in network_info.interfaces:
+            if iface.interface_type == interface_type:
+                ips.add(iface.ip_address)
+    return ips

src/exo/shared/tracing.py

@@ -0,0 +1,238 @@
+from __future__ import annotations
+
+import json
+import time
+from collections import defaultdict
+from collections.abc import Generator
+from contextlib import contextmanager
+from contextvars import ContextVar
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import cast, final
+
+from exo.shared.constants import EXO_TRACING_ENABLED
+from exo.worker.runner.bootstrap import logger
+
+# Context variable to track the current trace category for hierarchical nesting
+_current_category: ContextVar[str | None] = ContextVar("current_category", default=None)
+
+
+@final
+@dataclass(frozen=True)
+class TraceEvent:
+    name: str
+    start_us: int
+    duration_us: int
+    rank: int
+    category: str
+
+
+@final
+@dataclass
+class CategoryStats:
+    total_us: int = 0
+    count: int = 0
+    min_us: int = 0
+    max_us: int = 0
+
+    def add(self, duration_us: int) -> None:
+        if self.count == 0:
+            self.min_us = duration_us
+            self.max_us = duration_us
+        else:
+            self.min_us = min(self.min_us, duration_us)
+            self.max_us = max(self.max_us, duration_us)
+        self.total_us += duration_us
+        self.count += 1
+
+    @property
+    def avg_us(self) -> float:
+        return self.total_us / self.count if self.count > 0 else 0.0
+
+
+@final
+@dataclass
+class TraceStats:
+    total_wall_time_us: int = 0
+    by_category: dict[str, CategoryStats] = field(default_factory=dict)
+    by_rank: dict[int, dict[str, CategoryStats]] = field(default_factory=dict)
+
+
+# Global trace buffer - each rank accumulates traces here
+_trace_buffer: list[TraceEvent] = []
+
+
+def _record_span(
+    name: str, start_us: int, duration_us: int, rank: int, category: str
+) -> None:
+    _trace_buffer.append(
+        TraceEvent(
+            name=name,
+            start_us=start_us,
+            duration_us=duration_us,
+            rank=rank,
+            category=category,
+        )
+    )
+
+
+@contextmanager
+def trace(
+    name: str,
+    rank: int,
+    category: str = "compute",
+) -> Generator[None, None, None]:
+    """Context manager to trace any operation.
+
+    Nested traces automatically inherit the parent category, creating hierarchical
+    categories like "sync/compute" or "async/comms".
+
+    Args:
+        name: Name of the operation (e.g., "recv 0", "send 1", "joint_blocks")
+        rank: This rank's ID
+        category: Category for grouping in trace viewer ("comm", "compute", "step")
+
+    Example:
+        with trace(f"sync {t}", rank, "sync"):
+            with trace("joint_blocks", rank, "compute"):
+                # Recorded with category "sync/compute"
+                hidden_states = some_computation(...)
+    """
+    if not EXO_TRACING_ENABLED:
+        yield
+        return
+
+    # Combine with parent category if nested
+    parent = _current_category.get()
+    full_category = f"{parent}/{category}" if parent else category
+
+    # Set as current for nested traces
+    token = _current_category.set(full_category)
+
+    try:
+        start_us = int(time.time() * 1_000_000)
+        start_perf = time.perf_counter()
+        yield
+        duration_us = int((time.perf_counter() - start_perf) * 1_000_000)
+        _record_span(name, start_us, duration_us, rank, full_category)
+    finally:
+        _current_category.reset(token)
+
+
+def get_trace_buffer() -> list[TraceEvent]:
+    return list(_trace_buffer)
+
+
+def clear_trace_buffer() -> None:
+    _trace_buffer.clear()
+
+
+def export_trace(traces: list[TraceEvent], output_path: Path) -> None:
+    trace_events: list[dict[str, object]] = []
+
+    for event in traces:
+        # Chrome trace format uses "X" for complete events (with duration)
+        chrome_event: dict[str, object] = {
+            "name": event.name,
+            "cat": event.category,
+            "ph": "X",
+            "ts": event.start_us,
+            "dur": event.duration_us,
+            "pid": 0,
+            "tid": event.rank,
+            "args": {"rank": event.rank},
+        }
+        trace_events.append(chrome_event)
+
+    ranks_seen = set(t.rank for t in traces)
+    for rank in ranks_seen:
+        trace_events.append(
+            {
+                "name": "thread_name",
+                "ph": "M",  # Metadata event
+                "pid": 0,
+                "tid": rank,
+                "args": {"name": f"Rank {rank}"},
+            }
+        )
+
+    chrome_trace = {"traceEvents": trace_events}
+
+    try:
+        output_path.parent.mkdir(parents=True, exist_ok=True)
+        with open(output_path, "w") as f:
+            json.dump(chrome_trace, f, indent=2)
+    except OSError as e:
+        logger.warning("Failed to export trace to %s: %s", output_path, e)
+
+
+def load_trace_file(path: Path) -> list[TraceEvent]:
+    with open(path) as f:
+        data = cast(dict[str, list[dict[str, object]]], json.load(f))
+
+    events = data.get("traceEvents", [])
+    traces: list[TraceEvent] = []
+
+    for event in events:
+        # Skip metadata events
+        if event.get("ph") == "M":
+            continue
+
+        name = str(event.get("name", ""))
+        category = str(event.get("cat", ""))
+        ts_value = event.get("ts", 0)
+        dur_value = event.get("dur", 0)
+        tid_value = event.get("tid", 0)
+        start_us = int(ts_value) if isinstance(ts_value, (int, float, str)) else 0
+        duration_us = int(dur_value) if isinstance(dur_value, (int, float, str)) else 0
+
+        # Get rank from tid or args
+        rank = int(tid_value) if isinstance(tid_value, (int, float, str)) else 0
+        args = event.get("args")
+        if isinstance(args, dict):
+            args_dict = cast(dict[str, object], args)
+            rank_from_args = args_dict.get("rank")
+            if isinstance(rank_from_args, (int, float, str)):
+                rank = int(rank_from_args)
+
+        traces.append(
+            TraceEvent(
+                name=name,
+                start_us=start_us,
+                duration_us=duration_us,
+                rank=rank,
+                category=category,
+            )
+        )
+
+    return traces
+
+
+def compute_stats(traces: list[TraceEvent]) -> TraceStats:
+    stats = TraceStats()
+
+    if not traces:
+        return stats
+
+    # Calculate wall time from earliest start to latest end
+    min_start = min(t.start_us for t in traces)
+    max_end = max(t.start_us + t.duration_us for t in traces)
+    stats.total_wall_time_us = max_end - min_start
+
+    # Initialize nested dicts
+    by_category: dict[str, CategoryStats] = defaultdict(CategoryStats)
+    by_rank: dict[int, dict[str, CategoryStats]] = defaultdict(
+        lambda: defaultdict(CategoryStats)
+    )
+
+    for event in traces:
+        # By category
+        by_category[event.category].add(event.duration_us)
+
+        # By rank and category
+        by_rank[event.rank][event.category].add(event.duration_us)
+
+    stats.by_category = dict(by_category)
+    stats.by_rank = {k: dict(v) for k, v in by_rank.items()}
+
+    return stats

src/exo/shared/types/api.py

@@ -2,15 +2,15 @@ import time
 from collections.abc import Generator
 from typing import Annotated, Any, Literal
 
-from fastapi import UploadFile
 from pydantic import BaseModel, Field, field_validator
 from pydantic_core import PydanticUseDefault
 
 from exo.shared.models.model_cards import ModelCard, ModelId
-from exo.shared.types.common import CommandId
+from exo.shared.types.common import CommandId, NodeId
 from exo.shared.types.memory import Memory
 from exo.shared.types.worker.instances import Instance, InstanceId, InstanceMeta
-from exo.shared.types.worker.shards import Sharding
+from exo.shared.types.worker.shards import Sharding, ShardMetadata
+from exo.utils.pydantic_ext import CamelCaseModel
 
 FinishReason = Literal[
     "stop", "length", "tool_calls", "content_filter", "function_call", "error"
@@ -54,6 +54,18 @@ class ChatCompletionMessageText(BaseModel):
     text: str
 
 
+class ToolCallItem(BaseModel):
+    name: str
+    arguments: str
+
+
+class ToolCall(BaseModel):
+    id: str
+    index: int | None = None
+    type: Literal["function"] = "function"
+    function: ToolCallItem
+
+
 class ChatCompletionMessage(BaseModel):
     role: Literal["system", "user", "assistant", "developer", "tool", "function"]
     content: (
@@ -61,7 +73,7 @@ class ChatCompletionMessage(BaseModel):
     ) = None
     thinking: str | None = None  # Added for GPT-OSS harmony format support
     name: str | None = None
-    tool_calls: list[dict[str, Any]] | None = None
+    tool_calls: list[ToolCall] | None = None
     tool_call_id: str | None = None
     function_call: dict[str, Any] | None = None
 
@@ -103,8 +115,8 @@ class Usage(BaseModel):
     prompt_tokens: int
     completion_tokens: int
     total_tokens: int
-    prompt_tokens_details: PromptTokensDetails | None = None
-    completion_tokens_details: CompletionTokensDetails | None = None
+    prompt_tokens_details: PromptTokensDetails
+    completion_tokens_details: CompletionTokensDetails
 
 
 class StreamingChoiceResponse(BaseModel):
@@ -157,8 +169,12 @@ class BenchChatCompletionResponse(ChatCompletionResponse):
     generation_stats: GenerationStats | None = None
 
 
-class ChatCompletionTaskParams(BaseModel):
-    model: str
+class StreamOptions(BaseModel):
+    include_usage: bool = False
+
+
+class ChatCompletionRequest(BaseModel):
+    model: ModelId
     frequency_penalty: float | None = None
     messages: list[ChatCompletionMessage]
     logit_bias: dict[str, int] | None = None
@@ -171,15 +187,17 @@ class ChatCompletionTaskParams(BaseModel):
     seed: int | None = None
     stop: str | list[str] | None = None
     stream: bool = False
+    stream_options: StreamOptions | None = None
     temperature: float | None = None
     top_p: float | None = None
+    top_k: int | None = None
     tools: list[dict[str, Any]] | None = None
     tool_choice: str | dict[str, Any] | None = None
     parallel_tool_calls: bool | None = None
     user: str | None = None
 
 
-class BenchChatCompletionTaskParams(ChatCompletionTaskParams):
+class BenchChatCompletionRequest(ChatCompletionRequest):
     pass
 
 
@@ -263,28 +281,7 @@ class BenchImageGenerationTaskParams(ImageGenerationTaskParams):
 
 
 class ImageEditsTaskParams(BaseModel):
-    image: UploadFile
-    prompt: str
-    background: str | None = None
-    input_fidelity: float | None = None
-    mask: UploadFile | None = None
-    model: str
-    n: int | None = 1
-    output_compression: int | None = None
-    output_format: Literal["png", "jpeg", "webp"] = "png"
-    partial_images: int | None = 0
-    quality: Literal["high", "medium", "low"] | None = "medium"
-    response_format: Literal["url", "b64_json"] | None = "b64_json"
-    size: str | None = "1024x1024"
-    stream: bool | None = False
-    user: str | None = None
-    advanced_params: AdvancedImageParams | None = None
-    # Internal flag for benchmark mode - set by API, preserved through serialization
-    bench: bool = False
-
-
-class ImageEditsInternalParams(BaseModel):
-    """Serializable version of ImageEditsTaskParams for distributed task execution."""
+    """Internal task params for image-editing requests."""
 
     image_data: str = ""  # Base64-encoded image (empty when using chunked transfer)
     total_input_chunks: int = 0
@@ -340,3 +337,58 @@ class ImageListItem(BaseModel, frozen=True):
 
 class ImageListResponse(BaseModel, frozen=True):
     data: list[ImageListItem]
+
+
+class StartDownloadParams(CamelCaseModel):
+    target_node_id: NodeId
+    shard_metadata: ShardMetadata
+
+
+class StartDownloadResponse(CamelCaseModel):
+    command_id: CommandId
+
+
+class DeleteDownloadResponse(CamelCaseModel):
+    command_id: CommandId
+
+
+class TraceEventResponse(CamelCaseModel):
+    name: str
+    start_us: int
+    duration_us: int
+    rank: int
+    category: str
+
+
+class TraceResponse(CamelCaseModel):
+    task_id: str
+    traces: list[TraceEventResponse]
+
+
+class TraceCategoryStats(CamelCaseModel):
+    total_us: int
+    count: int
+    min_us: int
+    max_us: int
+    avg_us: float
+
+
+class TraceRankStats(CamelCaseModel):
+    by_category: dict[str, TraceCategoryStats]
+
+
+class TraceStatsResponse(CamelCaseModel):
+    task_id: str
+    total_wall_time_us: int
+    by_category: dict[str, TraceCategoryStats]
+    by_rank: dict[int, TraceRankStats]
+
+
+class TraceListItem(CamelCaseModel):
+    task_id: str
+    created_at: str
+    file_size: int
+
+
+class TraceListResponse(CamelCaseModel):
+    traces: list[TraceListItem]

src/exo/shared/types/chunks.py

@@ -1,31 +1,37 @@
 from collections.abc import Generator
-from enum import Enum
 from typing import Any, Literal
 
 from exo.shared.models.model_cards import ModelId
-from exo.shared.types.api import GenerationStats, ImageGenerationStats
+from exo.shared.types.api import GenerationStats, ImageGenerationStats, Usage
 from exo.utils.pydantic_ext import TaggedModel
 
 from .api import FinishReason
 from .common import CommandId
-
-
-class ChunkType(str, Enum):
-    Token = "Token"
-    Image = "Image"
+from .worker.runner_response import ToolCallItem
 
 
 class BaseChunk(TaggedModel):
-    idx: int
     model: ModelId
 
 
 class TokenChunk(BaseChunk):
     text: str
     token_id: int
-    finish_reason: FinishReason | None = None
+    usage: Usage | None
+    finish_reason: Literal["stop", "length", "content_filter"] | None = None
+    stats: GenerationStats | None = None
+
+
+class ErrorChunk(BaseChunk):
+    error_message: str
+    finish_reason: Literal["error"] = "error"
+
+
+class ToolCallChunk(BaseChunk):
+    tool_calls: list[ToolCallItem]
+    usage: Usage | None
+    finish_reason: Literal["tool_calls"] = "tool_calls"
     stats: GenerationStats | None = None
-    error_message: str | None = None
 
 
 class ImageChunk(BaseChunk):
@@ -63,4 +69,4 @@ class InputImageChunk(BaseChunk):
                 yield name, value
 
 
-GenerationChunk = TokenChunk | ImageChunk
+GenerationChunk = TokenChunk | ImageChunk | ToolCallChunk | ErrorChunk

src/exo/shared/types/claude_api.py

@@ -0,0 +1,214 @@
+"""Claude Messages API types for request/response conversion."""
+
+from typing import Any, Literal
+
+from pydantic import BaseModel, Field
+
+from exo.shared.types.common import ModelId
+
+# Tool definition types
+ClaudeToolInputSchema = dict[str, Any]
+
+
+class ClaudeToolDefinition(BaseModel, frozen=True):
+    """Tool definition in Claude Messages API request."""
+
+    name: str
+    description: str | None = None
+    input_schema: ClaudeToolInputSchema
+
+
+# Type aliases
+ClaudeRole = Literal["user", "assistant"]
+ClaudeStopReason = Literal["end_turn", "max_tokens", "stop_sequence", "tool_use"]
+
+
+# Content block types
+class ClaudeTextBlock(BaseModel, frozen=True):
+    """Text content block in Claude Messages API."""
+
+    type: Literal["text"] = "text"
+    text: str
+
+
+class ClaudeImageSource(BaseModel, frozen=True):
+    """Image source for Claude image blocks."""
+
+    type: Literal["base64", "url"]
+    media_type: str | None = None
+    data: str | None = None
+    url: str | None = None
+
+
+class ClaudeImageBlock(BaseModel, frozen=True):
+    """Image content block in Claude Messages API."""
+
+    type: Literal["image"] = "image"
+    source: ClaudeImageSource
+
+
+class ClaudeToolUseBlock(BaseModel, frozen=True):
+    """Tool use content block in Claude Messages API."""
+
+    type: Literal["tool_use"] = "tool_use"
+    id: str
+    name: str
+    input: dict[str, Any]
+
+
+class ClaudeToolResultBlock(BaseModel, frozen=True):
+    """Tool result content block in Claude Messages API request."""
+
+    type: Literal["tool_result"] = "tool_result"
+    tool_use_id: str
+    content: str | list[ClaudeTextBlock] | None = None
+    is_error: bool | None = None
+    cache_control: dict[str, str] | None = None
+
+
+ClaudeContentBlock = ClaudeTextBlock | ClaudeImageBlock | ClaudeToolUseBlock
+
+# Input content blocks can also include tool_result (sent by user after tool_use)
+ClaudeInputContentBlock = (
+    ClaudeTextBlock | ClaudeImageBlock | ClaudeToolUseBlock | ClaudeToolResultBlock
+)
+
+
+# Request types
+class ClaudeMessage(BaseModel, frozen=True):
+    """Message in Claude Messages API request."""
+
+    role: ClaudeRole
+    content: str | list[ClaudeInputContentBlock]
+
+
+class ClaudeMessagesRequest(BaseModel):
+    """Request body for Claude Messages API."""
+
+    model: ModelId
+    max_tokens: int
+    messages: list[ClaudeMessage]
+    system: str | list[ClaudeTextBlock] | None = None
+    stop_sequences: list[str] | None = None
+    stream: bool = False
+    temperature: float | None = None
+    top_p: float | None = None
+    top_k: int | None = None
+    tools: list[ClaudeToolDefinition] | None = None
+    metadata: dict[str, str] | None = None
+
+
+# Response types
+class ClaudeUsage(BaseModel, frozen=True):
+    """Token usage in Claude Messages API response."""
+
+    input_tokens: int
+    output_tokens: int
+
+
+class ClaudeMessagesResponse(BaseModel, frozen=True):
+    """Response body for Claude Messages API."""
+
+    id: str
+    type: Literal["message"] = "message"
+    role: Literal["assistant"] = "assistant"
+    content: list[ClaudeContentBlock]
+    model: str
+    stop_reason: ClaudeStopReason | None = None
+    stop_sequence: str | None = None
+    usage: ClaudeUsage
+
+
+# Streaming event types
+class ClaudeMessageStart(BaseModel, frozen=True):
+    """Partial message in message_start event."""
+
+    id: str
+    type: Literal["message"] = "message"
+    role: Literal["assistant"] = "assistant"
+    content: list[ClaudeTextBlock] = Field(default_factory=list)
+    model: str
+    stop_reason: ClaudeStopReason | None = None
+    stop_sequence: str | None = None
+    usage: ClaudeUsage
+
+
+class ClaudeMessageStartEvent(BaseModel, frozen=True):
+    """Event sent at start of message stream."""
+
+    type: Literal["message_start"] = "message_start"
+    message: ClaudeMessageStart
+
+
+class ClaudeContentBlockStartEvent(BaseModel, frozen=True):
+    """Event sent at start of a content block."""
+
+    type: Literal["content_block_start"] = "content_block_start"
+    index: int
+    content_block: ClaudeTextBlock | ClaudeToolUseBlock
+
+
+class ClaudeTextDelta(BaseModel, frozen=True):
+    """Delta for text content block."""
+
+    type: Literal["text_delta"] = "text_delta"
+    text: str
+
+
+class ClaudeInputJsonDelta(BaseModel, frozen=True):
+    """Delta for tool use input JSON content block."""
+
+    type: Literal["input_json_delta"] = "input_json_delta"
+    partial_json: str
+
+
+class ClaudeContentBlockDeltaEvent(BaseModel, frozen=True):
+    """Event sent for content block delta."""
+
+    type: Literal["content_block_delta"] = "content_block_delta"
+    index: int
+    delta: ClaudeTextDelta | ClaudeInputJsonDelta
+
+
+class ClaudeContentBlockStopEvent(BaseModel, frozen=True):
+    """Event sent at end of a content block."""
+
+    type: Literal["content_block_stop"] = "content_block_stop"
+    index: int
+
+
+class ClaudeMessageDeltaUsage(BaseModel, frozen=True):
+    """Usage in message_delta event."""
+
+    output_tokens: int
+
+
+class ClaudeMessageDelta(BaseModel, frozen=True):
+    """Delta in message_delta event."""
+
+    stop_reason: ClaudeStopReason | None = None
+    stop_sequence: str | None = None
+
+
+class ClaudeMessageDeltaEvent(BaseModel, frozen=True):
+    """Event sent with final message delta."""
+
+    type: Literal["message_delta"] = "message_delta"
+    delta: ClaudeMessageDelta
+    usage: ClaudeMessageDeltaUsage
+
+
+class ClaudeMessageStopEvent(BaseModel, frozen=True):
+    """Event sent at end of message stream."""
+
+    type: Literal["message_stop"] = "message_stop"
+
+
+ClaudeStreamEvent = (
+    ClaudeMessageStartEvent
+    | ClaudeContentBlockStartEvent
+    | ClaudeContentBlockDeltaEvent
+    | ClaudeContentBlockStopEvent
+    | ClaudeMessageDeltaEvent
+    | ClaudeMessageStopEvent
+)

src/exo/shared/types/commands.py

@@ -1,15 +1,15 @@
 from pydantic import Field
 
-from exo.shared.models.model_cards import ModelCard
+from exo.shared.models.model_cards import ModelCard, ModelId
 from exo.shared.types.api import (
-    ChatCompletionTaskParams,
-    ImageEditsInternalParams,
+    ImageEditsTaskParams,
     ImageGenerationTaskParams,
 )
 from exo.shared.types.chunks import InputImageChunk
 from exo.shared.types.common import CommandId, NodeId
+from exo.shared.types.text_generation import TextGenerationTaskParams
 from exo.shared.types.worker.instances import Instance, InstanceId, InstanceMeta
-from exo.shared.types.worker.shards import Sharding
+from exo.shared.types.worker.shards import Sharding, ShardMetadata
 from exo.utils.pydantic_ext import CamelCaseModel, TaggedModel
 
 
@@ -21,16 +21,16 @@ class TestCommand(BaseCommand):
     __test__ = False
 
 
-class ChatCompletion(BaseCommand):
-    request_params: ChatCompletionTaskParams
+class TextGeneration(BaseCommand):
+    task_params: TextGenerationTaskParams
 
 
 class ImageGeneration(BaseCommand):
-    request_params: ImageGenerationTaskParams
+    task_params: ImageGenerationTaskParams
 
 
 class ImageEdits(BaseCommand):
-    request_params: ImageEditsInternalParams
+    task_params: ImageEditsTaskParams
 
 
 class PlaceInstance(BaseCommand):
@@ -62,10 +62,23 @@ class RequestEventLog(BaseCommand):
     since_idx: int
 
 
+class StartDownload(BaseCommand):
+    target_node_id: NodeId
+    shard_metadata: ShardMetadata
+
+
+class DeleteDownload(BaseCommand):
+    target_node_id: NodeId
+    model_id: ModelId
+
+
+DownloadCommand = StartDownload | DeleteDownload
+
+
 Command = (
     TestCommand
     | RequestEventLog
-    | ChatCompletion
+    | TextGeneration
     | ImageGeneration
     | ImageEdits
     | PlaceInstance
@@ -79,3 +92,8 @@ Command = (
 class ForwarderCommand(CamelCaseModel):
     origin: NodeId
     command: Command
+
+
+class ForwarderDownloadCommand(CamelCaseModel):
+    origin: NodeId
+    command: DownloadCommand

src/exo/shared/types/events.py

@@ -1,4 +1,5 @@
 from datetime import datetime
+from typing import final
 
 from pydantic import Field
 
@@ -10,7 +11,7 @@ from exo.shared.types.worker.downloads import DownloadProgress
 from exo.shared.types.worker.instances import Instance, InstanceId
 from exo.shared.types.worker.runners import RunnerId, RunnerStatus
 from exo.utils.info_gatherer.info_gatherer import GatheredInfo
-from exo.utils.pydantic_ext import CamelCaseModel, TaggedModel
+from exo.utils.pydantic_ext import CamelCaseModel, FrozenModel, TaggedModel
 
 
 class EventId(Id):
@@ -109,6 +110,28 @@ class TopologyEdgeDeleted(BaseEvent):
     conn: Connection
 
 
+@final
+class TraceEventData(FrozenModel):
+    name: str
+    start_us: int
+    duration_us: int
+    rank: int
+    category: str
+
+
+@final
+class TracesCollected(BaseEvent):
+    task_id: TaskId
+    rank: int
+    traces: list[TraceEventData]
+
+
+@final
+class TracesMerged(BaseEvent):
+    task_id: TaskId
+    traces: list[TraceEventData]
+
+
 Event = (
     TestEvent
     | TaskCreated
@@ -127,6 +150,8 @@ Event = (
     | InputChunkReceived
     | TopologyEdgeCreated
     | TopologyEdgeDeleted
+    | TracesCollected
+    | TracesMerged
 )
 
 

src/exo/shared/types/mlx.py

@@ -0,0 +1,12 @@
+"""Shared types for MLX-related functionality."""
+
+from collections.abc import Sequence
+
+from mlx_lm.models.cache import (
+    KVCache,
+    QuantizedKVCache,
+    RotatingKVCache,
+)
+
+# This list contains one cache entry per transformer layer
+KVCacheType = Sequence[KVCache | RotatingKVCache | QuantizedKVCache]

src/exo/shared/types/openai_responses.py

@@ -0,0 +1,296 @@
+"""OpenAI Responses API wire types.
+
+These types model the OpenAI Responses API request/response format.
+ResponsesRequest is the API-level wire type; for the canonical internal
+task params type used by the inference pipeline, see
+``exo.shared.types.text_generation.TextGenerationTaskParams``.
+"""
+
+import time
+from typing import Any, Literal
+
+from pydantic import BaseModel, Field
+
+from exo.shared.types.common import ModelId
+
+# Type aliases
+ResponseStatus = Literal["completed", "failed", "in_progress", "incomplete"]
+ResponseRole = Literal["user", "assistant", "system", "developer"]
+
+
+# Request input content part types
+class ResponseInputTextPart(BaseModel, frozen=True):
+    """Text content part in a Responses API input message."""
+
+    type: Literal["input_text"] = "input_text"
+    text: str
+
+
+class ResponseOutputTextPart(BaseModel, frozen=True):
+    """Output text content part (used when replaying assistant messages in input)."""
+
+    type: Literal["output_text"] = "output_text"
+    text: str
+
+
+ResponseContentPart = ResponseInputTextPart | ResponseOutputTextPart
+
+
+# Request input item types
+class ResponseInputMessage(BaseModel, frozen=True):
+    """Input message for Responses API."""
+
+    role: ResponseRole
+    content: str | list[ResponseContentPart]
+    type: Literal["message"] = "message"
+
+
+class FunctionCallInputItem(BaseModel, frozen=True):
+    """Function call item replayed in input (from a previous assistant response)."""
+
+    type: Literal["function_call"] = "function_call"
+    id: str | None = None
+    call_id: str
+    name: str
+    arguments: str
+    status: ResponseStatus | None = None
+
+
+class FunctionCallOutputInputItem(BaseModel, frozen=True):
+    """Function call output item in input (user providing tool results)."""
+
+    type: Literal["function_call_output"] = "function_call_output"
+    call_id: str
+    output: str
+    id: str | None = None
+    status: ResponseStatus | None = None
+
+
+ResponseInputItem = (
+    ResponseInputMessage | FunctionCallInputItem | FunctionCallOutputInputItem
+)
+
+
+class ResponsesRequest(BaseModel, frozen=True):
+    """Request body for OpenAI Responses API.
+
+    This is the API wire type for the Responses endpoint. The canonical
+    internal task params type is ``TextGenerationTaskParams``; see the
+    ``responses_request_to_text_generation`` adapter for conversion.
+    """
+
+    # --- OpenAI Responses API standard fields ---
+    model: ModelId
+    input: str | list[ResponseInputItem]
+    instructions: str | None = None
+    max_output_tokens: int | None = None
+    temperature: float | None = None
+    top_p: float | None = None
+    stream: bool = False
+    tools: list[dict[str, Any]] | None = None
+    metadata: dict[str, str] | None = None
+
+    # --- exo extensions (not in OpenAI Responses API spec) ---
+    top_k: int | None = Field(
+        default=None,
+        description="[exo extension] Top-k sampling parameter. Not part of the OpenAI Responses API.",
+        json_schema_extra={"x-exo-extension": True},
+    )
+    stop: str | list[str] | None = Field(
+        default=None,
+        description="[exo extension] Stop sequence(s). Not part of the OpenAI Responses API.",
+        json_schema_extra={"x-exo-extension": True},
+    )
+    seed: int | None = Field(
+        default=None,
+        description="[exo extension] Seed for deterministic sampling. Not part of the OpenAI Responses API.",
+        json_schema_extra={"x-exo-extension": True},
+    )
+
+    # --- Internal fields (preserved during serialization, hidden from OpenAPI schema) ---
+    chat_template_messages: list[dict[str, Any]] | None = Field(
+        default=None,
+        description="Internal: pre-formatted messages for tokenizer chat template. Not part of the OpenAI Responses API.",
+        json_schema_extra={"x-exo-internal": True},
+    )
+
+
+# Response types
+class ResponseOutputText(BaseModel, frozen=True):
+    """Text content in response output."""
+
+    type: Literal["output_text"] = "output_text"
+    text: str
+    annotations: list[dict[str, str]] = Field(default_factory=list)
+
+
+class ResponseMessageItem(BaseModel, frozen=True):
+    """Message item in response output array."""
+
+    type: Literal["message"] = "message"
+    id: str
+    role: Literal["assistant"] = "assistant"
+    content: list[ResponseOutputText]
+    status: ResponseStatus = "completed"
+
+
+class ResponseFunctionCallItem(BaseModel, frozen=True):
+    """Function call item in response output array."""
+
+    type: Literal["function_call"] = "function_call"
+    id: str
+    call_id: str
+    name: str
+    arguments: str
+    status: ResponseStatus = "completed"
+
+
+ResponseItem = ResponseMessageItem | ResponseFunctionCallItem
+
+
+class ResponseUsage(BaseModel, frozen=True):
+    """Token usage in Responses API response."""
+
+    input_tokens: int
+    output_tokens: int
+    total_tokens: int
+
+
+class ResponsesResponse(BaseModel, frozen=True):
+    """Response body for OpenAI Responses API."""
+
+    id: str
+    object: Literal["response"] = "response"
+    created_at: int = Field(default_factory=lambda: int(time.time()))
+    status: ResponseStatus = "completed"
+    model: str
+    output: list[ResponseItem]
+    output_text: str
+    usage: ResponseUsage | None = None
+
+
+# Streaming event types
+class ResponseCreatedEvent(BaseModel, frozen=True):
+    """Event sent when response is created."""
+
+    type: Literal["response.created"] = "response.created"
+    sequence_number: int
+    response: ResponsesResponse
+
+
+class ResponseInProgressEvent(BaseModel, frozen=True):
+    """Event sent when response starts processing."""
+
+    type: Literal["response.in_progress"] = "response.in_progress"
+    sequence_number: int
+    response: ResponsesResponse
+
+
+class ResponseOutputItemAddedEvent(BaseModel, frozen=True):
+    """Event sent when an output item is added."""
+
+    type: Literal["response.output_item.added"] = "response.output_item.added"
+    sequence_number: int
+    output_index: int
+    item: ResponseItem
+
+
+class ResponseContentPartAddedEvent(BaseModel, frozen=True):
+    """Event sent when a content part is added."""
+
+    type: Literal["response.content_part.added"] = "response.content_part.added"
+    sequence_number: int
+    item_id: str
+    output_index: int
+    content_index: int
+    part: ResponseOutputText
+
+
+class ResponseTextDeltaEvent(BaseModel, frozen=True):
+    """Event sent for text delta during streaming."""
+
+    type: Literal["response.output_text.delta"] = "response.output_text.delta"
+    sequence_number: int
+    item_id: str
+    output_index: int
+    content_index: int
+    delta: str
+
+
+class ResponseTextDoneEvent(BaseModel, frozen=True):
+    """Event sent when text content is done."""
+
+    type: Literal["response.output_text.done"] = "response.output_text.done"
+    sequence_number: int
+    item_id: str
+    output_index: int
+    content_index: int
+    text: str
+
+
+class ResponseContentPartDoneEvent(BaseModel, frozen=True):
+    """Event sent when a content part is done."""
+
+    type: Literal["response.content_part.done"] = "response.content_part.done"
+    sequence_number: int
+    item_id: str
+    output_index: int
+    content_index: int
+    part: ResponseOutputText
+
+
+class ResponseOutputItemDoneEvent(BaseModel, frozen=True):
+    """Event sent when an output item is done."""
+
+    type: Literal["response.output_item.done"] = "response.output_item.done"
+    sequence_number: int
+    output_index: int
+    item: ResponseItem
+
+
+class ResponseFunctionCallArgumentsDeltaEvent(BaseModel, frozen=True):
+    """Event sent for function call arguments delta during streaming."""
+
+    type: Literal["response.function_call_arguments.delta"] = (
+        "response.function_call_arguments.delta"
+    )
+    sequence_number: int
+    item_id: str
+    output_index: int
+    delta: str
+
+
+class ResponseFunctionCallArgumentsDoneEvent(BaseModel, frozen=True):
+    """Event sent when function call arguments are complete."""
+
+    type: Literal["response.function_call_arguments.done"] = (
+        "response.function_call_arguments.done"
+    )
+    sequence_number: int
+    item_id: str
+    output_index: int
+    name: str
+    arguments: str
+
+
+class ResponseCompletedEvent(BaseModel, frozen=True):
+    """Event sent when response is completed."""
+
+    type: Literal["response.completed"] = "response.completed"
+    sequence_number: int
+    response: ResponsesResponse
+
+
+ResponsesStreamEvent = (
+    ResponseCreatedEvent
+    | ResponseInProgressEvent
+    | ResponseOutputItemAddedEvent
+    | ResponseContentPartAddedEvent
+    | ResponseTextDeltaEvent
+    | ResponseTextDoneEvent
+    | ResponseContentPartDoneEvent
+    | ResponseOutputItemDoneEvent
+    | ResponseFunctionCallArgumentsDeltaEvent
+    | ResponseFunctionCallArgumentsDoneEvent
+    | ResponseCompletedEvent
+)

src/exo/shared/types/profiling.py

@@ -1,5 +1,5 @@
 from collections.abc import Sequence
-from typing import Self
+from typing import Literal, Self
 
 import psutil
 
@@ -48,9 +48,13 @@ class SystemPerformanceProfile(CamelCaseModel):
     ecpu_usage: float = 0.0
 
 
+InterfaceType = Literal["wifi", "ethernet", "maybe_ethernet", "thunderbolt", "unknown"]
+
+
 class NetworkInterfaceInfo(CamelCaseModel):
     name: str
     ip_address: str
+    interface_type: InterfaceType = "unknown"
 
 
 class NodeIdentity(CamelCaseModel):
@@ -71,3 +75,11 @@ class NodeThunderboltInfo(CamelCaseModel):
     """Thunderbolt interface identifiers for a node."""
 
     interfaces: Sequence[ThunderboltIdentifier] = []
+
+
+class ThunderboltBridgeStatus(CamelCaseModel):
+    """Whether the Thunderbolt Bridge network service is enabled on this node."""
+
+    enabled: bool
+    exists: bool
+    service_name: str | None = None

src/exo/shared/types/state.py

@@ -13,6 +13,7 @@ from exo.shared.types.profiling import (
     NodeNetworkInfo,
     NodeThunderboltInfo,
     SystemPerformanceProfile,
+    ThunderboltBridgeStatus,
 )
 from exo.shared.types.tasks import Task, TaskId
 from exo.shared.types.worker.downloads import DownloadProgress
@@ -51,6 +52,10 @@ class State(CamelCaseModel):
     node_system: Mapping[NodeId, SystemPerformanceProfile] = {}
     node_network: Mapping[NodeId, NodeNetworkInfo] = {}
     node_thunderbolt: Mapping[NodeId, NodeThunderboltInfo] = {}
+    node_thunderbolt_bridge: Mapping[NodeId, ThunderboltBridgeStatus] = {}
+
+    # Detected cycles where all nodes have Thunderbolt bridge enabled (>2 nodes)
+    thunderbolt_bridge_cycles: Sequence[Sequence[NodeId]] = []
 
     @field_serializer("topology", mode="plain")
     def _encode_topology(self, value: Topology) -> TopologySnapshot:

src/exo/shared/types/tasks.py

@@ -3,11 +3,11 @@ from enum import Enum
 from pydantic import Field
 
 from exo.shared.types.api import (
-    ChatCompletionTaskParams,
-    ImageEditsInternalParams,
+    ImageEditsTaskParams,
     ImageGenerationTaskParams,
 )
 from exo.shared.types.common import CommandId, Id
+from exo.shared.types.text_generation import TextGenerationTaskParams
 from exo.shared.types.worker.instances import BoundInstance, InstanceId
 from exo.shared.types.worker.runners import RunnerId
 from exo.shared.types.worker.shards import ShardMetadata
@@ -52,9 +52,9 @@ class StartWarmup(BaseTask):  # emitted by Worker
     pass
 
 
-class ChatCompletion(BaseTask):  # emitted by Master
+class TextGeneration(BaseTask):  # emitted by Master
     command_id: CommandId
-    task_params: ChatCompletionTaskParams
+    task_params: TextGenerationTaskParams
 
     error_type: str | None = Field(default=None)
     error_message: str | None = Field(default=None)
@@ -70,7 +70,7 @@ class ImageGeneration(BaseTask):  # emitted by Master
 
 class ImageEdits(BaseTask):  # emitted by Master
     command_id: CommandId
-    task_params: ImageEditsInternalParams
+    task_params: ImageEditsTaskParams
 
     error_type: str | None = Field(default=None)
     error_message: str | None = Field(default=None)
@@ -86,7 +86,7 @@ Task = (
     | ConnectToGroup
     | LoadModel
     | StartWarmup
-    | ChatCompletion
+    | TextGeneration
     | ImageGeneration
     | ImageEdits
     | Shutdown

src/exo/shared/types/text_generation.py

@@ -0,0 +1,42 @@
+"""Canonical internal type for text generation task parameters.
+
+All external API formats (Chat Completions, Claude Messages, OpenAI Responses)
+are converted to TextGenerationTaskParams at the API boundary via adapters.
+"""
+
+from typing import Any, Literal
+
+from pydantic import BaseModel
+
+from exo.shared.types.common import ModelId
+
+MessageRole = Literal["user", "assistant", "system", "developer"]
+
+
+class InputMessage(BaseModel, frozen=True):
+    """Internal message for text generation pipelines."""
+
+    role: MessageRole
+    content: str
+
+
+class TextGenerationTaskParams(BaseModel, frozen=True):
+    """Canonical internal task params for text generation.
+
+    Every API adapter converts its wire type into this before handing
+    off to the master/worker pipeline.
+    """
+
+    model: ModelId
+    input: str | list[InputMessage]
+    instructions: str | None = None
+    max_output_tokens: int | None = None
+    temperature: float | None = None
+    top_p: float | None = None
+    stream: bool = False
+    tools: list[dict[str, Any]] | None = None
+    bench: bool = False
+    top_k: int | None = None
+    stop: str | list[str] | None = None
+    seed: int | None = None
+    chat_template_messages: list[dict[str, Any]] | None = None

src/exo/shared/types/topology.py

@@ -21,9 +21,6 @@ class RDMAConnection(FrozenModel):
     source_rdma_iface: str
     sink_rdma_iface: str
 
-    def is_thunderbolt(self) -> bool:
-        return True
-
 
 class SocketConnection(FrozenModel):
     sink_multiaddr: Multiaddr
@@ -31,9 +28,6 @@ class SocketConnection(FrozenModel):
     def __hash__(self):
         return hash(self.sink_multiaddr.ip_address)
 
-    def is_thunderbolt(self) -> bool:
-        return str(self.sink_multiaddr.ipv4_address).startswith("169.254")
-
 
 class Connection(FrozenModel):
     source: NodeId

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

@@ -1,7 +1,13 @@
 from collections.abc import Generator
 from typing import Any, Literal
 
-from exo.shared.types.api import FinishReason, GenerationStats, ImageGenerationStats
+from exo.shared.types.api import (
+    FinishReason,
+    GenerationStats,
+    ImageGenerationStats,
+    ToolCallItem,
+    Usage,
+)
 from exo.utils.pydantic_ext import TaggedModel
 
 
@@ -19,12 +25,14 @@ class GenerationResponse(BaseRunnerResponse):
     # logprobs: list[float] | None = None # too big. we can change to be top-k
     finish_reason: FinishReason | None = None
     stats: GenerationStats | None = None
+    usage: Usage | None
 
 
 class ImageGenerationResponse(BaseRunnerResponse):
     image_data: bytes
     format: Literal["png", "jpeg", "webp"] = "png"
     stats: ImageGenerationStats | None = None
+    image_index: int = 0
 
     def __repr_args__(self) -> Generator[tuple[str, Any], None, None]:
         for name, value in super().__repr_args__():  # pyright: ignore[reportAny]
@@ -39,6 +47,7 @@ class PartialImageResponse(BaseRunnerResponse):
     format: Literal["png", "jpeg", "webp"] = "png"
     partial_index: int
     total_partials: int
+    image_index: int = 0
 
     def __repr_args__(self) -> Generator[tuple[str, Any], None, None]:
         for name, value in super().__repr_args__():  # pyright: ignore[reportAny]
@@ -48,5 +57,10 @@ class PartialImageResponse(BaseRunnerResponse):
                 yield name, value
 
 
+class ToolCallResponse(BaseRunnerResponse):
+    tool_calls: list[ToolCallItem]
+    usage: Usage | None
+
+
 class FinishedResponse(BaseRunnerResponse):
     pass

src/exo/utils/info_gatherer/info_gatherer.py

@@ -13,12 +13,14 @@ from anyio.abc import TaskGroup
 from anyio.streams.buffered import BufferedByteReceiveStream
 from anyio.streams.text import TextReceiveStream
 from loguru import logger
+from pydantic import ValidationError
 
 from exo.shared.constants import EXO_CONFIG_FILE
 from exo.shared.types.memory import Memory
 from exo.shared.types.profiling import (
     MemoryUsage,
     NetworkInterfaceInfo,
+    ThunderboltBridgeStatus,
 )
 from exo.shared.types.thunderbolt import (
     ThunderboltConnection,
@@ -34,6 +36,142 @@ from .system_info import get_friendly_name, get_model_and_chip, get_network_inte
 IS_DARWIN = sys.platform == "darwin"
 
 
+async def _get_thunderbolt_devices() -> set[str] | None:
+    """Get Thunderbolt interface device names (e.g., en2, en3) from hardware ports.
+
+    Returns None if the networksetup command fails.
+    """
+    result = await anyio.run_process(
+        ["networksetup", "-listallhardwareports"],
+        check=False,
+    )
+    if result.returncode != 0:
+        logger.warning(
+            f"networksetup -listallhardwareports failed with code "
+            f"{result.returncode}: {result.stderr.decode()}"
+        )
+        return None
+
+    output = result.stdout.decode()
+    thunderbolt_devices: set[str] = set()
+    current_port: str | None = None
+
+    for line in output.splitlines():
+        line = line.strip()
+        if line.startswith("Hardware Port:"):
+            current_port = line.split(":", 1)[1].strip()
+        elif line.startswith("Device:") and current_port:
+            device = line.split(":", 1)[1].strip()
+            if "thunderbolt" in current_port.lower():
+                thunderbolt_devices.add(device)
+            current_port = None
+
+    return thunderbolt_devices
+
+
+async def _get_bridge_services() -> dict[str, str] | None:
+    """Get mapping of bridge device -> service name from network service order.
+
+    Returns None if the networksetup command fails.
+    """
+    result = await anyio.run_process(
+        ["networksetup", "-listnetworkserviceorder"],
+        check=False,
+    )
+    if result.returncode != 0:
+        logger.warning(
+            f"networksetup -listnetworkserviceorder failed with code "
+            f"{result.returncode}: {result.stderr.decode()}"
+        )
+        return None
+
+    # Parse service order to find bridge devices and their service names
+    # Format: "(1) Service Name\n(Hardware Port: ..., Device: bridge0)\n"
+    service_order_output = result.stdout.decode()
+    bridge_services: dict[str, str] = {}  # device -> service name
+    current_service: str | None = None
+
+    for line in service_order_output.splitlines():
+        line = line.strip()
+        # Match "(N) Service Name" or "(*) Service Name" (disabled)
+        # but NOT "(Hardware Port: ...)" lines
+        if (
+            line
+            and line.startswith("(")
+            and ")" in line
+            and not line.startswith("(Hardware Port:")
+        ):
+            paren_end = line.index(")")
+            if paren_end + 2 <= len(line):
+                current_service = line[paren_end + 2 :]
+        # Match "(Hardware Port: ..., Device: bridgeX)"
+        elif current_service and "Device: bridge" in line:
+            # Extract device name from "..., Device: bridge0)"
+            device_start = line.find("Device: ") + len("Device: ")
+            device_end = line.find(")", device_start)
+            if device_end > device_start:
+                device = line[device_start:device_end]
+                bridge_services[device] = current_service
+
+    return bridge_services
+
+
+async def _get_bridge_members(bridge_device: str) -> set[str]:
+    """Get member interfaces of a bridge device via ifconfig."""
+    result = await anyio.run_process(
+        ["ifconfig", bridge_device],
+        check=False,
+    )
+    if result.returncode != 0:
+        logger.debug(f"ifconfig {bridge_device} failed with code {result.returncode}")
+        return set()
+
+    members: set[str] = set()
+    ifconfig_output = result.stdout.decode()
+    for line in ifconfig_output.splitlines():
+        line = line.strip()
+        if line.startswith("member:"):
+            parts = line.split()
+            if len(parts) > 1:
+                members.add(parts[1])
+
+    return members
+
+
+async def _find_thunderbolt_bridge(
+    bridge_services: dict[str, str], thunderbolt_devices: set[str]
+) -> str | None:
+    """Find the service name of a bridge containing Thunderbolt interfaces.
+
+    Returns the service name if found, None otherwise.
+    """
+    for bridge_device, service_name in bridge_services.items():
+        members = await _get_bridge_members(bridge_device)
+        if members & thunderbolt_devices:  # intersection is non-empty
+            return service_name
+    return None
+
+
+async def _is_service_enabled(service_name: str) -> bool | None:
+    """Check if a network service is enabled.
+
+    Returns True if enabled, False if disabled, None on error.
+    """
+    result = await anyio.run_process(
+        ["networksetup", "-getnetworkserviceenabled", service_name],
+        check=False,
+    )
+    if result.returncode != 0:
+        logger.warning(
+            f"networksetup -getnetworkserviceenabled '{service_name}' "
+            f"failed with code {result.returncode}: {result.stderr.decode()}"
+        )
+        return None
+
+    stdout = result.stdout.decode().strip().lower()
+    return stdout == "enabled"
+
+
 class StaticNodeInformation(TaggedModel):
     """Node information that should NEVER change, to be gathered once at startup"""
 
@@ -58,19 +196,80 @@ class MacThunderboltConnections(TaggedModel):
     conns: Sequence[ThunderboltConnection]
 
 
+class ThunderboltBridgeInfo(TaggedModel):
+    status: ThunderboltBridgeStatus
+
+    @classmethod
+    async def gather(cls) -> Self | None:
+        """Check if a Thunderbolt Bridge network service is enabled on this node.
+
+        Detection approach:
+        1. Find all Thunderbolt interface devices (en2, en3, etc.) from hardware ports
+        2. Find bridge devices from network service order (not hardware ports, as
+           bridges may not appear there)
+        3. Check each bridge's members via ifconfig
+        4. If a bridge contains Thunderbolt interfaces, it's a Thunderbolt Bridge
+        5. Check if that network service is enabled
+        """
+        if not IS_DARWIN:
+            return None
+
+        def _no_bridge_status() -> Self:
+            return cls(
+                status=ThunderboltBridgeStatus(
+                    enabled=False, exists=False, service_name=None
+                )
+            )
+
+        try:
+            tb_devices = await _get_thunderbolt_devices()
+            if tb_devices is None:
+                return _no_bridge_status()
+
+            bridge_services = await _get_bridge_services()
+            if not bridge_services:
+                return _no_bridge_status()
+
+            tb_service_name = await _find_thunderbolt_bridge(
+                bridge_services, tb_devices
+            )
+            if not tb_service_name:
+                return _no_bridge_status()
+
+            enabled = await _is_service_enabled(tb_service_name)
+            if enabled is None:
+                return cls(
+                    status=ThunderboltBridgeStatus(
+                        enabled=False, exists=True, service_name=tb_service_name
+                    )
+                )
+
+            return cls(
+                status=ThunderboltBridgeStatus(
+                    enabled=enabled,
+                    exists=True,
+                    service_name=tb_service_name,
+                )
+            )
+        except Exception as e:
+            logger.warning(f"Failed to gather Thunderbolt Bridge info: {e}")
+            return None
+
+
 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"""
 
     @classmethod
     async def gather(cls) -> Self | None:
         cfg_file = anyio.Path(EXO_CONFIG_FILE)
+        await cfg_file.parent.mkdir(parents=True, exist_ok=True)
         await cfg_file.touch(exist_ok=True)
         async with await cfg_file.open("rb") as f:
             try:
                 contents = (await f.read()).decode("utf-8")
                 data = tomllib.loads(contents)
                 return cls.model_validate(data)
-            except (tomllib.TOMLDecodeError, UnicodeDecodeError):
+            except (tomllib.TOMLDecodeError, UnicodeDecodeError, ValidationError):
                 logger.warning("Invalid config file, skipping...")
                 return None
 
@@ -111,6 +310,7 @@ GatheredInfo = (
     | NodeNetworkInterfaces
     | MacThunderboltIdentifiers
     | MacThunderboltConnections
+    | ThunderboltBridgeInfo
     | NodeConfig
     | MiscData
     | StaticNodeInformation
@@ -125,6 +325,7 @@ class InfoGatherer:
     system_profiler_interval: float | None = 5 if IS_DARWIN else None
     memory_poll_rate: float | None = None if IS_DARWIN else 1
     macmon_interval: float | None = 1 if IS_DARWIN else None
+    thunderbolt_bridge_poll_interval: float | None = 10 if IS_DARWIN else None
     _tg: TaskGroup = field(init=False, default_factory=create_task_group)
 
     async def run(self):
@@ -133,6 +334,7 @@ class InfoGatherer:
                 if (macmon_path := shutil.which("macmon")) is not None:
                     tg.start_soon(self._monitor_macmon, macmon_path)
                 tg.start_soon(self._monitor_system_profiler_thunderbolt_data)
+                tg.start_soon(self._monitor_thunderbolt_bridge_status)
             tg.start_soon(self._watch_system_info)
             tg.start_soon(self._monitor_memory_usage)
             tg.start_soon(self._monitor_misc)
@@ -149,13 +351,8 @@ class InfoGatherer:
     async def _monitor_misc(self):
         if self.misc_poll_interval is None:
             return
-        prev = await MiscData.gather()
-        await self.info_sender.send(prev)
         while True:
-            curr = await MiscData.gather()
-            if prev != curr:
-                prev = curr
-                await self.info_sender.send(curr)
+            await self.info_sender.send(await MiscData.gather())
             await anyio.sleep(self.misc_poll_interval)
 
     async def _monitor_system_profiler_thunderbolt_data(self):
@@ -165,15 +362,12 @@ class InfoGatherer:
         if iface_map is None:
             return
 
-        old_idents = []
         while True:
             data = await ThunderboltConnectivity.gather()
             assert data is not None
 
             idents = [it for i in data if (it := i.ident(iface_map)) is not None]
-            if idents != old_idents:
-                await self.info_sender.send(MacThunderboltIdentifiers(idents=idents))
-            old_idents = idents
+            await self.info_sender.send(MacThunderboltIdentifiers(idents=idents))
 
             conns = [it for i in data if (it := i.conn()) is not None]
             await self.info_sender.send(MacThunderboltConnections(conns=conns))
@@ -198,14 +392,20 @@ class InfoGatherer:
     async def _watch_system_info(self):
         if self.interface_watcher_interval is None:
             return
-        old_nics = []
         while True:
-            nics = get_network_interfaces()
-            if nics != old_nics:
-                old_nics = nics
-                await self.info_sender.send(NodeNetworkInterfaces(ifaces=nics))
+            nics = await get_network_interfaces()
+            await self.info_sender.send(NodeNetworkInterfaces(ifaces=nics))
             await anyio.sleep(self.interface_watcher_interval)
 
+    async def _monitor_thunderbolt_bridge_status(self):
+        if self.thunderbolt_bridge_poll_interval is None:
+            return
+        while True:
+            curr = await ThunderboltBridgeInfo.gather()
+            if curr is not None:
+                await self.info_sender.send(curr)
+            await anyio.sleep(self.thunderbolt_bridge_poll_interval)
+
     async def _monitor_macmon(self, macmon_path: str):
         if self.macmon_interval is None:
             return

src/exo/utils/info_gatherer/system_info.py

@@ -5,7 +5,7 @@ from subprocess import CalledProcessError
 import psutil
 from anyio import run_process
 
-from exo.shared.types.profiling import NetworkInterfaceInfo
+from exo.shared.types.profiling import InterfaceType, NetworkInterfaceInfo
 
 
 async def get_friendly_name() -> str:
@@ -16,8 +16,7 @@ async def get_friendly_name() -> str:
     """
     hostname = socket.gethostname()
 
-    # TODO: better non mac support
-    if sys.platform != "darwin":  # 'darwin' is the platform name for macOS
+    if sys.platform != "darwin":
         return hostname
 
     try:
@@ -28,7 +27,41 @@ async def get_friendly_name() -> str:
     return process.stdout.decode("utf-8", errors="replace").strip() or hostname
 
 
-def get_network_interfaces() -> list[NetworkInterfaceInfo]:
+async def _get_interface_types_from_networksetup() -> dict[str, InterfaceType]:
+    """Parse networksetup -listallhardwareports to get interface types."""
+    if sys.platform != "darwin":
+        return {}
+
+    try:
+        result = await run_process(["networksetup", "-listallhardwareports"])
+    except CalledProcessError:
+        return {}
+
+    types: dict[str, InterfaceType] = {}
+    current_type: InterfaceType = "unknown"
+
+    for line in result.stdout.decode().splitlines():
+        if line.startswith("Hardware Port:"):
+            port_name = line.split(":", 1)[1].strip()
+            if "Wi-Fi" in port_name:
+                current_type = "wifi"
+            elif "Ethernet" in port_name or "LAN" in port_name:
+                current_type = "ethernet"
+            elif port_name.startswith("Thunderbolt"):
+                current_type = "thunderbolt"
+            else:
+                current_type = "unknown"
+        elif line.startswith("Device:"):
+            device = line.split(":", 1)[1].strip()
+            # enX is ethernet adapters or thunderbolt - these must be deprioritised
+            if device.startswith("en") and device not in ["en0", "en1"]:
+                current_type = "maybe_ethernet"
+            types[device] = current_type
+
+    return types
+
+
+async def get_network_interfaces() -> list[NetworkInterfaceInfo]:
     """
     Retrieves detailed network interface information on macOS.
     Parses output from 'networksetup -listallhardwareports' and 'ifconfig'
@@ -36,13 +69,18 @@ def get_network_interfaces() -> list[NetworkInterfaceInfo]:
     Returns a list of NetworkInterfaceInfo objects.
     """
     interfaces_info: list[NetworkInterfaceInfo] = []
+    interface_types = await _get_interface_types_from_networksetup()
 
     for iface, services in psutil.net_if_addrs().items():
         for service in services:
             match service.family:
                 case socket.AF_INET | socket.AF_INET6:
                     interfaces_info.append(
-                        NetworkInterfaceInfo(name=iface, ip_address=service.address)
+                        NetworkInterfaceInfo(
+                            name=iface,
+                            ip_address=service.address,
+                            interface_type=interface_types.get(iface, "unknown"),
+                        )
                     )
                 case _:
                     pass

src/exo/utils/keyed_backoff.py

@@ -0,0 +1,32 @@
+import time
+from typing import Generic, TypeVar
+
+K = TypeVar("K")
+
+
+class KeyedBackoff(Generic[K]):
+    """Tracks exponential backoff state per key."""
+
+    def __init__(self, base: float = 0.5, cap: float = 10.0):
+        self._base = base
+        self._cap = cap
+        self._attempts: dict[K, int] = {}
+        self._last_time: dict[K, float] = {}
+
+    def should_proceed(self, key: K) -> bool:
+        """Returns True if enough time has elapsed since last attempt."""
+        now = time.monotonic()
+        last = self._last_time.get(key, 0.0)
+        attempts = self._attempts.get(key, 0)
+        delay = min(self._cap, self._base * (2.0**attempts))
+        return now - last >= delay
+
+    def record_attempt(self, key: K) -> None:
+        """Record that an attempt was made for this key."""
+        self._last_time[key] = time.monotonic()
+        self._attempts[key] = self._attempts.get(key, 0) + 1
+
+    def reset(self, key: K) -> None:
+        """Reset backoff state for a key (e.g., on success)."""
+        self._attempts.pop(key, None)
+        self._last_time.pop(key, None)

src/exo/worker/engines/image/distributed_model.py

@@ -6,10 +6,10 @@ import mlx.core as mx
 from mflux.models.common.config.config import Config
 from PIL import Image
 
+from exo.download.download_utils import build_model_path
 from exo.shared.types.api import AdvancedImageParams
 from exo.shared.types.worker.instances import BoundInstance
 from exo.shared.types.worker.shards import PipelineShardMetadata
-from exo.worker.download.download_utils import build_model_path
 from exo.worker.engines.image.config import ImageModelConfig
 from exo.worker.engines.image.models import (
     create_adapter_for_model,
@@ -140,6 +140,7 @@ class DistributedImageModel:
             width=width,
             image_path=image_path,
             model_config=self._adapter.model.model_config,  # pyright: ignore[reportAny]
+            guidance=guidance_override if guidance_override is not None else 4.0,
         )
 
         num_sync_steps = self._config.get_num_sync_steps(steps)

src/exo/worker/engines/image/generate.py

@@ -11,7 +11,7 @@ from PIL import Image
 
 from exo.shared.types.api import (
     AdvancedImageParams,
-    ImageEditsInternalParams,
+    ImageEditsTaskParams,
     ImageGenerationStats,
     ImageGenerationTaskParams,
 )
@@ -67,7 +67,7 @@ def warmup_image_generator(model: DistributedImageModel) -> Image.Image | None:
 
 def generate_image(
     model: DistributedImageModel,
-    task: ImageGenerationTaskParams | ImageEditsInternalParams,
+    task: ImageGenerationTaskParams | ImageEditsTaskParams,
 ) -> Generator[ImageGenerationResponse | PartialImageResponse, None, None]:
     """Generate image(s), optionally yielding partial results.
 
@@ -75,19 +75,20 @@ def generate_image(
     intermediate images, then ImageGenerationResponse for the final image.
 
     Yields:
-        PartialImageResponse for intermediate images (if partial_images > 0)
-        ImageGenerationResponse for the final complete image
+        PartialImageResponse for intermediate images (if partial_images > 0, first image only)
+        ImageGenerationResponse for final complete images
     """
     width, height = parse_size(task.size)
     quality: Literal["low", "medium", "high"] = task.quality or "medium"
 
     advanced_params = task.advanced_params
     if advanced_params is not None and advanced_params.seed is not None:
-        seed = advanced_params.seed
+        base_seed = advanced_params.seed
     else:
-        seed = random.randint(0, 2**32 - 1)
+        base_seed = random.randint(0, 2**32 - 1)
 
     is_bench = getattr(task, "bench", False)
+    num_images = task.n or 1
 
     generation_start_time: float = 0.0
 
@@ -95,82 +96,95 @@ def generate_image(
         mx.reset_peak_memory()
         generation_start_time = time.perf_counter()
 
-    partial_images = task.partial_images or (3 if task.stream else 0)
+    partial_images = (
+        task.partial_images
+        if task.partial_images is not None and task.stream is not None and task.stream
+        else 0
+    )
 
     image_path: Path | None = None
 
     with tempfile.TemporaryDirectory() as tmpdir:
-        if isinstance(task, ImageEditsInternalParams):
+        if isinstance(task, ImageEditsTaskParams):
             # Decode base64 image data and save to temp file
             image_path = Path(tmpdir) / "input.png"
             image_path.write_bytes(base64.b64decode(task.image_data))
 
-        # Iterate over generator results
-        for result in model.generate(
-            prompt=task.prompt,
-            height=height,
-            width=width,
-            quality=quality,
-            seed=seed,
-            image_path=image_path,
-            partial_images=partial_images,
-            advanced_params=advanced_params,
-        ):
-            if isinstance(result, tuple):
-                # Partial image: (Image, partial_index, total_partials)
-                image, partial_idx, total_partials = result
-                buffer = io.BytesIO()
-                image_format = task.output_format.upper()
-                if image_format == "JPG":
-                    image_format = "JPEG"
-                if image_format == "JPEG" and image.mode == "RGBA":
-                    image = image.convert("RGB")
-                image.save(buffer, format=image_format)
+        for image_num in range(num_images):
+            # Increment seed for each image to ensure unique results
+            current_seed = base_seed + image_num
 
-                yield PartialImageResponse(
-                    image_data=buffer.getvalue(),
-                    format=task.output_format,
-                    partial_index=partial_idx,
-                    total_partials=total_partials,
-                )
-            else:
-                image = result
+            for result in model.generate(
+                prompt=task.prompt,
+                height=height,
+                width=width,
+                quality=quality,
+                seed=current_seed,
+                image_path=image_path,
+                partial_images=partial_images,
+                advanced_params=advanced_params,
+            ):
+                if isinstance(result, tuple):
+                    # Partial image: (Image, partial_index, total_partials)
+                    image, partial_idx, total_partials = result
+                    buffer = io.BytesIO()
+                    image_format = task.output_format.upper()
+                    if image_format == "JPG":
+                        image_format = "JPEG"
+                    if image_format == "JPEG" and image.mode == "RGBA":
+                        image = image.convert("RGB")
+                    image.save(buffer, format=image_format)
 
-                stats: ImageGenerationStats | None = None
-                if is_bench:
-                    generation_end_time = time.perf_counter()
-                    total_generation_time = generation_end_time - generation_start_time
-
-                    num_inference_steps = model.get_steps_for_quality(quality)
-
-                    seconds_per_step = (
-                        total_generation_time / num_inference_steps
-                        if num_inference_steps > 0
-                        else 0.0
+                    yield PartialImageResponse(
+                        image_data=buffer.getvalue(),
+                        format=task.output_format,
+                        partial_index=partial_idx,
+                        total_partials=total_partials,
+                        image_index=image_num,
                     )
+                else:
+                    image = result
 
-                    peak_memory_gb = mx.get_peak_memory() / (1024**3)
+                    # Only include stats on the final image
+                    stats: ImageGenerationStats | None = None
+                    if is_bench and image_num == num_images - 1:
+                        generation_end_time = time.perf_counter()
+                        total_generation_time = (
+                            generation_end_time - generation_start_time
+                        )
 
-                    stats = ImageGenerationStats(
-                        seconds_per_step=seconds_per_step,
-                        total_generation_time=total_generation_time,
-                        num_inference_steps=num_inference_steps,
-                        num_images=task.n or 1,
-                        image_width=width,
-                        image_height=height,
-                        peak_memory_usage=Memory.from_gb(peak_memory_gb),
+                        num_inference_steps = model.get_steps_for_quality(quality)
+                        total_steps = num_inference_steps * num_images
+
+                        seconds_per_step = (
+                            total_generation_time / total_steps
+                            if total_steps > 0
+                            else 0.0
+                        )
+
+                        peak_memory_gb = mx.get_peak_memory() / (1024**3)
+
+                        stats = ImageGenerationStats(
+                            seconds_per_step=seconds_per_step,
+                            total_generation_time=total_generation_time,
+                            num_inference_steps=num_inference_steps,
+                            num_images=num_images,
+                            image_width=width,
+                            image_height=height,
+                            peak_memory_usage=Memory.from_gb(peak_memory_gb),
+                        )
+
+                    buffer = io.BytesIO()
+                    image_format = task.output_format.upper()
+                    if image_format == "JPG":
+                        image_format = "JPEG"
+                    if image_format == "JPEG" and image.mode == "RGBA":
+                        image = image.convert("RGB")
+                    image.save(buffer, format=image_format)
+
+                    yield ImageGenerationResponse(
+                        image_data=buffer.getvalue(),
+                        format=task.output_format,
+                        stats=stats,
+                        image_index=image_num,
                     )
-
-                buffer = io.BytesIO()
-                image_format = task.output_format.upper()
-                if image_format == "JPG":
-                    image_format = "JPEG"
-                if image_format == "JPEG" and image.mode == "RGBA":
-                    image = image.convert("RGB")
-                image.save(buffer, format=image_format)
-
-                yield ImageGenerationResponse(
-                    image_data=buffer.getvalue(),
-                    format=task.output_format,
-                    stats=stats,
-                )

src/exo/worker/engines/image/models/__init__.py

@@ -33,6 +33,7 @@ _ADAPTER_REGISTRY: dict[str, AdapterFactory] = {
 # Config registry: maps model ID patterns to configs
 _CONFIG_REGISTRY: dict[str, ImageModelConfig] = {
     "flux.1-schnell": FLUX_SCHNELL_CONFIG,
+    "flux.1-krea-dev": FLUX_DEV_CONFIG,  # Must come before "flux.1-dev" for pattern matching
     "flux.1-dev": FLUX_DEV_CONFIG,
     "qwen-image-edit": QWEN_IMAGE_EDIT_CONFIG,  # Must come before "qwen-image" for pattern matching
     "qwen-image": QWEN_IMAGE_CONFIG,

src/exo/worker/engines/image/pipeline/runner.py

@@ -6,6 +6,11 @@ from mflux.models.common.config.config import Config
 from mflux.utils.exceptions import StopImageGenerationException
 from tqdm import tqdm
 
+from exo.shared.constants import EXO_TRACING_ENABLED
+from exo.shared.tracing import (
+    clear_trace_buffer,
+    trace,
+)
 from exo.shared.types.worker.shards import PipelineShardMetadata
 from exo.worker.engines.image.config import ImageModelConfig
 from exo.worker.engines.image.models.base import (
@@ -324,6 +329,7 @@ class DiffusionRunner:
             capture_steps = set()
 
         self._reset_all_caches()
+        clear_trace_buffer()
 
         time_steps = tqdm(range(runtime_config.num_inference_steps))
 
@@ -348,6 +354,7 @@ class DiffusionRunner:
                 ctx.in_loop(  # pyright: ignore[reportAny]
                     t=t,
                     latents=latents,
+                    time_steps=time_steps,
                 )
 
                 mx.eval(latents)
@@ -464,20 +471,22 @@ class DiffusionRunner:
         if self.group is None:
             return self._single_node_step(t, config, latents, prompt_data)
         elif t < config.init_time_step + num_sync_steps:
-            return self._sync_pipeline_step(
-                t,
-                config,
-                latents,
-                prompt_data,
-            )
+            with trace(name=f"sync {t}", rank=self.rank, category="sync"):
+                return self._sync_pipeline_step(
+                    t,
+                    config,
+                    latents,
+                    prompt_data,
+                )
         else:
-            return self._async_pipeline_step(
-                t,
-                config,
-                latents,
-                prompt_data,
-                is_first_async_step=t == config.init_time_step + num_sync_steps,
-            )
+            with trace(name=f"async {t}", rank=self.rank, category="async"):
+                return self._async_pipeline_step(
+                    t,
+                    config,
+                    latents,
+                    prompt_data,
+                    is_first_async_step=t == config.init_time_step + num_sync_steps,
+                )
 
     def _single_node_step(
         self,
@@ -585,30 +594,41 @@ class DiffusionRunner:
 
         if self.has_joint_blocks:
             if not self.is_first_stage:
-                hidden_states = mx.distributed.recv(
-                    (batch_size, num_img_tokens, hidden_dim),
-                    dtype,
-                    self.prev_rank,
-                    group=self.group,
-                )
-                encoder_hidden_states = mx.distributed.recv(
-                    (batch_size, text_seq_len, hidden_dim),
-                    dtype,
-                    self.prev_rank,
-                    group=self.group,
-                )
-                mx.eval(hidden_states, encoder_hidden_states)
+                with trace(
+                    name=f"recv {self.prev_rank}", rank=self.rank, category="comms"
+                ):
+                    hidden_states = mx.distributed.recv(
+                        (batch_size, num_img_tokens, hidden_dim),
+                        dtype,
+                        self.prev_rank,
+                        group=self.group,
+                    )
+                    encoder_hidden_states = mx.distributed.recv(
+                        (batch_size, text_seq_len, hidden_dim),
+                        dtype,
+                        self.prev_rank,
+                        group=self.group,
+                    )
+                    mx.eval(hidden_states, encoder_hidden_states)
 
             assert self.joint_block_wrappers is not None
             assert encoder_hidden_states is not None
-            for wrapper in self.joint_block_wrappers:
-                wrapper.set_patch(BlockWrapperMode.CACHING)
-                encoder_hidden_states, hidden_states = wrapper(
-                    hidden_states=hidden_states,
-                    encoder_hidden_states=encoder_hidden_states,
-                    text_embeddings=text_embeddings,
-                    rotary_embeddings=image_rotary_embeddings,
-                )
+            with trace(
+                name="joint_blocks",
+                rank=self.rank,
+                category="compute",
+            ):
+                for wrapper in self.joint_block_wrappers:
+                    wrapper.set_patch(BlockWrapperMode.CACHING)
+                    encoder_hidden_states, hidden_states = wrapper(
+                        hidden_states=hidden_states,
+                        encoder_hidden_states=encoder_hidden_states,
+                        text_embeddings=text_embeddings,
+                        rotary_embeddings=image_rotary_embeddings,
+                    )
+
+                if EXO_TRACING_ENABLED:
+                    mx.eval(encoder_hidden_states, hidden_states)
 
         if self.owns_concat_stage:
             assert encoder_hidden_states is not None
@@ -619,45 +639,63 @@ class DiffusionRunner:
             if self.has_single_blocks or self.is_last_stage:
                 hidden_states = concatenated
             else:
-                concatenated = mx.distributed.send(
-                    concatenated, self.next_rank, group=self.group
-                )
-                mx.async_eval(concatenated)
+                with trace(
+                    name=f"send {self.next_rank}", rank=self.rank, category="comms"
+                ):
+                    concatenated = mx.distributed.send(
+                        concatenated, self.next_rank, group=self.group
+                    )
+                    mx.async_eval(concatenated)
 
         elif self.has_joint_blocks and not self.is_last_stage:
             assert encoder_hidden_states is not None
-            hidden_states = mx.distributed.send(
-                hidden_states, self.next_rank, group=self.group
-            )
-            encoder_hidden_states = mx.distributed.send(
-                encoder_hidden_states, self.next_rank, group=self.group
-            )
-            mx.async_eval(hidden_states, encoder_hidden_states)
-
-        if self.has_single_blocks:
-            if not self.owns_concat_stage and not self.is_first_stage:
-                hidden_states = mx.distributed.recv(
-                    (batch_size, text_seq_len + num_img_tokens, hidden_dim),
-                    dtype,
-                    self.prev_rank,
-                    group=self.group,
-                )
-                mx.eval(hidden_states)
-
-            assert self.single_block_wrappers is not None
-            for wrapper in self.single_block_wrappers:
-                wrapper.set_patch(BlockWrapperMode.CACHING)
-                hidden_states = wrapper(
-                    hidden_states=hidden_states,
-                    text_embeddings=text_embeddings,
-                    rotary_embeddings=image_rotary_embeddings,
-                )
-
-            if not self.is_last_stage:
+            with trace(name=f"send {self.next_rank}", rank=self.rank, category="comms"):
                 hidden_states = mx.distributed.send(
                     hidden_states, self.next_rank, group=self.group
                 )
-                mx.async_eval(hidden_states)
+                encoder_hidden_states = mx.distributed.send(
+                    encoder_hidden_states, self.next_rank, group=self.group
+                )
+                mx.async_eval(hidden_states, encoder_hidden_states)
+
+        if self.has_single_blocks:
+            if not self.owns_concat_stage and not self.is_first_stage:
+                with trace(
+                    name=f"recv {self.prev_rank}", rank=self.rank, category="comms"
+                ):
+                    hidden_states = mx.distributed.recv(
+                        (batch_size, text_seq_len + num_img_tokens, hidden_dim),
+                        dtype,
+                        self.prev_rank,
+                        group=self.group,
+                    )
+                    mx.eval(hidden_states)
+
+            assert self.single_block_wrappers is not None
+            with trace(
+                name="single blocks",
+                rank=self.rank,
+                category="compute",
+            ):
+                for wrapper in self.single_block_wrappers:
+                    wrapper.set_patch(BlockWrapperMode.CACHING)
+                    hidden_states = wrapper(
+                        hidden_states=hidden_states,
+                        text_embeddings=text_embeddings,
+                        rotary_embeddings=image_rotary_embeddings,
+                    )
+
+                if EXO_TRACING_ENABLED:
+                    mx.eval(hidden_states)
+
+            if not self.is_last_stage:
+                with trace(
+                    name=f"send {self.next_rank}", rank=self.rank, category="comms"
+                ):
+                    hidden_states = mx.distributed.send(
+                        hidden_states, self.next_rank, group=self.group
+                    )
+                    mx.async_eval(hidden_states)
 
         hidden_states = hidden_states[:, text_seq_len:, ...]
 
@@ -741,14 +779,20 @@ class DiffusionRunner:
             )
 
             if not self.is_first_stage:
-                hidden_states = mx.distributed.send(hidden_states, 0, group=self.group)
-                mx.async_eval(hidden_states)
+                with trace(name="send 0", rank=self.rank, category="comms"):
+                    hidden_states = mx.distributed.send(
+                        hidden_states, 0, group=self.group
+                    )
+                    mx.async_eval(hidden_states)
 
         elif self.is_first_stage:
-            hidden_states = mx.distributed.recv_like(
-                prev_latents, src=self.world_size - 1, group=self.group
-            )
-            mx.eval(hidden_states)
+            with trace(
+                name=f"recv {self.world_size - 1}", rank=self.rank, category="comms"
+            ):
+                hidden_states = mx.distributed.recv_like(
+                    prev_latents, src=self.world_size - 1, group=self.group
+                )
+                mx.eval(hidden_states)
 
         else:
             hidden_states = prev_latents
@@ -808,10 +852,13 @@ class DiffusionRunner:
                 and not self.is_last_stage
                 and not is_first_async_step
             ):
-                patch = mx.distributed.recv_like(
-                    patch, src=self.prev_rank, group=self.group
-                )
-                mx.eval(patch)
+                with trace(
+                    name=f"recv {self.prev_rank}", rank=self.rank, category="comms"
+                ):
+                    patch = mx.distributed.recv_like(
+                        patch, src=self.prev_rank, group=self.group
+                    )
+                    mx.eval(patch)
 
             step_patch = mx.concatenate([patch, patch], axis=0) if needs_cfg else patch
 
@@ -842,10 +889,13 @@ class DiffusionRunner:
                 )
 
                 if not self.is_first_stage and t != config.num_inference_steps - 1:
-                    patch_latents[patch_idx] = mx.distributed.send(
-                        patch_latents[patch_idx], self.next_rank, group=self.group
-                    )
-                    mx.async_eval(patch_latents[patch_idx])
+                    with trace(
+                        name=f"send {self.next_rank}", rank=self.rank, category="comms"
+                    ):
+                        patch_latents[patch_idx] = mx.distributed.send(
+                            patch_latents[patch_idx], self.next_rank, group=self.group
+                        )
+                        mx.async_eval(patch_latents[patch_idx])
 
         return mx.concatenate(patch_latents, axis=1)
 
@@ -884,22 +934,28 @@ class DiffusionRunner:
         if self.has_joint_blocks:
             if not self.is_first_stage:
                 patch_len = patch.shape[1]
-                patch = mx.distributed.recv(
-                    (batch_size, patch_len, hidden_dim),
-                    patch.dtype,
-                    self.prev_rank,
-                    group=self.group,
-                )
-                mx.eval(patch)
-
-                if patch_idx == 0:
-                    encoder_hidden_states = mx.distributed.recv(
-                        (batch_size, text_seq_len, hidden_dim),
+                with trace(
+                    name=f"recv {self.prev_rank}", rank=self.rank, category="comms"
+                ):
+                    patch = mx.distributed.recv(
+                        (batch_size, patch_len, hidden_dim),
                         patch.dtype,
                         self.prev_rank,
                         group=self.group,
                     )
-                    mx.eval(encoder_hidden_states)
+                    mx.eval(patch)
+
+                if patch_idx == 0:
+                    with trace(
+                        name=f"recv {self.prev_rank}", rank=self.rank, category="comms"
+                    ):
+                        encoder_hidden_states = mx.distributed.recv(
+                            (batch_size, text_seq_len, hidden_dim),
+                            patch.dtype,
+                            self.prev_rank,
+                            group=self.group,
+                        )
+                        mx.eval(encoder_hidden_states)
 
             if self.is_first_stage:
                 patch, encoder_hidden_states = self.adapter.compute_embeddings(
@@ -908,14 +964,22 @@ class DiffusionRunner:
 
             assert self.joint_block_wrappers is not None
             assert encoder_hidden_states is not None
-            for wrapper in self.joint_block_wrappers:
-                wrapper.set_patch(BlockWrapperMode.PATCHED, start_token, end_token)
-                encoder_hidden_states, patch = wrapper(
-                    hidden_states=patch,
-                    encoder_hidden_states=encoder_hidden_states,
-                    text_embeddings=text_embeddings,
-                    rotary_embeddings=image_rotary_embeddings,
-                )
+            with trace(
+                name=f"joint patch {patch_idx}",
+                rank=self.rank,
+                category="compute",
+            ):
+                for wrapper in self.joint_block_wrappers:
+                    wrapper.set_patch(BlockWrapperMode.PATCHED, start_token, end_token)
+                    encoder_hidden_states, patch = wrapper(
+                        hidden_states=patch,
+                        encoder_hidden_states=encoder_hidden_states,
+                        text_embeddings=text_embeddings,
+                        rotary_embeddings=image_rotary_embeddings,
+                    )
+
+                if EXO_TRACING_ENABLED:
+                    mx.eval(encoder_hidden_states, patch)
 
         if self.owns_concat_stage:
             assert encoder_hidden_states is not None
@@ -924,49 +988,70 @@ class DiffusionRunner:
             if self.has_single_blocks or self.is_last_stage:
                 patch = patch_concat
             else:
-                patch_concat = mx.distributed.send(
-                    patch_concat, self.next_rank, group=self.group
-                )
-                mx.async_eval(patch_concat)
+                with trace(
+                    name=f"send {self.next_rank}", rank=self.rank, category="comms"
+                ):
+                    patch_concat = mx.distributed.send(
+                        patch_concat, self.next_rank, group=self.group
+                    )
+                    mx.async_eval(patch_concat)
 
         elif self.has_joint_blocks and not self.is_last_stage:
-            patch = mx.distributed.send(patch, self.next_rank, group=self.group)
-            mx.async_eval(patch)
+            with trace(name=f"send {self.next_rank}", rank=self.rank, category="comms"):
+                patch = mx.distributed.send(patch, self.next_rank, group=self.group)
+                mx.async_eval(patch)
 
             if patch_idx == 0:
                 assert encoder_hidden_states is not None
-                encoder_hidden_states = mx.distributed.send(
-                    encoder_hidden_states, self.next_rank, group=self.group
-                )
-                mx.async_eval(encoder_hidden_states)
+                with trace(
+                    name=f"send {self.next_rank}", rank=self.rank, category="comms"
+                ):
+                    encoder_hidden_states = mx.distributed.send(
+                        encoder_hidden_states, self.next_rank, group=self.group
+                    )
+                    mx.async_eval(encoder_hidden_states)
 
         if self.has_single_blocks:
             if not self.owns_concat_stage and not self.is_first_stage:
                 patch_len = patch.shape[1]
-                patch = mx.distributed.recv(
-                    (batch_size, text_seq_len + patch_len, hidden_dim),
-                    patch.dtype,
-                    self.prev_rank,
-                    group=self.group,
-                )
-                mx.eval(patch)
+                with trace(
+                    name=f"recv {self.prev_rank}", rank=self.rank, category="comms"
+                ):
+                    patch = mx.distributed.recv(
+                        (batch_size, text_seq_len + patch_len, hidden_dim),
+                        patch.dtype,
+                        self.prev_rank,
+                        group=self.group,
+                    )
+                    mx.eval(patch)
 
             assert self.single_block_wrappers is not None
-            for wrapper in self.single_block_wrappers:
-                wrapper.set_patch(BlockWrapperMode.PATCHED, start_token, end_token)
-                patch = wrapper(
-                    hidden_states=patch,
-                    text_embeddings=text_embeddings,
-                    rotary_embeddings=image_rotary_embeddings,
-                )
+            with trace(
+                name=f"single patch {patch_idx}",
+                rank=self.rank,
+                category="compute",
+            ):
+                for wrapper in self.single_block_wrappers:
+                    wrapper.set_patch(BlockWrapperMode.PATCHED, start_token, end_token)
+                    patch = wrapper(
+                        hidden_states=patch,
+                        text_embeddings=text_embeddings,
+                        rotary_embeddings=image_rotary_embeddings,
+                    )
+
+                if EXO_TRACING_ENABLED:
+                    mx.eval(patch)
 
             if not self.is_last_stage:
-                patch = mx.distributed.send(patch, self.next_rank, group=self.group)
-                mx.async_eval(patch)
+                with trace(
+                    name=f"send {self.next_rank}", rank=self.rank, category="comms"
+                ):
+                    patch = mx.distributed.send(patch, self.next_rank, group=self.group)
+                    mx.async_eval(patch)
 
         noise: mx.array | None = None
         if self.is_last_stage:
-            patch_img_only = patch[:, text_seq_len:, :]
-            noise = self.adapter.final_projection(patch_img_only, text_embeddings)
+            patch = patch[:, text_seq_len:, :]
+            noise = self.adapter.final_projection(patch, text_embeddings)
 
         return noise, encoder_hidden_states

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

@@ -19,8 +19,11 @@ from mlx_lm.models.deepseek_v32 import DeepseekV32MLP
 from mlx_lm.models.deepseek_v32 import Model as DeepseekV32Model
 from mlx_lm.models.glm4_moe import Model as Glm4MoeModel
 from mlx_lm.models.glm4_moe import MoE
+from mlx_lm.models.glm4_moe_lite import Glm4MoeLiteDecoderLayer, Glm4MoeLiteMLP
+from mlx_lm.models.glm4_moe_lite import Model as GLM4MoeLiteModel
 from mlx_lm.models.gpt_oss import GptOssMoeModel
 from mlx_lm.models.gpt_oss import Model as GptOssModel
+from mlx_lm.models.kimi_k25 import Model as KimiK25Model
 from mlx_lm.models.llama import Model as LlamaModel
 from mlx_lm.models.minimax import Model as MiniMaxModel
 from mlx_lm.models.ministral3 import Model as Ministral3Model
@@ -145,6 +148,10 @@ class PipelineLastLayer(CustomMlxLayer):
             if cache is not None:
                 cache.keys = mx.depends(cache.keys, output)  # type: ignore[reportUnknownMemberType]
 
+        output = mx.distributed.all_gather(output, group=self.group)[
+            -output.shape[0] :
+        ]  # type :ignore
+
         return output
 
 
@@ -194,6 +201,9 @@ def pipeline_auto_parallel(
     device_rank, world_size = model_shard_meta.device_rank, model_shard_meta.world_size
 
     layers = layers[start_layer:end_layer]
+    for layer in layers:
+        mx.eval(layer)  # type: ignore
+
     layers[0] = PipelineFirstLayer(layers[0], device_rank, group=group)
     layers[-1] = PipelineLastLayer(
         layers[-1],
@@ -252,10 +262,6 @@ def patch_pipeline_model[T](model: T, group: mx.distributed.Group) -> T:
         if cache is not None:
             cache[-1].state = mx.depends(cache[-1].state, logits)  # type: ignore
 
-        logits = mx.distributed.all_gather(logits, group=group)[
-            -logits.shape[0] :
-        ]  # type :ignore
-
         return logits
 
     cls.__call__ = patched_call
@@ -334,15 +340,7 @@ def tensor_auto_parallel(
         group=group,
     )
 
-    if hasattr(model, "shard") and not isinstance(model, GptOssModel):
-        try:
-            model.shard(group)  # type: ignore
-            return patch_tensor_model(model)
-        except (AttributeError, TypeError, NameError):
-            pass
-
     if isinstance(model, (LlamaModel, Ministral3Model)):
-        logger.warning("shouldn't be hit - upstream sharding exists")
         tensor_parallel_sharding_strategy = LlamaShardingStrategy(
             group,
             all_to_sharded_linear,
@@ -350,8 +348,7 @@ def tensor_auto_parallel(
             all_to_sharded_linear_in_place,
             sharded_to_all_linear_in_place,
         )
-    elif isinstance(model, (DeepseekV3Model, DeepseekV32Model)):
-        logger.warning("shouldn't be hit - upstream sharding exists")
+    elif isinstance(model, (DeepseekV3Model, DeepseekV32Model, KimiK25Model)):
         tensor_parallel_sharding_strategy = DeepSeekShardingStrategy(
             group,
             all_to_sharded_linear,
@@ -367,6 +364,14 @@ def tensor_auto_parallel(
             all_to_sharded_linear_in_place,
             sharded_to_all_linear_in_place,
         )
+    elif isinstance(model, GLM4MoeLiteModel):
+        tensor_parallel_sharding_strategy = GLM4MoeLiteShardingStrategy(
+            group,
+            all_to_sharded_linear,
+            sharded_to_all_linear,
+            all_to_sharded_linear_in_place,
+            sharded_to_all_linear_in_place,
+        )
     elif isinstance(model, (Qwen3MoeModel, Glm4MoeModel, Qwen3NextModel)):
         tensor_parallel_sharding_strategy = QwenShardingStrategy(
             group,
@@ -441,7 +446,7 @@ class LlamaShardingStrategy(TensorParallelShardingStrategy):
             layer.mlp.gate_proj = self.all_to_sharded_linear(layer.mlp.gate_proj)
             layer.mlp.down_proj = self.sharded_to_all_linear(layer.mlp.down_proj)
             layer.mlp.up_proj = self.all_to_sharded_linear(layer.mlp.up_proj)
-
+            mx.eval(layer)
         return model
 
 
@@ -452,7 +457,7 @@ def _set_layers(model: nn.Module, layers: list[_LayerCallable]) -> None:
 
         # Update DeepSeek V3 specific parameters when layers are shrunk
         if isinstance(
-            model, (DeepseekV3Model, DeepseekV32Model, Glm4MoeModel)
+            model, (DeepseekV3Model, DeepseekV32Model, Glm4MoeModel, KimiK25Model)
         ) and hasattr(inner_model_instance, "num_layers"):
             logger.info(
                 f"Setting num_layers to {len(layers)} for model {model.model.__class__.__name__}"
@@ -516,6 +521,8 @@ class DeepSeekShardingStrategy(TensorParallelShardingStrategy):
                 layer.mlp = ShardedDeepseekV3MoE(layer.mlp)  # type: ignore
                 layer.mlp.sharding_group = self.group
 
+            mx.eval(layer)
+
         return model
 
 
@@ -533,6 +540,84 @@ class ShardedDeepseekV3MoE(CustomMlxLayer):
         return y
 
 
+class GLM4MoeLiteShardingStrategy(TensorParallelShardingStrategy):
+    def shard_model(
+        self,
+        model: nn.Module,
+        timeout_seconds: float,
+        on_timeout: TimeoutCallback | None,
+    ) -> nn.Module:
+        model = cast(GLM4MoeLiteModel, model)
+        for layer in model.layers:  # type: ignore
+            layer = cast(Glm4MoeLiteDecoderLayer, layer)
+            eval_with_timeout(
+                layer.parameters(),
+                timeout_seconds / len(model.layers),  # type: ignore
+                on_timeout,
+            )
+            if layer.self_attn.q_lora_rank is None:  # type: ignore
+                layer.self_attn.q_proj = self.all_to_sharded_linear(
+                    layer.self_attn.q_proj
+                )
+            else:
+                layer.self_attn.q_b_proj = self.all_to_sharded_linear(
+                    layer.self_attn.q_b_proj
+                )
+
+            layer.self_attn.o_proj = self.sharded_to_all_linear(layer.self_attn.o_proj)
+            layer.self_attn.num_heads //= self.N
+
+            # Logic from upstream mlx
+            num_heads = layer.self_attn.num_heads
+            sh = self.group.rank() * num_heads
+            eh = sh + num_heads
+
+            def shard_heads(w: mx.array, sh: int = sh, eh: int = eh) -> mx.array:
+                return w[sh:eh]
+
+            layer.self_attn.embed_q.apply(shard_heads)
+            layer.self_attn.unembed_out.apply(shard_heads)
+
+            if isinstance(layer.mlp, Glm4MoeLiteMLP):
+                layer.mlp.gate_proj = self.all_to_sharded_linear(layer.mlp.gate_proj)
+                layer.mlp.down_proj = self.sharded_to_all_linear(layer.mlp.down_proj)
+                layer.mlp.up_proj = self.all_to_sharded_linear(layer.mlp.up_proj)
+
+            else:
+                if getattr(layer.mlp, "shared_experts", None) is not None:
+                    self.all_to_sharded_linear_in_place(
+                        layer.mlp.shared_experts.gate_proj
+                    )
+                    self.sharded_to_all_linear_in_place(
+                        layer.mlp.shared_experts.down_proj
+                    )
+                    self.all_to_sharded_linear_in_place(
+                        layer.mlp.shared_experts.up_proj
+                    )
+                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)
+                layer.mlp = ShardedGLM4MoeLiteMoE(layer.mlp)  # type: ignore
+                layer.mlp.sharding_group = self.group  # type: ignore
+            mx.eval(layer)
+
+        return model
+
+
+class ShardedGLM4MoeLiteMoE(CustomMlxLayer):
+    def __init__(self, layer: _LayerCallable):
+        super().__init__(layer)
+        self.sharding_group: mx.distributed.Group | None = None
+
+    def __call__(self, x: mx.array) -> mx.array:
+        if self.sharding_group is not None:
+            x = sum_gradients(self.sharding_group)(x)
+        y = self.original_layer.__call__(x)
+        if self.sharding_group is not None:
+            y = mx.distributed.all_sum(y, group=self.sharding_group)
+        return y
+
+
 class MiniMaxShardingStrategy(TensorParallelShardingStrategy):
     def shard_model(
         self,
@@ -541,6 +626,7 @@ class MiniMaxShardingStrategy(TensorParallelShardingStrategy):
         on_timeout: TimeoutCallback | None,
     ) -> nn.Module:
         model = cast(MiniMaxModel, model)
+        rank = self.group.rank()
         for layer in model.layers:
             eval_with_timeout(
                 layer.parameters(), timeout_seconds / len(model.layers), on_timeout
@@ -550,6 +636,16 @@ class MiniMaxShardingStrategy(TensorParallelShardingStrategy):
             layer.self_attn.k_proj = self.all_to_sharded_linear(layer.self_attn.k_proj)
             layer.self_attn.v_proj = self.all_to_sharded_linear(layer.self_attn.v_proj)
             layer.self_attn.o_proj = self.sharded_to_all_linear(layer.self_attn.o_proj)
+
+            # Shard qk_norm weights if present (must match sharded head count)
+            if getattr(layer.self_attn, "use_qk_norm", False):
+                layer.self_attn.q_norm.weight = layer.self_attn.q_norm.weight.split(  # type: ignore
+                    self.N, axis=-1
+                )[rank]
+                layer.self_attn.k_norm.weight = layer.self_attn.k_norm.weight.split(  # type: ignore
+                    self.N, axis=-1
+                )[rank]
+
             layer.self_attn.num_attention_heads //= self.N
             layer.self_attn.num_key_value_heads //= self.N
 
@@ -566,7 +662,7 @@ class MiniMaxShardingStrategy(TensorParallelShardingStrategy):
             )
             layer.block_sparse_moe = ShardedQwenMoE(layer.block_sparse_moe)  # pyright: ignore[reportAttributeAccessIssue, reportArgumentType]
             layer.block_sparse_moe.sharding_group = self.group  # pyright: ignore[reportAttributeAccessIssue]
-
+            mx.eval(layer)
         return model
 
 
@@ -607,6 +703,7 @@ class QwenShardingStrategy(TensorParallelShardingStrategy):
                 layer.mlp.down_proj = self.sharded_to_all_linear(layer.mlp.down_proj)
                 layer.mlp.up_proj = self.all_to_sharded_linear(layer.mlp.up_proj)
 
+            mx.eval(layer)
         return model
 
 
@@ -661,7 +758,7 @@ class GptOssShardingStrategy(TensorParallelShardingStrategy):
 
             layer.mlp = ShardedGptOssMoE(layer.mlp)  # type: ignore
             layer.mlp.sharding_group = self.group  # pyright: ignore[reportAttributeAccessIssue]
-
+            mx.eval(layer)
         return model
 
 

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

@@ -1,104 +1,234 @@
-# type: ignore
-# TODO: Fix this file, including types!
+import os
 from copy import deepcopy
-from typing import Callable
+from typing import Any, cast
 
 import mlx.core as mx
-from mlx_lm import stream_generate
-from mlx_lm.models.cache import _BaseCache, trim_prompt_cache
+import psutil
+from mlx_lm.models.cache import (
+    KVCache,
+    QuantizedKVCache,
+    RotatingKVCache,
+    trim_prompt_cache,
+)
+from mlx_lm.models.gpt_oss import Model as GptOssModel
 from mlx_lm.tokenizer_utils import TokenizerWrapper
 
+from exo.shared.types.memory import Memory
+from exo.shared.types.mlx import KVCacheType
 from exo.worker.engines.mlx import Model
-from exo.worker.engines.mlx.constants import KEEP_KV_SIZE, KV_BITS, KV_GROUP_SIZE
-from exo.worker.engines.mlx.utils_mlx import make_kv_cache
+from exo.worker.engines.mlx.constants import CACHE_GROUP_SIZE, KV_CACHE_BITS
+from exo.worker.runner.bootstrap import logger
+
+# Fraction of device memory above which LRU eviction kicks in
+_DEFAULT_MEMORY_THRESHOLD = 0.9
+_MEMORY_THRESHOLD = float(
+    os.environ.get("EXO_MEMORY_THRESHOLD", _DEFAULT_MEMORY_THRESHOLD)
+)
 
 
 class KVPrefixCache:
-    def __init__(self):
-        # Only one prefix cache per runner.
-        self.prompts: list[mx.array] = []  # mx array of tokens (ints)
-        self.caches: list[list[_BaseCache]] = []
-
-    def add_kv_cache(
-        self, tokenizer: TokenizerWrapper, prompt: str, cache: list[_BaseCache]
+    def __init__(
+        self, tokenizer: TokenizerWrapper, group: mx.distributed.Group | None = None
     ):
-        tokenized_prompt = self.encode_prompt(tokenizer, prompt)
+        self.prompts: list[mx.array] = []  # mx array of tokens (ints)
+        self.caches: list[KVCacheType] = []
+        self._last_used: list[int] = []  # monotonic counter of last access per entry
+        self._access_counter: int = 0
+        self._tokenizer: TokenizerWrapper = tokenizer
+        self._group = group
+
+    def clear(self):
+        """Clear all cached prompts and caches."""
+        self.prompts.clear()
+        self.caches.clear()
+        self._last_used.clear()
+
+    def add_kv_cache(self, prompt: str, cache: KVCacheType):
+        """Add a new cache entry. Evicts LRU entries if memory is high."""
+        self._evict_if_needed()
+        tokenized_prompt = encode_prompt(self._tokenizer, prompt)
         self.prompts.append(tokenized_prompt)
         self.caches.append(deepcopy(cache))
+        self._access_counter += 1
+        self._last_used.append(self._access_counter)
+        logger.info(f"KV cache added: {len(tokenized_prompt)} tokens")
+
+    def update_kv_cache(
+        self,
+        index: int,
+        prompt: str,
+        cache: KVCacheType,
+    ):
+        """Update an existing cache entry in-place."""
+        tokenized_prompt = encode_prompt(self._tokenizer, prompt)
+        self.prompts[index] = tokenized_prompt
+        self.caches[index] = deepcopy(cache)
+        self._access_counter += 1
+        self._last_used[index] = self._access_counter
+        logger.info(f"KV cache updated (index {index}): {len(tokenized_prompt)} tokens")
 
     def get_kv_cache(
         self,
         model: Model,
-        tokenizer: TokenizerWrapper,
-        sampler: Callable[[mx.array], mx.array],
         prompt: str,
-    ) -> list[_BaseCache]:
-        tokenized_prompt = self.encode_prompt(tokenizer, prompt)
+    ) -> tuple[KVCacheType, mx.array, int | None]:
+        """Get KV cache for prompt, returning remaining tokens to prefill.
+
+        Returns:
+            Tuple of (cache, remaining_tokens, matched_index) where:
+            - cache: KV cache to use for generation
+            - remaining_tokens: tokens that still need prefilling
+            - matched_index: index of the matched entry (None if no match)
+        """
+        tokenized_prompt = encode_prompt(self._tokenizer, prompt)
         max_length = len(tokenized_prompt)
 
         best_snapshot_index, best_snapshot_length = None, 0
 
         for i, cached_prompt in enumerate(self.prompts):
-            length = _get_prefix_length(tokenized_prompt, cached_prompt)
+            length = get_prefix_length(tokenized_prompt, cached_prompt)
 
             if length == max_length:
-                return self.caches[i]
+                # Exact match - cached prompt starts with our entire prompt
+                # Trim cache to prompt length - 1, return last token for stream_generate
+                prompt_cache = deepcopy(self.caches[i])
+                cached_length = cache_length(self.caches[i])
+                tokens_to_trim = cached_length - (max_length - 1)
+                if tokens_to_trim > 0:
+                    trim_prompt_cache(cast(list[Any], prompt_cache), tokens_to_trim)
+                self._access_counter += 1
+                self._last_used[i] = self._access_counter
+                logger.info(f"KV cache exact match: {max_length} tokens (instant)")
+                return prompt_cache, tokenized_prompt[-1:], i
 
             if length > best_snapshot_length:
                 best_snapshot_index, best_snapshot_length = i, length
 
         if best_snapshot_index is not None:
-            prompt_cache = deepcopy(self.caches[best_snapshot_index])
-            trim_prompt_cache(prompt_cache, max_length - best_snapshot_length)
-            tokenized_prompt = tokenized_prompt[best_snapshot_index:]
-
-        else:
-            prompt_cache = make_kv_cache(
-                model,
-                # max_kv_size=MAX_KV_SIZE,
-                # keep=KEEP_KV_SIZE
+            new_tokens = max_length - best_snapshot_length
+            logger.info(
+                f"KV cache prefix match: {best_snapshot_length}/{max_length} tokens "
+                f"(reusing {best_snapshot_length}, need to prefill {new_tokens})"
             )
 
-        prefill(model, tokenizer, sampler, tokenized_prompt, prompt_cache)
+            prompt_cache = deepcopy(self.caches[best_snapshot_index])
 
-        return prompt_cache
+            # Trim removes tokens from the end, so we trim (cached_length - prefix_length) to keep the prefix
+            cached_length = cache_length(self.caches[best_snapshot_index])
+            tokens_to_trim = cached_length - best_snapshot_length
+            if tokens_to_trim > 0:
+                trim_prompt_cache(cast(list[Any], prompt_cache), tokens_to_trim)
 
-    def encode_prompt(self, tokenizer: TokenizerWrapper, prompt: str) -> mx.array:
-        add_special_tokens = tokenizer.bos_token is None or not prompt.startswith(
-            tokenizer.bos_token
+            self._access_counter += 1
+            self._last_used[best_snapshot_index] = self._access_counter
+            remaining_tokens = tokenized_prompt[best_snapshot_length:]
+            return prompt_cache, remaining_tokens, best_snapshot_index
+
+        else:
+            prompt_cache = make_kv_cache(model)
+            if len(self.prompts) == 0:
+                logger.info(f"KV cache empty, need to prefill {max_length} tokens")
+            else:
+                logger.info(
+                    f"KV cache no prefix match, need to prefill {max_length} tokens"
+                )
+
+            return prompt_cache, tokenized_prompt, None
+
+    def _evict_if_needed(self):
+        """Evict least recently used entries while memory usage is high."""
+        if len(self.caches) == 0:
+            return
+
+        # Evict LRU entries until below threshold or only one entry left
+        while (
+            len(self.caches) > 1
+            and self.get_memory_used_percentage() > _MEMORY_THRESHOLD
+        ):
+            lru_index = self._last_used.index(min(self._last_used))
+            evicted_tokens = len(self.prompts[lru_index])
+            self.prompts.pop(lru_index)
+            self.caches.pop(lru_index)
+            self._last_used.pop(lru_index)
+            logger.info(
+                f"KV cache evicted LRU entry ({evicted_tokens} tokens) due to memory usage"
+            )
+
+    def get_memory_used_percentage(self) -> float:
+        local_pressure: float = get_memory_used_percentage()
+
+        if self._group is None:
+            return local_pressure
+
+        all_pressure = mx.distributed.all_gather(
+            mx.array([local_pressure], dtype=mx.float32),
+            group=self._group,
         )
-        tokenized_prompt = tokenizer.encode(
-            prompt, add_special_tokens=add_special_tokens
-        )
-        return mx.array(tokenized_prompt)
+        # .item() evals.
+        max_pressure = float(mx.max(all_pressure).item())
+        return max_pressure
 
 
-def _get_prefix_length(prompt: mx.array, cached_prompt: mx.array) -> int:
-    n = min(int(prompt.shape[0]), int(cached_prompt.shape[0]), KEEP_KV_SIZE)
+def encode_prompt(tokenizer: TokenizerWrapper, prompt: str) -> mx.array:
+    """Encode a prompt string to token array.
+
+    For chat-templated prompts (which have their own structure markers like
+    <|im_user|>, <|im_middle|>, etc.), we should NOT add BOS/EOS tokens as
+    that would corrupt the prompt structure.
+    """
+    # Chat templates define their own structure - don't add BOS/EOS
+    tokenized_prompt = tokenizer.encode(prompt, add_special_tokens=False)
+    return mx.array(tokenized_prompt)
+
+
+def cache_length(cache: KVCacheType) -> int:
+    """Get the number of tokens in a KV cache."""
+    # Use .offset attribute which all cache types have (len() not implemented in older QuantizedKVCache)
+    return max(c.offset for c in cache)  # type: ignore
+
+
+def get_prefix_length(prompt: mx.array, cached_prompt: mx.array) -> int:
+    """Find the length of the common prefix between two token arrays."""
+    n = min(int(prompt.shape[0]), int(cached_prompt.shape[0]))
     if n == 0:
         return 0
 
-    equal = (prompt[:n] == cached_prompt[:n]).astype(mx.int32)
+    equal = mx.equal(prompt[:n], cached_prompt[:n]).astype(mx.int32)
     prefix_mask = mx.cumprod(equal)  # stays 1 until first mismatch, then 0 forever
     return int(mx.sum(prefix_mask).item())
 
 
-def prefill(
-    model: Model,
-    tokenizer: TokenizerWrapper,
-    sampler: Callable[[mx.array], mx.array],
-    prompt: mx.array,
-    cache: list[_BaseCache],
-) -> None:
-    for _ in stream_generate(
-        model=model,
-        tokenizer=tokenizer,
-        prompt=prompt,
-        max_tokens=0,
-        sampler=sampler,
-        prompt_cache=cache,
-        prefill_step_size=2048,
-        kv_group_size=KV_GROUP_SIZE,
-        kv_bits=KV_BITS,
-    ):
-        pass
+def get_available_memory() -> Memory:
+    mem: int = psutil.virtual_memory().available
+    return Memory.from_bytes(mem)
+
+
+def get_memory_used_percentage() -> float:
+    mem = psutil.virtual_memory()
+    # percent is 0-100
+    return float(mem.percent / 100)
+
+
+def make_kv_cache(
+    model: Model, max_kv_size: int | None = None, keep: int = 0
+) -> KVCacheType:
+    assert hasattr(model, "layers")
+
+    # TODO: Do this for all models
+    if hasattr(model, "make_cache") and isinstance(model, GptOssModel):
+        logger.info("Using MLX LM's make cache")
+        return model.make_cache()  # type: ignore
+
+    if max_kv_size is None:
+        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)
+                for _ in model.layers
+            ]
+    else:
+        logger.info(f"Using rotating KV cache with {max_kv_size=} with {keep=}")
+        return [RotatingKVCache(max_size=max_kv_size, keep=keep) for _ in model.layers]

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

@@ -4,7 +4,7 @@
 KV_GROUP_SIZE: int | None = 32
 KV_BITS: int | None = None
 ATTENTION_KV_BITS: int | None = 4
-MAX_TOKENS: int = 8192
+MAX_TOKENS: int = 32168
 MAX_KV_SIZE: int | None = 3200
 KEEP_KV_SIZE: int | None = 1600
 QUANTIZE_MODEL_MODE: str | None = "affine"

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

@@ -1,48 +1,93 @@
+import time
 from typing import Any, Callable, Generator, cast, get_args
 
 import mlx.core as mx
 from mlx_lm.generate import stream_generate
-from mlx_lm.models.cache import KVCache
+from mlx_lm.models.cache import trim_prompt_cache
 from mlx_lm.sample_utils import make_sampler
 from mlx_lm.tokenizer_utils import TokenizerWrapper
 
-# from exo.engines.mlx.cache import KVPrefixCache
 from exo.shared.types.api import (
-    BenchChatCompletionTaskParams,
-    ChatCompletionMessage,
+    CompletionTokensDetails,
     FinishReason,
     GenerationStats,
+    PromptTokensDetails,
+    Usage,
 )
+from exo.shared.types.common import ModelId
 from exo.shared.types.memory import Memory
-from exo.shared.types.tasks import ChatCompletionTaskParams
+from exo.shared.types.mlx import KVCacheType
+from exo.shared.types.text_generation import TextGenerationTaskParams
 from exo.shared.types.worker.runner_response import (
     GenerationResponse,
 )
 from exo.worker.engines.mlx import Model
+from exo.worker.engines.mlx.cache import KVPrefixCache, encode_prompt, make_kv_cache
 from exo.worker.engines.mlx.constants import KV_BITS, KV_GROUP_SIZE, MAX_TOKENS
 from exo.worker.engines.mlx.utils_mlx import (
     apply_chat_template,
-    make_kv_cache,
     mx_barrier,
 )
 from exo.worker.runner.bootstrap import logger
 
 generation_stream = mx.new_stream(mx.default_device())
 
+_MIN_PREFIX_HIT_TO_UPDATE = 1000
 
-def maybe_quantize_kv_cache(
-    prompt_cache: list[KVCache | Any],
-    quantized_kv_start: int,
-    kv_group_size: int,
-    kv_bits: int | None,
-) -> None:
-    if kv_bits is None:
-        return
-    for e, c in enumerate(prompt_cache):
-        if (
-            hasattr(c, "to_quantized") and c.offset >= quantized_kv_start  # type: ignore
-        ):
-            prompt_cache[e] = c.to_quantized(group_size=kv_group_size, bits=kv_bits)
+
+def prefill(
+    model: Model,
+    tokenizer: TokenizerWrapper,
+    sampler: Callable[[mx.array], mx.array],
+    prompt_tokens: mx.array,
+    cache: KVCacheType,
+) -> tuple[float, int]:
+    """Prefill the KV cache with prompt tokens.
+
+    This runs the model over the prompt tokens to populate the cache,
+    then trims off the extra generated token.
+
+    Returns:
+        tokens_per_sec
+    """
+    num_tokens = len(prompt_tokens)
+    if num_tokens == 0:
+        return 0.0, 0
+
+    logger.debug(f"Prefilling {num_tokens} tokens...")
+    start_time = time.perf_counter()
+
+    def progress_callback(processed: int, total: int) -> None:
+        elapsed = time.time() - start_time
+        tok_per_sec = processed / elapsed if elapsed > 0 else 0
+        logger.debug(
+            f"Prefill progress: {processed}/{total} tokens ({tok_per_sec:.1f} tok/s)"
+        )
+
+    # 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=2048,
+        kv_group_size=KV_GROUP_SIZE,
+        kv_bits=KV_BITS,
+        prompt_progress_callback=progress_callback,
+    ):
+        break  # Stop after first iteration - cache is now filled
+    trim_prompt_cache(cast(list[Any], cache), 1)
+
+    elapsed = time.perf_counter() - start_time
+    tokens_per_sec = num_tokens / elapsed if elapsed > 0 else 0.0
+    logger.debug(
+        f"Prefill complete: {num_tokens} tokens in {elapsed:.2f}s "
+        f"({tokens_per_sec:.1f} tok/s)"
+    )
+    return tokens_per_sec, num_tokens
 
 
 def warmup_inference(
@@ -53,14 +98,9 @@ def warmup_inference(
 
     warmup_prompt = apply_chat_template(
         tokenizer=tokenizer,
-        chat_task_data=ChatCompletionTaskParams(
-            model="",
-            messages=[
-                ChatCompletionMessage(
-                    role="user",
-                    content=content,
-                )
-            ],
+        task_params=TextGenerationTaskParams(
+            model=ModelId(""),
+            input=content,
         ),
     )
 
@@ -118,20 +158,32 @@ def eos_ids_from_tokenizer(tokenizer: TokenizerWrapper) -> list[int]:
 def mlx_generate(
     model: Model,
     tokenizer: TokenizerWrapper,
-    task: ChatCompletionTaskParams,
+    task: TextGenerationTaskParams,
     prompt: str,
+    kv_prefix_cache: KVPrefixCache | None = None,
 ) -> Generator[GenerationResponse]:
     # Ensure that generation stats only contains peak memory for this generation
     mx.reset_peak_memory()
-    is_bench: bool = isinstance(task, BenchChatCompletionTaskParams)
-
-    # Currently we support chat-completion tasks only.
-    logger.debug(f"task_params: {task}")
-
     if task.seed is not None:
         mx.random.seed(task.seed)
 
-    caches = make_kv_cache(model=model)
+    # Do not use the prefix cache if we are trying to do benchmarks.
+    is_bench = task.bench
+    if is_bench:
+        kv_prefix_cache = None
+
+    # Use prefix cache if available, otherwise create fresh cache
+    prefix_hit_length = 0
+    matched_index: int | None = None
+    if kv_prefix_cache is None:
+        caches = make_kv_cache(model=model)
+        prompt_tokens = encode_prompt(tokenizer, prompt)
+    else:
+        caches, prompt_tokens, matched_index = kv_prefix_cache.get_kv_cache(
+            model, prompt
+        )
+        all_prompt_tokens = encode_prompt(tokenizer, prompt)
+        prefix_hit_length = len(all_prompt_tokens) - len(prompt_tokens)
 
     logits_processors: list[Callable[[mx.array, mx.array], mx.array]] = []
     if is_bench:
@@ -142,49 +194,141 @@ def mlx_generate(
     sampler = make_sampler(
         temp=task.temperature if task.temperature is not None else 0.7,
         top_p=task.top_p if task.top_p is not None else 1.0,
+        top_k=task.top_k if task.top_k is not None else 0,
     )
 
-    max_tokens = task.max_tokens or MAX_TOKENS
-    for out in stream_generate(
-        model=model,
-        tokenizer=tokenizer,
-        prompt=prompt,
-        max_tokens=max_tokens,
-        sampler=sampler,
-        logits_processors=logits_processors,
-        prompt_cache=caches,
-        # TODO: Dynamically change prefill step size to be the maximum possible without timing out.
-        prefill_step_size=2048,
-        kv_group_size=KV_GROUP_SIZE,
-        kv_bits=KV_BITS,
+    # Normalize stop sequences to a list
+    stop_sequences: list[str] = (
+        ([task.stop] if isinstance(task.stop, str) else task.stop)
+        if task.stop is not None
+        else []
+    )
+    max_stop_len = max((len(s) for s in stop_sequences), default=0)
+
+    # Prefill cache with all tokens except the last one
+    prefill_tps, prefill_tokens = prefill(
+        model, tokenizer, sampler, prompt_tokens[:-1], caches
+    )
+
+    # stream_generate starts from the last token
+    last_token = prompt_tokens[-1:]
+
+    max_tokens = task.max_output_tokens or MAX_TOKENS
+    accumulated_text = ""
+    generated_text_parts: list[str] = []
+    generation_start_time = time.perf_counter()
+    usage: Usage | None = None
+    in_thinking = False
+    reasoning_tokens = 0
+    think_start = tokenizer.think_start
+    think_end = tokenizer.think_end
+    for completion_tokens, out in enumerate(
+        stream_generate(
+            model=model,
+            tokenizer=tokenizer,
+            prompt=last_token,
+            max_tokens=max_tokens,
+            sampler=sampler,
+            logits_processors=logits_processors,
+            prompt_cache=caches,
+            # TODO: Dynamically change prefill step size to be the maximum possible without timing out.
+            prefill_step_size=2048,
+            kv_group_size=KV_GROUP_SIZE,
+            kv_bits=KV_BITS,
+        ),
+        start=1,
     ):
+        generated_text_parts.append(out.text)
         logger.info(out.text)
+        accumulated_text += out.text
+
+        if think_start is not None and out.text == think_start:
+            in_thinking = True
+        elif think_end is not None and out.text == think_end:
+            in_thinking = False
+        if in_thinking:
+            reasoning_tokens += 1
+
+        # Check for stop sequences
+        text = out.text
+        finish_reason: FinishReason | None = cast(
+            FinishReason | None, out.finish_reason
+        )
+        stop_matched = False
+
+        if stop_sequences:
+            for stop_seq in stop_sequences:
+                if stop_seq in accumulated_text:
+                    # Trim text to just before the stop sequence
+                    stop_index = accumulated_text.find(stop_seq)
+                    text_before_stop = accumulated_text[:stop_index]
+                    chunk_start = len(accumulated_text) - len(out.text)
+                    text = text_before_stop[chunk_start:]
+                    finish_reason = "stop"
+                    stop_matched = True
+                    break
+
+        is_done = finish_reason is not None
 
         stats: GenerationStats | None = None
-        if out.finish_reason is not None:
+        if is_done:
             stats = GenerationStats(
-                prompt_tps=float(out.prompt_tps),
+                prompt_tps=float(prefill_tps or out.prompt_tps),
                 generation_tps=float(out.generation_tps),
-                prompt_tokens=int(out.prompt_tokens),
+                prompt_tokens=int(prefill_tokens + out.prompt_tokens),
                 generation_tokens=int(out.generation_tokens),
                 peak_memory_usage=Memory.from_gb(out.peak_memory),
             )
-
-            if out.finish_reason not in get_args(FinishReason):
-                # We don't throw here as this failure case is really not all that bad
-                # Just log the error and move on
+            if not stop_matched and out.finish_reason not in get_args(FinishReason):
                 logger.warning(
                     f"Model generated unexpected finish_reason: {out.finish_reason}"
                 )
 
+            usage = Usage(
+                prompt_tokens=int(out.prompt_tokens),
+                completion_tokens=completion_tokens,
+                total_tokens=int(out.prompt_tokens) + completion_tokens,
+                prompt_tokens_details=PromptTokensDetails(
+                    cached_tokens=prefix_hit_length
+                ),
+                completion_tokens_details=CompletionTokensDetails(
+                    reasoning_tokens=reasoning_tokens
+                ),
+            )
+
         yield GenerationResponse(
-            text=out.text,
+            text=text,
             token=out.token,
-            finish_reason=cast(FinishReason | None, out.finish_reason),
+            finish_reason=finish_reason,
             stats=stats,
+            usage=usage,
         )
 
-        if out.finish_reason is not None:
+        if is_done:
+            # Log generation stats
+            generation_elapsed = time.perf_counter() - generation_start_time
+            generated_tokens = len(generated_text_parts)
+            generation_tps = (
+                generated_tokens / generation_elapsed if generation_elapsed > 0 else 0.0
+            )
+            logger.debug(
+                f"Generation complete: prefill {prompt_tokens} tokens @ "
+                f"{prefill_tps:.1f} tok/s, generated {generated_tokens} tokens @ "
+                f"{generation_tps:.1f} tok/s"
+            )
+            if kv_prefix_cache is not None:
+                full_prompt = prompt + "".join(generated_text_parts)
+                if (
+                    matched_index is not None
+                    and prefix_hit_length >= _MIN_PREFIX_HIT_TO_UPDATE
+                ):
+                    kv_prefix_cache.update_kv_cache(matched_index, full_prompt, caches)
+                else:
+                    kv_prefix_cache.add_kv_cache(full_prompt, caches)
             break
 
+        # Limit accumulated_text to what's needed for stop sequence detection
+        if max_stop_len > 0 and len(accumulated_text) > max_stop_len:
+            accumulated_text = accumulated_text[-max_stop_len:]
+
         # TODO: Do we want an mx_barrier?

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

@@ -18,15 +18,12 @@ try:
 except ImportError:
     pass  # transformers < 5.0 or bytes_to_unicode not available
 
-from mlx_lm.models.cache import KVCache, QuantizedKVCache, RotatingKVCache
+from mlx_lm.models.cache import KVCache
 from mlx_lm.models.deepseek_v3 import DeepseekV3Model
-from mlx_lm.models.gpt_oss import Model as GptOssModel
 from mlx_lm.tokenizer_utils import TokenizerWrapper
 
 from exo.shared.models.model_cards import ModelId
 from exo.worker.engines.mlx.constants import (
-    CACHE_GROUP_SIZE,
-    KV_CACHE_BITS,
     TRUST_REMOTE_CODE,
 )
 
@@ -41,10 +38,10 @@ import mlx.nn as nn
 from mlx_lm.utils import load_model
 from pydantic import RootModel
 
-from exo.shared.types.api import ChatCompletionMessageText
+from exo.download.download_utils import build_model_path
 from exo.shared.types.common import Host
 from exo.shared.types.memory import Memory
-from exo.shared.types.tasks import ChatCompletionTaskParams
+from exo.shared.types.text_generation import TextGenerationTaskParams
 from exo.shared.types.worker.instances import (
     BoundInstance,
     MlxJacclInstance,
@@ -55,7 +52,6 @@ from exo.shared.types.worker.shards import (
     ShardMetadata,
     TensorShardMetadata,
 )
-from exo.worker.download.download_utils import build_model_path
 from exo.worker.engines.mlx import Model
 from exo.worker.engines.mlx.auto_parallel import (
     TimeoutCallback,
@@ -168,12 +164,11 @@ def mlx_distributed_init(
 
                 jaccl_coordinator = jaccl_coordinators[bound_instance.bound_node_id]
 
-                # TODO: update once upstream fixes
                 logger.info(
-                    f"rank {rank} MLX_JACCL_DEVICES: {coordination_file} with devices: {jaccl_devices_json}"
+                    f"rank {rank} MLX_IBV_DEVICES: {coordination_file} with devices: {jaccl_devices_json}"
                 )
                 logger.info(f"rank {rank} MLX_JACCL_COORDINATOR: {jaccl_coordinator}")
-                os.environ["MLX_JACCL_DEVICES"] = coordination_file
+                os.environ["MLX_IBV_DEVICES"] = coordination_file
                 os.environ["MLX_RANK"] = str(rank)
                 os.environ["MLX_JACCL_COORDINATOR"] = jaccl_coordinator
                 group = mx.distributed.init(backend="jaccl", strict=True)
@@ -262,10 +257,10 @@ def shard_and_load(
 
     logger.info(f"Group size: {group.size()}, group rank: {group.rank()}")
 
-    # Estimate timeout based on model size
-    base_timeout = float(os.environ.get("EXO_MODEL_LOAD_TIMEOUT", "60"))
+    # Estimate timeout based on model size (5x default for large queued workloads)
+    base_timeout = float(os.environ.get("EXO_MODEL_LOAD_TIMEOUT", "300"))
     model_size_gb = get_weights_size(shard_metadata).in_bytes / (1024**3)
-    timeout_seconds = base_timeout + model_size_gb / 5
+    timeout_seconds = base_timeout + model_size_gb
     logger.info(
         f"Evaluating model parameters with timeout of {timeout_seconds:.0f}s "
         f"(model size: {model_size_gb:.1f}GB)"
@@ -342,8 +337,35 @@ def load_tokenizer_for_model_id(
 
     # Kimi uses a custom TikTokenTokenizer that transformers 5.x can't load via AutoTokenizer
     if "kimi-k2" in model_id_lower:
+        import importlib.util
+        import types
+
         sys.path.insert(0, str(model_path))
-        from tokenization_kimi import TikTokenTokenizer  # type: ignore[import-not-found]  # noqa: I001
+
+        # Load tool_declaration_ts first (tokenization_kimi imports it with relative import)
+        tool_decl_path = model_path / "tool_declaration_ts.py"
+        if tool_decl_path.exists():
+            spec = importlib.util.spec_from_file_location(
+                "tool_declaration_ts", tool_decl_path
+            )
+            if spec and spec.loader:
+                tool_decl_module = importlib.util.module_from_spec(spec)
+                sys.modules["tool_declaration_ts"] = tool_decl_module
+                spec.loader.exec_module(tool_decl_module)
+
+        # Load tokenization_kimi with patched source (convert relative to absolute import)
+        tok_path = model_path / "tokenization_kimi.py"
+        source = tok_path.read_text()
+        source = source.replace("from .tool_declaration_ts", "from tool_declaration_ts")
+        spec = importlib.util.spec_from_file_location("tokenization_kimi", tok_path)
+        if spec:
+            tok_module = types.ModuleType("tokenization_kimi")
+            tok_module.__file__ = str(tok_path)
+            sys.modules["tokenization_kimi"] = tok_module
+            exec(compile(source, tok_path, "exec"), tok_module.__dict__)  # noqa: S102
+            TikTokenTokenizer = tok_module.TikTokenTokenizer  # type: ignore[attr-defined]  # noqa: N806
+        else:
+            from tokenization_kimi import TikTokenTokenizer  # type: ignore[import-not-found]  # noqa: I001
 
         hf_tokenizer: Any = TikTokenTokenizer.from_pretrained(model_path)  # pyright: ignore[reportUnknownVariableType,reportUnknownMemberType]
 
@@ -365,36 +387,71 @@ def load_tokenizer_for_model_id(
     return tokenizer
 
 
+def _normalize_tool_calls(msg_dict: dict[str, Any]) -> None:
+    """Normalize tool_calls in a message dict.
+
+    OpenAI format has tool_calls[].function.arguments as a JSON string,
+    but some chat templates (e.g., GLM) expect it as a dict.
+    """
+    tool_calls = msg_dict.get("tool_calls")
+    if not tool_calls or not isinstance(tool_calls, list):
+        return
+
+    for tc in tool_calls:  # pyright: ignore[reportUnknownVariableType]
+        if not isinstance(tc, dict):
+            continue
+        func = tc.get("function")  # pyright: ignore[reportUnknownMemberType,reportUnknownVariableType]
+        if not isinstance(func, dict):
+            continue
+        args = func.get("arguments")  # pyright: ignore[reportUnknownMemberType,reportUnknownVariableType]
+        if isinstance(args, str):
+            with contextlib.suppress(json.JSONDecodeError):
+                func["arguments"] = json.loads(args)
+
+
 def apply_chat_template(
     tokenizer: TokenizerWrapper,
-    chat_task_data: ChatCompletionTaskParams,
+    task_params: TextGenerationTaskParams,
 ) -> str:
-    # Now we can properly access the messages
-    messages = chat_task_data.messages
+    """Convert TextGenerationTaskParams to a chat template prompt.
 
+    Converts the internal format (input + instructions) to a messages list
+    that can be processed by the tokenizer's chat template.
+
+    When chat_template_messages is available (from Chat Completions API),
+    uses those directly to preserve tool_calls, thinking, and other fields.
+    Otherwise builds messages from the task params input/instructions.
+    """
     formatted_messages: list[dict[str, Any]] = []
-    for message in messages:
-        if isinstance(message.content, ChatCompletionMessageText):
-            message.content = message.content.text
-        if isinstance(message.content, list):
-            if len(message.content) == 0:
-                logger.warning("Received prompt with no content, skipping")
-                continue
+    if task_params.chat_template_messages is not None:
+        # Use pre-formatted messages that preserve tool_calls, thinking, etc.
+        formatted_messages = list(task_params.chat_template_messages)
+        for msg in formatted_messages:
+            _normalize_tool_calls(msg)
+    else:
+        # Add system message (instructions) if present
+        if task_params.instructions:
+            formatted_messages.append(
+                {"role": "system", "content": task_params.instructions}
+            )
 
-            message.content = "\n".join(c.text for c in message.content).strip()
-        if message.content is None and message.thinking is None:
-            continue
-
-        # Null values are not valid when applying templates in tokenizer
-        formatted_messages.append(
-            {k: v for k, v in message.model_dump().items() if v is not None}  # type: ignore
-        )
+        # Convert input to messages
+        if isinstance(task_params.input, str):
+            # Simple string input becomes a single user message
+            formatted_messages.append({"role": "user", "content": task_params.input})
+        else:
+            # List of InputMessage
+            for msg in task_params.input:
+                if not msg.content:
+                    logger.warning("Received message with empty content, skipping")
+                    continue
+                formatted_messages.append({"role": msg.role, "content": msg.content})
 
     prompt: str = tokenizer.apply_chat_template(
         formatted_messages,
         tokenize=False,
         add_generation_prompt=True,
-        tools=chat_task_data.tools,
+        tools=task_params.tools,
     )
 
     logger.info(prompt)
@@ -435,31 +492,6 @@ class NullKVCache(KVCache):
         raise NotImplementedError("We should not be setting a NullKVCache.")
 
 
-def make_kv_cache(
-    model: Model, max_kv_size: int | None = None, keep: int = 0
-) -> list[KVCache | RotatingKVCache | QuantizedKVCache]:
-    assert hasattr(model, "layers")
-
-    # TODO: Do this for all models
-    if hasattr(model, "make_cache") and isinstance(model, GptOssModel):
-        logger.info("Using MLX LM's make cache")
-        return model.make_cache()  # type: ignore
-
-    if max_kv_size is None:
-        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)
-                for _ in model.layers
-            ]
-    else:
-        logger.info(f"Using rotating KV cache with {max_kv_size=} with {keep=}")
-        return [RotatingKVCache(max_size=max_kv_size, keep=keep) for _ in model.layers]
-
-
 def mlx_force_oom(size: int = 40000) -> None:
     """
     Force an Out-Of-Memory (OOM) error in MLX by performing large tensor operations.

src/exo/worker/main.py

@@ -1,16 +1,21 @@
 from datetime import datetime, timezone
 from random import random
+from typing import Iterator
 
 import anyio
-from anyio import CancelScope, create_task_group, current_time, fail_after
+from anyio import CancelScope, create_task_group, fail_after
 from anyio.abc import TaskGroup
 from loguru import logger
 
-from exo.routing.connection_message import ConnectionMessage, ConnectionMessageType
 from exo.shared.apply import apply
 from exo.shared.models.model_cards import ModelId
-from exo.shared.types.api import ImageEditsInternalParams
-from exo.shared.types.commands import ForwarderCommand, RequestEventLog
+from exo.shared.types.api import ImageEditsTaskParams
+from exo.shared.types.commands import (
+    ForwarderCommand,
+    ForwarderDownloadCommand,
+    RequestEventLog,
+    StartDownload,
+)
 from exo.shared.types.common import CommandId, NodeId, SessionId
 from exo.shared.types.events import (
     Event,
@@ -18,7 +23,6 @@ from exo.shared.types.events import (
     ForwarderEvent,
     IndexedEvent,
     InputChunkReceived,
-    NodeDownloadProgress,
     NodeGatheredInfo,
     TaskCreated,
     TaskStatusUpdated,
@@ -36,22 +40,12 @@ from exo.shared.types.tasks import (
     TaskStatus,
 )
 from exo.shared.types.topology import Connection, SocketConnection
-from exo.shared.types.worker.downloads import (
-    DownloadCompleted,
-    DownloadOngoing,
-    DownloadPending,
-    DownloadProgress,
-)
 from exo.shared.types.worker.runners import RunnerId
-from exo.shared.types.worker.shards import ShardMetadata
 from exo.utils.channels import Receiver, Sender, channel
 from exo.utils.event_buffer import OrderedBuffer
 from exo.utils.info_gatherer.info_gatherer import GatheredInfo, InfoGatherer
 from exo.utils.info_gatherer.net_profile import check_reachable
-from exo.worker.download.download_utils import (
-    map_repo_download_progress_to_download_progress_data,
-)
-from exo.worker.download.shard_downloader import RepoDownloadProgress, ShardDownloader
+from exo.utils.keyed_backoff import KeyedBackoff
 from exo.worker.plan import plan
 from exo.worker.runner.runner_supervisor import RunnerSupervisor
 
@@ -61,31 +55,27 @@ class Worker:
         self,
         node_id: NodeId,
         session_id: SessionId,
-        shard_downloader: ShardDownloader,
         *,
-        connection_message_receiver: Receiver[ConnectionMessage],
         global_event_receiver: Receiver[ForwarderEvent],
         local_event_sender: Sender[ForwarderEvent],
         # This is for requesting updates. It doesn't need to be a general command sender right now,
         # but I think it's the correct way to be thinking about commands
         command_sender: Sender[ForwarderCommand],
+        download_command_sender: Sender[ForwarderDownloadCommand],
+        event_index_counter: Iterator[int],
     ):
         self.node_id: NodeId = node_id
         self.session_id: SessionId = session_id
 
-        self.shard_downloader: ShardDownloader = shard_downloader
-        self._pending_downloads: dict[RunnerId, ShardMetadata] = {}
-
         self.global_event_receiver = global_event_receiver
         self.local_event_sender = local_event_sender
-        self.local_event_index = 0
+        self.event_index_counter = event_index_counter
         self.command_sender = command_sender
-        self.connection_message_receiver = connection_message_receiver
+        self.download_command_sender = download_command_sender
         self.event_buffer = OrderedBuffer[Event]()
         self.out_for_delivery: dict[EventId, ForwarderEvent] = {}
 
         self.state: State = State()
-        self.download_status: dict[ModelId, DownloadProgress] = {}
         self.runners: dict[RunnerId, RunnerSupervisor] = {}
         self._tg: TaskGroup = create_task_group()
 
@@ -100,6 +90,8 @@ class Worker:
         self.input_chunk_buffer: dict[CommandId, dict[int, str]] = {}
         self.input_chunk_counts: dict[CommandId, int] = {}
 
+        self._download_backoff: KeyedBackoff[ModelId] = KeyedBackoff(base=0.5, cap=10.0)
+
     async def run(self):
         logger.info("Starting Worker")
 
@@ -110,8 +102,6 @@ class Worker:
             tg.start_soon(info_gatherer.run)
             tg.start_soon(self._forward_info, info_recv)
             tg.start_soon(self.plan_step)
-            tg.start_soon(self._emit_existing_download_progress)
-            tg.start_soon(self._connection_message_event_writer)
             tg.start_soon(self._resend_out_for_delivery)
             tg.start_soon(self._event_applier)
             tg.start_soon(self._forward_events)
@@ -120,6 +110,7 @@ class Worker:
         # Actual shutdown code - waits for all tasks to complete before executing.
         self.local_event_sender.close()
         self.command_sender.close()
+        self.download_command_sender.close()
         for runner in self.runners.values():
             runner.shutdown()
 
@@ -178,11 +169,9 @@ class Worker:
     async def plan_step(self):
         while True:
             await anyio.sleep(0.1)
-            # 3. based on the updated state, we plan & execute an operation.
             task: Task | None = plan(
                 self.node_id,
                 self.runners,
-                self.download_status,
                 self.state.downloads,
                 self.state.instances,
                 self.state.runners,
@@ -206,42 +195,26 @@ class Worker:
                         )
                     )
                 case DownloadModel(shard_metadata=shard):
-                    if shard.model_card.model_id not in self.download_status:
-                        progress = DownloadPending(
-                            shard_metadata=shard, node_id=self.node_id
-                        )
-                        self.download_status[shard.model_card.model_id] = progress
-                        await self.event_sender.send(
-                            NodeDownloadProgress(download_progress=progress)
-                        )
-                    initial_progress = (
-                        await self.shard_downloader.get_shard_download_status_for_shard(
-                            shard
+                    model_id = shard.model_card.model_id
+                    if not self._download_backoff.should_proceed(model_id):
+                        continue
+
+                    self._download_backoff.record_attempt(model_id)
+
+                    await self.download_command_sender.send(
+                        ForwarderDownloadCommand(
+                            origin=self.node_id,
+                            command=StartDownload(
+                                target_node_id=self.node_id,
+                                shard_metadata=shard,
+                            ),
                         )
                     )
-                    if initial_progress.status == "complete":
-                        progress = DownloadCompleted(
-                            shard_metadata=shard,
-                            node_id=self.node_id,
-                            total_bytes=initial_progress.total_bytes,
+                    await self.event_sender.send(
+                        TaskStatusUpdated(
+                            task_id=task.task_id, task_status=TaskStatus.Running
                         )
-                        self.download_status[shard.model_card.model_id] = progress
-                        await self.event_sender.send(
-                            NodeDownloadProgress(download_progress=progress)
-                        )
-                        await self.event_sender.send(
-                            TaskStatusUpdated(
-                                task_id=task.task_id,
-                                task_status=TaskStatus.Complete,
-                            )
-                        )
-                    else:
-                        await self.event_sender.send(
-                            TaskStatusUpdated(
-                                task_id=task.task_id, task_status=TaskStatus.Running
-                            )
-                        )
-                        self._handle_shard_download_process(task, initial_progress)
+                    )
                 case Shutdown(runner_id=runner_id):
                     try:
                         with fail_after(3):
@@ -267,7 +240,7 @@ class Worker:
                         command_id=task.command_id,
                         instance_id=task.instance_id,
                         task_status=task.task_status,
-                        task_params=ImageEditsInternalParams(
+                        task_params=ImageEditsTaskParams(
                             image_data=assembled,
                             total_input_chunks=task.task_params.total_input_chunks,
                             prompt=task.task_params.prompt,
@@ -302,41 +275,6 @@ class Worker:
         instance = self.state.instances[task.instance_id]
         return instance.shard_assignments.node_to_runner[self.node_id]
 
-    async def _connection_message_event_writer(self):
-        with self.connection_message_receiver as connection_messages:
-            async for msg in connection_messages:
-                await self.event_sender.send(
-                    self._convert_connection_message_to_event(msg)
-                )
-
-    def _convert_connection_message_to_event(self, msg: ConnectionMessage):
-        match msg.connection_type:
-            case ConnectionMessageType.Connected:
-                return TopologyEdgeCreated(
-                    conn=Connection(
-                        source=self.node_id,
-                        sink=msg.node_id,
-                        edge=SocketConnection(
-                            sink_multiaddr=Multiaddr(
-                                address=f"/ip4/{msg.remote_ipv4}/tcp/{msg.remote_tcp_port}"
-                            ),
-                        ),
-                    ),
-                )
-
-            case ConnectionMessageType.Disconnected:
-                return TopologyEdgeDeleted(
-                    conn=Connection(
-                        source=self.node_id,
-                        sink=msg.node_id,
-                        edge=SocketConnection(
-                            sink_multiaddr=Multiaddr(
-                                address=f"/ip4/{msg.remote_ipv4}/tcp/{msg.remote_tcp_port}"
-                            ),
-                        ),
-                    ),
-                )
-
     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.
@@ -386,78 +324,17 @@ class Worker:
         self._tg.start_soon(runner.run)
         return runner
 
-    def _handle_shard_download_process(
-        self,
-        task: DownloadModel,
-        initial_progress: RepoDownloadProgress,
-    ):
-        """Manages the shard download process with progress tracking."""
-        status = DownloadOngoing(
-            node_id=self.node_id,
-            shard_metadata=task.shard_metadata,
-            download_progress=map_repo_download_progress_to_download_progress_data(
-                initial_progress
-            ),
-        )
-        self.download_status[task.shard_metadata.model_card.model_id] = status
-        self.event_sender.send_nowait(NodeDownloadProgress(download_progress=status))
-
-        last_progress_time = 0.0
-        throttle_interval_secs = 1.0
-
-        async def download_progress_callback(
-            shard: ShardMetadata, progress: RepoDownloadProgress
-        ) -> None:
-            nonlocal self
-            nonlocal last_progress_time
-            if progress.status == "complete":
-                status = DownloadCompleted(
-                    shard_metadata=shard,
-                    node_id=self.node_id,
-                    total_bytes=progress.total_bytes,
-                )
-                self.download_status[shard.model_card.model_id] = status
-                await self.event_sender.send(
-                    NodeDownloadProgress(download_progress=status)
-                )
-                await self.event_sender.send(
-                    TaskStatusUpdated(
-                        task_id=task.task_id, task_status=TaskStatus.Complete
-                    )
-                )
-            elif (
-                progress.status == "in_progress"
-                and current_time() - last_progress_time > throttle_interval_secs
-            ):
-                status = DownloadOngoing(
-                    node_id=self.node_id,
-                    shard_metadata=shard,
-                    download_progress=map_repo_download_progress_to_download_progress_data(
-                        progress
-                    ),
-                )
-                self.download_status[shard.model_card.model_id] = status
-                await self.event_sender.send(
-                    NodeDownloadProgress(download_progress=status)
-                )
-                last_progress_time = current_time()
-
-        self.shard_downloader.on_progress(download_progress_callback)
-        self._tg.start_soon(self.shard_downloader.ensure_shard, task.shard_metadata)
-
     async def _forward_events(self) -> None:
         with self.event_receiver as events:
             async for event in events:
+                idx = next(self.event_index_counter)
                 fe = ForwarderEvent(
-                    origin_idx=self.local_event_index,
+                    origin_idx=idx,
                     origin=self.node_id,
                     session=self.session_id,
                     event=event,
                 )
-                logger.debug(
-                    f"Worker published event {self.local_event_index}: {str(event)[:100]}"
-                )
-                self.local_event_index += 1
+                logger.debug(f"Worker published event {idx}: {str(event)[:100]}")
                 await self.local_event_sender.send(fe)
                 self.out_for_delivery[event.event_id] = fe
 
@@ -505,42 +382,3 @@ class Worker:
                     await self.event_sender.send(TopologyEdgeDeleted(conn=conn))
 
             await anyio.sleep(10)
-
-    async def _emit_existing_download_progress(self) -> None:
-        try:
-            while True:
-                logger.debug("Fetching and emitting existing download progress...")
-                async for (
-                    _,
-                    progress,
-                ) in self.shard_downloader.get_shard_download_status():
-                    if progress.status == "complete":
-                        status = DownloadCompleted(
-                            node_id=self.node_id,
-                            shard_metadata=progress.shard,
-                            total_bytes=progress.total_bytes,
-                        )
-                    elif progress.status in ["in_progress", "not_started"]:
-                        if progress.downloaded_bytes_this_session.in_bytes == 0:
-                            status = DownloadPending(
-                                node_id=self.node_id, shard_metadata=progress.shard
-                            )
-                        else:
-                            status = DownloadOngoing(
-                                node_id=self.node_id,
-                                shard_metadata=progress.shard,
-                                download_progress=map_repo_download_progress_to_download_progress_data(
-                                    progress
-                                ),
-                            )
-                    else:
-                        continue
-
-                    self.download_status[progress.shard.model_card.model_id] = status
-                    await self.event_sender.send(
-                        NodeDownloadProgress(download_progress=status)
-                    )
-                logger.debug("Done emitting existing download progress.")
-                await anyio.sleep(5 * 60)  # 5 minutes
-        except Exception as e:
-            logger.error(f"Error emitting existing download progress: {e}")

src/exo/worker/plan.py

@@ -2,10 +2,8 @@
 
 from collections.abc import Mapping, Sequence
 
-from exo.shared.models.model_cards import ModelId
 from exo.shared.types.common import CommandId, NodeId
 from exo.shared.types.tasks import (
-    ChatCompletion,
     ConnectToGroup,
     CreateRunner,
     DownloadModel,
@@ -17,9 +15,11 @@ from exo.shared.types.tasks import (
     Task,
     TaskId,
     TaskStatus,
+    TextGeneration,
 )
 from exo.shared.types.worker.downloads import (
     DownloadCompleted,
+    DownloadFailed,
     DownloadOngoing,
     DownloadProgress,
 )
@@ -44,9 +44,6 @@ def plan(
     node_id: NodeId,
     # Runners is expected to be FRESH and so should not come from state
     runners: Mapping[RunnerId, RunnerSupervisor],
-    # DL_status is expected to be FRESH and so should not come from state
-    download_status: Mapping[ModelId, DownloadProgress],
-    # gdls is not expected to be fresh
     global_download_status: Mapping[NodeId, Sequence[DownloadProgress]],
     instances: Mapping[InstanceId, Instance],
     all_runners: Mapping[RunnerId, RunnerStatus],  # all global
@@ -58,7 +55,7 @@ def plan(
     return (
         _kill_runner(runners, all_runners, instances)
         or _create_runner(node_id, runners, instances)
-        or _model_needs_download(runners, download_status)
+        or _model_needs_download(node_id, runners, global_download_status)
         or _init_distributed_backend(runners, all_runners)
         or _load_model(runners, all_runners, global_download_status)
         or _ready_to_warmup(runners, all_runners)
@@ -114,15 +111,22 @@ def _create_runner(
 
 
 def _model_needs_download(
+    node_id: NodeId,
     runners: Mapping[RunnerId, RunnerSupervisor],
-    download_status: Mapping[ModelId, DownloadProgress],
+    global_download_status: Mapping[NodeId, Sequence[DownloadProgress]],
 ) -> DownloadModel | None:
+    local_downloads = global_download_status.get(node_id, [])
+    download_status = {
+        dp.shard_metadata.model_card.model_id: dp for dp in local_downloads
+    }
+
     for runner in runners.values():
         model_id = runner.bound_instance.bound_shard.model_card.model_id
         if isinstance(runner.status, RunnerIdle) and (
             model_id not in download_status
             or not isinstance(
-                download_status[model_id], (DownloadOngoing, DownloadCompleted)
+                download_status[model_id],
+                (DownloadOngoing, DownloadCompleted, DownloadFailed),
             )
         ):
             # We don't invalidate download_status randomly in case a file gets deleted on disk
@@ -271,7 +275,7 @@ def _pending_tasks(
     for task in tasks.values():
         # for now, just forward chat completions
         # TODO(ciaran): do this better!
-        if not isinstance(task, (ChatCompletion, ImageGeneration, ImageEdits)):
+        if not isinstance(task, (TextGeneration, ImageGeneration, ImageEdits)):
             continue
         if task.task_status not in (TaskStatus.Pending, TaskStatus.Running):
             continue

src/exo/worker/runner/runner.py

@@ -1,8 +1,9 @@
 import base64
+import json
 import time
 from collections.abc import Generator
 from functools import cache
-from typing import Literal
+from typing import Any, Callable, Literal
 
 import mlx.core as mx
 from mlx_lm.models.gpt_oss import Model as GptOssModel
@@ -13,11 +14,13 @@ from openai_harmony import (  # pyright: ignore[reportMissingTypeStubs]
     StreamableParser,
     load_harmony_encoding,
 )
+from pydantic import ValidationError
 
-from exo.shared.constants import EXO_MAX_CHUNK_SIZE
+from exo.shared.constants import EXO_MAX_CHUNK_SIZE, EXO_TRACING_ENABLED
 from exo.shared.models.model_cards import ModelId, ModelTask
-from exo.shared.types.api import ChatCompletionMessageText, ImageGenerationStats
-from exo.shared.types.chunks import ImageChunk, TokenChunk
+from exo.shared.tracing import clear_trace_buffer, get_trace_buffer
+from exo.shared.types.api import ImageGenerationStats
+from exo.shared.types.chunks import ErrorChunk, ImageChunk, TokenChunk, ToolCallChunk
 from exo.shared.types.common import CommandId
 from exo.shared.types.events import (
     ChunkGenerated,
@@ -25,9 +28,10 @@ from exo.shared.types.events import (
     RunnerStatusUpdated,
     TaskAcknowledged,
     TaskStatusUpdated,
+    TraceEventData,
+    TracesCollected,
 )
 from exo.shared.types.tasks import (
-    ChatCompletion,
     ConnectToGroup,
     ImageEdits,
     ImageGeneration,
@@ -35,13 +39,18 @@ from exo.shared.types.tasks import (
     Shutdown,
     StartWarmup,
     Task,
+    TaskId,
     TaskStatus,
+    TextGeneration,
 )
+from exo.shared.types.text_generation import TextGenerationTaskParams
 from exo.shared.types.worker.instances import BoundInstance
 from exo.shared.types.worker.runner_response import (
     GenerationResponse,
     ImageGenerationResponse,
     PartialImageResponse,
+    ToolCallItem,
+    ToolCallResponse,
 )
 from exo.shared.types.worker.runners import (
     RunnerConnected,
@@ -66,6 +75,7 @@ from exo.worker.engines.image import (
     warmup_image_generator,
 )
 from exo.worker.engines.mlx import Model
+from exo.worker.engines.mlx.cache import KVPrefixCache
 from exo.worker.engines.mlx.generator.generate import mlx_generate, warmup_inference
 from exo.worker.engines.mlx.utils_mlx import (
     apply_chat_template,
@@ -99,14 +109,19 @@ def main(
     model: Model | DistributedImageModel | None = None
     tokenizer = None
     group = None
+    kv_prefix_cache: KVPrefixCache | None = None
 
     current_status: RunnerStatus = RunnerIdle()
     logger.info("runner created")
     event_sender.send(
         RunnerStatusUpdated(runner_id=runner_id, runner_status=current_status)
     )
+    seen = set[TaskId]()
     with task_receiver as tasks:
         for task in tasks:
+            if task.task_id in seen:
+                logger.warning("repeat task - potential error")
+            seen.add(task.task_id)
             event_sender.send(
                 TaskStatusUpdated(task_id=task.task_id, task_status=TaskStatus.Running)
             )
@@ -154,6 +169,11 @@ def main(
                         model, tokenizer = load_mlx_items(
                             bound_instance, group, on_timeout=on_model_load_timeout
                         )
+                        logger.info(
+                            f"model has_tool_calling={tokenizer.has_tool_calling}"
+                        )
+                        kv_prefix_cache = KVPrefixCache(tokenizer, group)
+
                     elif (
                         ModelTask.TextToImage in shard_metadata.model_card.tasks
                         or ModelTask.ImageToImage in shard_metadata.model_card.tasks
@@ -163,7 +183,6 @@ def main(
                         raise ValueError(
                             f"Unknown model task(s): {shard_metadata.model_card.tasks}"
                         )
-
                     current_status = RunnerLoaded()
                     logger.info("runner loaded")
                 case StartWarmup() if isinstance(current_status, RunnerLoaded):
@@ -204,7 +223,7 @@ def main(
 
                     current_status = RunnerReady()
                     logger.info("runner ready")
-                case ChatCompletion(task_params=task_params, command_id=command_id) if (
+                case TextGeneration(task_params=task_params, command_id=command_id) if (
                     isinstance(current_status, RunnerReady)
                 ):
                     logger.info(f"received chat request: {task}")
@@ -217,10 +236,9 @@ def main(
                     )
                     assert model and not isinstance(model, DistributedImageModel)
                     assert tokenizer
-                    assert task_params.messages[0].content is not None
 
                     try:
-                        _check_for_debug_prompts(task_params.messages[0].content)
+                        _check_for_debug_prompts(task_params)
 
                         # Build prompt once - used for both generation and thinking detection
                         prompt = apply_chat_template(tokenizer, task_params)
@@ -231,12 +249,9 @@ def main(
                             tokenizer=tokenizer,
                             task=task_params,
                             prompt=prompt,
+                            kv_prefix_cache=kv_prefix_cache,
                         )
 
-                        # GPT-OSS specific parsing to match other model formats.
-                        if isinstance(model, GptOssModel):
-                            mlx_generator = parse_gpt_oss(mlx_generator)
-
                         # For other thinking models (GLM, etc.), check if we need to
                         # prepend the thinking tag that was consumed by the chat template
                         if detect_thinking_prompt_suffix(prompt, tokenizer):
@@ -244,25 +259,82 @@ def main(
                                 mlx_generator, tokenizer
                             )
 
-                        # TODO: Add tool call parser here
+                        # Kimi-K2 has tool call sections - we don't care about them
+                        if "kimi" in shard_metadata.model_card.model_id.lower():
+                            mlx_generator = filter_kimi_tokens(mlx_generator)
+                            patch_kimi_tokenizer(tokenizer)
 
+                        # GLM models need patched parser (upstream has bug with None regex match)
+                        elif "glm" in shard_metadata.model_card.model_id.lower():
+                            patch_glm_tokenizer(tokenizer)
+
+                        # GPT-OSS specific parsing to match other model formats.
+                        elif isinstance(model, GptOssModel):
+                            mlx_generator = parse_gpt_oss(mlx_generator)
+
+                        if tokenizer.has_tool_calling and not isinstance(
+                            model, GptOssModel
+                        ):
+                            assert tokenizer.tool_call_start
+                            assert tokenizer.tool_call_end
+                            assert tokenizer.tool_parser  # pyright: ignore[reportAny]
+                            mlx_generator = parse_tool_calls(
+                                mlx_generator,
+                                tokenizer.tool_call_start,
+                                tokenizer.tool_call_end,
+                                tokenizer.tool_parser,  # pyright: ignore[reportAny]
+                            )
+
+                        completion_tokens = 0
                         for response in mlx_generator:
                             match response:
                                 case GenerationResponse():
-                                    if device_rank == 0:
+                                    completion_tokens += 1
+                                    if (
+                                        device_rank == 0
+                                        and response.finish_reason == "error"
+                                    ):
+                                        event_sender.send(
+                                            ChunkGenerated(
+                                                command_id=command_id,
+                                                chunk=ErrorChunk(
+                                                    error_message=response.text,
+                                                    model=shard_metadata.model_card.model_id,
+                                                ),
+                                            )
+                                        )
+
+                                    elif device_rank == 0:
+                                        assert response.finish_reason not in (
+                                            "error",
+                                            "tool_calls",
+                                            "function_call",
+                                        )
                                         event_sender.send(
                                             ChunkGenerated(
                                                 command_id=command_id,
                                                 chunk=TokenChunk(
-                                                    idx=response.token,
                                                     model=shard_metadata.model_card.model_id,
                                                     text=response.text,
                                                     token_id=response.token,
+                                                    usage=response.usage,
                                                     finish_reason=response.finish_reason,
                                                     stats=response.stats,
                                                 ),
                                             )
                                         )
+                                case ToolCallResponse():
+                                    if device_rank == 0:
+                                        event_sender.send(
+                                            ChunkGenerated(
+                                                command_id=command_id,
+                                                chunk=ToolCallChunk(
+                                                    tool_calls=response.tool_calls,
+                                                    model=shard_metadata.model_card.model_id,
+                                                    usage=response.usage,
+                                                ),
+                                            )
+                                        )
 
                     # can we make this more explicit?
                     except Exception as e:
@@ -270,11 +342,8 @@ def main(
                             event_sender.send(
                                 ChunkGenerated(
                                     command_id=command_id,
-                                    chunk=TokenChunk(
-                                        idx=0,
+                                    chunk=ErrorChunk(
                                         model=shard_metadata.model_card.model_id,
-                                        text="",
-                                        token_id=0,
                                         finish_reason="error",
                                         error_message=str(e),
                                     ),
@@ -328,24 +397,24 @@ def main(
                                             image_index,
                                         )
                                         image_index += 1
+                    # can we make this more explicit?
                     except Exception as e:
                         if shard_metadata.device_rank == shard_metadata.world_size - 1:
                             event_sender.send(
                                 ChunkGenerated(
                                     command_id=command_id,
-                                    chunk=ImageChunk(
-                                        idx=0,
+                                    chunk=ErrorChunk(
                                         model=shard_metadata.model_card.model_id,
-                                        data="",
-                                        chunk_index=0,
-                                        total_chunks=1,
-                                        image_index=0,
                                         finish_reason="error",
                                         error_message=str(e),
                                     ),
                                 )
                             )
                         raise
+                    finally:
+                        _send_traces_if_enabled(
+                            event_sender, task.task_id, shard_metadata.device_rank
+                        )
 
                     current_status = RunnerReady()
                     logger.info("runner ready")
@@ -396,19 +465,18 @@ def main(
                             event_sender.send(
                                 ChunkGenerated(
                                     command_id=command_id,
-                                    chunk=ImageChunk(
-                                        idx=0,
+                                    chunk=ErrorChunk(
                                         model=shard_metadata.model_card.model_id,
-                                        data="",
-                                        chunk_index=0,
-                                        total_chunks=1,
-                                        image_index=0,
                                         finish_reason="error",
                                         error_message=str(e),
                                     ),
                                 )
                             )
                         raise
+                    finally:
+                        _send_traces_if_enabled(
+                            event_sender, task.task_id, shard_metadata.device_rank
+                        )
 
                     current_status = RunnerReady()
                     logger.info("runner ready")
@@ -446,18 +514,58 @@ def get_gpt_oss_encoding():
     return encoding
 
 
-def parse_gpt_oss(
-    responses: Generator[GenerationResponse],
+def filter_kimi_tokens(
+    responses: Generator[GenerationResponse | ToolCallResponse],
 ) -> Generator[GenerationResponse]:
+    for resp in responses:
+        assert isinstance(resp, GenerationResponse)
+        if (
+            resp.text == "<|tool_calls_section_begin|>"
+            or resp.text == "<|tool_calls_section_end|>"
+        ):
+            continue
+        yield resp
+
+
+def parse_gpt_oss(
+    responses: Generator[GenerationResponse | ToolCallResponse],
+) -> Generator[GenerationResponse | ToolCallResponse]:
     encoding = get_gpt_oss_encoding()
     stream = StreamableParser(encoding, role=Role.ASSISTANT)
     thinking = False
+    current_tool_name: str | None = None
+    tool_arg_parts: list[str] = []
 
     for response in responses:
+        assert isinstance(response, GenerationResponse)
         stream.process(response.token)
 
         delta = stream.last_content_delta
         ch = stream.current_channel
+        recipient = stream.current_recipient
+
+        if recipient != current_tool_name:
+            if current_tool_name is not None:
+                prefix = "functions."
+                if current_tool_name.startswith(prefix):
+                    current_tool_name = current_tool_name[len(prefix) :]
+                yield ToolCallResponse(
+                    tool_calls=[
+                        ToolCallItem(
+                            name=current_tool_name,
+                            arguments="".join(tool_arg_parts).strip(),
+                        )
+                    ],
+                    usage=response.usage,
+                )
+                tool_arg_parts = []
+            current_tool_name = recipient
+
+        # If inside a tool call, accumulate arguments
+        if current_tool_name is not None:
+            if delta:
+                tool_arg_parts.append(delta)
+            continue
 
         if ch == "analysis" and not thinking:
             thinking = True
@@ -474,13 +582,12 @@ def parse_gpt_oss(
             if thinking:
                 yield response.model_copy(update={"text": "</think>"})
             yield response
-            break
 
 
 def parse_thinking_models(
-    responses: Generator[GenerationResponse],
+    responses: Generator[GenerationResponse | ToolCallResponse],
     tokenizer: TokenizerWrapper,
-) -> Generator[GenerationResponse]:
+) -> Generator[GenerationResponse | ToolCallResponse]:
     """
     For models that inject thinking tags in the prompt (like GLM-4.7),
     prepend the thinking tag to the output stream so the frontend
@@ -488,6 +595,9 @@ def parse_thinking_models(
     """
     first = True
     for response in responses:
+        if isinstance(response, ToolCallResponse):
+            yield response
+            continue
         if first:
             first = False
             yield response.model_copy(
@@ -526,7 +636,6 @@ def _send_image_chunk(
             ChunkGenerated(
                 command_id=command_id,
                 chunk=ImageChunk(
-                    idx=chunk_index,
                     model=model_id,
                     data=chunk_data,
                     chunk_index=chunk_index,
@@ -542,6 +651,36 @@ def _send_image_chunk(
         )
 
 
+def _send_traces_if_enabled(
+    event_sender: MpSender[Event],
+    task_id: TaskId,
+    rank: int,
+) -> None:
+    if not EXO_TRACING_ENABLED:
+        return
+
+    traces = get_trace_buffer()
+    if traces:
+        trace_data = [
+            TraceEventData(
+                name=t.name,
+                start_us=t.start_us,
+                duration_us=t.duration_us,
+                rank=t.rank,
+                category=t.category,
+            )
+            for t in traces
+        ]
+        event_sender.send(
+            TracesCollected(
+                task_id=task_id,
+                rank=rank,
+                traces=trace_data,
+            )
+        )
+    clear_trace_buffer()
+
+
 def _process_image_response(
     response: ImageGenerationResponse | PartialImageResponse,
     command_id: CommandId,
@@ -559,7 +698,7 @@ def _process_image_response(
         command_id=command_id,
         model_id=shard_metadata.model_card.model_id,
         event_sender=event_sender,
-        image_index=response.partial_index if is_partial else image_index,
+        image_index=response.image_index,
         is_partial=is_partial,
         partial_index=response.partial_index if is_partial else None,
         total_partials=response.total_partials if is_partial else None,
@@ -568,22 +707,229 @@ def _process_image_response(
     )
 
 
+def parse_tool_calls(
+    responses: Generator[GenerationResponse | ToolCallResponse],
+    tool_call_start: str,
+    tool_call_end: str,
+    tool_parser: Callable[[str], dict[str, Any] | list[dict[str, Any]]],
+) -> Generator[GenerationResponse | ToolCallResponse]:
+    in_tool_call = False
+    tool_call_text_parts: list[str] = []
+    for response in responses:
+        assert isinstance(response, GenerationResponse)
+        # assumption: the tool call start is one token
+        if response.text == tool_call_start:
+            in_tool_call = True
+            continue
+        # assumption: the tool call end is one token
+        if in_tool_call and response.text == tool_call_end:
+            try:
+                # tool_parser returns an arbitrarily nested python dictionary
+                # we actually don't want the python dictionary, we just want to
+                # parse the top level { function: ..., arguments: ... } structure
+                # as we're just gonna hand it back to the api anyway
+                parsed = tool_parser("".join(tool_call_text_parts).strip())
+                logger.info(f"parsed {tool_call_text_parts=} into {parsed=}")
+                if isinstance(parsed, list):
+                    tools = [_validate_single_tool(tool) for tool in parsed]
+                else:
+                    tools = [_validate_single_tool(parsed)]
+                yield ToolCallResponse(tool_calls=tools, usage=response.usage)
+
+            except (
+                json.JSONDecodeError,
+                ValidationError,
+                ValueError,
+                AttributeError,
+            ) as e:
+                # ValueError: our parsers raise this for malformed tool calls
+                # AttributeError: upstream parsers (e.g. glm47) may raise this when regex doesn't match
+                logger.opt(exception=e).warning("tool call parsing failed")
+                # assumption: talking about tool calls, not making a tool call
+                response.text = (
+                    tool_call_start + "".join(tool_call_text_parts) + tool_call_end
+                )
+                yield response
+
+            in_tool_call = False
+            tool_call_text_parts = []
+            continue
+
+        if in_tool_call:
+            tool_call_text_parts.append(response.text)
+            if response.finish_reason is not None:
+                logger.info(
+                    "toll call parsing interrupted, yield partial tool call as text"
+                )
+                yield GenerationResponse(
+                    text=tool_call_start + "".join(tool_call_text_parts),
+                    token=0,
+                    finish_reason=response.finish_reason,
+                    usage=None,
+                )
+            continue
+        # fallthrough
+        yield response
+
+
+def patch_kimi_tokenizer(tokenizer: TokenizerWrapper):
+    """
+    Version of to-be-upstreamed kimi-k2 tool parser
+    """
+    import ast
+    import json
+    from typing import Any
+
+    import regex as re
+
+    # kimi has a fixed function naming scheme, with a json formatted arg
+    #   functions.multiply:0 <|tool_call_argument_begin|> {"a": 2, "b": 3}
+    _func_name_regex = re.compile(
+        r"^\s*(.+):\d+\s*<\|tool_call_argument_begin\|>", re.DOTALL
+    )
+    _func_arg_regex = re.compile(r"<\|tool_call_argument_begin\|>\s*(.*)\s*", re.DOTALL)
+
+    # kimi has a tool_calls_section - we're leaving this up to the caller to handle
+    tool_call_start = "<|tool_call_begin|>"
+    tool_call_end = "<|tool_call_end|>"
+
+    def _deserialize(value: str) -> Any:  # pyright: ignore[reportAny]
+        try:
+            return json.loads(value)  # pyright: ignore[reportAny]
+        except Exception:
+            pass
+
+        try:
+            return ast.literal_eval(value)  # pyright: ignore[reportAny]
+        except Exception:
+            pass
+        return value
+
+    def parse_tool_call(text: str, tools: Any | None = None):
+        func_name_match = _func_name_regex.search(text)
+        if func_name_match is None:
+            raise ValueError(f"Could not parse function name from tool call: {text!r}")
+        func_name = func_name_match.group(1)
+        # strip off the `functions.` prefix, if it exists.
+        func_name = func_name[func_name.find(".") + 1 :]
+
+        func_args_match = _func_arg_regex.search(text)
+        if func_args_match is None:
+            raise ValueError(f"Could not parse function args from tool call: {text!r}")
+        func_args = func_args_match.group(1)
+        # the args should be valid json - no need to check against our tools to deserialize
+        arg_dct = _deserialize(func_args)  # pyright: ignore[reportAny]
+
+        return dict(name=func_name, arguments=arg_dct)  # pyright: ignore[reportAny]
+
+    tokenizer._tool_call_start = tool_call_start
+    tokenizer._tool_call_end = tool_call_end
+    tokenizer._tool_parser = parse_tool_call
+
+
+def patch_glm_tokenizer(tokenizer: TokenizerWrapper):
+    """
+    Fixed version of mlx_lm's glm47 tool parser that handles regex match failures.
+    """
+    import ast
+    import json
+    from typing import Any
+
+    import regex as re
+
+    _func_name_regex = re.compile(r"^(.*?)<arg_key>", re.DOTALL)
+    _func_arg_regex = re.compile(
+        r"<arg_key>(.*?)</arg_key>(?:\n|\s)*<arg_value>(.*?)(?:</arg_value>|(?=<arg_key>)|$)",
+        re.DOTALL,
+    )
+
+    tool_call_start = "<tool_call>"
+    tool_call_end = "</tool_call>"
+
+    def _is_string_type(
+        tool_name: str,
+        arg_name: str,
+        tools: list[Any] | None,
+    ) -> bool:
+        if tools is None:
+            return False
+        for tool in tools:  # pyright: ignore[reportAny]
+            func = tool["function"]  # pyright: ignore[reportAny]
+            if func["name"] == tool_name:
+                params = func["parameters"]  # pyright: ignore[reportAny]
+                if params is None:
+                    return False
+                props = params.get("properties", {})  # pyright: ignore[reportAny]
+                arg_props = props.get(arg_name, {})  # pyright: ignore[reportAny]
+                arg_type = arg_props.get("type", None)  # pyright: ignore[reportAny]
+                return arg_type == "string"  # pyright: ignore[reportAny]
+        return False
+
+    def _deserialize(value: str) -> Any:  # pyright: ignore[reportAny]
+        try:
+            return json.loads(value)  # pyright: ignore[reportAny]
+        except Exception:
+            pass
+        try:
+            return ast.literal_eval(value)  # pyright: ignore[reportAny]
+        except Exception:
+            pass
+        return value
+
+    def parse_tool_call(text: str, tools: list[Any] | None = None):
+        func_name_match = _func_name_regex.search(text)
+        if func_name_match is None:
+            raise ValueError(f"Could not parse function name from tool call: {text!r}")
+        func_name = func_name_match.group(1)
+
+        pairs = _func_arg_regex.findall(text)
+        arg_dct: dict[str, Any] = {}
+        for key, value in pairs:  # pyright: ignore[reportAny]
+            arg_key = key.strip()  # pyright: ignore[reportAny]
+            arg_val = value.strip()  # pyright: ignore[reportAny]
+            if not _is_string_type(func_name, arg_key, tools):  # pyright: ignore[reportAny]
+                arg_val = _deserialize(arg_val)  # pyright: ignore[reportAny]
+            arg_dct[arg_key] = arg_val
+        return dict(name=func_name, arguments=arg_dct)
+
+    tokenizer._tool_call_start = tool_call_start
+    tokenizer._tool_call_end = tool_call_end
+    tokenizer._tool_parser = parse_tool_call
+
+
+def _validate_single_tool(obj: dict[str, Any]) -> ToolCallItem:
+    if (
+        ((name := obj.get("name")) is not None)
+        and ((args := obj.get("arguments")) is not None)
+        and isinstance(name, str)
+    ):
+        return ToolCallItem(name=name, arguments=json.dumps(args))
+    else:
+        raise ValidationError
+
+
 EXO_RUNNER_MUST_FAIL = "EXO RUNNER MUST FAIL"
 EXO_RUNNER_MUST_OOM = "EXO RUNNER MUST OOM"
 EXO_RUNNER_MUST_TIMEOUT = "EXO RUNNER MUST TIMEOUT"
 
 
-def _check_for_debug_prompts(
-    prompt: str | ChatCompletionMessageText | list[ChatCompletionMessageText],
-):
-    if isinstance(prompt, list):
-        if len(prompt) == 0:
-            logger.debug("Empty message prompt received in debug prompt")
-            return
-        prompt = prompt[0]
+def _check_for_debug_prompts(task_params: TextGenerationTaskParams) -> None:
+    """Check for debug prompt triggers in the input.
 
-    if isinstance(prompt, ChatCompletionMessageText):
-        prompt = prompt.text
+    Extracts the first user input text and checks for debug triggers.
+    """
+    prompt: str
+    if isinstance(task_params.input, str):
+        prompt = task_params.input
+    else:
+        # List of InputMessage - get first message content
+        if len(task_params.input) == 0:
+            logger.debug("Empty message list in debug prompt check")
+            return
+        prompt = task_params.input[0].content
+
+    if not prompt:
+        return
 
     if EXO_RUNNER_MUST_FAIL in prompt:
         logger.info("raising exception")

src/exo/worker/runner/runner_supervisor.py

@@ -127,20 +127,25 @@ class RunnerSupervisor:
         self._tg.cancel_scope.cancel()
 
     async def start_task(self, task: Task):
+        if task.task_id in self.pending:
+            logger.warning(
+                f"Skipping invalid task {task} as it has already been submitted"
+            )
+            return
         if task.task_id in self.completed:
-            logger.info(
+            logger.warning(
                 f"Skipping invalid task {task} as it has already been completed"
             )
+            return
         logger.info(f"Starting task {task}")
         event = anyio.Event()
         self.pending[task.task_id] = event
         try:
-            self._task_sender.send(task)
+            await self._task_sender.send_async(task)
         except ClosedResourceError:
             logger.warning(f"Task {task} dropped, runner closed communication.")
             return
         await event.wait()
-        logger.info(f"Finished task {task}")
 
     async def _forward_events(self):
         with self._ev_recv as events:

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

@@ -12,10 +12,9 @@ import mlx.nn as nn
 
 from exo.shared.constants import EXO_MODELS_DIR
 from exo.shared.models.model_cards import ModelCard, ModelTask
-from exo.shared.types.api import ChatCompletionMessage
 from exo.shared.types.common import ModelId
 from exo.shared.types.memory import Memory
-from exo.shared.types.tasks import ChatCompletionTaskParams
+from exo.shared.types.text_generation import TextGenerationTaskParams
 from exo.shared.types.worker.shards import PipelineShardMetadata, TensorShardMetadata
 from exo.worker.engines.mlx import Model
 from exo.worker.engines.mlx.generator.generate import mlx_generate
@@ -113,10 +112,10 @@ def run_gpt_oss_pipeline_device(
         tokens = tokens[:prompt_tokens]
         prompt_text = tokenizer.decode(tokens)
 
-        task = ChatCompletionTaskParams(
+        task = TextGenerationTaskParams(
             model=DEFAULT_GPT_OSS_MODEL_ID,
-            messages=[ChatCompletionMessage(role="user", content=prompt_text)],
-            max_tokens=max_tokens,
+            input=prompt_text,
+            max_output_tokens=max_tokens,
         )
 
         prompt = apply_chat_template(tokenizer, task)
@@ -181,10 +180,10 @@ def run_gpt_oss_tensor_parallel_device(
         tokens = tokens[:prompt_tokens]
         prompt_text = tokenizer.decode(tokens)
 
-        task = ChatCompletionTaskParams(
+        task = TextGenerationTaskParams(
             model=DEFAULT_GPT_OSS_MODEL_ID,
-            messages=[ChatCompletionMessage(role="user", content=prompt_text)],
-            max_tokens=max_tokens,
+            input=prompt_text,
+            max_output_tokens=max_tokens,
         )
 
         prompt = apply_chat_template(tokenizer, task)

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

@@ -0,0 +1,542 @@
+# type: ignore
+import time
+from typing import cast
+from unittest.mock import patch
+
+import mlx.core as mx
+import pytest
+from mlx_lm.models.cache import KVCache
+from mlx_lm.sample_utils import make_sampler
+
+from exo.shared.types.common import ModelId
+from exo.shared.types.text_generation import InputMessage, TextGenerationTaskParams
+from exo.worker.engines.mlx import Model
+from exo.worker.engines.mlx.cache import (
+    KVPrefixCache,
+    cache_length,
+    encode_prompt,
+    get_prefix_length,
+    make_kv_cache,
+)
+from exo.worker.engines.mlx.generator.generate import mlx_generate, prefill
+from exo.worker.engines.mlx.utils_mlx import apply_chat_template
+from exo.worker.tests.unittests.test_mlx.conftest import (
+    DEFAULT_GPT_OSS_CONFIG,
+    DEFAULT_GPT_OSS_MODEL_ID,
+)
+
+
+def _check_model_exists() -> bool:
+    return DEFAULT_GPT_OSS_CONFIG.model_path.exists()
+
+
+class TestGetPrefixLength:
+    def test_identical_arrays(self):
+        a = mx.array([1, 2, 3, 4, 5])
+        b = mx.array([1, 2, 3, 4, 5])
+        assert get_prefix_length(a, b) == 5
+
+    def test_no_common_prefix(self):
+        a = mx.array([1, 2, 3])
+        b = mx.array([4, 5, 6])
+        assert get_prefix_length(a, b) == 0
+
+    def test_partial_prefix(self):
+        a = mx.array([1, 2, 3, 4, 5])
+        b = mx.array([1, 2, 3, 7, 8])
+        assert get_prefix_length(a, b) == 3
+
+    def test_prompt_longer_than_cached(self):
+        a = mx.array([1, 2, 3, 4, 5])
+        b = mx.array([1, 2, 3])
+        assert get_prefix_length(a, b) == 3
+
+    def test_cached_longer_than_prompt(self):
+        a = mx.array([1, 2, 3])
+        b = mx.array([1, 2, 3, 4, 5])
+        assert get_prefix_length(a, b) == 3
+
+    def test_single_token_match(self):
+        a = mx.array([1, 2, 3])
+        b = mx.array([1, 5, 6])
+        assert get_prefix_length(a, b) == 1
+
+    def test_empty_prompt(self):
+        a = mx.array([]).astype(mx.int32)
+        b = mx.array([1, 2, 3])
+        assert get_prefix_length(a, b) == 0
+
+    def test_empty_cached(self):
+        a = mx.array([1, 2, 3])
+        b = mx.array([]).astype(mx.int32)
+        assert get_prefix_length(a, b) == 0
+
+    def test_both_empty(self):
+        a = mx.array([]).astype(mx.int32)
+        b = mx.array([]).astype(mx.int32)
+        assert get_prefix_length(a, b) == 0
+
+
+class TestKVPrefix:
+    @pytest.fixture
+    def mock_tokenizer(self):
+        """Create a minimal mock tokenizer for tests that don't need real tokenization."""
+        from unittest.mock import MagicMock
+
+        tokenizer = MagicMock()
+        tokenizer.encode.return_value = [1, 2, 3]
+        return tokenizer
+
+    def test_starts_empty(self, mock_tokenizer):
+        cache = KVPrefixCache(mock_tokenizer)
+        assert len(cache.prompts) == 0
+        assert len(cache.caches) == 0
+
+    def test_clear_empties_cache(self, mock_tokenizer):
+        cache = KVPrefixCache(mock_tokenizer)
+        cache.prompts.append(mx.array([1, 2, 3]))
+        cache.caches.append([KVCache()])
+        cache.clear()
+        assert len(cache.prompts) == 0
+        assert len(cache.caches) == 0
+
+    def test_clear_on_empty_cache(self, mock_tokenizer):
+        cache = KVPrefixCache(mock_tokenizer)
+        cache.clear()
+        assert len(cache.prompts) == 0
+
+
+def _load_gpt_oss() -> tuple[Model, object]:
+    from mlx_lm.utils import load_model
+
+    from exo.worker.engines.mlx.utils_mlx import load_tokenizer_for_model_id
+
+    model_path = DEFAULT_GPT_OSS_CONFIG.model_path
+    model_id = ModelId(DEFAULT_GPT_OSS_MODEL_ID)
+
+    model, _ = load_model(model_path, lazy=False)
+    tokenizer = load_tokenizer_for_model_id(model_id, model_path)
+    return cast(Model, model), tokenizer
+
+
+@pytest.mark.slow
+@pytest.mark.skipif(
+    not _check_model_exists(),
+    reason=f"GPT-OSS model not found at {DEFAULT_GPT_OSS_CONFIG.model_path}",
+)
+class TestKVPrefixCacheWithModel:
+    @pytest.fixture(scope="class")
+    def model_and_tokenizer(self):
+        model, tokenizer = _load_gpt_oss()
+        return model, tokenizer
+
+    def test_prefill_populates_cache(self, model_and_tokenizer):
+        model, tokenizer = model_and_tokenizer
+
+        task = TextGenerationTaskParams(
+            model=DEFAULT_GPT_OSS_MODEL_ID,
+            input=[InputMessage(role="user", content="Hello!!")],
+            max_output_tokens=1,
+        )
+        prompt = apply_chat_template(tokenizer, task)
+        tokens = encode_prompt(tokenizer, prompt)
+        cache = make_kv_cache(model)
+
+        prefill(model, tokenizer, make_sampler(0.0), tokens, cache)
+
+        # Cache should now hold the prompt tokens
+        assert cache_length(cache) == len(tokens)
+
+    def test_add_and_get_exact_match(self, model_and_tokenizer):
+        model, tokenizer = model_and_tokenizer
+
+        task = TextGenerationTaskParams(
+            model=DEFAULT_GPT_OSS_MODEL_ID,
+            input=[InputMessage(role="user", content="Test exact")],
+            max_output_tokens=1,
+        )
+        prompt = apply_chat_template(tokenizer, task)
+        tokens = encode_prompt(tokenizer, prompt)
+        cache = make_kv_cache(model)
+
+        prefill(model, tokenizer, make_sampler(0.0), tokens, cache)
+
+        kv_prefix_cache = KVPrefixCache(tokenizer)
+        kv_prefix_cache.add_kv_cache(prompt, cache)
+
+        assert len(kv_prefix_cache.prompts) == 1
+        stored_length = cache_length(kv_prefix_cache.caches[0])
+        assert stored_length > 0
+
+        # Retrieve with same prompt: exact match
+        result_cache, remaining_tokens, matched_index = kv_prefix_cache.get_kv_cache(
+            model, prompt
+        )
+        assert matched_index == 0
+
+        # Exact match returns only last token
+        assert len(remaining_tokens) == 1
+        assert mx.array_equal(remaining_tokens, tokens[-1:])
+
+    def test_add_and_get_prefix_match(self, model_and_tokenizer):
+        """get_kv_cache with a longer prompt sharing prefix should return partial match."""
+        model, tokenizer = model_and_tokenizer
+
+        short_task = TextGenerationTaskParams(
+            model=DEFAULT_GPT_OSS_MODEL_ID,
+            input=[InputMessage(role="user", content="Hi")],
+            max_output_tokens=1,
+        )
+        short_prompt = apply_chat_template(tokenizer, short_task)
+        short_tokens = encode_prompt(tokenizer, short_prompt)
+        cache = make_kv_cache(model)
+
+        prefill(model, tokenizer, make_sampler(0.0), short_tokens, cache)
+
+        kv_prefix_cache = KVPrefixCache(tokenizer)
+        kv_prefix_cache.add_kv_cache(short_prompt, cache)
+
+        # Query with longer prompt that shares the chat template prefix
+        long_task = TextGenerationTaskParams(
+            model=DEFAULT_GPT_OSS_MODEL_ID,
+            input=[InputMessage(role="user", content="Hi there, how are you?")],
+            max_output_tokens=1,
+        )
+        long_prompt = apply_chat_template(tokenizer, long_task)
+        long_tokens = encode_prompt(tokenizer, long_prompt)
+
+        # The prompts share a prefix (chat template preamble + "Hi")
+        expected_prefix = get_prefix_length(long_tokens, short_tokens)
+        assert expected_prefix > 0, (
+            "Prompts should share a prefix from the chat template"
+        )
+
+        result_cache, remaining_tokens, matched_index = kv_prefix_cache.get_kv_cache(
+            model, long_prompt
+        )
+        assert matched_index == 0
+
+        # remaining_tokens should be the suffix after the shared prefix
+        assert len(remaining_tokens) == len(long_tokens) - expected_prefix
+        assert mx.array_equal(remaining_tokens, long_tokens[expected_prefix:])
+
+    def test_stored_cache_not_mutated_after_get_and_generation(
+        self, model_and_tokenizer
+    ):
+        """Getting a cache and then mutating it (as generation does) must not corrupt stored cache."""
+        model, tokenizer = model_and_tokenizer
+
+        task = TextGenerationTaskParams(
+            model=DEFAULT_GPT_OSS_MODEL_ID,
+            input=[InputMessage(role="user", content="Mutation test")],
+            max_output_tokens=1,
+        )
+        prompt = apply_chat_template(tokenizer, task)
+        tokens = encode_prompt(tokenizer, prompt)
+        cache = make_kv_cache(model)
+
+        prefill(model, tokenizer, make_sampler(0.0), tokens, cache)
+
+        kv_prefix_cache = KVPrefixCache(tokenizer)
+        kv_prefix_cache.add_kv_cache(prompt, cache)
+
+        stored_length = cache_length(kv_prefix_cache.caches[0])
+
+        # Get cache and mutate it (simulating what generation does)
+        result_cache, _, matched_index = kv_prefix_cache.get_kv_cache(model, prompt)
+        assert matched_index == 0
+
+        # Simulate generation: feed many additional tokens through the cache
+        head_dim = result_cache[0].keys.shape[-1]
+        num_heads = result_cache[0].keys.shape[1]
+        extra_keys = mx.random.normal((1, num_heads, 50, head_dim))
+        extra_values = mx.random.normal((1, num_heads, 50, head_dim))
+        for layer_cache in result_cache:
+            layer_cache.update_and_fetch(extra_keys, extra_values)
+        mx.eval([c.keys for c in result_cache])
+
+        # Stored cache must be unchanged
+        assert cache_length(kv_prefix_cache.caches[0]) == stored_length
+
+    def test_stored_cache_survives_repeated_get_mutate_cycles(
+        self, model_and_tokenizer
+    ):
+        """Multiple get+mutate cycles (like repeated user requests) must not corrupt cache."""
+        model, tokenizer = model_and_tokenizer
+
+        task = TextGenerationTaskParams(
+            model=DEFAULT_GPT_OSS_MODEL_ID,
+            input=[InputMessage(role="user", content="Repeat test")],
+            max_output_tokens=1,
+        )
+        prompt = apply_chat_template(tokenizer, task)
+        tokens = encode_prompt(tokenizer, prompt)
+        cache = make_kv_cache(model)
+
+        prefill(model, tokenizer, make_sampler(0.0), tokens, cache)
+
+        kv_prefix_cache = KVPrefixCache(tokenizer)
+        kv_prefix_cache.add_kv_cache(prompt, cache)
+
+        stored_length = cache_length(kv_prefix_cache.caches[0])
+
+        for i in range(3):
+            result_cache, _, _ = kv_prefix_cache.get_kv_cache(model, prompt)
+
+            head_dim = result_cache[0].keys.shape[-1]
+            num_heads = result_cache[0].keys.shape[1]
+            extra = mx.random.normal((1, num_heads, 30, head_dim))
+            for layer_cache in result_cache:
+                layer_cache.update_and_fetch(extra, extra)
+            mx.eval([c.keys for c in result_cache])
+
+            assert cache_length(kv_prefix_cache.caches[0]) == stored_length, (
+                f"Failed on loop {i}"
+            )
+
+    def test_mlx_generate_populates_cache(self, model_and_tokenizer):
+        """mlx_generate should save the cache after generation completes."""
+        model, tokenizer = model_and_tokenizer
+
+        kv_prefix_cache = KVPrefixCache(tokenizer)
+        task = TextGenerationTaskParams(
+            model=DEFAULT_GPT_OSS_MODEL_ID,
+            input=[InputMessage(role="user", content="Hello")],
+            max_output_tokens=5,
+        )
+        prompt = apply_chat_template(tokenizer, task)
+        prompt_tokens = encode_prompt(tokenizer, prompt)
+
+        # Consume the entire generator so the cache-saving code after yield runs
+        generated_tokens = 0
+        for _response in mlx_generate(
+            model=model,
+            tokenizer=tokenizer,
+            task=task,
+            prompt=prompt,
+            kv_prefix_cache=kv_prefix_cache,
+        ):
+            generated_tokens += 1
+
+        assert len(kv_prefix_cache.prompts) == 1
+        assert len(kv_prefix_cache.caches) == 1
+        # Cache should contain prompt + generated tokens
+        expected_length = len(prompt_tokens) + generated_tokens
+        assert cache_length(kv_prefix_cache.caches[0]) == expected_length
+
+    def test_mlx_generate_second_call_gets_prefix_hit(self, model_and_tokenizer):
+        """Second mlx_generate call with same prompt should get a prefix hit from stored cache."""
+        model, tokenizer = model_and_tokenizer
+
+        kv_prefix_cache = KVPrefixCache(tokenizer)
+        task = TextGenerationTaskParams(
+            model=DEFAULT_GPT_OSS_MODEL_ID,
+            input=[InputMessage(role="user", content="Reuse test")],
+            max_output_tokens=5,
+        )
+        prompt = apply_chat_template(tokenizer, task)
+        prompt_tokens = encode_prompt(tokenizer, prompt)
+
+        # First generation populates cache
+        for _response in mlx_generate(
+            model=model,
+            tokenizer=tokenizer,
+            task=task,
+            prompt=prompt,
+            kv_prefix_cache=kv_prefix_cache,
+        ):
+            pass
+
+        assert len(kv_prefix_cache.prompts) == 1
+
+        # Second call should find a prefix match (the stored cache contains
+        # prompt + generated tokens, which shares the prompt prefix)
+        result_cache, remaining_tokens, matched_index = kv_prefix_cache.get_kv_cache(
+            model, prompt
+        )
+        # The stored cache is longer than the prompt (it includes generated tokens),
+        # so this is a prefix match where our prompt is fully contained
+        assert matched_index == 0
+        # Exact match: remaining_tokens is just the last token
+        assert len(remaining_tokens) == 1
+        assert mx.array_equal(remaining_tokens, prompt_tokens[-1:])
+
+    def test_mlx_generate_long_prompt_updates_cache_in_place(self, model_and_tokenizer):
+        """With a prompt > 1000 tokens, second generation should update the cache entry in-place."""
+        model, tokenizer = model_and_tokenizer
+
+        kv_prefix_cache = KVPrefixCache(tokenizer)
+
+        # Build a long user message (> 1000 tokens) to exceed _MIN_PREFIX_HIT_TO_UPDATE
+        base_text = "The quick brown fox jumps over the lazy dog. "
+        base_tokens = tokenizer.encode(base_text)
+        repeats = (1200 // len(base_tokens)) + 2
+        long_content = base_text * repeats
+
+        task1 = TextGenerationTaskParams(
+            model=DEFAULT_GPT_OSS_MODEL_ID,
+            input=[InputMessage(role="user", content=long_content)],
+            max_output_tokens=5,
+        )
+        prompt1 = apply_chat_template(tokenizer, task1)
+        prompt1_tokens = encode_prompt(tokenizer, prompt1)
+        assert len(prompt1_tokens) > 1000, (
+            "Prompt must exceed _MIN_PREFIX_HIT_TO_UPDATE"
+        )
+
+        # First generation populates the cache (must prefill all tokens)
+        t0 = time.perf_counter()
+        for _response in mlx_generate(
+            model=model,
+            tokenizer=tokenizer,
+            task=task1,
+            prompt=prompt1,
+            kv_prefix_cache=kv_prefix_cache,
+        ):
+            pass
+        first_gen_time = time.perf_counter() - t0
+
+        assert len(kv_prefix_cache.prompts) == 1
+        first_cache_length = cache_length(kv_prefix_cache.caches[0])
+
+        # Second generation: same long prompt + extra content (simulating multi-turn)
+        task2 = TextGenerationTaskParams(
+            model=DEFAULT_GPT_OSS_MODEL_ID,
+            input=[
+                InputMessage(role="user", content=long_content),
+                InputMessage(role="assistant", content="Sure, I can help."),
+                InputMessage(role="user", content="Tell me more."),
+            ],
+            max_output_tokens=5,
+        )
+        prompt2 = apply_chat_template(tokenizer, task2)
+        prompt2_tokens = encode_prompt(tokenizer, prompt2)
+
+        # Verify the prompts share a long prefix
+        prefix_len = get_prefix_length(prompt2_tokens, prompt1_tokens)
+        assert prefix_len > 1000, "Prompts must share > 1000 token prefix"
+
+        # Second generation should reuse the cached prefix (only prefill new tokens)
+        t0 = time.perf_counter()
+        for _response in mlx_generate(
+            model=model,
+            tokenizer=tokenizer,
+            task=task2,
+            prompt=prompt2,
+            kv_prefix_cache=kv_prefix_cache,
+        ):
+            pass
+        second_gen_time = time.perf_counter() - t0
+
+        # Second generation should be significantly faster due to prefix cache hit - hopefully not flaky
+        assert second_gen_time < first_gen_time * 0.5, (
+            f"Expected prefix cache speedup: "
+            f"first={first_gen_time:.2f}s, second={second_gen_time:.2f}s"
+        )
+
+        # With prefix_hit > 1000, should update in-place (not add a second entry)
+        assert len(kv_prefix_cache.prompts) == 1
+        # Updated cache should be longer (prompt2 + generated > prompt1 + generated)
+        updated_cache_length = cache_length(kv_prefix_cache.caches[0])
+        assert updated_cache_length > first_cache_length
+
+    def test_mlx_generate_stored_cache_not_mutated(self, model_and_tokenizer):
+        """After mlx_generate saves a cache, a second generation must not corrupt the stored copy."""
+        model, tokenizer = model_and_tokenizer
+
+        kv_prefix_cache = KVPrefixCache(tokenizer)
+        task = TextGenerationTaskParams(
+            model=DEFAULT_GPT_OSS_MODEL_ID,
+            input=[InputMessage(role="user", content="Immutable test")],
+            max_output_tokens=5,
+        )
+        prompt = apply_chat_template(tokenizer, task)
+
+        # First generation populates cache
+        for _response in mlx_generate(
+            model=model,
+            tokenizer=tokenizer,
+            task=task,
+            prompt=prompt,
+            kv_prefix_cache=kv_prefix_cache,
+        ):
+            pass
+
+        firstcache_length = cache_length(kv_prefix_cache.caches[0])
+
+        # Second generation gets the cache and mutates it during generation
+        for _response in mlx_generate(
+            model=model,
+            tokenizer=tokenizer,
+            task=task,
+            prompt=prompt,
+            kv_prefix_cache=kv_prefix_cache,
+        ):
+            pass
+
+        # The first stored cache must not have been mutated by the second generation
+        assert cache_length(kv_prefix_cache.caches[0]) == firstcache_length
+
+    def test_evicts_lru_entry_under_memory_pressure(self, model_and_tokenizer):
+        """Under memory pressure, adding a new cache entry evicts the least recently used one."""
+        model, tokenizer = model_and_tokenizer
+
+        kv_prefix_cache = KVPrefixCache(tokenizer)
+
+        # Add three cache entries with different prompts
+        prompts = ["First entry", "Second entry", "Third entry"]
+        for i, content in enumerate(prompts):
+            task = TextGenerationTaskParams(
+                model=DEFAULT_GPT_OSS_MODEL_ID,
+                input=[InputMessage(role="user", content=content)],
+                max_output_tokens=1,
+            )
+            prompt = apply_chat_template(tokenizer, task)
+            tokens = encode_prompt(tokenizer, prompt)
+            cache = make_kv_cache(model)
+            prefill(model, tokenizer, make_sampler(0.0), tokens, cache)
+            kv_prefix_cache.add_kv_cache(prompt, cache)
+            # Stagger _last_used so LRU order is deterministic
+            kv_prefix_cache._last_used[i] = float(i)
+
+        assert len(kv_prefix_cache.prompts) == 3
+
+        # Access the third entry to make it most recently used
+        kv_prefix_cache._last_used[2] = 100.0
+        # Entry 0 (_last_used=0.0) is LRU, entry 1 (_last_used=1.0) is next
+
+        # Simulate memory pressure: active memory exceeds threshold
+        fake_limit = 1000
+        fake_active = int(fake_limit * 0.90)  # Above _MEMORY_THRESHOLD (0.85)
+
+        with (
+            patch(
+                "exo.worker.engines.mlx.cache.mx.metal.get_active_memory",
+                return_value=fake_active,
+            ),
+            patch(
+                "exo.worker.engines.mlx.cache.mx.metal.device_info",
+                return_value={"max_recommended_working_set_size": fake_limit},
+            ),
+        ):
+            # Trigger eviction by adding a new entry
+            task = TextGenerationTaskParams(
+                model=DEFAULT_GPT_OSS_MODEL_ID,
+                input=[InputMessage(role="user", content="New entry")],
+                max_output_tokens=1,
+            )
+            prompt = apply_chat_template(tokenizer, task)
+            tokens = encode_prompt(tokenizer, prompt)
+            cache = make_kv_cache(model)
+            prefill(model, tokenizer, make_sampler(0.0), tokens, cache)
+            kv_prefix_cache.add_kv_cache(prompt, cache)
+
+        # LRU entries should have been evicted (entries 0, 1, 2 in order of _last_used)
+        # Since fake_active stays above threshold after each eviction (we don't change it),
+        # all old entries get evicted, leaving only the newly added one
+        assert len(kv_prefix_cache.prompts) == 1
+        # The surviving entry should be the newly added one
+        new_tokens = encode_prompt(tokenizer, prompt)
+        assert get_prefix_length(kv_prefix_cache.prompts[0], new_tokens) == len(
+            new_tokens
+        )

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

@@ -11,12 +11,12 @@ from pathlib import Path
 
 import pytest
 
-from exo.shared.models.model_cards import MODEL_CARDS, ModelCard, ModelId
-from exo.worker.download.download_utils import (
+from exo.download.download_utils import (
     download_file_with_retry,
     ensure_models_dir,
     fetch_file_list_with_cache,
 )
+from exo.shared.models.model_cards import MODEL_CARDS, ModelCard, ModelId
 from exo.worker.engines.mlx.utils_mlx import (
     get_eos_token_ids_for_model,
     load_tokenizer_for_model_id,

src/exo/worker/tests/unittests/test_plan/test_download_and_loading.py

@@ -1,5 +1,5 @@
 import exo.worker.plan as plan_mod
-from exo.shared.types.common import ModelId, NodeId
+from exo.shared.types.common import NodeId
 from exo.shared.types.memory import Memory
 from exo.shared.types.tasks import LoadModel
 from exo.shared.types.worker.downloads import DownloadCompleted, DownloadProgress
@@ -45,13 +45,9 @@ def test_plan_requests_download_when_waiting_and_shard_not_downloaded():
     instances = {INSTANCE_1_ID: instance}
     all_runners = {RUNNER_1_ID: RunnerIdle()}
 
-    # No entry for this shard -> should trigger DownloadModel
-    download_status: dict[ModelId, DownloadProgress] = {}
-
     result = plan_mod.plan(
         node_id=NODE_A,
         runners=runners,  # type: ignore
-        download_status=download_status,
         global_download_status={NODE_A: []},
         instances=instances,
         all_runners=all_runners,
@@ -92,14 +88,6 @@ def test_plan_loads_model_when_all_shards_downloaded_and_waiting():
         RUNNER_2_ID: RunnerConnected(),
     }
 
-    # Local node has already marked its shard as downloaded (not actually used by _load_model)
-    local_download_status = {
-        MODEL_A_ID: DownloadCompleted(
-            shard_metadata=shard1, node_id=NODE_A, total_bytes=Memory()
-        )
-    }
-
-    # Global view has completed downloads for both nodes
     global_download_status = {
         NODE_A: [
             DownloadCompleted(
@@ -116,7 +104,6 @@ def test_plan_loads_model_when_all_shards_downloaded_and_waiting():
     result = plan_mod.plan(
         node_id=NODE_A,
         runners=runners,  # type: ignore
-        download_status=local_download_status,
         global_download_status=global_download_status,
         instances=instances,
         all_runners=all_runners,
@@ -148,30 +135,26 @@ def test_plan_does_not_request_download_when_shard_already_downloaded():
     instances = {INSTANCE_1_ID: instance}
     all_runners = {RUNNER_1_ID: RunnerIdle()}
 
-    # Local status claims the shard is downloaded already
-    local_download_status = {
-        MODEL_A_ID: DownloadCompleted(
-            shard_metadata=shard, node_id=NODE_A, total_bytes=Memory()
-        )
-    }
-
-    # Global view hasn't caught up yet (no completed shards recorded for NODE_A)
+    # Global state shows shard is downloaded for NODE_A
     global_download_status: dict[NodeId, list[DownloadProgress]] = {
-        NODE_A: [],
+        NODE_A: [
+            DownloadCompleted(
+                shard_metadata=shard, node_id=NODE_A, total_bytes=Memory()
+            )
+        ],
         NODE_B: [],
     }
 
     result = plan_mod.plan(
         node_id=NODE_A,
         runners=runners,  # type: ignore
-        download_status=local_download_status,
         global_download_status=global_download_status,
         instances=instances,
         all_runners=all_runners,
         tasks={},
     )
 
-    assert result is None
+    assert not isinstance(result, plan_mod.DownloadModel)
 
 
 def test_plan_does_not_load_model_until_all_shards_downloaded_globally():
@@ -202,12 +185,6 @@ def test_plan_does_not_load_model_until_all_shards_downloaded_globally():
         RUNNER_2_ID: RunnerConnected(),
     }
 
-    # Only NODE_A's shard is recorded as downloaded globally
-    local_download_status = {
-        MODEL_A_ID: DownloadCompleted(
-            shard_metadata=shard1, node_id=NODE_A, total_bytes=Memory()
-        )
-    }
     global_download_status = {
         NODE_A: [
             DownloadCompleted(
@@ -220,7 +197,6 @@ def test_plan_does_not_load_model_until_all_shards_downloaded_globally():
     result = plan_mod.plan(
         node_id=NODE_A,
         runners=runners,  # type: ignore
-        download_status=local_download_status,
         global_download_status=global_download_status,
         instances=instances,
         all_runners=all_runners,
@@ -245,7 +221,6 @@ def test_plan_does_not_load_model_until_all_shards_downloaded_globally():
     result = plan_mod.plan(
         node_id=NODE_A,
         runners=runners,  # type: ignore
-        download_status=local_download_status,
         global_download_status=global_download_status,
         instances=instances,
         all_runners=all_runners,

src/exo/worker/tests/unittests/test_plan/test_runner_lifecycle.py

@@ -47,8 +47,7 @@ def test_plan_kills_runner_when_instance_missing():
 
     result = plan_mod.plan(
         node_id=NODE_A,
-        runners=runners,  # type: ignore
-        download_status={},
+        runners=runners,  # type: ignore[arg-type]
         global_download_status={NODE_A: []},
         instances=instances,
         all_runners=all_runners,
@@ -87,8 +86,7 @@ def test_plan_kills_runner_when_sibling_failed():
 
     result = plan_mod.plan(
         node_id=NODE_A,
-        runners=runners,  # type: ignore
-        download_status={},
+        runners=runners,  # type: ignore[arg-type]
         global_download_status={NODE_A: []},
         instances=instances,
         all_runners=all_runners,
@@ -120,7 +118,6 @@ def test_plan_creates_runner_when_missing_for_node():
     result = plan_mod.plan(
         node_id=NODE_A,
         runners=runners,
-        download_status={},
         global_download_status={NODE_A: []},
         instances=instances,
         all_runners=all_runners,
@@ -158,8 +155,7 @@ def test_plan_does_not_create_runner_when_supervisor_already_present():
 
     result = plan_mod.plan(
         node_id=NODE_A,
-        runners=runners,  # type: ignore
-        download_status={},
+        runners=runners,  # type: ignore[arg-type]
         global_download_status={NODE_A: []},
         instances=instances,
         all_runners=all_runners,
@@ -189,7 +185,6 @@ def test_plan_does_not_create_runner_for_unassigned_node():
     result = plan_mod.plan(
         node_id=NODE_A,
         runners=runners,  # type: ignore
-        download_status={},
         global_download_status={NODE_A: []},
         instances=instances,
         all_runners=all_runners,