formatting

.github/actions/typecheck/action.yml

@@ -0,0 +1,12 @@
+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,14 +26,73 @@ jobs:
           name: exo
           authToken: "${{ secrets.CACHIX_AUTH_TOKEN }}"
 
-      - name: Load nix develop environment
-        run: nix run github:nicknovitski/nix-develop/v1
+      - 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: Sync dependencies
-        run: uv sync --all-packages
+      - name: Pull LFS files
+        run: |
+          echo "Pulling Git LFS files..."
+          git lfs pull
+        shell: bash
 
-      - name: Run type checker
-        run: uv run basedpyright --project pyproject.toml
+      - 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
 
   nix:
     name: Build and check (${{ matrix.system }})
@@ -64,63 +123,6 @@ 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 '
@@ -132,14 +134,3 @@ 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,6 +28,3 @@ 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 frontier AI locally. Maintained by [exo labs](https://x.com/exolabs).
+exo: Run your own AI cluster at home with everyday devices. 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,10 +107,6 @@ 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:**
@@ -234,7 +230,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).
 
-Please refer to the caveats for immediate troubleshooting.
+Note that on Mac Studio, you cannot use the Thunderbolt 5 port next to the Ethernet port.
 
 To enable RDMA on macOS, follow these steps:
 
@@ -251,14 +247,6 @@ 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,8 +342,6 @@
 				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;
 		};
@@ -399,8 +397,6 @@
 				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,7 +14,6 @@ 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
@@ -25,8 +24,6 @@ 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) {
@@ -306,49 +303,6 @@ 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()
                     }
@@ -469,44 +423,6 @@ 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) {
@@ -549,7 +465,6 @@ struct ContentView: View {
                     Text(thunderboltStatusText)
                         .font(.caption2)
                         .foregroundColor(thunderboltStatusColor)
-                    clusterThunderboltBridgeView
                     interfaceIpList
                     rdmaStatusView
                     sendBugReportButton

app/EXO/EXO/EXOApp.swift

@@ -21,7 +21,6 @@ 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)
 
@@ -42,13 +41,10 @@ struct EXOApp: App {
         let localNetwork = LocalNetworkChecker()
         _localNetworkChecker = StateObject(wrappedValue: localNetwork)
         _updater = StateObject(wrappedValue: updater)
-        let thunderboltBridge = ThunderboltBridgeService(clusterStateService: service)
-        _thunderboltBridgeService = StateObject(wrappedValue: thunderboltBridge)
         enableLaunchAtLoginIfNeeded()
-        // 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)
+        NetworkSetupHelper.ensureLaunchDaemonInstalled()
+        // Check local network access BEFORE launching exo
+        localNetwork.check()
         controller.scheduleLaunch(after: 15)
         service.startPolling()
         networkStatus.startPolling()
@@ -62,7 +58,6 @@ struct EXOApp: App {
                 .environmentObject(networkStatusService)
                 .environmentObject(localNetworkChecker)
                 .environmentObject(updater)
-                .environmentObject(thunderboltBridgeService)
         } label: {
             menuBarIcon
         }
@@ -135,7 +130,6 @@ struct EXOApp: App {
                 "Failed to register EXO for launch at login: \(error.localizedDescription)")
         }
     }
-
 }
 
 /// Helper for managing EXO's launch-at-login registration
@@ -225,7 +219,7 @@ private final class ExoUpdaterDelegate: NSObject, SPUUpdaterDelegate {
         }
     }
 
-    nonisolated private func showNotification(title: String, body: String) {
+    private func showNotification(title: String, body: String) {
         let center = UNUserNotificationCenter.current()
         let content = UNMutableNotificationContent()
         content.title = title

app/EXO/EXO/ExoProcessController.swift

@@ -3,8 +3,6 @@ import Combine
 import Foundation
 
 private let customNamespaceKey = "EXOCustomNamespace"
-private let hfTokenKey = "EXOHFToken"
-private let enableImageModelsKey = "EXOEnableImageModels"
 
 @MainActor
 final class ExoProcessController: ObservableObject {
@@ -39,22 +37,6 @@ 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?
@@ -209,12 +191,6 @@ 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,43 +5,17 @@ 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)
@@ -50,34 +24,15 @@ 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
     }
 }
 
@@ -147,18 +102,6 @@ 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?
@@ -177,51 +120,16 @@ struct SystemInfo: Decodable {
 }
 
 struct Topology: Decodable {
-    /// 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
-    }
+    let nodes: [TopologyNode]
+    let connections: [TopologyConnection]?
 }
 
-/// 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 TopologyNode: Decodable {
+    let nodeId: String
+    let nodeProfile: NodeProfile
 }
 
-struct TopologyConnection {
+struct TopologyConnection: Decodable {
     let localNodeId: String
     let sendBackNodeId: String
 }
@@ -293,7 +201,7 @@ struct ClusterTask {
     let modelName: String?
     let promptPreview: String?
     let errorMessage: String?
-    let parameters: TextGenerationTaskParameters?
+    let parameters: ChatCompletionTaskParameters?
 
     var sortPriority: Int {
         switch status {
@@ -330,12 +238,12 @@ struct ClusterTaskPayload: Decodable {
     let taskStatus: TaskStatus?
     let instanceId: String?
     let commandId: String?
-    let taskParams: TextGenerationTaskParameters?
+    let taskParams: ChatCompletionTaskParameters?
     let errorType: String?
     let errorMessage: String?
 }
 
-struct TextGenerationTaskParameters: Decodable, Equatable {
+struct ChatCompletionTaskParameters: Decodable, Equatable {
     let model: String?
     let messages: [ChatCompletionMessage]?
     let maxTokens: Int?
@@ -374,7 +282,7 @@ extension ClusterTask {
         guard let id = payload.taskId else { return nil }
         let status = payload.taskStatus ?? .unknown
         switch kindKey {
-        case "TextGeneration":
+        case "ChatCompletion":
             self.init(
                 id: id,
                 status: status,

app/EXO/EXO/Services/BugReportService.swift

@@ -55,16 +55,12 @@ 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,
-            clusterTbBridgeStatus: clusterTbBridgeStatus
+            isManual: isManual
         )
 
         let uploads: [(path: String, data: Data?)] = [
@@ -182,19 +178,18 @@ struct BugReportService {
     }
 
     private func readThunderboltBridgeDisabled() -> Bool? {
-        // 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
+        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
         }
-
-        guard let isEnabled = ThunderboltBridgeDetector.isServiceEnabled(serviceName: serviceName)
-        else {
-            return nil
+        if output.contains("disabled") {
+            return true
         }
-
-        // Return true if disabled, false if enabled
-        return !isEnabled
+        return nil
     }
 
     private func readInterfaces() -> [DebugInfo.InterfaceStatus] {
@@ -273,12 +268,11 @@ struct BugReportService {
         hostName: String,
         ifconfig: String,
         debugInfo: DebugInfo,
-        isManual: Bool,
-        clusterTbBridgeStatus: [[String: Any]]?
+        isManual: Bool
     ) -> Data? {
         let system = readSystemMetadata()
         let exo = readExoMetadata()
-        var payload: [String: Any] = [
+        let payload: [String: Any] = [
             "timestamp": timestamp,
             "host": hostName,
             "ifconfig": ifconfig,
@@ -288,38 +282,9 @@ 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,7 +41,6 @@ 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 {
@@ -49,39 +48,10 @@ final class LocalNetworkChecker: ObservableObject {
         set { UserDefaults.standard.set(newValue, forKey: Self.hasCompletedInitialCheckKey) }
     }
 
-    /// Checks if local network access is working (one-time check).
+    /// Checks if local network access is working.
     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()
-        // Only show "checking" status on first check to avoid UI flicker
-        if status == .unknown {
-            status = .checking
-        }
+        status = .checking
 
         // Use longer timeout on first launch to allow time for permission prompt
         let isFirstCheck = !hasCompletedInitialCheck
@@ -90,15 +60,12 @@ final class LocalNetworkChecker: ObservableObject {
         checkTask = Task { [weak self] in
             guard let self else { return }
 
-            Self.logger.debug("Checking local network connectivity (first check: \(isFirstCheck))")
+            Self.logger.info("Checking local network connectivity (first check: \(isFirstCheck))")
             let result = await self.checkConnectivity(timeout: timeout)
             self.status = result
             self.hasCompletedInitialCheck = true
 
-            // Only log on state changes or first check to reduce noise
-            if isFirstCheck || result != self.status {
-                Self.logger.info("Local network check: \(result.displayText)")
-            }
+            Self.logger.info("Local network check complete: \(result.displayText)")
         }
     }
 
@@ -174,7 +141,6 @@ final class LocalNetworkChecker: ObservableObject {
     }
 
     func stop() {
-        stopPeriodicChecking()
         checkTask?.cancel()
         checkTask = nil
         connection?.cancel()

app/EXO/EXO/Services/NetworkSetupHelper.swift

@@ -7,20 +7,14 @@ 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 = 1786
+    private static let requiredStartInterval: Int = 1791
 
     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
@@ -34,69 +28,19 @@ 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
         """
 
-    /// 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() {
+    static func ensureLaunchDaemonInstalled() {
         // 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 {
-                try installLaunchDaemon()
+                if daemonAlreadyInstalled() {
+                    return
+                }
+                try await installLaunchDaemon()
                 logger.info("Network setup launch daemon installed and started")
             } catch {
                 logger.error(
@@ -119,9 +63,48 @@ 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 || legacyScriptExists || plistExists
+        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"
+        """
     }
 
     private static func daemonAlreadyInstalled() -> Bool {
@@ -157,7 +140,7 @@ enum NetworkSetupHelper {
         return true
     }
 
-    private static func installLaunchDaemon() throws {
+    private static func installLaunchDaemon() async throws {
         let installerScript = makeInstallerScript()
         try runShellAsAdmin(installerScript)
     }
@@ -168,19 +151,8 @@ 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'
@@ -212,82 +184,13 @@ 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,18 +153,22 @@ private struct NetworkStatusFetcher {
     }
 
     private func readThunderboltBridgeState() -> ThunderboltState? {
-        // 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 {
+        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
+            }
             return nil
         }
-
-        return isEnabled ? .enabled : .disabled
+        let output = result.output.lowercased()
+        if output.contains("enabled") {
+            return .enabled
+        }
+        if output.contains("disabled") {
+            return .disabled
+        }
+        return nil
     }
 
     private func readBridgeInactive() -> Bool? {

app/EXO/EXO/Services/ThunderboltBridgeDetector.swift

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

@@ -1,261 +0,0 @@
-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: TextGenerationTaskParameters?
+    let parameters: ChatCompletionTaskParameters?
 
     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 ?? [])
+        let topologyNodeIds = Set(topology?.nodes.map(\.nodeId) ?? [])
         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 nodeId in topologyNodes {
-                if let viewModel = nodesById[nodeId] {
+            for topoNode in topologyNodes {
+                if let viewModel = nodesById[topoNode.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
+    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
-    echo_warn "EXO.app found at: $app_path"
-    echo_warn "You may want to move it to Trash manually."
-  fi
 done
 
 echo ""
@@ -151,3 +151,4 @@ echo ""
 echo "Manual step required:"
 echo "  Remove EXO from Login Items in System Settings → General → Login Items"
 echo ""
+

bench/exo_bench.py

@@ -5,13 +5,10 @@ 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
@@ -19,84 +16,6 @@ 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):
@@ -105,7 +24,7 @@ class ExoHttpError(RuntimeError):
 
 
 class ExoClient:
-    def __init__(self, host: str, port: int, timeout_s: float = 7200.0):
+    def __init__(self, host: str, port: int, timeout_s: float = 600.0):
         self.host = host
         self.port = port
         self.timeout_s = timeout_s
@@ -261,7 +180,14 @@ def parse_int_list(values: list[str]) -> list[int]:
             part = part.strip()
             if part:
                 items.append(int(part))
-    return items
+
+    seen: set[int] = set()
+    out: list[int] = []
+    for x in items:
+        if x not in seen:
+            out.append(x)
+            seen.add(x)
+    return out
 
 
 def resolve_model_short_id(client: ExoClient, model_arg: str) -> tuple[str, str]:
@@ -314,11 +240,7 @@ def run_one_completion(
 
     stats = out.get("generation_stats")
 
-    # 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 ""
+    preview = (out.get("choices") or [{}])[0]["message"]["content"][:200]
 
     return {
         "elapsed_s": elapsed,
@@ -355,29 +277,12 @@ class PromptSizer:
                 f"Target ({target}) is smaller than template overhead ({self.base_tokens})."
             )
 
-        # 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
-
-        # 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
+        content = ""
         tok = self.count_fn(content)
-        logger.info(f"{tok=}")
+
+        while tok < target:
+            content += self.atom
+            tok = self.count_fn(content)
 
         if tok != target:
             raise RuntimeError(
@@ -443,7 +348,7 @@ def main() -> int:
         help="Warmup runs per placement (uses first pp/tg).",
     )
     ap.add_argument(
-        "--timeout", type=float, default=7200.0, help="HTTP timeout (seconds)."
+        "--timeout", type=float, default=600.0, help="HTTP timeout (seconds)."
     )
     ap.add_argument(
         "--json-out",
@@ -453,11 +358,6 @@ 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)
@@ -469,15 +369,6 @@ 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)
 
@@ -486,7 +377,10 @@ def main() -> int:
     )
     previews = previews_resp.get("previews") or []
 
-    tokenizer = load_tokenizer_for_bench(full_model_id)
+    tokenizer = AutoTokenizer.from_pretrained(
+        full_model_id,
+        trust_remote_code=True,
+    )
     if tokenizer is None:
         raise RuntimeError("[exo-bench] tokenizer load failed")
 
@@ -592,55 +486,60 @@ def main() -> int:
                 )
                 logger.debug(f"  warmup {i + 1}/{args.warmup} done")
 
-            # 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))
-
-            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
+            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,
+                            }
                         )
-                    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)
+                        runs.append(row)
+                        all_rows.append(row)
 
-                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
-                    )
+                    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)
+                        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,61 +3,6 @@
   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;
@@ -97,12 +42,11 @@
       '';
 
       # 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="${dashboardDeps}/lib/node_modules/exo-dashboard/node_modules"
+        export NODE_PATH="${dashboardFull}/lib/node_modules/exo-dashboard/node_modules"
         exec ${pkgs.nodejs}/bin/node \
-          ${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" \
+          ${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" \
           "$@"
       '';
     };

dashboard/src/lib/components/ChatForm.svelte

@@ -89,10 +89,7 @@
 
   const isImageModel = $derived(() => {
     if (!currentModel) return false;
-    return (
-      modelSupportsTextToImage(currentModel) ||
-      modelSupportsImageEditing(currentModel)
-    );
+    return modelSupportsTextToImage(currentModel);
   });
 
   const isEditOnlyWithoutImage = $derived(
@@ -649,23 +646,6 @@
             </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,36 +110,6 @@
     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 });
   }
@@ -164,92 +134,90 @@
 <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 (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
+    <!-- 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'
+            : ''}"
         >
-        <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'
-              : ''}"
+          {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"
           >
-            {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>
+            <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}
           </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}
-            </div>
-          </div>
-        {/if}
-      </div>
-    {/if}
+      {/if}
+    </div>
 
     <!-- Quality -->
     <div class="flex items-center gap-1.5">
@@ -357,59 +325,6 @@
       </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,32 +5,22 @@
     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(),
-    filteredNodes = new Set(),
-    onNodeClick,
-  }: Props = $props();
+  let { class: className = "", highlightedNodes = new Set() }: 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];
@@ -532,72 +522,10 @@
         }
       }
 
-      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", 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;
-          }
-        });
+        .style("cursor", "pointer");
 
       // Add tooltip
       nodeG
@@ -606,6 +534,27 @@
           `${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;
@@ -630,7 +579,6 @@
         // Main body (uniform color)
         nodeG
           .append("rect")
-          .attr("class", "node-outline")
           .attr("x", x)
           .attr("y", y)
           .attr("width", iconBaseWidth)
@@ -713,7 +661,6 @@
         // Main body (uniform color)
         nodeG
           .append("rect")
-          .attr("class", "node-outline")
           .attr("x", x)
           .attr("y", y)
           .attr("width", iconBaseWidth)
@@ -791,7 +738,6 @@
         // Screen outer frame
         nodeG
           .append("rect")
-          .attr("class", "node-outline")
           .attr("x", screenX)
           .attr("y", y)
           .attr("width", screenWidth)
@@ -900,7 +846,6 @@
         // Main shape
         nodeG
           .append("polygon")
-          .attr("class", "node-outline")
           .attr("points", hexPoints)
           .attr("fill", fillColor)
           .attr("stroke", wireColor)
@@ -1119,41 +1064,11 @@
           .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(() => {
-    // Track all reactive dependencies that affect rendering
-    const _data = data;
-    const _hoveredNodeId = hoveredNodeId;
-    const _filteredNodes = filteredNodes;
-    const _highlightedNodes = highlightedNodes;
-    if (_data) {
+    if (data) {
       renderGraph();
     }
   });
@@ -1176,8 +1091,12 @@
 
 <style>
   :global(.graph-node) {
-    /* Only transition opacity for filtered-out nodes, no transition on hover stroke changes */
-    transition: opacity 0.2s ease;
+    transition:
+      transform 0.2s ease,
+      opacity 0.2s ease;
+  }
+  :global(.graph-node:hover) {
+    filter: brightness(1.1);
   }
   :global(.graph-link) {
     stroke: var(--exo-light-gray, #b3b3b3);

dashboard/src/routes/+page.svelte

@@ -19,9 +19,6 @@
     selectedPreviewModelId,
     isLoadingPreviews,
     selectPreviewModel,
-    togglePreviewNodeFilter,
-    clearPreviewNodeFilter,
-    previewNodeFilter,
     createConversation,
     setSelectedChatModel,
     selectedChatModel,
@@ -31,8 +28,6 @@
     toggleTopologyOnlyMode,
     chatSidebarVisible,
     toggleChatSidebarVisible,
-    thunderboltBridgeCycles,
-    nodeThunderboltBridge,
     type DownloadProgress,
     type PlacementPreview,
   } from "$lib/stores/app.svelte";
@@ -54,41 +49,6 @@
   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);
 
@@ -130,15 +90,6 @@
       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";
 
@@ -230,9 +181,6 @@
   // 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")
@@ -784,8 +732,6 @@
     instanceWrapped: unknown,
   ): {
     isDownloading: boolean;
-    isFailed: boolean;
-    errorMessage: string | null;
     progress: DownloadProgress | null;
     statusText: string;
     perNode: Array<{
@@ -797,8 +743,6 @@
     if (!downloadsData || Object.keys(downloadsData).length === 0) {
       return {
         isDownloading: false,
-        isFailed: false,
-        errorMessage: null,
         progress: null,
         statusText: "RUNNING",
         perNode: [],
@@ -810,8 +754,6 @@
     if (!instance || typeof instance !== "object") {
       return {
         isDownloading: false,
-        isFailed: false,
-        errorMessage: null,
         progress: null,
         statusText: "PREPARING",
         perNode: [],
@@ -867,26 +809,6 @@
           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;
 
@@ -922,8 +844,6 @@
       const statusInfo = deriveInstanceStatus(instanceWrapped);
       return {
         isDownloading: false,
-        isFailed: statusInfo.statusText === "FAILED",
-        errorMessage: null,
         progress: null,
         statusText: statusInfo.statusText,
         perNode: [],
@@ -936,8 +856,6 @@
 
     return {
       isDownloading: true,
-      isFailed: false,
-      errorMessage: null,
       progress: {
         totalBytes,
         downloadedBytes,
@@ -1533,7 +1451,7 @@
 
   // Get ALL filtered previews based on current settings (matching minimum nodes)
   // Note: previewsData already contains previews for the selected model (fetched via API)
-  // Backend handles node_ids filtering, we filter by sharding/instance type and min nodes
+  // We filter by sharding/instance type and min nodes, returning ALL eligible previews
   const filteredPreviews = $derived(() => {
     if (!selectedModelId || previewsData.length === 0) return [];
 
@@ -1666,86 +1584,7 @@
           <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"
@@ -1785,111 +1624,7 @@
             <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 -->
@@ -2326,13 +2061,6 @@
                           >
                             {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>
@@ -2378,9 +2106,6 @@
                     {@const isImageModel = modelSupportsImageGeneration(
                       foundModel.id,
                     )}
-                    {@const isImageEditModel = modelSupportsImageEditing(
-                      foundModel.id,
-                    )}
                     <span
                       class="flex items-center justify-between gap-2 w-full pr-4"
                     >
@@ -2407,22 +2132,6 @@
                             <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
                         >
@@ -2495,9 +2204,6 @@
                       {@const isImageModel = modelSupportsImageGeneration(
                         model.id,
                       )}
-                      {@const isImageEditModel = modelSupportsImageEditing(
-                        model.id,
-                      )}
                       <button
                         type="button"
                         onclick={() => {
@@ -2538,23 +2244,6 @@
                               <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
@@ -2875,36 +2564,7 @@
               <div
                 class="relative aspect-square bg-exo-dark-gray rounded-lg overflow-hidden"
               >
-                <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}
+                <TopologyGraph highlightedNodes={highlightedNodes()} />
               </div>
             </button>
 
@@ -3333,13 +2993,6 @@
                             >
                               {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,8 +6,6 @@
     type DownloadProgress,
     refreshState,
     lastUpdate as lastUpdateStore,
-    startDownload,
-    deleteDownload,
   } from "$lib/stores/app.svelte";
   import HeaderNav from "$lib/components/HeaderNav.svelte";
 
@@ -30,7 +28,6 @@
     etaMs: number;
     status: "completed" | "downloading";
     files: FileProgress[];
-    shardMetadata?: Record<string, unknown>;
   };
 
   type NodeEntry = {
@@ -175,6 +172,33 @@
   }
 
   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 {
@@ -272,12 +296,6 @@
             }
           }
 
-          // 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,
@@ -294,7 +312,6 @@
                 ? "completed"
                 : "downloading",
             files,
-            shardMetadata,
           };
 
           const existing = modelMap.get(modelId);
@@ -356,6 +373,7 @@
   onMount(() => {
     // Ensure we fetch at least once when visiting downloads directly
     refreshState();
+    fetchModels();
   });
 </script>
 
@@ -464,7 +482,7 @@
                     {#if model.status !== "completed"}
                       <div class="text-[11px] text-exo-light-gray font-mono">
                         {formatBytes(model.downloadedBytes)} / {formatBytes(
-                          model.totalBytes,
+                          getModelTotalBytes(model.modelId, model.totalBytes),
                         )}
                       </div>
                     {/if}
@@ -479,52 +497,6 @@
                     >
                       {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

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

@@ -1,367 +0,0 @@
-<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,9 +21,7 @@
           "nixpkgs"
         ],
         "purescript-overlay": "purescript-overlay",
-        "pyproject-nix": [
-          "pyproject-nix"
-        ]
+        "pyproject-nix": "pyproject-nix"
       },
       "locked": {
         "lastModified": 1765953015,
@@ -151,44 +149,19 @@
         "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": 1764134915,
-        "narHash": "sha256-xaKvtPx6YAnA3HQVp5LwyYG1MaN4LLehpQI8xEdBvBY=",
+        "lastModified": 1763017646,
+        "narHash": "sha256-Z+R2lveIp6Skn1VPH3taQIuMhABg1IizJd8oVdmdHsQ=",
         "owner": "pyproject-nix",
         "repo": "pyproject.nix",
-        "rev": "2c8df1383b32e5443c921f61224b198a2282a657",
+        "rev": "47bd6f296502842643078d66128f7b5e5370790c",
         "type": "github"
       },
       "original": {
@@ -205,10 +178,7 @@
         "flake-parts": "flake-parts",
         "nixpkgs": "nixpkgs",
         "nixpkgs-swift": "nixpkgs-swift",
-        "pyproject-build-systems": "pyproject-build-systems",
-        "pyproject-nix": "pyproject-nix",
-        "treefmt-nix": "treefmt-nix",
-        "uv2nix": "uv2nix"
+        "treefmt-nix": "treefmt-nix"
       }
     },
     "rust-analyzer-src": {
@@ -269,29 +239,6 @@
         "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,26 +24,6 @@
     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)
@@ -68,7 +48,6 @@
         inputs.treefmt-nix.flakeModule
         ./dashboard/parts.nix
         ./rust/parts.nix
-        ./python/parts.nix
       ];
 
       perSystem =
@@ -79,11 +58,6 @@
           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 = {
@@ -105,24 +79,14 @@
                 enable = true;
                 package = pkgsSwift.swiftPackages.swift-format;
               };
-              shfmt.enable = true;
             };
           };
 
-          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;
-              };
-            }
-          );
+          checks.lint = pkgs.runCommand "lint-check" { } ''
+            export RUFF_CACHE_DIR="$TMPDIR/ruff-cache"
+            ${pkgs.ruff}/bin/ruff check ${inputs.self}/
+            touch $out
+          '';
 
           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:
-    treefmt || nix fmt
+    nix fmt
 
 lint:
     uv run ruff check --fix

nix/darwin-build-fixes.patch

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

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

@@ -1,158 +0,0 @@
-{ 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.4; sys_platform == 'darwin'",
-    "mlx[cpu]==0.30.4; sys_platform == 'linux'",
-    "mlx-lm",
+    "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",
     "tiktoken>=0.12.0", # required for kimi k2 tokenizer
     "hypercorn>=0.18.0",
     "openai-harmony>=0.0.8",
     "httpx>=0.28.1",
     "tomlkit>=0.14.0",
     "pillow>=11.0,<12.0", # compatibility with mflux
-    "mflux==0.15.4",
+    "mflux>=0.14.2",
     "python-multipart>=0.0.21",
 ]
 
@@ -34,6 +34,7 @@ dependencies = [
 exo-master = "exo.master.main:main"
 exo-worker = "exo.worker.main:main"
 exo = "exo.main:main"
+exo-rsh = "exo.rsh.client:main"
 
 # dependencies only required for development
 [dependency-groups]
@@ -63,7 +64,6 @@ 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

@@ -1,94 +0,0 @@
-{ 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/cli/__init__.py

@@ -0,0 +1,32 @@
+"""Exo CLI - SLURM-compatible job management commands."""
+
+
+def run_subcommand(command: str, args: list[str]) -> int:
+    """Route to the appropriate subcommand handler.
+
+    Args:
+        command: The subcommand name (sbatch, squeue, scancel, salloc)
+        args: Command line arguments for the subcommand
+
+    Returns:
+        Exit code from the subcommand
+    """
+    if command == "sbatch":
+        from exo.cli.sbatch import main
+
+        return main(args)
+    elif command == "squeue":
+        from exo.cli.squeue import main
+
+        return main(args)
+    elif command == "scancel":
+        from exo.cli.scancel import main
+
+        return main(args)
+    elif command == "salloc":
+        from exo.cli.salloc import main
+
+        return main(args)
+    else:
+        print(f"Unknown subcommand: {command}")
+        return 1

src/exo/cli/common.py

@@ -0,0 +1,118 @@
+"""Common utilities for Exo CLI commands."""
+
+import json
+import os
+import urllib.request
+from typing import Any
+from urllib.error import HTTPError, URLError
+
+# Default API endpoint
+DEFAULT_API_HOST = "localhost"
+DEFAULT_API_PORT = 52415
+
+
+def get_api_base() -> str:
+    """Get the API base URL from environment or defaults."""
+    host = os.environ.get("EXO_API_HOST", DEFAULT_API_HOST)
+    port = os.environ.get("EXO_API_PORT", str(DEFAULT_API_PORT))
+    return f"http://{host}:{port}"
+
+
+def api_request(
+    method: str,
+    path: str,
+    data: dict[str, Any] | None = None,
+) -> dict[str, Any] | list[Any]:
+    """Make an API request to the Exo server.
+
+    Args:
+        method: HTTP method (GET, POST, DELETE, etc.)
+        path: API path (e.g., "/flash/instances")
+        data: Optional JSON data for POST/PUT requests
+
+    Returns:
+        Parsed JSON response
+
+    Raises:
+        SystemExit: On connection or HTTP errors
+    """
+    url = f"{get_api_base()}{path}"
+
+    request_data = None
+    if data is not None:
+        request_data = json.dumps(data).encode("utf-8")
+
+    req = urllib.request.Request(
+        url,
+        data=request_data,
+        method=method,
+    )
+    req.add_header("Content-Type", "application/json")
+
+    try:
+        with urllib.request.urlopen(req, timeout=30) as response:  # pyright: ignore[reportAny]
+            body: str = response.read().decode("utf-8")  # pyright: ignore[reportAny]
+            if body:
+                return json.loads(body)  # pyright: ignore[reportAny]
+            return {}
+    except HTTPError as e:
+        error_body = e.read().decode("utf-8") if e.fp else ""
+        print(f"API error: {e.code} {e.reason}")
+        if error_body:
+            try:
+                error_json: dict[str, str] = json.loads(error_body)  # pyright: ignore[reportAny]
+                if "detail" in error_json:
+                    print(f"  {error_json['detail']}")
+            except json.JSONDecodeError:
+                print(f"  {error_body}")
+        raise SystemExit(1) from None
+    except URLError as e:
+        print(f"Connection error: {e.reason}")
+        print(f"Is Exo running at {get_api_base()}?")
+        raise SystemExit(1) from None
+
+
+def truncate_id(instance_id: str, length: int = 8) -> str:
+    """Truncate a UUID for display.
+
+    Args:
+        instance_id: Full UUID string
+        length: Number of characters to keep
+
+    Returns:
+        Truncated ID without hyphens
+    """
+    return instance_id.replace("-", "")[:length]
+
+
+def format_table(headers: list[str], rows: list[list[str]]) -> str:
+    """Format data as a simple text table.
+
+    Args:
+        headers: Column headers
+        rows: List of rows, each row is a list of column values
+
+    Returns:
+        Formatted table string
+    """
+    if not rows:
+        return "  ".join(f"{h:<10}" for h in headers)
+
+    # Calculate column widths
+    widths = [len(h) for h in headers]
+    for row in rows:
+        for i, cell in enumerate(row):
+            if i < len(widths):
+                widths[i] = max(widths[i], len(cell))
+
+    # Build format string
+    fmt = "  ".join(f"{{:<{w}}}" for w in widths)
+
+    # Format output
+    lines = [fmt.format(*headers)]
+    for row in rows:
+        # Pad row if needed
+        padded = row + [""] * (len(headers) - len(row))
+        lines.append(fmt.format(*padded[: len(headers)]))
+
+    return "\n".join(lines)

src/exo/cli/salloc.py

@@ -0,0 +1,100 @@
+"""salloc - Allocate nodes for interactive use.
+
+Usage:
+    exo salloc [options] [-- command [args...]]
+
+Options:
+    -N, --nodes N       Number of nodes to allocate (default: 1)
+    --hosts HOSTS       Comma-separated list of hostnames
+
+If a command is provided after --, it will be executed with
+SLURM-like environment variables set:
+    SLURM_JOB_NODELIST  - Comma-separated list of allocated nodes
+    SLURM_NNODES        - Number of allocated nodes
+
+Examples:
+    exo salloc --nodes=2 --hosts=node1,node2 -- mpirun ./my_program
+    exo salloc --hosts=localhost -- bash
+"""
+
+import argparse
+import os
+import subprocess
+import sys
+
+
+def main(args: list[str]) -> int:
+    """Main entry point for salloc command."""
+    # Split args at -- if present
+    cmd_args: list[str] = []
+    salloc_args = args
+
+    if "--" in args:
+        idx = args.index("--")
+        salloc_args = args[:idx]
+        cmd_args = args[idx + 1 :]
+
+    parser = argparse.ArgumentParser(
+        prog="exo salloc",
+        description="Allocate nodes for interactive use",
+    )
+    parser.add_argument(
+        "-N",
+        "--nodes",
+        type=int,
+        default=1,
+        help="Number of nodes to allocate (default: 1)",
+    )
+    parser.add_argument(
+        "--hosts",
+        help="Comma-separated list of hostnames (required)",
+    )
+
+    parsed = parser.parse_args(salloc_args)
+
+    nodes: int = parsed.nodes  # pyright: ignore[reportAny]
+    hosts: str | None = parsed.hosts  # pyright: ignore[reportAny]
+
+    # Require explicit hosts since we can't discover them from topology
+    if not hosts:
+        print("Error: --hosts is required (e.g., --hosts=node1,node2)", file=sys.stderr)
+        print("       The Exo topology doesn't expose hostnames.", file=sys.stderr)
+        return 1
+
+    host_list = [h.strip() for h in hosts.split(",") if h.strip()]
+
+    if len(host_list) < nodes:
+        print(
+            f"Error: Requested {nodes} nodes but only {len(host_list)} hosts provided",
+            file=sys.stderr,
+        )
+        return 1
+
+    # Use first N hosts
+    allocated_hosts = host_list[:nodes]
+    nodelist = ",".join(allocated_hosts)
+
+    # Set environment variables
+    env = os.environ.copy()
+    env["SLURM_JOB_NODELIST"] = nodelist
+    env["SLURM_NNODES"] = str(nodes)
+
+    print(f"salloc: Granted job allocation on {nodes} node(s)")
+    print(f"salloc: Nodes: {nodelist}")
+
+    if cmd_args:
+        # Run the command
+        print(f"salloc: Running: {' '.join(cmd_args)}")
+        result = subprocess.run(cmd_args, env=env)
+        return result.returncode
+    else:
+        # Start interactive shell
+        shell = os.environ.get("SHELL", "/bin/bash")
+        print(f"salloc: Starting shell {shell}")
+        print("salloc: Use 'exit' to release allocation")
+        result = subprocess.run([shell], env=env)
+        return result.returncode
+
+
+if __name__ == "__main__":
+    sys.exit(main(sys.argv[1:]))

src/exo/cli/sbatch.py

@@ -0,0 +1,233 @@
+"""sbatch - Submit a batch job to Exo.
+
+Usage:
+    exo sbatch [options] <script|executable>
+    exo sbatch --job-name=NAME --nodes=N <executable>
+
+Options:
+    -J, --job-name NAME         Job name
+    -N, --nodes N               Number of nodes (default: 1)
+    --ntasks-per-node N         Tasks per node (default: 1)
+    -D, --chdir DIR             Working directory
+    --hosts HOSTS               Comma-separated list of hostnames
+
+Job scripts can contain #SBATCH directives:
+    #!/bin/bash
+    #SBATCH --job-name=Sod2D
+    #SBATCH --nodes=2
+    #SBATCH --chdir=/path/to/workdir
+
+    /path/to/flash4
+"""
+
+import argparse
+import os
+import re
+import sys
+
+from exo.cli.common import api_request, truncate_id
+
+
+def parse_job_script(script_path: str) -> tuple[dict[str, str], str | None]:
+    """Parse a job script for #SBATCH directives and executable.
+
+    Args:
+        script_path: Path to the job script
+
+    Returns:
+        Tuple of (directives dict, executable path or None)
+    """
+    directives: dict[str, str] = {}
+    executable: str | None = None
+
+    with open(script_path, "r") as f:
+        for line in f:
+            line = line.strip()
+
+            # Parse #SBATCH directives
+            if line.startswith("#SBATCH"):
+                # Handle both --option=value and --option value formats
+                match = re.match(r"#SBATCH\s+(-\w|--[\w-]+)(?:=|\s+)(.+)", line)
+                if match:
+                    opt, val = match.groups()
+                    directives[opt.lstrip("-")] = val.strip()
+                continue
+
+            # Skip comments and empty lines
+            if line.startswith("#") or not line:
+                continue
+
+            # First non-comment, non-directive line is the executable
+            if executable is None:
+                # Handle lines like "/path/to/flash4" or "srun /path/to/flash4"
+                parts = line.split()
+                if parts:
+                    # Skip srun/mpirun prefixes if present
+                    for part in parts:
+                        if not part.startswith("-") and "/" in part:
+                            executable = part
+                            break
+                    if executable is None and parts:
+                        executable = parts[-1]  # Last token
+
+    return directives, executable
+
+
+def main(args: list[str]) -> int:
+    """Main entry point for sbatch command."""
+    parser = argparse.ArgumentParser(
+        prog="exo sbatch",
+        description="Submit a batch job to Exo",
+    )
+    parser.add_argument(
+        "script",
+        help="Job script or executable path",
+    )
+    parser.add_argument(
+        "-J",
+        "--job-name",
+        dest="job_name",
+        help="Job name",
+    )
+    parser.add_argument(
+        "-N",
+        "--nodes",
+        type=int,
+        default=1,
+        help="Number of nodes (default: 1)",
+    )
+    parser.add_argument(
+        "--ntasks-per-node",
+        type=int,
+        default=1,
+        help="Tasks per node (default: 1)",
+    )
+    parser.add_argument(
+        "-D",
+        "--chdir",
+        help="Working directory",
+    )
+    parser.add_argument(
+        "--hosts",
+        help="Comma-separated list of hostnames",
+    )
+
+    parsed = parser.parse_args(args)
+
+    # Extract typed values from namespace
+    script_path: str = parsed.script  # pyright: ignore[reportAny]
+    arg_job_name: str | None = parsed.job_name  # pyright: ignore[reportAny]
+    arg_nodes: int = parsed.nodes  # pyright: ignore[reportAny]
+    arg_ntasks: int = parsed.ntasks_per_node  # pyright: ignore[reportAny]
+    arg_chdir: str | None = parsed.chdir  # pyright: ignore[reportAny]
+    arg_hosts: str | None = parsed.hosts  # pyright: ignore[reportAny]
+
+    # Determine if input is a script or direct executable
+    executable: str | None = None
+    directives: dict[str, str] = {}
+
+    if os.path.isfile(script_path):
+        # Check if it's a binary file (executable) or text script
+        is_binary = False
+        try:
+            with open(script_path, "rb") as f:
+                chunk = f.read(512)
+                # Binary files typically contain null bytes
+                is_binary = b"\x00" in chunk
+        except OSError:
+            pass
+
+        if is_binary:
+            # It's a binary executable
+            executable = script_path
+        else:
+            # Try to read as text
+            try:
+                with open(script_path, "r") as f:
+                    first_line = f.readline()
+                    f.seek(0)
+                    content = f.read(1024)
+
+                if first_line.startswith("#!") or "#SBATCH" in content:
+                    # It's a job script - parse it
+                    directives, executable = parse_job_script(script_path)
+                else:
+                    # It's an executable (text but no shebang/directives)
+                    executable = script_path
+            except UnicodeDecodeError:
+                # Can't read as text - treat as binary executable
+                executable = script_path
+    else:
+        # Not a file - treat as executable path
+        executable = script_path
+
+    if executable is None:
+        print("Error: No executable found in job script", file=sys.stderr)
+        return 1
+
+    # Build job parameters - CLI args override script directives
+    job_name = arg_job_name or directives.get("job-name") or directives.get("J")
+    if not job_name:
+        # Generate name from executable
+        job_name = os.path.basename(executable).replace(".", "_")
+
+    nodes = arg_nodes
+    if "nodes" in directives:
+        nodes = int(directives["nodes"])
+    if "N" in directives:
+        nodes = int(directives["N"])
+    if arg_nodes != 1:  # CLI override
+        nodes = arg_nodes
+
+    ntasks = arg_ntasks
+    if "ntasks-per-node" in directives:
+        ntasks = int(directives["ntasks-per-node"])
+    if arg_ntasks != 1:  # CLI override
+        ntasks = arg_ntasks
+
+    workdir = arg_chdir or directives.get("chdir") or directives.get("D")
+    if not workdir:
+        workdir = os.getcwd()
+
+    hosts = arg_hosts or directives.get("hosts") or ""
+
+    # Resolve executable to absolute path
+    if not os.path.isabs(executable):
+        executable = os.path.abspath(os.path.join(workdir, executable))
+
+    # Submit job via API using query parameters
+    from urllib.parse import urlencode
+
+    params = {
+        "simulation_name": job_name,
+        "flash_executable_path": executable,
+        "parameter_file_path": "",  # FLASH par file - use default
+        "working_directory": workdir,
+        "ranks_per_node": str(ntasks),
+        "min_nodes": str(nodes),
+        "hosts": hosts,
+    }
+
+    query_string = urlencode(params)
+    result = api_request("POST", f"/flash/launch?{query_string}")
+
+    # Print job submission confirmation
+    if isinstance(result, dict):
+        instance_id_val = result.get("instance_id")
+
+        if instance_id_val is not None:
+            job_id = truncate_id(str(instance_id_val))  # pyright: ignore[reportAny]
+            print(f"Submitted batch job {job_id}")
+        else:
+            # Instance created asynchronously - user should check squeue
+            print("Job submitted successfully")
+            print("Use 'exo squeue' to view job ID")
+    else:
+        print("Job submitted successfully")
+        print("Use 'exo squeue' to view job ID")
+
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main(sys.argv[1:]))

src/exo/cli/scancel.py

@@ -0,0 +1,95 @@
+"""scancel - Cancel jobs in the Exo queue.
+
+Usage:
+    exo scancel <jobid> [<jobid>...]
+
+Arguments:
+    jobid   Job ID (or prefix) to cancel. Can specify multiple.
+
+Examples:
+    exo scancel abc123          # Cancel job starting with abc123
+    exo scancel abc123 def456   # Cancel multiple jobs
+"""
+
+import argparse
+import sys
+from typing import Any, cast
+
+from exo.cli.common import api_request, truncate_id
+
+
+def main(args: list[str]) -> int:
+    """Main entry point for scancel command."""
+    parser = argparse.ArgumentParser(
+        prog="exo scancel",
+        description="Cancel jobs in the Exo queue",
+    )
+    parser.add_argument(
+        "jobids",
+        nargs="+",
+        help="Job ID(s) to cancel",
+    )
+
+    parsed = parser.parse_args(args)
+    jobids: list[str] = parsed.jobids  # pyright: ignore[reportAny]
+
+    # Fetch current jobs to resolve partial IDs
+    result = api_request("GET", "/flash/instances")
+    if isinstance(result, list):
+        instances = cast(list[dict[str, Any]], result)
+    else:
+        instances = cast(list[dict[str, Any]], result.get("instances", []))
+
+    # Build lookup of full IDs
+    id_map: dict[str, str] = {}
+    for inst in instances:
+        iid = inst.get("instance_id", "")  # pyright: ignore[reportAny]
+        full_id = str(iid) if iid else ""  # pyright: ignore[reportAny]
+        if full_id:
+            # Map both full ID and truncated versions
+            normalized = full_id.replace("-", "").lower()
+            id_map[normalized] = full_id
+            # Also map prefixes
+            for length in range(4, len(normalized) + 1):
+                prefix = normalized[:length]
+                if prefix not in id_map:
+                    id_map[prefix] = full_id
+
+    cancelled = 0
+    errors = 0
+
+    for jobid in jobids:
+        search = jobid.lower().replace("-", "")
+
+        # Find matching full ID
+        full_id = id_map.get(search)
+        if not full_id:
+            # Try prefix match
+            matches = [fid for key, fid in id_map.items() if key.startswith(search)]
+            if len(matches) == 1:
+                full_id = matches[0]
+            elif len(matches) > 1:
+                print(f"Ambiguous job ID: {jobid} matches multiple jobs")
+                errors += 1
+                continue
+            else:
+                print(f"Job not found: {jobid}")
+                errors += 1
+                continue
+
+        # Cancel the job
+        try:
+            api_request("DELETE", f"/flash/{full_id}")
+            print(f"Job {truncate_id(full_id)} cancelled")
+            cancelled += 1
+        except SystemExit:
+            print(f"Failed to cancel job {truncate_id(full_id)}")
+            errors += 1
+
+    if errors > 0 and cancelled == 0:
+        return 1
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main(sys.argv[1:]))

src/exo/cli/squeue.py

@@ -0,0 +1,165 @@
+"""squeue - View the Exo job queue.
+
+Usage:
+    exo squeue [options]
+
+Options:
+    -l, --long      Show detailed output
+    -j, --job ID    Show only this job
+
+Output columns:
+    JOBID   - Job identifier (truncated UUID)
+    NAME    - Job name
+    NODES   - Number of nodes
+    STATE   - Job state (PENDING, RUNNING, FAILED, etc.)
+"""
+
+import argparse
+import sys
+from typing import Any, cast
+
+from exo.cli.common import api_request, format_table, truncate_id
+
+# Map Exo runner statuses to SLURM-like states
+STATUS_MAP: dict[str, str] = {
+    "RunnerIdle": "PENDING",
+    "RunnerConnecting": "CONFIGURING",
+    "RunnerConnected": "CONFIGURING",
+    "RunnerLoading": "CONFIGURING",
+    "RunnerLoaded": "CONFIGURING",
+    "RunnerWarmingUp": "CONFIGURING",
+    "RunnerReady": "COMPLETING",
+    "RunnerRunning": "RUNNING",
+    "RunnerShuttingDown": "COMPLETING",
+    "RunnerShutdown": "COMPLETED",
+    "RunnerFailed": "FAILED",
+}
+
+
+def get_job_state(runner_statuses: dict[str, Any]) -> str:
+    """Determine overall job state from runner statuses."""
+    if not runner_statuses:
+        return "PENDING"
+
+    states: set[str] = set()
+    for status_val in runner_statuses.values():  # pyright: ignore[reportAny]
+        if isinstance(status_val, dict):
+            # Extract status type from discriminated union
+            type_val = status_val.get("type", "RunnerIdle")  # pyright: ignore[reportUnknownMemberType,reportUnknownVariableType]
+            status_type = str(type_val) if type_val else "RunnerIdle"  # pyright: ignore[reportUnknownArgumentType]
+        elif isinstance(status_val, str):
+            status_type = status_val
+        else:
+            status_type = "RunnerIdle"
+        # Strip parentheses from status strings like "RunnerRunning()"
+        if status_type.endswith("()"):
+            status_type = status_type[:-2]
+        states.add(STATUS_MAP.get(status_type, "UNKNOWN"))
+
+    # Priority order for overall state
+    if "FAILED" in states:
+        return "FAILED"
+    if "RUNNING" in states:
+        return "RUNNING"
+    if "CONFIGURING" in states:
+        return "CONFIGURING"
+    if "COMPLETING" in states:
+        return "COMPLETING"
+    if "COMPLETED" in states:
+        return "COMPLETED"
+    if "PENDING" in states:
+        return "PENDING"
+    return "UNKNOWN"
+
+
+def main(args: list[str]) -> int:
+    """Main entry point for squeue command."""
+    parser = argparse.ArgumentParser(
+        prog="exo squeue",
+        description="View the Exo job queue",
+    )
+    parser.add_argument(
+        "-l",
+        "--long",
+        action="store_true",
+        help="Show detailed output",
+    )
+    parser.add_argument(
+        "-j",
+        "--job",
+        help="Show only this job ID",
+    )
+
+    parsed = parser.parse_args(args)
+
+    # Extract typed values
+    long_format: bool = parsed.long  # pyright: ignore[reportAny]
+    job_filter: str | None = parsed.job  # pyright: ignore[reportAny]
+
+    # Fetch jobs from API - returns list directly
+    result = api_request("GET", "/flash/instances")
+    # API returns list directly, not {"instances": [...]}
+    if isinstance(result, list):
+        instances = cast(list[dict[str, Any]], result)
+    else:
+        instances = cast(list[dict[str, Any]], result.get("instances", []))
+
+    if not instances:
+        # No jobs - just print header
+        if long_format:
+            print("JOBID           NAME            NODES  RANKS  STATE        WORKDIR")
+        else:
+            print("JOBID       NAME            NODES  STATE")
+        return 0
+
+    # Filter by job ID if specified
+    if job_filter:
+        search = job_filter.lower()
+        filtered: list[dict[str, Any]] = []
+        for i in instances:
+            iid = i.get("instance_id", "")  # pyright: ignore[reportAny]
+            if search in str(iid).lower().replace("-", ""):  # pyright: ignore[reportAny]
+                filtered.append(i)
+        instances = filtered
+
+    # Build table
+    rows: list[list[str]] = []
+
+    if long_format:
+        headers = ["JOBID", "NAME", "NODES", "RANKS", "STATE", "WORKDIR"]
+        for inst in instances:
+            iid_val = inst.get("instance_id", "")  # pyright: ignore[reportAny]
+            instance_id = str(iid_val) if iid_val else ""  # pyright: ignore[reportAny]
+            job_id = truncate_id(instance_id, 12)
+            name_val = inst.get("simulation_name", "")  # pyright: ignore[reportAny]
+            name = (str(name_val) if name_val else "")[:15]  # pyright: ignore[reportAny]
+            runner_statuses = cast(dict[str, Any], inst.get("runner_statuses", {}))
+            nodes = str(len(runner_statuses))
+            ranks_val = inst.get("total_ranks", 0)  # pyright: ignore[reportAny]
+            ranks = str(ranks_val) if ranks_val else "0"  # pyright: ignore[reportAny]
+            state = get_job_state(runner_statuses)
+            workdir_val = inst.get("working_directory", "")  # pyright: ignore[reportAny]
+            workdir = str(workdir_val) if workdir_val else ""  # pyright: ignore[reportAny]
+            # Truncate workdir for display
+            if len(workdir) > 30:
+                workdir = "..." + workdir[-27:]
+            rows.append([job_id, name, nodes, ranks, state, workdir])
+    else:
+        headers = ["JOBID", "NAME", "NODES", "STATE"]
+        for inst in instances:
+            iid_val = inst.get("instance_id", "")  # pyright: ignore[reportAny]
+            instance_id = str(iid_val) if iid_val else ""  # pyright: ignore[reportAny]
+            job_id = truncate_id(instance_id, 8)
+            name_val = inst.get("simulation_name", "")  # pyright: ignore[reportAny]
+            name = (str(name_val) if name_val else "")[:15]  # pyright: ignore[reportAny]
+            runner_statuses = cast(dict[str, Any], inst.get("runner_statuses", {}))
+            nodes = str(len(runner_statuses))
+            state = get_job_state(runner_statuses)
+            rows.append([job_id, name, nodes, state])
+
+    print(format_table(headers, rows))
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main(sys.argv[1:]))

src/exo/download/coordinator.py

@@ -1,284 +0,0 @@
-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/tests/test_download_verification.py

@@ -1,451 +0,0 @@
-"""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,11 +1,10 @@
 import argparse
-import itertools
 import multiprocessing as mp
 import os
 import resource
 import signal
 from dataclasses import dataclass, field
-from typing import Iterator, Self
+from typing import Self
 
 import anyio
 from anyio.abc import TaskGroup
@@ -13,8 +12,6 @@ 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
@@ -24,6 +21,7 @@ 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
 
 
@@ -31,7 +29,6 @@ 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]
@@ -39,7 +36,6 @@ 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
@@ -53,26 +49,8 @@ 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,
@@ -80,7 +58,6 @@ 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:
@@ -90,11 +67,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
@@ -122,25 +99,13 @@ class Node:
             election_result_sender=er_send,
         )
 
-        return cls(
-            router,
-            download_coordinator,
-            worker,
-            election,
-            er_recv,
-            master,
-            api,
-            node_id,
-            event_index_counter,
-        )
+        return cls(router, worker, election, er_recv, master, api, node_id)
 
     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:
@@ -205,36 +170,21 @@ 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:
@@ -245,6 +195,14 @@ class Node:
 
 
 def main():
+    # Check for SLURM-compatible subcommands first
+    import sys
+
+    if len(sys.argv) > 1 and sys.argv[1] in ("sbatch", "squeue", "scancel", "salloc"):
+        from exo.cli import run_subcommand
+
+        sys.exit(run_subcommand(sys.argv[1], sys.argv[2:]))
+
     args = Args.parse()
     soft, hard = resource.getrlimit(resource.RLIMIT_NOFILE)
     resource.setrlimit(resource.RLIMIT_NOFILE, (max(soft, 65535), hard))
@@ -255,6 +213,11 @@ def main():
     logger.info("Starting EXO")
     logger.info(f"EXO_LIBP2P_NAMESPACE: {os.getenv('EXO_LIBP2P_NAMESPACE')}")
 
+    # Discover and register plugins
+    from exo.plugins.registry import discover_plugins
+
+    discover_plugins()
+
     # Set FAST_SYNCH override env var for runner subprocesses
     if args.fast_synch is True:
         os.environ["EXO_FAST_SYNCH"] = "on"
@@ -276,7 +239,6 @@ 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
@@ -319,11 +281,6 @@ 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

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

src/exo/master/adapters/chat_completions.py

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

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

@@ -1,369 +0,0 @@
-"""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/main.py

@@ -10,9 +10,10 @@ from exo.master.placement import (
     get_transition_events,
     place_instance,
 )
+from exo.plugins.registry import PluginRegistry
 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,10 +24,10 @@ 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 (
+    BaseEvent,
     Event,
     ForwarderEvent,
     IndexedEvent,
@@ -36,11 +37,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,
 )
@@ -51,9 +52,6 @@ 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
@@ -87,11 +85,9 @@ class Master:
         self._loopback_event_sender: Sender[ForwarderEvent] = (
             local_event_receiver.clone_sender()
         )
-        self._multi_buffer = MultiSourceBuffer[NodeId, Event]()
+        self._multi_buffer = MultiSourceBuffer[NodeId, BaseEvent]()
         # 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]] = {}
+        self._event_log: list[BaseEvent] = []
 
     async def run(self):
         logger.info("Starting Master")
@@ -123,11 +119,11 @@ class Master:
                     match command:
                         case TestCommand():
                             pass
-                        case TextGeneration():
+                        case ChatCompletion():
                             for instance in self.state.instances.values():
                                 if (
                                     instance.shard_assignments.model_id
-                                    == command.task_params.model
+                                    == command.request_params.model
                                 ):
                                     task_count = sum(
                                         1
@@ -140,7 +136,7 @@ class Master:
 
                             if not instance_task_counts:
                                 raise ValueError(
-                                    f"No instance found for model {command.task_params.model}"
+                                    f"No instance found for model {command.request_params.model}"
                                 )
 
                             available_instance_ids = sorted(
@@ -154,12 +150,12 @@ class Master:
                             generated_events.append(
                                 TaskCreated(
                                     task_id=task_id,
-                                    task=TextGenerationTask(
+                                    task=ChatCompletionTask(
                                         task_id=task_id,
                                         command_id=command.command_id,
                                         instance_id=available_instance_ids[0],
                                         task_status=TaskStatus.Pending,
-                                        task_params=command.task_params,
+                                        task_params=command.request_params,
                                     ),
                                 )
                             )
@@ -169,7 +165,7 @@ class Master:
                             for instance in self.state.instances.values():
                                 if (
                                     instance.shard_assignments.model_id
-                                    == command.task_params.model
+                                    == command.request_params.model
                                 ):
                                     task_count = sum(
                                         1
@@ -182,7 +178,7 @@ class Master:
 
                             if not instance_task_counts:
                                 raise ValueError(
-                                    f"No instance found for model {command.task_params.model}"
+                                    f"No instance found for model {command.request_params.model}"
                                 )
 
                             available_instance_ids = sorted(
@@ -193,37 +189,25 @@ 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=selected_instance_id,
+                                        instance_id=available_instance_ids[0],
                                         task_status=TaskStatus.Pending,
-                                        task_params=command.task_params,
+                                        task_params=command.request_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.task_params.model
+                                    == command.request_params.model
                                 ):
                                     task_count = sum(
                                         1
@@ -236,7 +220,7 @@ class Master:
 
                             if not instance_task_counts:
                                 raise ValueError(
-                                    f"No instance found for model {command.task_params.model}"
+                                    f"No instance found for model {command.request_params.model}"
                                 )
 
                             available_instance_ids = sorted(
@@ -247,32 +231,20 @@ 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=selected_instance_id,
+                                        instance_id=available_instance_ids[0],
                                         task_status=TaskStatus.Pending,
-                                        task_params=command.task_params,
+                                        task_params=command.request_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(
@@ -326,6 +298,17 @@ class Master:
                                 await self._send_event(
                                     IndexedEvent(idx=i, event=self._event_log[i])
                                 )
+                        case _:
+                            # Check if a plugin handles this command
+                            registry = PluginRegistry.get()
+                            plugin = registry.get_plugin_for_command(command)
+                            if plugin is not None:
+                                events = plugin.process_command(
+                                    command,
+                                    self.state.topology,
+                                    self.state.instances,
+                                )
+                                generated_events.extend(events)
                     for event in generated_events:
                         await self.event_sender.send(event)
                 except ValueError as e:
@@ -365,10 +348,6 @@ 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)
@@ -407,29 +386,3 @@ 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

@@ -24,7 +24,7 @@ from exo.shared.types.events import Event, InstanceCreated, InstanceDeleted
 from exo.shared.types.memory import Memory
 from exo.shared.types.profiling import MemoryUsage, NodeNetworkInfo
 from exo.shared.types.worker.instances import (
-    Instance,
+    BaseInstance,
     InstanceId,
     InstanceMeta,
     MlxJacclInstance,
@@ -35,14 +35,14 @@ 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 port
+    return port - 1 if port <= 52415 else 52414
 
 
 def add_instance_to_placements(
     command: CreateInstance,
     topology: Topology,
-    current_instances: Mapping[InstanceId, Instance],
-) -> Mapping[InstanceId, Instance]:
+    current_instances: Mapping[InstanceId, BaseInstance],
+) -> Mapping[InstanceId, BaseInstance]:
     # TODO: validate against topology
 
     return {**current_instances, command.instance.instance_id: command.instance}
@@ -51,21 +51,12 @@ def add_instance_to_placements(
 def place_instance(
     command: PlaceInstance,
     topology: Topology,
-    current_instances: Mapping[InstanceId, Instance],
+    current_instances: Mapping[InstanceId, BaseInstance],
     node_memory: Mapping[NodeId, MemoryUsage],
     node_network: Mapping[NodeId, NodeNetworkInfo],
-    required_nodes: set[NodeId] | None = None,
-) -> dict[InstanceId, Instance]:
+) -> dict[InstanceId, BaseInstance]:
     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
     )
@@ -96,12 +87,12 @@ def place_instance(
 
     smallest_cycles = get_smallest_cycles(cycles_with_sufficient_memory)
 
-    smallest_rdma_cycles = [
-        cycle for cycle in smallest_cycles if topology.is_rdma_cycle(cycle)
+    smallest_tb_cycles = [
+        cycle for cycle in smallest_cycles if topology.is_thunderbolt_cycle(cycle)
     ]
 
-    if command.instance_meta == InstanceMeta.MlxJaccl and smallest_rdma_cycles != []:
-        smallest_cycles = smallest_rdma_cycles
+    if smallest_tb_cycles != []:
+        smallest_cycles = smallest_tb_cycles
 
     cycles_with_leaf_nodes: list[Cycle] = [
         cycle
@@ -168,8 +159,8 @@ def place_instance(
 
 def delete_instance(
     command: DeleteInstance,
-    current_instances: Mapping[InstanceId, Instance],
-) -> dict[InstanceId, Instance]:
+    current_instances: Mapping[InstanceId, BaseInstance],
+) -> dict[InstanceId, BaseInstance]:
     target_instances = dict(deepcopy(current_instances))
     if command.instance_id in target_instances:
         del target_instances[command.instance_id]
@@ -178,8 +169,8 @@ def delete_instance(
 
 
 def get_transition_events(
-    current_instances: Mapping[InstanceId, Instance],
-    target_instances: Mapping[InstanceId, Instance],
+    current_instances: Mapping[InstanceId, BaseInstance],
+    target_instances: Mapping[InstanceId, BaseInstance],
 ) -> Sequence[Event]:
     events: list[Event] = []
 

src/exo/master/placement_utils.py

@@ -197,6 +197,49 @@ 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,
@@ -222,6 +265,9 @@ 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"
                 )
@@ -233,11 +279,22 @@ def _find_connection_ip(
     node_i: NodeId,
     node_j: NodeId,
     cycle_digraph: Topology,
-) -> Generator[str, None, None]:
+) -> Generator[tuple[str, bool]]:
     """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
+            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
 
 
 def _find_ip_prioritised(
@@ -246,25 +303,43 @@ 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: ethernet > wifi > unknown > thunderbolt
+    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
     """
     ips = list(_find_connection_ip(node_id, other_node_id, cycle_digraph))
-    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
+    # 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
     }
-    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))
+
+    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
 
 
 def get_mlx_ring_hosts_by_node(
@@ -306,6 +381,9 @@ 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"
                 )
@@ -338,6 +416,9 @@ 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,9 +1,13 @@
 # pyright: reportUnusedFunction=false, reportAny=false
-from typing import Any
+from typing import Any, get_args
 
 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."""
@@ -44,3 +48,95 @@ 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

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

@@ -1,265 +0,0 @@
-"""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,14 +7,15 @@ 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, ModelTask
+from exo.shared.models.model_cards import ModelCard, ModelId, ModelTask
+from exo.shared.types.api import ChatCompletionMessage, ChatCompletionTaskParams
 from exo.shared.types.commands import (
+    ChatCompletion,
     CommandId,
     ForwarderCommand,
     PlaceInstance,
-    TextGeneration,
 )
-from exo.shared.types.common import ModelId, NodeId, SessionId
+from exo.shared.types.common import NodeId, SessionId
 from exo.shared.types.events import (
     ForwarderEvent,
     IndexedEvent,
@@ -26,9 +27,8 @@ 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,16 +127,20 @@ 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 TextGeneration Command")
+        logger.info("inject a ChatCompletion Command")
         await command_sender.send(
             ForwarderCommand(
                 origin=node_id,
                 command=(
-                    TextGeneration(
+                    ChatCompletion(
                         command_id=CommandId(),
-                        task_params=TextGenerationTaskParams(
-                            model=ModelId("llama-3.2-1b"),
-                            input="Hello, how are you?",
+                        request_params=ChatCompletionTaskParams(
+                            model="llama-3.2-1b",
+                            messages=[
+                                ChatCompletionMessage(
+                                    role="user", content="Hello, how are you?"
+                                )
+                            ],
                         ),
                     )
                 ),
@@ -186,10 +190,12 @@ 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, TextGenerationTask)
-        assert events[2].event.task.task_params == TextGenerationTaskParams(
-            model=ModelId("llama-3.2-1b"),
-            input="Hello, how are you?",
+        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?")
+            ],
         )
 
         await master.shutdown()

src/exo/master/tests/test_openai_responses_api.py

@@ -1,48 +0,0 @@
-"""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,6 +3,7 @@ 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,
@@ -13,7 +14,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 NodeId
+from exo.shared.types.common import Host, NodeId
 from exo.shared.types.memory import Memory
 from exo.shared.types.profiling import (
     NetworkInterfaceInfo,
@@ -272,6 +273,45 @@ 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/plugins/__init__.py

@@ -0,0 +1,24 @@
+"""Exo Plugin System.
+
+This module provides the plugin architecture for extending exo with custom
+workload types (simulations, ML frameworks, etc.) without modifying core code.
+"""
+
+from exo.plugins.base import EXOPlugin, PluginCommand, PluginInstance
+from exo.plugins.registry import PluginRegistry, discover_plugins
+from exo.plugins.type_registry import (
+    command_registry,
+    event_registry,
+    instance_registry,
+)
+
+__all__ = [
+    "EXOPlugin",
+    "PluginCommand",
+    "PluginInstance",
+    "PluginRegistry",
+    "discover_plugins",
+    "command_registry",
+    "event_registry",
+    "instance_registry",
+]

src/exo/plugins/base.py

@@ -0,0 +1,171 @@
+"""Base classes and protocols for Exo plugins."""
+
+from abc import ABC, abstractmethod
+from collections.abc import Callable, Mapping, Sequence
+from typing import TYPE_CHECKING, Any
+
+from pydantic import Field
+
+from exo.shared.types.common import CommandId
+from exo.shared.types.events import Event
+from exo.shared.types.tasks import Task
+from exo.shared.types.worker.instances import InstanceId
+from exo.shared.types.worker.runners import RunnerId
+from exo.utils.pydantic_ext import TaggedModel
+
+if TYPE_CHECKING:
+    from exo.shared.topology import Topology
+    from exo.shared.types.worker.instances import BaseInstance, BoundInstance
+    from exo.utils.channels import MpReceiver, MpSender
+    from exo.worker.runner.runner_supervisor import RunnerSupervisor
+
+
+class PluginCommand(TaggedModel):
+    """Base class for plugin-defined commands.
+
+    All plugin commands must inherit from this class. Commands are serialized
+    with their class name as a tag for routing.
+    """
+
+    command_id: CommandId = Field(default_factory=CommandId)
+
+
+class PluginInstance(TaggedModel):
+    """Base class for plugin-defined instances.
+
+    All plugin instances must inherit from this class. Plugins are expected
+    to define their own instance type with workload-specific fields.
+    """
+
+    instance_id: InstanceId
+
+
+class EXOPlugin(ABC):
+    """Protocol that all exo plugins must implement.
+
+    A plugin provides:
+    - Custom command types for API -> Master communication
+    - Custom instance types representing running workloads
+    - Placement logic for distributing work across nodes
+    - Planning logic for local task scheduling
+    - Runner implementation for executing work
+    """
+
+    @property
+    @abstractmethod
+    def name(self) -> str:
+        """Unique identifier for this plugin (e.g., 'flash', 'pytorch', 'mpi')."""
+        ...
+
+    @property
+    @abstractmethod
+    def version(self) -> str:
+        """Semantic version string (e.g., '1.0.0')."""
+        ...
+
+    # ========== Type Registration ==========
+
+    @abstractmethod
+    def get_command_types(self) -> Sequence[type]:
+        """Return command types this plugin handles.
+
+        These commands are routed to this plugin's process_command method.
+        Can return core BaseCommand types or PluginCommand types.
+        """
+        ...
+
+    @abstractmethod
+    def get_instance_type(self) -> type:
+        """Return the instance type this plugin creates.
+
+        This instance type is used for routing in planning and runner bootstrap.
+        Can return core Instance types or PluginInstance types.
+        """
+        ...
+
+    # ========== API Routes ==========
+
+    @abstractmethod
+    def get_api_routes(
+        self,
+    ) -> Sequence[tuple[str, str, Callable[..., Any]]]:
+        """Return FastAPI routes to register.
+
+        Each tuple: (method, path, handler)
+        Example: [('post', '/flash/launch', self.launch_handler)]
+
+        Handlers receive a PluginContext with access to:
+        - state: Current State object
+        - send_command: Async function to send commands
+        - node_id: Current node's ID
+        """
+        ...
+
+    # ========== Master Command Handling ==========
+
+    @abstractmethod
+    def handles_command(self, command: Any) -> bool:  # pyright: ignore[reportAny]
+        """Return True if this plugin handles the given command type."""
+        ...
+
+    @abstractmethod
+    def process_command(
+        self,
+        command: Any,  # pyright: ignore[reportAny]
+        topology: "Topology",
+        current_instances: Mapping[InstanceId, "BaseInstance"],
+    ) -> Sequence[Event]:
+        """Process a command and return events to emit.
+
+        Typically creates placement and returns InstanceCreated/InstanceDeleted events.
+
+        Args:
+            command: The command to process
+            topology: Current cluster topology
+            current_instances: Currently running instances
+
+        Returns:
+            Sequence of events to emit (e.g., InstanceCreated, InstanceDeleted)
+        """
+        ...
+
+    # ========== Worker Planning ==========
+
+    @abstractmethod
+    def handles_instance(self, instance: object) -> bool:
+        """Return True if this plugin manages the given instance type."""
+        ...
+
+    @abstractmethod
+    def plan_task(
+        self,
+        runners: Mapping[RunnerId, "RunnerSupervisor"],
+        instances: Mapping[InstanceId, "BaseInstance"],
+    ) -> Task | None:
+        """Plan the next task for plugin instances.
+
+        Called during each planning cycle.
+        Return None if no task is needed.
+        """
+        ...
+
+    @abstractmethod
+    def should_skip_download(self, instance: object) -> bool:
+        """Return True if this instance type doesn't need model downloads."""
+        ...
+
+    # ========== Runner Bootstrap ==========
+
+    @abstractmethod
+    def create_runner(
+        self,
+        bound_instance: "BoundInstance",
+        event_sender: "MpSender[Event]",
+        task_receiver: "MpReceiver[Task]",
+    ) -> None:
+        """Entry point for the runner process.
+
+        Called in a subprocess to execute the actual workload.
+        This function should block until the workload completes.
+        """
+        ...

src/exo/plugins/context.py

@@ -0,0 +1,21 @@
+"""Context objects passed to plugin handlers."""
+
+from collections.abc import Awaitable, Callable
+from dataclasses import dataclass
+
+from exo.shared.types.commands import BaseCommand
+from exo.shared.types.common import NodeId
+from exo.shared.types.state import State
+
+
+@dataclass
+class PluginContext:
+    """Context provided to plugin API handlers.
+
+    This gives plugins access to the current state and the ability to send
+    commands without direct access to internal API components.
+    """
+
+    state: State
+    send_command: Callable[[BaseCommand], Awaitable[None]]
+    node_id: NodeId

src/exo/plugins/implementations/__init__.py

@@ -0,0 +1,5 @@
+"""Plugin implementations directory.
+
+Each subdirectory should contain a plugin with a register() function
+that returns an EXOPlugin instance.
+"""

src/exo/plugins/implementations/flash/__init__.py

@@ -0,0 +1,15 @@
+"""FLASH Plugin - MPI-based simulation support for Exo."""
+
+from exo.plugins.implementations.flash.plugin import FLASHPlugin
+from exo.plugins.implementations.flash.types import (
+    FLASHInstance,
+    LaunchFLASH,
+    StopFLASH,
+)
+
+__all__ = ["FLASHPlugin", "FLASHInstance", "LaunchFLASH", "StopFLASH", "register"]
+
+
+def register() -> FLASHPlugin:
+    """Entry point for plugin discovery."""
+    return FLASHPlugin()

src/exo/plugins/implementations/flash/api_handlers.py

@@ -0,0 +1,109 @@
+"""FLASH plugin API handlers."""
+
+from typing import Any
+
+from fastapi import HTTPException
+
+from exo.plugins.context import PluginContext
+from exo.plugins.implementations.flash.types import (
+    FLASHInstance,
+    LaunchFLASH,
+    StopFLASH,
+)
+
+
+async def handle_launch_flash(
+    ctx: PluginContext,
+    simulation_name: str,
+    flash_executable_path: str,
+    working_directory: str,
+    parameter_file_path: str = "",
+    ranks_per_node: int = 1,
+    min_nodes: int = 1,
+    hosts: str = "",
+) -> dict[str, str]:
+    """Launch a FLASH MPI simulation across the cluster.
+
+    Args:
+        ctx: Plugin context with state and send_command
+        simulation_name: Name of the simulation
+        flash_executable_path: Path to the FLASH executable
+        working_directory: Working directory for the simulation
+        parameter_file_path: Path to parameter file (optional)
+        ranks_per_node: Number of MPI ranks per node
+        min_nodes: Minimum number of nodes required
+        hosts: Optional comma-separated hostnames (e.g., "s14,james21-1").
+               If not provided, IPs are discovered from topology edges.
+    """
+    command = LaunchFLASH(
+        simulation_name=simulation_name,
+        flash_executable_path=flash_executable_path,
+        parameter_file_path=parameter_file_path,
+        working_directory=working_directory,
+        ranks_per_node=ranks_per_node,
+        min_nodes=min_nodes,
+        hosts=hosts,
+    )
+    await ctx.send_command(command)
+
+    return {
+        "message": "FLASH launch command received",
+        "command_id": str(command.command_id),
+        "simulation_name": simulation_name,
+    }
+
+
+async def handle_stop_flash(
+    ctx: PluginContext,
+    instance_id: str,
+) -> dict[str, str]:
+    """Stop a running FLASH simulation."""
+    from exo.shared.types.worker.instances import InstanceId
+
+    inst_id = InstanceId(instance_id)
+
+    if inst_id not in ctx.state.instances:
+        raise HTTPException(status_code=404, detail="Instance not found")
+
+    instance = ctx.state.instances[inst_id]
+    if not isinstance(instance, FLASHInstance):
+        raise HTTPException(
+            status_code=400, detail="Instance is not a FLASH simulation"
+        )
+
+    command = StopFLASH(instance_id=inst_id)
+    await ctx.send_command(command)
+
+    return {
+        "message": "Stop command received",
+        "command_id": str(command.command_id),
+        "instance_id": str(instance_id),
+    }
+
+
+async def handle_list_flash_instances(ctx: PluginContext) -> list[dict[str, Any]]:
+    """List all FLASH simulation instances."""
+    flash_instances: list[dict[str, Any]] = []
+    for instance_id, instance in ctx.state.instances.items():
+        if isinstance(instance, FLASHInstance):
+            # Get runner statuses for this instance
+            runner_statuses: dict[str, str | None] = {}
+            for (
+                node_id,
+                runner_id,
+            ) in instance.shard_assignments.node_to_runner.items():
+                runner_status = ctx.state.runners.get(runner_id)
+                runner_statuses[str(node_id)] = (
+                    str(runner_status) if runner_status else None
+                )
+
+            flash_instances.append(
+                {
+                    "instance_id": str(instance_id),
+                    "simulation_name": instance.simulation_name,
+                    "total_ranks": instance.total_ranks,
+                    "working_directory": instance.working_directory,
+                    "runner_statuses": runner_statuses,
+                }
+            )
+    return flash_instances

src/exo/plugins/implementations/flash/placement.py

@@ -0,0 +1,167 @@
+"""FLASH plugin placement logic."""
+
+from collections.abc import Mapping
+from copy import deepcopy
+
+from loguru import logger
+
+from exo.plugins.implementations.flash.types import FLASHInstance, LaunchFLASH
+from exo.shared.models.model_cards import ModelCard
+from exo.shared.topology import Topology
+from exo.shared.types.common import Host, ModelId, NodeId
+from exo.shared.types.memory import Memory
+from exo.shared.types.topology import SocketConnection
+from exo.shared.types.worker.instances import BaseInstance, InstanceId
+from exo.shared.types.worker.runners import (
+    RunnerId,
+    ShardAssignments,
+)
+from exo.shared.types.worker.shards import PipelineShardMetadata
+
+
+def place_flash_instance(
+    command: LaunchFLASH,
+    topology: Topology,
+    current_instances: Mapping[InstanceId, BaseInstance],
+) -> dict[InstanceId, BaseInstance]:
+    """Place a FLASH simulation instance across available nodes.
+
+    Unlike MLX instances which use ring/JACCL topology for tensor parallelism,
+    FLASH instances use MPI for communication. We just need to provide the
+    node IPs so the runner can generate an MPI hostfile.
+    """
+    instance_id = InstanceId()
+    target_instances: dict[InstanceId, BaseInstance] = dict(deepcopy(current_instances))
+
+    all_nodes = list(topology.list_nodes())
+
+    if len(all_nodes) < command.min_nodes:
+        raise ValueError(
+            f"Not enough nodes: need {command.min_nodes}, have {len(all_nodes)}"
+        )
+
+    # Select nodes (take the first min_nodes)
+    selected_nodes = all_nodes[: command.min_nodes]
+
+    logger.info(
+        f"Placing FLASH instance '{command.simulation_name}' on {len(selected_nodes)} nodes"
+    )
+
+    # Build shard assignments (one runner per node for FLASH)
+    runner_to_shard: dict[RunnerId, PipelineShardMetadata] = {}
+    node_to_runner: dict[NodeId, RunnerId] = {}
+
+    # Create a dummy ModelCard for FLASH (required by ShardMetadata interface)
+    flash_model_card = ModelCard(
+        model_id=ModelId(command.simulation_name),
+        storage_size=Memory(in_bytes=0),
+        n_layers=1,
+        hidden_size=1,
+        supports_tensor=False,
+        tasks=[],
+    )
+
+    for i, node_id in enumerate(selected_nodes):
+        runner_id = RunnerId()
+        node_to_runner[node_id] = runner_id
+        runner_to_shard[runner_id] = PipelineShardMetadata(
+            device_rank=i,
+            world_size=len(selected_nodes),
+            model_card=flash_model_card,
+            start_layer=0,
+            end_layer=1,
+            n_layers=1,
+        )
+
+    shard_assignments = ShardAssignments(
+        model_id=ModelId(command.simulation_name),
+        runner_to_shard=runner_to_shard,
+        node_to_runner=node_to_runner,
+    )
+
+    # Build hosts_by_node - get hostnames/IPs for MPI hostfile generation
+    hosts_by_node: dict[NodeId, list[Host]] = {}
+
+    # If explicit hosts are provided, use them directly
+    if command.hosts:
+        explicit_hosts = [h.strip() for h in command.hosts.split(",") if h.strip()]
+        logger.info(f"FLASH placement: explicit hosts provided: {explicit_hosts}")
+        for i, node_id in enumerate(selected_nodes):
+            if i < len(explicit_hosts):
+                hosts_by_node[node_id] = [Host(ip=explicit_hosts[i], port=0)]
+                logger.info(
+                    f"FLASH placement: node {node_id} (rank {i}) -> IP {explicit_hosts[i]}"
+                )
+            else:
+                logger.warning(
+                    f"Not enough hosts provided for node {i}, using localhost"
+                )
+                hosts_by_node[node_id] = [Host(ip="127.0.0.1", port=0)]
+        logger.info(
+            f"FLASH placement: coordinator will be rank 0 at IP {explicit_hosts[0]}"
+        )
+    else:
+        # Get each node's own IP by looking at how OTHER nodes see it.
+        # out_edges(A) gives edges A->B with B's IP in sink_multiaddr,
+        # so to find A's IP we look at edges B->A from any other node.
+        for node_id in selected_nodes:
+            candidate_ips: set[str] = set()
+
+            for other_node in all_nodes:
+                if other_node == node_id:
+                    continue
+                for conn in topology.out_edges(other_node):
+                    if conn.sink == node_id and isinstance(conn.edge, SocketConnection):
+                        ip = conn.edge.sink_multiaddr.ip_address
+                        # Skip link-local and localhost addresses
+                        if not ip.startswith("169.254.") and not ip.startswith("127."):
+                            candidate_ips.add(ip)
+
+            # Prefer private network IPs (10.x, 192.168.x) over Tailscale CGNAT (100.64-127.x)
+            chosen_ip: str | None = None
+            for ip in candidate_ips:
+                if ip.startswith(("10.", "192.168.")):
+                    chosen_ip = ip
+                    break
+            if chosen_ip is None and candidate_ips:
+                chosen_ip = next(iter(candidate_ips))
+
+            if chosen_ip:
+                hosts_by_node[node_id] = [Host(ip=chosen_ip, port=0)]
+            else:
+                logger.warning(
+                    f"Could not determine IP for node {node_id}, using localhost"
+                )
+                hosts_by_node[node_id] = [Host(ip="127.0.0.1", port=0)]
+
+            logger.info(
+                f"FLASH placement: node {node_id} -> IP {hosts_by_node[node_id][0].ip}"
+                f" (candidates: {candidate_ips})"
+            )
+
+    total_ranks = len(selected_nodes) * command.ranks_per_node
+
+    # Determine coordinator IP - first node's first host IP
+    first_node_id: NodeId = next(iter(hosts_by_node.keys()))
+    coordinator_ip: str = (
+        hosts_by_node[first_node_id][0].ip
+        if hosts_by_node[first_node_id]
+        else "127.0.0.1"
+    )
+
+    target_instances[instance_id] = FLASHInstance(
+        instance_id=instance_id,
+        shard_assignments=shard_assignments,
+        hosts_by_node=hosts_by_node,
+        flash_executable_path=command.flash_executable_path,
+        parameter_file_path=command.parameter_file_path,
+        working_directory=command.working_directory,
+        ranks_per_node=command.ranks_per_node,
+        total_ranks=total_ranks,
+        simulation_name=command.simulation_name,
+        coordinator_ip=coordinator_ip,
+    )
+
+    logger.info(f"Created FLASH instance {instance_id} with {total_ranks} total ranks")
+
+    return target_instances

src/exo/plugins/implementations/flash/planning.py

@@ -0,0 +1,37 @@
+"""FLASH plugin planning logic."""
+
+from collections.abc import Mapping
+
+from exo.plugins.implementations.flash.types import FLASHInstance
+from exo.shared.types.tasks import LoadModel, Task
+from exo.shared.types.worker.instances import BaseInstance, InstanceId
+from exo.shared.types.worker.runners import RunnerId, RunnerIdle
+from exo.worker.runner.runner_supervisor import RunnerSupervisor
+
+
+def plan_flash(
+    runners: Mapping[RunnerId, RunnerSupervisor],
+    instances: Mapping[InstanceId, BaseInstance],
+) -> Task | None:
+    """Plan tasks specifically for FLASH instances.
+
+    FLASH instances have a simpler lifecycle:
+    - CreateRunner (handled by core _create_runner)
+    - LoadModel (starts the simulation immediately)
+    - Shutdown (handled by core _kill_runner)
+
+    This function handles the LoadModel step for FLASH instances,
+    skipping the MLX-specific download/init/warmup steps.
+    """
+    for runner in runners.values():
+        instance = runner.bound_instance.instance
+
+        # Only handle FLASH instances
+        if not isinstance(instance, FLASHInstance):
+            continue
+
+        # If runner is idle, emit LoadModel to start the simulation
+        if isinstance(runner.status, RunnerIdle):
+            return LoadModel(instance_id=instance.instance_id)
+
+    return None

src/exo/plugins/implementations/flash/plugin.py

@@ -0,0 +1,98 @@
+"""FLASH Plugin - Main plugin class."""
+
+from collections.abc import Callable, Mapping, Sequence
+from typing import Any
+
+from exo.plugins.base import EXOPlugin
+from exo.plugins.implementations.flash.api_handlers import (
+    handle_launch_flash,
+    handle_list_flash_instances,
+    handle_stop_flash,
+)
+from exo.plugins.implementations.flash.placement import place_flash_instance
+from exo.plugins.implementations.flash.planning import plan_flash
+from exo.plugins.implementations.flash.runner import main as flash_runner_main
+from exo.plugins.implementations.flash.types import (
+    FLASHInstance,
+    LaunchFLASH,
+    StopFLASH,
+)
+from exo.shared.topology import Topology
+from exo.shared.types.commands import DeleteInstance
+from exo.shared.types.events import Event
+from exo.shared.types.tasks import Task
+from exo.shared.types.worker.instances import BaseInstance, BoundInstance, InstanceId
+from exo.shared.types.worker.runners import RunnerId
+from exo.utils.channels import MpReceiver, MpSender
+from exo.worker.runner.runner_supervisor import RunnerSupervisor
+
+
+class FLASHPlugin(EXOPlugin):
+    """Plugin for FLASH MPI simulations."""
+
+    @property
+    def name(self) -> str:
+        return "flash"
+
+    @property
+    def version(self) -> str:
+        return "1.0.0"
+
+    def get_command_types(self) -> Sequence[type]:
+        return [LaunchFLASH, StopFLASH]
+
+    def get_instance_type(self) -> type:
+        return FLASHInstance
+
+    def get_api_routes(
+        self,
+    ) -> Sequence[tuple[str, str, Callable[..., Any]]]:
+        return [
+            ("post", "/flash/launch", handle_launch_flash),
+            ("delete", "/flash/{instance_id}", handle_stop_flash),
+            ("get", "/flash/instances", handle_list_flash_instances),
+        ]
+
+    def handles_command(self, command: Any) -> bool:  # pyright: ignore[reportAny]
+        return isinstance(command, (LaunchFLASH, StopFLASH))
+
+    def process_command(
+        self,
+        command: Any,  # pyright: ignore[reportAny]
+        topology: Topology,
+        current_instances: Mapping[InstanceId, BaseInstance],
+    ) -> Sequence[Event]:
+        from exo.master.placement import delete_instance, get_transition_events
+
+        if isinstance(command, LaunchFLASH):
+            placement = place_flash_instance(command, topology, current_instances)
+            return list(get_transition_events(current_instances, placement))
+        elif isinstance(command, StopFLASH):
+            placement = delete_instance(
+                DeleteInstance(instance_id=command.instance_id),
+                current_instances,
+            )
+            return list(get_transition_events(current_instances, placement))
+        return []
+
+    def handles_instance(self, instance: object) -> bool:
+        return isinstance(instance, FLASHInstance)
+
+    def plan_task(
+        self,
+        runners: Mapping[RunnerId, RunnerSupervisor],
+        instances: Mapping[InstanceId, BaseInstance],
+    ) -> Task | None:
+        return plan_flash(runners, instances)
+
+    def should_skip_download(self, instance: object) -> bool:
+        # FLASH instances don't need model downloads
+        return True
+
+    def create_runner(
+        self,
+        bound_instance: BoundInstance,
+        event_sender: MpSender[Event],
+        task_receiver: MpReceiver[Task],
+    ) -> None:
+        flash_runner_main(bound_instance, event_sender, task_receiver)

src/exo/plugins/implementations/flash/runner.py

@@ -0,0 +1,305 @@
+"""FLASH MPI Runner - spawns and monitors FLASH simulations.
+
+Exo-native distributed MPI:
+- Exo handles node discovery and coordination
+- Coordinator generates hostfile from Exo topology
+- mpirun uses exo-rsh (no SSH required) to spawn on remote nodes
+- exo-rsh connects to each node's Exo API (/execute endpoint) for remote execution
+- Workers just report ready and wait
+"""
+# ruff: noqa: I001 - Import order intentional (plugin types before shared types)
+
+import os
+import shutil
+import socket
+import subprocess
+import threading
+
+from loguru import logger
+
+from exo.shared.types.events import (
+    Event,
+    RunnerStatusUpdated,
+    TaskAcknowledged,
+    TaskStatusUpdated,
+)
+from exo.shared.types.tasks import (
+    LoadModel,
+    Shutdown,
+    Task,
+    TaskStatus,
+)
+from exo.plugins.implementations.flash.types import FLASHInstance
+from exo.shared.types.worker.instances import BoundInstance
+from exo.shared.types.worker.runners import (
+    RunnerFailed,
+    RunnerIdle,
+    RunnerLoading,
+    RunnerReady,
+    RunnerRunning,
+    RunnerShutdown,
+    RunnerShuttingDown,
+    RunnerStatus,
+)
+from exo.utils.channels import MpReceiver, MpSender
+
+# Find mpirun in PATH, fallback to common locations
+MPIRUN_PATH = shutil.which("mpirun") or "/opt/homebrew/bin/mpirun"
+
+# exo-rsh is installed as console script by exo package
+_exo_rsh_path = shutil.which("exo-rsh")
+if not _exo_rsh_path:
+    raise RuntimeError("exo-rsh not found in PATH - this should be installed with exo")
+EXO_RSH_PATH: str = _exo_rsh_path
+
+
+def get_my_rank(instance: FLASHInstance, my_node_id: str) -> int:
+    """Determine this node's rank based on position in hosts_by_node."""
+    for i, node_id in enumerate(instance.hosts_by_node.keys()):
+        if str(node_id) == str(my_node_id):
+            return i
+    return -1
+
+
+def get_coordinator_host(instance: FLASHInstance) -> str:
+    """Get the IP of the coordinator node."""
+    return instance.coordinator_ip
+
+
+def resolve_host(host: str) -> str:
+    """Resolve host string to a usable hostname for MPI hostfile.
+
+    Accepts either an IP address or hostname. For IPs, attempts to resolve
+    to a hostname via DNS/mDNS. Hostnames are returned as-is after validation.
+    """
+    # Check if input is already a hostname (not an IP)
+    try:
+        socket.inet_aton(host)
+        is_ip = True
+    except socket.error:
+        is_ip = False
+
+    if not is_ip:
+        # Already a hostname, verify it resolves and return as-is
+        try:
+            socket.gethostbyname(host)
+            return host
+        except socket.gaierror:
+            logger.warning(f"Hostname {host} does not resolve, using anyway")
+            return host
+
+    # It's an IP address, try to resolve to hostname
+    try:
+        hostname, _, _ = socket.gethostbyaddr(host)
+        hostname = hostname.split(".")[0]
+        logger.info(f"Resolved {host} to {hostname}")
+        return hostname
+    except socket.herror:
+        pass
+
+    # Fall back to IP
+    logger.warning(f"Could not resolve {host} to hostname, using IP directly")
+    return host
+
+
+def generate_hostfile(instance: FLASHInstance, working_dir: str) -> str:
+    """Generate MPI hostfile from instance topology."""
+    hostfile_path = os.path.join(working_dir, "flash_hosts.txt")
+    with open(hostfile_path, "w") as f:
+        for _node_id, hosts in instance.hosts_by_node.items():
+            if hosts:
+                host = resolve_host(hosts[0].ip)
+                f.write(f"{host} slots={instance.ranks_per_node}\n")
+    logger.info(f"Generated hostfile at {hostfile_path}")
+    with open(hostfile_path, "r") as f:
+        logger.info(f"Hostfile contents:\n{f.read()}")
+    return hostfile_path
+
+
+def main(
+    bound_instance: BoundInstance,
+    event_sender: MpSender[Event],
+    task_receiver: MpReceiver[Task],
+) -> None:
+    """Main FLASH runner loop.
+
+    Coordinator: generates hostfile and runs mpirun (uses exo-rsh instead of SSH)
+    Workers: just report ready and wait for mpirun to spawn processes on them
+    """
+    assert isinstance(bound_instance.instance, FLASHInstance)
+    instance = bound_instance.instance
+    runner_id = bound_instance.bound_runner_id
+    my_node_id = str(bound_instance.bound_node_id)
+
+    logger.info(f"FLASH runner starting for simulation: {instance.simulation_name}")
+
+    my_rank = get_my_rank(instance, my_node_id)
+    world_size = len(instance.hosts_by_node)
+    is_coordinator = my_rank == 0
+    coordinator_ip = get_coordinator_host(instance)
+
+    logger.info(
+        f"FLASH node: rank={my_rank}, world_size={world_size}, coordinator={is_coordinator}"
+    )
+    logger.info(f"FLASH coordinator IP: {coordinator_ip}")
+
+    process: subprocess.Popen[bytes] | None = None
+    current_status: RunnerStatus = RunnerIdle()
+    shutdown_requested = False
+
+    event_sender.send(
+        RunnerStatusUpdated(runner_id=runner_id, runner_status=current_status)
+    )
+
+    def monitor_output(proc: subprocess.Popen[bytes]) -> None:
+        """Monitor FLASH stdout for progress updates."""
+        if proc.stdout is None:
+            return
+        for line in iter(proc.stdout.readline, b""):
+            if shutdown_requested:
+                break
+            try:
+                decoded: str = line.decode("utf-8", errors="replace").strip()
+                if decoded:
+                    logger.info(f"[FLASH] {decoded}")
+            except Exception as e:
+                logger.warning(f"Error parsing FLASH output: {e}")
+
+    with task_receiver as tasks:
+        for task in tasks:
+            event_sender.send(
+                TaskStatusUpdated(task_id=task.task_id, task_status=TaskStatus.Running)
+            )
+            event_sender.send(TaskAcknowledged(task_id=task.task_id))
+
+            match task:
+                case LoadModel() if isinstance(current_status, RunnerIdle):
+                    current_status = RunnerLoading()
+                    logger.info("Starting FLASH simulation")
+                    event_sender.send(
+                        RunnerStatusUpdated(
+                            runner_id=runner_id, runner_status=current_status
+                        )
+                    )
+
+                    try:
+                        if is_coordinator:
+                            # Coordinator: generate hostfile and run mpirun
+                            hostfile = generate_hostfile(
+                                instance, instance.working_directory
+                            )
+
+                            iface = instance.network_interface
+                            cmd = [
+                                MPIRUN_PATH,
+                                "-np",
+                                str(instance.total_ranks),
+                                "--hostfile",
+                                hostfile,
+                                "--wdir",
+                                instance.working_directory,
+                                "--oversubscribe",
+                                "--allow-run-as-root",
+                                "--mca",
+                                "btl",
+                                "tcp,self",
+                                "--mca",
+                                "btl_tcp_if_include",
+                                iface,
+                                "--mca",
+                                "oob_tcp_if_include",
+                                iface,
+                                "--mca",
+                                "plm_rsh_no_tree_spawn",
+                                "1",
+                            ]
+
+                            # Use exo-rsh for remote execution (no SSH needed)
+                            cmd.extend(["--mca", "plm_rsh_agent", EXO_RSH_PATH])
+
+                            cmd.append(instance.flash_executable_path)
+
+                            logger.info(f"FLASH distributed launch: {' '.join(cmd)}")
+
+                            process = subprocess.Popen(
+                                cmd,
+                                cwd=instance.working_directory,
+                                stdout=subprocess.PIPE,
+                                stderr=subprocess.STDOUT,
+                            )
+
+                            monitor_thread = threading.Thread(
+                                target=monitor_output, args=(process,), daemon=True
+                            )
+                            monitor_thread.start()
+
+                            current_status = RunnerRunning()
+                            logger.info(
+                                f"FLASH running on {world_size} nodes with {instance.total_ranks} ranks"
+                            )
+
+                        else:
+                            # Worker: mpirun on coordinator will use exo-rsh to spawn processes here
+                            logger.info(
+                                f"Worker {my_rank}: Ready for mpirun to spawn processes via exo-rsh"
+                            )
+                            current_status = RunnerRunning()
+
+                    except Exception as e:
+                        logger.error(f"Failed to start FLASH: {e}")
+                        import traceback
+
+                        logger.error(traceback.format_exc())
+                        current_status = RunnerFailed(error_message=str(e))
+
+                case Shutdown():
+                    shutdown_requested = True
+                    current_status = RunnerShuttingDown()
+                    logger.info("FLASH runner shutting down")
+                    event_sender.send(
+                        RunnerStatusUpdated(
+                            runner_id=runner_id, runner_status=current_status
+                        )
+                    )
+
+                    if process and process.poll() is None:
+                        logger.info("Terminating FLASH simulation")
+                        process.terminate()
+                        try:
+                            process.wait(timeout=10)
+                        except subprocess.TimeoutExpired:
+                            logger.warning("FLASH didn't terminate, killing")
+                            process.kill()
+                            process.wait()
+
+                    current_status = RunnerShutdown()
+
+                case _:
+                    if process and process.poll() is not None:
+                        exit_code = process.returncode
+                        if exit_code == 0:
+                            logger.info("FLASH simulation completed successfully")
+                            current_status = RunnerReady()
+                        else:
+                            logger.error(
+                                f"FLASH simulation failed with code {exit_code}"
+                            )
+                            current_status = RunnerFailed(
+                                error_message=f"Exit code {exit_code}"
+                            )
+
+            event_sender.send(
+                TaskStatusUpdated(task_id=task.task_id, task_status=TaskStatus.Complete)
+            )
+            event_sender.send(
+                RunnerStatusUpdated(runner_id=runner_id, runner_status=current_status)
+            )
+
+            if isinstance(current_status, RunnerShutdown):
+                break
+
+    if process and process.poll() is None:
+        process.terminate()
+        process.wait(timeout=5)
+
+    logger.info("FLASH runner exiting")

src/exo/plugins/implementations/flash/types.py

@@ -0,0 +1,62 @@
+"""FLASH plugin types - commands and instances."""
+
+from exo.plugins.type_registry import command_registry, instance_registry
+from exo.shared.types.commands import BaseCommand
+from exo.shared.types.common import Host, NodeId
+from exo.shared.types.worker.instances import BaseInstance, InstanceId
+from exo.shared.types.worker.runners import RunnerId
+from exo.shared.types.worker.shards import ShardMetadata
+
+# ============================================================================
+# Commands
+# ============================================================================
+
+
+@command_registry.register
+class LaunchFLASH(BaseCommand):
+    """Command to launch a FLASH MPI simulation."""
+
+    simulation_name: str
+    flash_executable_path: str
+    parameter_file_path: str
+    working_directory: str
+    ranks_per_node: int = 1
+    min_nodes: int = 1
+    # Optional: explicit hostnames for MPI (e.g., "s14,james21-1")
+    # Used when topology edges don't contain IP addresses
+    hosts: str = ""
+
+
+@command_registry.register
+class StopFLASH(BaseCommand):
+    """Command to stop a running FLASH simulation."""
+
+    instance_id: InstanceId
+
+
+# ============================================================================
+# Instances
+# ============================================================================
+
+
+@instance_registry.register
+class FLASHInstance(BaseInstance):
+    """Instance for FLASH MPI simulation.
+
+    Unlike MLX instances which do tensor parallelism, FLASH instances
+    coordinate MPI processes across nodes. Each node runs one or more
+    MPI ranks of the FLASH simulation.
+    """
+
+    hosts_by_node: dict[NodeId, list[Host]]
+    flash_executable_path: str
+    parameter_file_path: str
+    working_directory: str
+    ranks_per_node: int = 1
+    total_ranks: int
+    simulation_name: str
+    coordinator_ip: str
+    network_interface: str = "en0"  # Network interface for MPI (e.g., en0, eth0)
+
+    def shard(self, runner_id: RunnerId) -> ShardMetadata | None:
+        return self.shard_assignments.runner_to_shard.get(runner_id, None)

src/exo/plugins/registry.py

@@ -0,0 +1,110 @@
+"""Plugin registry for discovering and managing plugins."""
+
+from collections.abc import Callable, Sequence
+from typing import Any
+
+from loguru import logger
+
+from exo.plugins.base import EXOPlugin
+
+
+class PluginRegistry:
+    """Central registry for all plugins."""
+
+    _instance: "PluginRegistry | None" = None
+
+    def __init__(self) -> None:
+        self._plugins: dict[str, EXOPlugin] = {}
+        self._command_handlers: dict[type, EXOPlugin] = {}
+        self._instance_handlers: dict[type, EXOPlugin] = {}
+
+    @classmethod
+    def get(cls) -> "PluginRegistry":
+        """Get the singleton registry instance."""
+        if cls._instance is None:
+            cls._instance = cls()
+        return cls._instance
+
+    @classmethod
+    def reset(cls) -> None:
+        """Reset the singleton instance (useful for testing)."""
+        cls._instance = None
+
+    def register(self, plugin: EXOPlugin) -> None:
+        """Register a plugin and its types."""
+        if plugin.name in self._plugins:
+            raise ValueError(f"Plugin '{plugin.name}' already registered")
+
+        logger.info(f"Registering plugin: {plugin.name} v{plugin.version}")
+
+        self._plugins[plugin.name] = plugin
+
+        # Register command handlers
+        for cmd_type in plugin.get_command_types():
+            self._command_handlers[cmd_type] = plugin
+            logger.debug(f"  Registered command: {cmd_type.__name__}")
+
+        # Register instance handler
+        instance_type = plugin.get_instance_type()
+        self._instance_handlers[instance_type] = plugin
+        logger.debug(f"  Registered instance: {instance_type.__name__}")
+
+    def get_plugin(self, name: str) -> EXOPlugin | None:
+        """Get a plugin by name."""
+        return self._plugins.get(name)
+
+    def get_plugin_for_command(self, command: object) -> EXOPlugin | None:
+        """Get the plugin that handles a command."""
+        for plugin in self._plugins.values():
+            if plugin.handles_command(command):
+                return plugin
+        return None
+
+    def get_plugin_for_instance(self, instance: object) -> EXOPlugin | None:
+        """Get the plugin that manages an instance."""
+        for plugin in self._plugins.values():
+            if plugin.handles_instance(instance):
+                return plugin
+        return None
+
+    def all_plugins(self) -> Sequence[EXOPlugin]:
+        """Get all registered plugins."""
+        return list(self._plugins.values())
+
+    def get_all_api_routes(
+        self,
+    ) -> Sequence[tuple[str, str, Callable[..., Any], EXOPlugin]]:
+        """Get all API routes from all plugins."""
+        routes: list[tuple[str, str, Callable[..., Any], EXOPlugin]] = []
+        for plugin in self._plugins.values():
+            for method, path, handler in plugin.get_api_routes():
+                routes.append((method, path, handler, plugin))
+        return routes
+
+
+def discover_plugins() -> None:
+    """Auto-discover and register plugins from the implementations directory.
+
+    Plugins should have a register() function that returns an EXOPlugin instance.
+    """
+    import importlib
+    import pkgutil
+
+    registry = PluginRegistry.get()
+
+    try:
+        import exo.plugins.implementations as impl_package
+
+        for _, module_name, _ in pkgutil.iter_modules(impl_package.__path__):
+            try:
+                module = importlib.import_module(
+                    f"exo.plugins.implementations.{module_name}"
+                )
+                if hasattr(module, "register"):
+                    plugin = module.register()  # pyright: ignore[reportAny]
+                    if plugin is not None:
+                        registry.register(plugin)  # pyright: ignore[reportAny]
+            except Exception as e:
+                logger.warning(f"Failed to load plugin {module_name}: {e}")
+    except ImportError:
+        logger.debug("No plugin implementations package found")

src/exo/plugins/type_registry.py

@@ -0,0 +1,84 @@
+"""Dynamic type registry for plugin types.
+
+This module provides a registry system that allows plugins to register their
+command and instance types dynamically, eliminating the need for static union
+types and avoiding circular imports.
+"""
+
+from typing import TypeVar
+
+from loguru import logger
+
+from exo.utils.pydantic_ext import CamelCaseModel
+
+# TypeVar for preserving exact types through the register decorator
+_TCls = TypeVar("_TCls", bound=type[CamelCaseModel])
+
+
+class TypeRegistry[T: CamelCaseModel]:
+    """Registry for dynamically registered Pydantic types.
+
+    Enables plugins to register their types at import time. Deserialization
+    uses the class name from the tagged JSON format to look up the correct type.
+    """
+
+    def __init__(self, name: str) -> None:
+        self._name = name
+        self._types: dict[str, type[T]] = {}
+
+    def register(self, cls: _TCls) -> _TCls:
+        """Decorator to register a type with this registry.
+
+        Preserves the exact type through the decorator for proper type checking.
+        """
+        self._types[cls.__name__] = cls  # type: ignore[assignment]
+        logger.debug(f"{self._name}: registered {cls.__name__}")
+        return cls
+
+    def get(self, name: str) -> type[T] | None:
+        """Look up a type by class name."""
+        return self._types.get(name)
+
+    def all_types(self) -> dict[str, type[T]]:
+        """Return all registered types."""
+        return dict(self._types)
+
+    def deserialize(self, data: dict[str, dict[str, object]] | CamelCaseModel) -> T:
+        """Deserialize dict to the appropriate registered type.
+
+        Supports two formats:
+        1. Tagged format: {"ClassName": {...fields...}} - used for network serialization
+        2. Flat format: {...fields...} - used for API requests, tries each type
+        """
+        # If already deserialized (e.g., from Pydantic), return as-is
+        if isinstance(data, CamelCaseModel):
+            return data  # type: ignore[return-value]
+
+        # Check for tagged format: single key that matches a registered type
+        if len(data) == 1:
+            class_name: str = next(iter(data.keys()))
+            cls = self._types.get(class_name)
+            if cls is not None:
+                return cls.model_validate(data[class_name], strict=False)
+
+        # Flat format: try each registered type, use first that validates
+        errors: list[str] = []
+        for type_name, cls in self._types.items():
+            try:
+                return cls.model_validate(data, strict=False)
+            except Exception as e:  # noqa: BLE001
+                errors.append(f"{type_name}: {e}")
+
+        # None matched - provide helpful error
+        available = ", ".join(self._types.keys())
+        raise ValueError(
+            f"{self._name}: could not deserialize data. "
+            f"Available types: {available}. Errors: {'; '.join(errors[:3])}"
+        )
+
+
+# Global registries for commands, instances, events, and tasks
+command_registry: TypeRegistry[CamelCaseModel] = TypeRegistry("CommandRegistry")
+instance_registry: TypeRegistry[CamelCaseModel] = TypeRegistry("InstanceRegistry")
+event_registry: TypeRegistry[CamelCaseModel] = TypeRegistry("EventRegistry")
+task_registry: TypeRegistry[CamelCaseModel] = TypeRegistry("TaskRegistry")

src/exo/routing/router.py

@@ -216,8 +216,6 @@ 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, ForwarderDownloadCommand
+from exo.shared.types.commands import ForwarderCommand
 from exo.shared.types.events import (
     ForwarderEvent,
 )
@@ -30,7 +30,7 @@ class TypedTopic[T: CamelCaseModel]:
 
     @staticmethod
     def serialize(t: T) -> bytes:
-        return t.model_dump_json().encode("utf-8")
+        return t.model_dump_json(by_alias=True, serialize_as_any=True).encode("utf-8")
 
     def deserialize(self, b: bytes) -> T:
         return self.model_type.model_validate_json(b.decode("utf-8"))
@@ -45,6 +45,3 @@ ELECTION_MESSAGES = TypedTopic(
 CONNECTION_MESSAGES = TypedTopic(
     "connection_messages", PublishPolicy.Never, ConnectionMessage
 )
-DOWNLOAD_COMMANDS = TypedTopic(
-    "download_commands", PublishPolicy.Always, ForwarderDownloadCommand
-)

src/exo/rsh/__init__.py

@@ -0,0 +1,13 @@
+"""Exo RSH - Remote Shell for MPI without SSH.
+
+This module provides a remote execution mechanism that allows mpirun to spawn
+processes on remote nodes without requiring SSH setup. It works by:
+
+1. Each Exo node runs an API server on port 52415 with an /execute endpoint
+2. The exo-rsh script acts as a drop-in replacement for ssh
+3. When mpirun calls "exo-rsh hostname command", it HTTP POSTs to the target's /execute
+4. The target executes the command and returns output
+
+Usage:
+    mpirun --mca plm_rsh_agent exo-rsh -np 4 --hostfile hosts.txt ./program
+"""

src/exo/rsh/client.py

@@ -0,0 +1,101 @@
+#!/usr/bin/env python3
+"""exo-rsh - Remote shell client for MPI.
+
+This script is called by mpirun as a replacement for ssh.
+Usage: exo-rsh [ssh-options...] hostname command [args...]
+
+It connects to the target node's Exo API (port 52415) and executes the command.
+"""
+
+import json
+import socket
+import sys
+from typing import Any, cast
+from urllib.error import URLError
+from urllib.request import Request, urlopen
+
+# Use the same port as Exo's API server
+EXO_API_PORT = 52415
+
+
+def resolve_hostname(hostname: str) -> str:
+    """Resolve hostname to IP address."""
+    try:
+        return socket.gethostbyname(hostname)
+    except socket.gaierror:
+        # If resolution fails, try using the hostname directly
+        return hostname
+
+
+def main():
+    # Parse arguments - mpirun calls us like: exo-rsh [options] hostname command [args...]
+    # SSH options we might see: -x (disable X11), -o options, etc.
+    args = sys.argv[1:]
+
+    # Skip SSH-style options
+    hostname = None
+    command_start = 0
+
+    i = 0
+    while i < len(args):
+        arg = args[i]
+        if arg.startswith("-"):
+            # Skip option and its value if needed
+            if arg in ("-o", "-i", "-l", "-p", "-F"):
+                i += 2  # Skip option and its argument
+                continue
+            i += 1
+            continue
+        else:
+            # First non-option is the hostname
+            hostname = arg
+            command_start = i + 1
+            break
+        i += 1
+
+    if hostname is None or command_start >= len(args):
+        print("Usage: exo-rsh [options] hostname command [args...]", file=sys.stderr)
+        sys.exit(1)
+
+    command = args[command_start:]
+
+    # Resolve hostname to IP
+    ip = resolve_hostname(hostname)
+
+    # Make request to Exo API
+    url = f"http://{ip}:{EXO_API_PORT}/execute"
+    data = json.dumps({"command": command}).encode("utf-8")
+
+    try:
+        req = Request(url, data=data, headers={"Content-Type": "application/json"})
+        with urlopen(req, timeout=300) as response:  # pyright: ignore[reportAny]
+            response_body: bytes = cast(bytes, response.read())  # pyright: ignore[reportAny]
+            result: dict[str, Any] = json.loads(response_body.decode("utf-8"))  # pyright: ignore[reportAny]
+
+        # Output stdout/stderr
+        stdout: str = cast(str, result.get("stdout", ""))
+        stderr: str = cast(str, result.get("stderr", ""))
+        exit_code: int = cast(int, result.get("exit_code", 0))
+
+        if stdout:
+            sys.stdout.write(stdout)
+            sys.stdout.flush()
+        if stderr:
+            sys.stderr.write(stderr)
+            sys.stderr.flush()
+
+        sys.exit(exit_code)
+
+    except URLError as e:
+        print(
+            f"exo-rsh: Failed to connect to {hostname}:{EXO_API_PORT}: {e}",
+            file=sys.stderr,
+        )
+        sys.exit(255)
+    except Exception as e:
+        print(f"exo-rsh: Error: {e}", file=sys.stderr)
+        sys.exit(1)
+
+
+if __name__ == "__main__":
+    main()

src/exo/shared/apply.py

@@ -6,8 +6,8 @@ from loguru import logger
 
 from exo.shared.types.common import NodeId
 from exo.shared.types.events import (
+    BaseEvent,
     ChunkGenerated,
-    Event,
     IndexedEvent,
     InputChunkReceived,
     InstanceCreated,
@@ -25,20 +25,17 @@ 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
+from exo.shared.types.tasks import BaseTask, TaskId, TaskStatus
 from exo.shared.types.topology import Connection, RDMAConnection
 from exo.shared.types.worker.downloads import DownloadProgress
-from exo.shared.types.worker.instances import Instance, InstanceId
+from exo.shared.types.worker.instances import BaseInstance, InstanceId
 from exo.shared.types.worker.runners import RunnerId, RunnerStatus
 from exo.utils.info_gatherer.info_gatherer import (
     MacmonMetrics,
@@ -49,20 +46,14 @@ from exo.utils.info_gatherer.info_gatherer import (
     NodeConfig,
     NodeNetworkInterfaces,
     StaticNodeInformation,
-    ThunderboltBridgeInfo,
 )
 
 
-def event_apply(event: Event, state: State) -> State:
+def event_apply(event: BaseEvent, state: State) -> State:
     """Apply an event to state."""
     match event:
         case (
-            TestEvent()
-            | ChunkGenerated()
-            | TaskAcknowledged()
-            | InputChunkReceived()
-            | TracesCollected()
-            | TracesMerged()
+            TestEvent() | ChunkGenerated() | TaskAcknowledged() | InputChunkReceived()
         ):  # Pass-through events that don't modify state
             return state
         case InstanceCreated():
@@ -91,6 +82,10 @@ def event_apply(event: Event, state: State) -> State:
             return apply_topology_edge_created(event, state)
         case TopologyEdgeDeleted():
             return apply_topology_edge_deleted(event, state)
+        case _:
+            # Unknown event types from plugins are ignored
+            logger.debug(f"Ignoring unknown event type: {type(event).__name__}")
+            return state
 
 
 def apply(state: State, event: IndexedEvent) -> State:
@@ -131,12 +126,12 @@ def apply_node_download_progress(event: NodeDownloadProgress, state: State) -> S
 
 
 def apply_task_created(event: TaskCreated, state: State) -> State:
-    new_tasks: Mapping[TaskId, Task] = {**state.tasks, event.task_id: event.task}
+    new_tasks: Mapping[TaskId, BaseTask] = {**state.tasks, event.task_id: event.task}
     return state.model_copy(update={"tasks": new_tasks})
 
 
 def apply_task_deleted(event: TaskDeleted, state: State) -> State:
-    new_tasks: Mapping[TaskId, Task] = {
+    new_tasks: Mapping[TaskId, BaseTask] = {
         tid: task for tid, task in state.tasks.items() if tid != event.task_id
     }
     return state.model_copy(update={"tasks": new_tasks})
@@ -155,7 +150,7 @@ def apply_task_status_updated(event: TaskStatusUpdated, state: State) -> State:
         update["error_message"] = None
 
     updated_task = state.tasks[event.task_id].model_copy(update=update)
-    new_tasks: Mapping[TaskId, Task] = {**state.tasks, event.task_id: updated_task}
+    new_tasks: Mapping[TaskId, BaseTask] = {**state.tasks, event.task_id: updated_task}
     return state.model_copy(update={"tasks": new_tasks})
 
 
@@ -167,13 +162,13 @@ def apply_task_failed(event: TaskFailed, state: State) -> State:
     updated_task = state.tasks[event.task_id].model_copy(
         update={"error_type": event.error_type, "error_message": event.error_message}
     )
-    new_tasks: Mapping[TaskId, Task] = {**state.tasks, event.task_id: updated_task}
+    new_tasks: Mapping[TaskId, BaseTask] = {**state.tasks, event.task_id: updated_task}
     return state.model_copy(update={"tasks": new_tasks})
 
 
 def apply_instance_created(event: InstanceCreated, state: State) -> State:
     instance = event.instance
-    new_instances: Mapping[InstanceId, Instance] = {
+    new_instances: Mapping[InstanceId, BaseInstance] = {
         **state.instances,
         instance.instance_id: instance,
     }
@@ -181,7 +176,7 @@ def apply_instance_created(event: InstanceCreated, state: State) -> State:
 
 
 def apply_instance_deleted(event: InstanceDeleted, state: State) -> State:
-    new_instances: Mapping[InstanceId, Instance] = {
+    new_instances: Mapping[InstanceId, BaseInstance] = {
         iid: inst for iid, inst in state.instances.items() if iid != event.instance_id
     }
     return state.model_copy(update={"instances": new_instances})
@@ -234,21 +229,6 @@ 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,
@@ -259,8 +239,6 @@ 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,
         }
     )
 
@@ -338,22 +316,6 @@ 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,10 +49,3 @@ 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, Any
+from typing import Annotated
 
 import aiofiles
 import aiofiles.os as aios
@@ -7,16 +7,8 @@ import tomlkit
 from anyio import Path, open_file
 from huggingface_hub import model_info
 from loguru import logger
-from pydantic import (
-    AliasChoices,
-    BaseModel,
-    Field,
-    PositiveInt,
-    field_validator,
-    model_validator,
-)
+from pydantic import BaseModel, Field, PositiveInt, field_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
@@ -128,14 +120,6 @@ 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"),
@@ -270,7 +254,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(44900),
+        storage_size=Memory.from_mb(84700),
         n_layers=48,
         hidden_size=2048,
         supports_tensor=True,
@@ -426,310 +410,174 @@ 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
 
-    architectures: list[str] | None = None
+    # 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
-    layer_count: int = Field(
-        validation_alias=AliasChoices(
-            "num_hidden_layers",
-            "num_layers",
-            "n_layer",
-            "n_layers",
-            "num_decoder_layers",
-            "decoder_layers",
-        )
-    )
+    architectures: list[str] | None = None
 
     @property
     def supports_tensor(self) -> bool:
@@ -744,32 +592,30 @@ class ConfigData(BaseModel):
             ["GptOssForCausalLM"],
         ]
 
-    @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
+    @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,
+        ]
 
-        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
+        for layer_count in layer_fields:
+            if layer_count is not None:
+                return layer_count
 
-        return data
+        raise ValueError(
+            f"No layer count found in config.json: {self.model_dump_json()}"
+        )
 
 
 async def get_config_data(model_id: ModelId) -> ConfigData:
     """Downloads and parses config.json for a model."""
-    from exo.download.download_utils import (
+    from exo.worker.download.download_utils import (
         download_file_with_retry,
         ensure_models_dir,
     )
@@ -781,7 +627,7 @@ async def get_config_data(model_id: ModelId) -> ConfigData:
         "main",
         "config.json",
         target_dir,
-        lambda curr_bytes, total_bytes, is_renamed: logger.debug(
+        lambda curr_bytes, total_bytes, is_renamed: logger.info(
             f"Downloading config.json for {model_id}: {curr_bytes}/{total_bytes} ({is_renamed=})"
         ),
     )
@@ -791,11 +637,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.download.download_utils import (
+    from exo.shared.types.worker.downloads import ModelSafetensorsIndex
+    from exo.worker.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)
@@ -804,7 +650,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.debug(
+        lambda curr_bytes, total_bytes, is_renamed: logger.info(
             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, mark
+from pytest import LogCaptureFixture
 
 from exo.routing.router import get_node_id_keypair
 from exo.shared.constants import EXO_NODE_ID_KEYPAIR
@@ -74,7 +74,6 @@ 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,11 +7,6 @@ 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,
@@ -193,25 +188,24 @@ class Topology:
             cycles.append(Cycle(node_ids=[node_id]))
         return cycles
 
-    def get_rdma_cycles(self) -> list[Cycle]:
-        rdma_edges = [
+    def get_cycles_tb(self) -> list[Cycle]:
+        tb_edges = [
             (u, v, conn)
             for u, v, conn in self._graph.weighted_edge_list()
-            if isinstance(conn, RDMAConnection)
+            if conn.is_thunderbolt()
         ]
 
-        rdma_graph: rx.PyDiGraph[NodeId, SocketConnection | RDMAConnection] = (
-            rx.PyDiGraph()
-        )
-        rdma_graph.add_nodes_from(self._graph.nodes())
+        tb_graph: rx.PyDiGraph[NodeId, SocketConnection] = rx.PyDiGraph()
+        tb_graph.add_nodes_from(self._graph.nodes())
 
-        for u, v, conn in rdma_edges:
-            rdma_graph.add_edge(u, v, conn)
+        for u, v, conn in tb_edges:
+            if isinstance(conn, SocketConnection):
+                tb_graph.add_edge(u, v, conn)
 
-        cycle_idxs = rx.simple_cycles(rdma_graph)
+        cycle_idxs = rx.simple_cycles(tb_graph)
         cycles: list[Cycle] = []
         for cycle_idx in cycle_idxs:
-            cycle = Cycle(node_ids=[rdma_graph[idx] for idx in cycle_idx])
+            cycle = Cycle(node_ids=[tb_graph[idx] for idx in cycle_idx])
             cycles.append(cycle)
 
         return cycles
@@ -225,83 +219,18 @@ class Topology:
                 topology.add_connection(connection)
         return topology
 
-    def is_rdma_cycle(self, cycle: Cycle) -> bool:
+    def is_thunderbolt_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_rdma = False
+                has_tb = False
                 for edge in self._graph.get_all_edge_data(rid, neighbor_rid):
-                    if isinstance(edge, RDMAConnection):
-                        has_rdma = True
+                    if edge.is_thunderbolt():
+                        has_tb = True
                         break
-                if not has_rdma:
+                if not has_tb:
                     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

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

@@ -1,16 +1,22 @@
 import time
 from collections.abc import Generator
-from typing import Annotated, Any, Literal
+from typing import Annotated, Any, Literal, cast
 
+from fastapi import UploadFile
 from pydantic import BaseModel, Field, field_validator
 from pydantic_core import PydanticUseDefault
 
+from exo.plugins.type_registry import instance_registry
 from exo.shared.models.model_cards import ModelCard, ModelId
-from exo.shared.types.common import CommandId, NodeId
+from exo.shared.types.common import CommandId
 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, ShardMetadata
-from exo.utils.pydantic_ext import CamelCaseModel
+from exo.shared.types.worker.instances import (
+    BaseInstance,
+    Instance,
+    InstanceId,
+    InstanceMeta,
+)
+from exo.shared.types.worker.shards import Sharding
 
 FinishReason = Literal[
     "stop", "length", "tool_calls", "content_filter", "function_call", "error"
@@ -54,18 +60,6 @@ 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: (
@@ -73,7 +67,7 @@ class ChatCompletionMessage(BaseModel):
     ) = None
     thinking: str | None = None  # Added for GPT-OSS harmony format support
     name: str | None = None
-    tool_calls: list[ToolCall] | None = None
+    tool_calls: list[dict[str, Any]] | None = None
     tool_call_id: str | None = None
     function_call: dict[str, Any] | None = None
 
@@ -115,8 +109,8 @@ class Usage(BaseModel):
     prompt_tokens: int
     completion_tokens: int
     total_tokens: int
-    prompt_tokens_details: PromptTokensDetails
-    completion_tokens_details: CompletionTokensDetails
+    prompt_tokens_details: PromptTokensDetails | None = None
+    completion_tokens_details: CompletionTokensDetails | None = None
 
 
 class StreamingChoiceResponse(BaseModel):
@@ -169,12 +163,8 @@ class BenchChatCompletionResponse(ChatCompletionResponse):
     generation_stats: GenerationStats | None = None
 
 
-class StreamOptions(BaseModel):
-    include_usage: bool = False
-
-
-class ChatCompletionRequest(BaseModel):
-    model: ModelId
+class ChatCompletionTaskParams(BaseModel):
+    model: str
     frequency_penalty: float | None = None
     messages: list[ChatCompletionMessage]
     logit_bias: dict[str, int] | None = None
@@ -187,17 +177,15 @@ class ChatCompletionRequest(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 BenchChatCompletionRequest(ChatCompletionRequest):
+class BenchChatCompletionTaskParams(ChatCompletionTaskParams):
     pass
 
 
@@ -218,6 +206,12 @@ class PlaceInstanceParams(BaseModel):
 class CreateInstanceParams(BaseModel):
     instance: Instance
 
+    @field_validator("instance", mode="before")
+    @classmethod
+    def validate_instance(cls, v: Any) -> BaseInstance:  # noqa: ANN401  # pyright: ignore[reportAny]
+        """Validate instance using registry to handle both tagged and flat formats."""
+        return cast(BaseInstance, instance_registry.deserialize(v))  # pyright: ignore[reportAny]
+
 
 class PlacementPreview(BaseModel):
     model_id: ModelId
@@ -281,7 +275,28 @@ class BenchImageGenerationTaskParams(ImageGenerationTaskParams):
 
 
 class ImageEditsTaskParams(BaseModel):
-    """Internal task params for image-editing requests."""
+    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."""
 
     image_data: str = ""  # Base64-encoded image (empty when using chunked transfer)
     total_input_chunks: int = 0
@@ -337,58 +352,3 @@ 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,37 +1,31 @@
 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, Usage
+from exo.shared.types.api import GenerationStats, ImageGenerationStats
 from exo.utils.pydantic_ext import TaggedModel
 
 from .api import FinishReason
 from .common import CommandId
-from .worker.runner_response import ToolCallItem
+
+
+class ChunkType(str, Enum):
+    Token = "Token"
+    Image = "Image"
 
 
 class BaseChunk(TaggedModel):
+    idx: int
     model: ModelId
 
 
 class TokenChunk(BaseChunk):
     text: str
     token_id: int
-    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"
+    finish_reason: FinishReason | None = None
     stats: GenerationStats | None = None
+    error_message: str | None = None
 
 
 class ImageChunk(BaseChunk):
@@ -69,4 +63,4 @@ class InputImageChunk(BaseChunk):
                 yield name, value
 
 
-GenerationChunk = TokenChunk | ImageChunk | ToolCallChunk | ErrorChunk
+GenerationChunk = TokenChunk | ImageChunk

src/exo/shared/types/claude_api.py

@@ -1,214 +0,0 @@
-"""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,38 +1,54 @@
-from pydantic import Field
+"""Command types for exo.
 
-from exo.shared.models.model_cards import ModelCard, ModelId
+Commands are registered dynamically via the command_registry, allowing plugins
+to add their own command types without modifying this file.
+"""
+
+from typing import Any, cast
+
+from pydantic import Field, field_validator
+
+from exo.plugins.type_registry import command_registry
+from exo.shared.models.model_cards import ModelCard
 from exo.shared.types.api import (
-    ImageEditsTaskParams,
+    ChatCompletionTaskParams,
+    ImageEditsInternalParams,
     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, ShardMetadata
+from exo.shared.types.worker.shards import Sharding
 from exo.utils.pydantic_ext import CamelCaseModel, TaggedModel
 
 
 class BaseCommand(TaggedModel):
+    """Base class for all commands."""
+
     command_id: CommandId = Field(default_factory=CommandId)
 
 
+@command_registry.register
 class TestCommand(BaseCommand):
     __test__ = False
 
 
-class TextGeneration(BaseCommand):
-    task_params: TextGenerationTaskParams
+@command_registry.register
+class ChatCompletion(BaseCommand):
+    request_params: ChatCompletionTaskParams
 
 
+@command_registry.register
 class ImageGeneration(BaseCommand):
-    task_params: ImageGenerationTaskParams
+    request_params: ImageGenerationTaskParams
 
 
+@command_registry.register
 class ImageEdits(BaseCommand):
-    task_params: ImageEditsTaskParams
+    request_params: ImageEditsInternalParams
 
 
+@command_registry.register
 class PlaceInstance(BaseCommand):
     model_card: ModelCard
     sharding: Sharding
@@ -40,45 +56,38 @@ class PlaceInstance(BaseCommand):
     min_nodes: int
 
 
+@command_registry.register
 class CreateInstance(BaseCommand):
     instance: Instance
 
 
+@command_registry.register
 class DeleteInstance(BaseCommand):
     instance_id: InstanceId
 
 
+@command_registry.register
 class TaskFinished(BaseCommand):
     finished_command_id: CommandId
 
 
+@command_registry.register
 class SendInputChunk(BaseCommand):
     """Command to send an input image chunk (converted to event by master)."""
 
     chunk: InputImageChunk
 
 
+@command_registry.register
 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
-
-
+# Union type for core commands - used by ForwarderCommand for network deserialization
 Command = (
     TestCommand
     | RequestEventLog
-    | TextGeneration
+    | ChatCompletion
     | ImageGeneration
     | ImageEdits
     | PlaceInstance
@@ -90,10 +99,14 @@ Command = (
 
 
 class ForwarderCommand(CamelCaseModel):
-    origin: NodeId
-    command: Command
+    """Wrapper for commands that includes origin node."""
 
-
-class ForwarderDownloadCommand(CamelCaseModel):
     origin: NodeId
-    command: DownloadCommand
+    command: BaseCommand
+
+    @field_validator("command", mode="before")
+    @classmethod
+    def validate_command(cls, v: Any) -> BaseCommand:  # noqa: ANN401  # pyright: ignore[reportAny]
+        """Validate command, using registry for plugin commands not in Command union."""
+        # First try the registry (handles both core and plugin commands)
+        return cast(BaseCommand, command_registry.deserialize(v))  # pyright: ignore[reportAny]

src/exo/shared/types/events.py

@@ -1,17 +1,18 @@
 from datetime import datetime
-from typing import final
+from typing import Any, cast
 
-from pydantic import Field
+from pydantic import Field, field_validator
 
+from exo.plugins.type_registry import event_registry, instance_registry, task_registry
 from exo.shared.topology import Connection
 from exo.shared.types.chunks import GenerationChunk, InputImageChunk
 from exo.shared.types.common import CommandId, Id, NodeId, SessionId
-from exo.shared.types.tasks import Task, TaskId, TaskStatus
+from exo.shared.types.tasks import BaseTask, TaskId, TaskStatus
 from exo.shared.types.worker.downloads import DownloadProgress
-from exo.shared.types.worker.instances import Instance, InstanceId
+from exo.shared.types.worker.instances import BaseInstance, 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, FrozenModel, TaggedModel
+from exo.utils.pydantic_ext import CamelCaseModel, TaggedModel
 
 
 class EventId(Id):
@@ -26,36 +27,53 @@ class BaseEvent(TaggedModel):
     _master_time_stamp: None | datetime = None
 
 
+@event_registry.register
 class TestEvent(BaseEvent):
     __test__ = False
 
 
+@event_registry.register
 class TaskCreated(BaseEvent):
     task_id: TaskId
-    task: Task
+    task: BaseTask
+
+    @field_validator("task", mode="before")
+    @classmethod
+    def validate_task(cls, v: Any) -> BaseTask:  # noqa: ANN401  # pyright: ignore[reportAny]
+        return cast(BaseTask, task_registry.deserialize(v))  # pyright: ignore[reportAny]
 
 
+@event_registry.register
 class TaskAcknowledged(BaseEvent):
     task_id: TaskId
 
 
+@event_registry.register
 class TaskDeleted(BaseEvent):
     task_id: TaskId
 
 
+@event_registry.register
 class TaskStatusUpdated(BaseEvent):
     task_id: TaskId
     task_status: TaskStatus
 
 
+@event_registry.register
 class TaskFailed(BaseEvent):
     task_id: TaskId
     error_type: str
     error_message: str
 
 
+@event_registry.register
 class InstanceCreated(BaseEvent):
-    instance: Instance
+    instance: BaseInstance
+
+    @field_validator("instance", mode="before")
+    @classmethod
+    def validate_instance(cls, v: Any) -> BaseInstance:  # noqa: ANN401  # pyright: ignore[reportAny]
+        return cast(BaseInstance, instance_registry.deserialize(v))  # pyright: ignore[reportAny]
 
     def __eq__(self, other: object) -> bool:
         if isinstance(other, InstanceCreated):
@@ -64,74 +82,63 @@ class InstanceCreated(BaseEvent):
         return False
 
 
+@event_registry.register
 class InstanceDeleted(BaseEvent):
     instance_id: InstanceId
 
 
+@event_registry.register
 class RunnerStatusUpdated(BaseEvent):
     runner_id: RunnerId
     runner_status: RunnerStatus
 
 
+@event_registry.register
 class RunnerDeleted(BaseEvent):
     runner_id: RunnerId
 
 
+@event_registry.register
 class NodeTimedOut(BaseEvent):
     node_id: NodeId
 
 
 # TODO: bikeshed this name
+@event_registry.register
 class NodeGatheredInfo(BaseEvent):
     node_id: NodeId
     when: str  # this is a manually cast datetime overrode by the master when the event is indexed, rather than the local time on the device
     info: GatheredInfo
 
 
+@event_registry.register
 class NodeDownloadProgress(BaseEvent):
     download_progress: DownloadProgress
 
 
+@event_registry.register
 class ChunkGenerated(BaseEvent):
     command_id: CommandId
     chunk: GenerationChunk
 
 
+@event_registry.register
 class InputChunkReceived(BaseEvent):
     command_id: CommandId
     chunk: InputImageChunk
 
 
+@event_registry.register
 class TopologyEdgeCreated(BaseEvent):
     conn: Connection
 
 
+@event_registry.register
 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]
-
-
+# Union type for Pydantic validation - tries each type in order
 Event = (
     TestEvent
     | TaskCreated
@@ -150,8 +157,6 @@ Event = (
     | InputChunkReceived
     | TopologyEdgeCreated
     | TopologyEdgeDeleted
-    | TracesCollected
-    | TracesMerged
 )
 
 
@@ -159,7 +164,12 @@ class IndexedEvent(CamelCaseModel):
     """An event indexed by the master, with a globally unique index"""
 
     idx: int = Field(ge=0)
-    event: Event
+    event: BaseEvent
+
+    @field_validator("event", mode="before")
+    @classmethod
+    def validate_event(cls, v: Any) -> BaseEvent:  # noqa: ANN401  # pyright: ignore[reportAny]
+        return cast(BaseEvent, event_registry.deserialize(v))  # pyright: ignore[reportAny]
 
 
 class ForwarderEvent(CamelCaseModel):
@@ -168,4 +178,9 @@ class ForwarderEvent(CamelCaseModel):
     origin_idx: int = Field(ge=0)
     origin: NodeId
     session: SessionId
-    event: Event
+    event: BaseEvent
+
+    @field_validator("event", mode="before")
+    @classmethod
+    def validate_event(cls, v: Any) -> BaseEvent:  # noqa: ANN401  # pyright: ignore[reportAny]
+        return cast(BaseEvent, event_registry.deserialize(v))  # pyright: ignore[reportAny]

src/exo/shared/types/mlx.py

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

@@ -1,296 +0,0 @@
-"""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 Literal, Self
+from typing import Self
 
 import psutil
 
@@ -48,13 +48,9 @@ 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):
@@ -75,11 +71,3 @@ 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,11 +13,10 @@ from exo.shared.types.profiling import (
     NodeNetworkInfo,
     NodeThunderboltInfo,
     SystemPerformanceProfile,
-    ThunderboltBridgeStatus,
 )
-from exo.shared.types.tasks import Task, TaskId
+from exo.shared.types.tasks import BaseTask, TaskId
 from exo.shared.types.worker.downloads import DownloadProgress
-from exo.shared.types.worker.instances import Instance, InstanceId
+from exo.shared.types.worker.instances import BaseInstance, InstanceId
 from exo.shared.types.worker.runners import RunnerId, RunnerStatus
 from exo.utils.pydantic_ext import CamelCaseModel
 
@@ -38,10 +37,10 @@ class State(CamelCaseModel):
         strict=True,
         arbitrary_types_allowed=True,
     )
-    instances: Mapping[InstanceId, Instance] = {}
+    instances: Mapping[InstanceId, BaseInstance] = {}
     runners: Mapping[RunnerId, RunnerStatus] = {}
     downloads: Mapping[NodeId, Sequence[DownloadProgress]] = {}
-    tasks: Mapping[TaskId, Task] = {}
+    tasks: Mapping[TaskId, BaseTask] = {}
     last_seen: Mapping[NodeId, datetime] = {}
     topology: Topology = Field(default_factory=Topology)
     last_event_applied_idx: int = Field(default=-1, ge=-1)
@@ -52,10 +51,16 @@ 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("instances", mode="plain")
+    def _encode_instances(
+        self, value: Mapping[InstanceId, BaseInstance]
+    ) -> dict[str, Any]:
+        """Serialize instances with full subclass fields."""
+        return {
+            str(k): v.model_dump(by_alias=True, serialize_as_any=True)
+            for k, v in value.items()
+        }
 
     @field_serializer("topology", mode="plain")
     def _encode_topology(self, value: Topology) -> TopologySnapshot:

src/exo/shared/types/tasks.py

@@ -2,12 +2,13 @@ from enum import Enum
 
 from pydantic import Field
 
+from exo.plugins.type_registry import task_registry
 from exo.shared.types.api import (
-    ImageEditsTaskParams,
+    ChatCompletionTaskParams,
+    ImageEditsInternalParams,
     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
@@ -32,34 +33,41 @@ class BaseTask(TaggedModel):
     instance_id: InstanceId
 
 
+@task_registry.register
 class CreateRunner(BaseTask):  # emitted by Worker
     bound_instance: BoundInstance
 
 
+@task_registry.register
 class DownloadModel(BaseTask):  # emitted by Worker
     shard_metadata: ShardMetadata
 
 
+@task_registry.register
 class LoadModel(BaseTask):  # emitted by Worker
     pass
 
 
+@task_registry.register
 class ConnectToGroup(BaseTask):  # emitted by Worker
     pass
 
 
+@task_registry.register
 class StartWarmup(BaseTask):  # emitted by Worker
     pass
 
 
-class TextGeneration(BaseTask):  # emitted by Master
+@task_registry.register
+class ChatCompletion(BaseTask):  # emitted by Master
     command_id: CommandId
-    task_params: TextGenerationTaskParams
+    task_params: ChatCompletionTaskParams
 
     error_type: str | None = Field(default=None)
     error_message: str | None = Field(default=None)
 
 
+@task_registry.register
 class ImageGeneration(BaseTask):  # emitted by Master
     command_id: CommandId
     task_params: ImageGenerationTaskParams
@@ -68,25 +76,28 @@ class ImageGeneration(BaseTask):  # emitted by Master
     error_message: str | None = Field(default=None)
 
 
+@task_registry.register
 class ImageEdits(BaseTask):  # emitted by Master
     command_id: CommandId
-    task_params: ImageEditsTaskParams
+    task_params: ImageEditsInternalParams
 
     error_type: str | None = Field(default=None)
     error_message: str | None = Field(default=None)
 
 
+@task_registry.register
 class Shutdown(BaseTask):  # emitted by Worker
     runner_id: RunnerId
 
 
+# Union type for Pydantic validation - tries each type in order
 Task = (
     CreateRunner
     | DownloadModel
     | ConnectToGroup
     | LoadModel
     | StartWarmup
-    | TextGeneration
+    | ChatCompletion
     | ImageGeneration
     | ImageEdits
     | Shutdown

src/exo/shared/types/text_generation.py

@@ -1,42 +0,0 @@
-"""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,6 +21,9 @@ class RDMAConnection(FrozenModel):
     source_rdma_iface: str
     sink_rdma_iface: str
 
+    def is_thunderbolt(self) -> bool:
+        return True
+
 
 class SocketConnection(FrozenModel):
     sink_multiaddr: Multiaddr
@@ -28,6 +31,9 @@ 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/instances.py

@@ -1,7 +1,15 @@
+"""Instance types for exo.
+
+Instances are registered dynamically via the instance_registry, allowing plugins
+to add their own instance types without modifying this file.
+"""
+
 from enum import Enum
+from typing import Any, cast
 
-from pydantic import model_validator
+from pydantic import field_validator, model_validator
 
+from exo.plugins.type_registry import instance_registry
 from exo.shared.types.common import Host, Id, NodeId
 from exo.shared.types.worker.runners import RunnerId, ShardAssignments, ShardMetadata
 from exo.utils.pydantic_ext import CamelCaseModel, TaggedModel
@@ -17,6 +25,8 @@ class InstanceMeta(str, Enum):
 
 
 class BaseInstance(TaggedModel):
+    """Base class for all instance types."""
+
     instance_id: InstanceId
     shard_assignments: ShardAssignments
 
@@ -24,25 +34,36 @@ class BaseInstance(TaggedModel):
         return self.shard_assignments.runner_to_shard.get(runner_id, None)
 
 
+@instance_registry.register
 class MlxRingInstance(BaseInstance):
     hosts_by_node: dict[NodeId, list[Host]]
     ephemeral_port: int
 
 
+@instance_registry.register
 class MlxJacclInstance(BaseInstance):
     jaccl_devices: list[list[str | None]]
     jaccl_coordinators: dict[NodeId, str]
 
 
-# TODO: Single node instance
+# Union type for Pydantic validation - tries each type in order
+# This is used by API endpoints (dashboard) which send flat format
 Instance = MlxRingInstance | MlxJacclInstance
 
 
 class BoundInstance(CamelCaseModel):
-    instance: Instance
+    """An instance bound to a specific runner on a specific node."""
+
+    instance: BaseInstance
     bound_runner_id: RunnerId
     bound_node_id: NodeId
 
+    @field_validator("instance", mode="before")
+    @classmethod
+    def validate_instance(cls, v: Any) -> BaseInstance:  # noqa: ANN401  # pyright: ignore[reportAny]
+        """Validate instance using registry to handle both tagged and flat formats."""
+        return cast(BaseInstance, instance_registry.deserialize(v))  # pyright: ignore[reportAny]
+
     @property
     def bound_shard(self) -> ShardMetadata:
         shard = self.instance.shard(self.bound_runner_id)

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

@@ -1,13 +1,7 @@
 from collections.abc import Generator
 from typing import Any, Literal
 
-from exo.shared.types.api import (
-    FinishReason,
-    GenerationStats,
-    ImageGenerationStats,
-    ToolCallItem,
-    Usage,
-)
+from exo.shared.types.api import FinishReason, GenerationStats, ImageGenerationStats
 from exo.utils.pydantic_ext import TaggedModel
 
 
@@ -25,14 +19,12 @@ 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]
@@ -47,7 +39,6 @@ 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]
@@ -57,10 +48,5 @@ 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,14 +13,12 @@ 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,
@@ -36,142 +34,6 @@ 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"""
 
@@ -196,80 +58,19 @@ 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, ValidationError):
+            except (tomllib.TOMLDecodeError, UnicodeDecodeError):
                 logger.warning("Invalid config file, skipping...")
                 return None
 
@@ -310,7 +111,6 @@ GatheredInfo = (
     | NodeNetworkInterfaces
     | MacThunderboltIdentifiers
     | MacThunderboltConnections
-    | ThunderboltBridgeInfo
     | NodeConfig
     | MiscData
     | StaticNodeInformation
@@ -325,7 +125,6 @@ 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):
@@ -334,7 +133,6 @@ 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)
@@ -351,8 +149,13 @@ 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:
-            await self.info_sender.send(await MiscData.gather())
+            curr = await MiscData.gather()
+            if prev != curr:
+                prev = curr
+                await self.info_sender.send(curr)
             await anyio.sleep(self.misc_poll_interval)
 
     async def _monitor_system_profiler_thunderbolt_data(self):
@@ -362,12 +165,15 @@ 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]
-            await self.info_sender.send(MacThunderboltIdentifiers(idents=idents))
+            if idents != old_idents:
+                await self.info_sender.send(MacThunderboltIdentifiers(idents=idents))
+            old_idents = idents
 
             conns = [it for i in data if (it := i.conn()) is not None]
             await self.info_sender.send(MacThunderboltConnections(conns=conns))
@@ -392,20 +198,14 @@ class InfoGatherer:
     async def _watch_system_info(self):
         if self.interface_watcher_interval is None:
             return
+        old_nics = []
         while True:
-            nics = await get_network_interfaces()
-            await self.info_sender.send(NodeNetworkInterfaces(ifaces=nics))
+            nics = get_network_interfaces()
+            if nics != old_nics:
+                old_nics = nics
+                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 InterfaceType, NetworkInterfaceInfo
+from exo.shared.types.profiling import NetworkInterfaceInfo
 
 
 async def get_friendly_name() -> str:
@@ -16,7 +16,8 @@ async def get_friendly_name() -> str:
     """
     hostname = socket.gethostname()
 
-    if sys.platform != "darwin":
+    # TODO: better non mac support
+    if sys.platform != "darwin":  # 'darwin' is the platform name for macOS
         return hostname
 
     try:
@@ -27,41 +28,7 @@ async def get_friendly_name() -> str:
     return process.stdout.decode("utf-8", errors="replace").strip() or hostname
 
 
-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]:
+def get_network_interfaces() -> list[NetworkInterfaceInfo]:
     """
     Retrieves detailed network interface information on macOS.
     Parses output from 'networksetup -listallhardwareports' and 'ifconfig'
@@ -69,18 +36,13 @@ async 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,
-                            interface_type=interface_types.get(iface, "unknown"),
-                        )
+                        NetworkInterfaceInfo(name=iface, ip_address=service.address)
                     )
                 case _:
                     pass

src/exo/utils/keyed_backoff.py

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