Use sliding window

AGENTS.md

@@ -4,7 +4,7 @@ This file provides guidance to AI coding agents when working with code in this r
 
 ## Project Overview
 
-exo is a distributed AI inference system that connects multiple devices into a cluster. It enables running large language models across multiple machines using MLX as the inference backend and zenoh for peer-to-peer networking.
+exo is a distributed AI inference system that connects multiple devices into a cluster. It enables running large language models across multiple machines using MLX as the inference backend and libp2p for peer-to-peer networking.
 
 ## Build & Run Commands
 
@@ -69,7 +69,7 @@ If `nix fmt` changes any files, stage them before committing. The CI runs `nix f
 
 ### Node Composition
 A single exo `Node` (src/exo/main.py) runs multiple components:
-- **Router**: zenoh-based pub/sub messaging via Rust bindings (exo_rs)
+- **Router**: libp2p-based pub/sub messaging via Rust bindings (exo_pyo3_bindings)
 - **Worker**: Handles inference tasks, downloads models, manages runner processes
 - **Master**: Coordinates cluster state, places model instances across nodes
 - **Election**: Bully algorithm for master election
@@ -81,7 +81,7 @@ Components communicate via typed pub/sub topics (src/exo/routing/topics.py):
 - `LOCAL_EVENTS`: Workers send events to master for indexing
 - `COMMANDS`: Workers/API send commands to master
 - `ELECTION_MESSAGES`: Election protocol messages
-- `CONNECTION_MESSAGES`: zenoh connection updates
+- `CONNECTION_MESSAGES`: libp2p connection updates
 
 ### Event Sourcing
 The system uses event sourcing for state management:
@@ -98,8 +98,8 @@ The system uses event sourcing for state management:
 
 ### Rust Components
 Rust code in `rust/` provides:
-- `networking`: zenoh networking (gossipsub, peer discovery)
-- `exo_rs`: PyO3 bindings exposing Rust to Python
+- `networking`: libp2p networking (gossipsub, peer discovery)
+- `exo_pyo3_bindings`: PyO3 bindings exposing Rust to Python
 - `system_custodian`: System-level operations
 
 ### Dashboard

Cargo.toml

@@ -1,6 +1,6 @@
 [workspace]
 resolver = "3"
-members = ["rust/exo_rs", "rust/networking"]
+members = ["rust/networking", "rust/exo_rs", "rust/util"]
 
 [workspace.package]
 version = "0.0.1"
@@ -20,72 +20,31 @@ opt-level = 3
 [workspace.dependencies]
 ## Crate members as common dependencies
 networking = { path = "rust/networking" }
+util = { path = "rust/util" }
 
-# pyo3
-pyo3 = "0.28.3"
-pyo3-async-runtimes = "0.28.0"
-pyo3-log = "0.13.2"
-pyo3-stub-gen = "0.22.2"
-
-# util
+# Macro dependecies
 extend = "1.2"
 delegate = "0.13"
 
 # Utility dependencies
+keccak-const = "0.2"
 nix = "0.31"
 
 # Async dependencies
+async-stream = "0.3"
 tokio = "1.46"
 futures-lite = "2.6.1"
-async-stream = "0.3.6"
-pin-project = "1.1.10"
-serde_json = "1.0.149"
-rand = "0.10.1"
-parking_lot = "0.12.5"
-pidfile-rs = "0.3.1"
+futures-timer = "3.0"
+
+# Data structures
+either = "1.15"
 
 # Tracing/logging
 log = "0.4"
-env_logger = "0.11.10"
 
 # networking
-zenoh = "=1.9.0"
-zenoh-ext = "=1.9.0"
-zenoh-plugin-storage-manager = { version = "=1.9.0", default-features = false }
-zenoh-plugin-trait = "=1.9.0"
-netwatcher = "0.6.0"
-bytemuck = "1.25.0"
-
-[patch.crates-io]
-zenoh = { git = "https://github.com/evanev7/zenoh.git", branch = "exo" }
-zenoh-ext = { git = "https://github.com/evanev7/zenoh.git", branch = "exo" }
-zenoh-buffers = { git = "https://github.com/evanev7/zenoh.git", branch = "exo" }
-zenoh-codec = { git = "https://github.com/evanev7/zenoh.git", branch = "exo" }
-zenoh-collections = { git = "https://github.com/evanev7/zenoh.git", branch = "exo" }
-zenoh-config = { git = "https://github.com/evanev7/zenoh.git", branch = "exo" }
-zenoh-core = { git = "https://github.com/evanev7/zenoh.git", branch = "exo" }
-zenoh-crypto = { git = "https://github.com/evanev7/zenoh.git", branch = "exo" }
-zenoh-keyexpr = { git = "https://github.com/evanev7/zenoh.git", branch = "exo" }
-zenoh-link = { git = "https://github.com/evanev7/zenoh.git", branch = "exo" }
-zenoh-link-commons = { git = "https://github.com/evanev7/zenoh.git", branch = "exo" }
-zenoh-link-quic = { git = "https://github.com/evanev7/zenoh.git", branch = "exo" }
-zenoh-link-quic_datagram = { git = "https://github.com/evanev7/zenoh.git", branch = "exo" }
-zenoh-link-tcp = { git = "https://github.com/evanev7/zenoh.git", branch = "exo" }
-zenoh-link-tls = { git = "https://github.com/evanev7/zenoh.git", branch = "exo" }
-zenoh-link-udp = { git = "https://github.com/evanev7/zenoh.git", branch = "exo" }
-zenoh-link-unixsock_stream = { git = "https://github.com/evanev7/zenoh.git", branch = "exo" }
-zenoh-link-ws = { git = "https://github.com/evanev7/zenoh.git", branch = "exo" }
-zenoh-macros = { git = "https://github.com/evanev7/zenoh.git", branch = "exo" }
-zenoh-plugin-trait = { git = "https://github.com/evanev7/zenoh.git", branch = "exo" }
-zenoh-protocol = { git = "https://github.com/evanev7/zenoh.git", branch = "exo" }
-zenoh-result = { git = "https://github.com/evanev7/zenoh.git", branch = "exo" }
-zenoh-runtime = { git = "https://github.com/evanev7/zenoh.git", branch = "exo" }
-zenoh-sync = { git = "https://github.com/evanev7/zenoh.git", branch = "exo" }
-zenoh-task = { git = "https://github.com/evanev7/zenoh.git", branch = "exo" }
-zenoh-transport = { git = "https://github.com/evanev7/zenoh.git", branch = "exo" }
-zenoh-util = { git = "https://github.com/evanev7/zenoh.git", branch = "exo" }
-zenoh-plugin-storage-manager = { git = "https://github.com/evanev7/zenoh.git", branch = "exo" }
-zenoh_backend_traits = { git = "https://github.com/evanev7/zenoh.git", branch = "exo" }
+libp2p = "0.56"
+libp2p-tcp = "0.44"
 
 [workspace.lints.rust]
 static_mut_refs = "warn"      # Or use "warn" instead of deny

MISSED_THINGS.md

@@ -0,0 +1,41 @@
+# Missed things
+[X] Log EXO_LIBP2P_NAMESPACE on start in exo/main.py
+[X] Ordering of warmup was changed, which is wrong. It was changed to rank < n-1, then rank=n-1. It should be rank!=0 then rank=0 (this matches the auto_parallel implementation. NOTE: we use a different convention to mlx-lm, our terminal rank is rank=n-1 whereas mlx-lm is rank=0 hence i can see why this was changed wrongly).
+[X] Downloads keying by model_id not shard_metadata (worker/plan.py, worker/main.py).
+[X] Fetching download status of all models on start
+[X] Deduplication of tasks in plan_step.
+[X] resolve_allow_patterns should just be wildcard now.
+[X] no mx_barrier in genreate.py mlx_generate at the end.
+[] cache assertion not needed in auto_parallel.py PipelineLastLayer.
+[X] GPTOSS support dropped in auto_parallel.py.
+[X] sharding changed "all-to-sharded" became _all_to_sharded in auto_parallel.py.
+[X] same as above with "sharded-to-all" became _sharded_to_all in auto_parallel.py.
+[X] Dropped support for Ministral3Model, DeepseekV32Model, Glm4MoeModel, Qwen3NextModel, GptOssMode in auto_parallel.py.
+[] Dropped prefill/decode code in auto_parallel.py and utils_mlx.py.
+[X] KV_CACHE_BITS should be None to disable quantized KV cache.
+[X] Dropped _set_nofile_limit in utils_mlx.py.
+[X] We have group optional in load_mlx_items in utils_mlx.py.
+[X] Dropped add_missing_chat_templates for GptOss in load_mlx_items in utils_mlx.py.
+[X] Dropped model.make_cache in make_kv_cache in utils_mlx.py.
+[X] We put cache limit back in utils_mlx.py.
+[X] topology.py remove_node removes the connections after checking if node is is in self._node_id_to_rx_id_map. on beta_1 it checks after, so would remove stale connections I guess?
+[X] Missing Glm 4.7 model cards (this isn't ready yet but should be picked up, probably create an issue... the blocker is transforemrs version doesn't support the tokenizer for Glm 4.7. rc-1 does but we can't upgrade as it breaks other things.)
+[] try-except in _command_processor only excepts ValueError. This was silently failing leading to un-debuggable errors (we had a KeyError that was happening ). Changed this to catch Exception instead of ValueError. See exo-v2 89ae38405e0052e3c22405daf094b065878aa873 and fb99fea69b5a39017efc90c5dad0072e677455f0.
+[X] In placement.py, place_instance no longer looks at model_meta.supports_tensor and check if this tensor parallel number of nodes is supported by the model's tensor dimensions.
+[X] In placement.py, place_instanec, we no longer have the special case to exclude DeepSeek v3.1 pipeline parallel (it doesn't work).
+[] logger.warning("You have likely selected ibv for a single node instance; falling back to MlxRing") was changed to debug. That will spam this warning since it happens every time we query instance previews.
+[X] In placement_utils.py, get_mlx_jaccl_coordinators, We no longer prioritise Jaccl Coordinator IP. Now it picks the first one, which is unstable (Jaccl coordinator over TB5 is unstable).
+
+
+
+[X] Downloads keying by model_id not shard_metadata (worker/plan.py, worker/main.py).
+[X] Fetching download status of all models on start
+[X] Deduplication of tasks in plan_step.
+[X] resolve_allow_patterns should just be wildcard now.
+[X] KV_CACHE_BITS should be None to disable quantized KV cache.
+[X] We put cache limit back in utils_mlx.py.
+[X] In placement.py, place_instance no longer looks at model_meta.supports_tensor and check if this tensor parallel number of nodes is supported by the model's tensor dimensions.
+[X] In placement.py, place_instanec, we no longer have the special case to exclude DeepSeek v3.1 pipeline parallel (it doesn't work).
+[X] In placement_utils.py, get_mlx_jaccl_coordinators, We no longer prioritise Jaccl Coordinator IP. Now it picks the first one, which is unstable (Jaccl coordinator over TB5 is unstable).
+
+

app/EXO/EXO/ExoProcessController.swift

@@ -352,7 +352,7 @@ final class ExoProcessController: ObservableObject {
     private func makeEnvironment(for runtimeURL: URL) -> [String: String] {
         var environment = ProcessInfo.processInfo.environment
         environment["EXO_RUNTIME_DIR"] = runtimeURL.path
-        environment["EXO_ZENOH_NAMESPACE"] = computeNamespace()
+        environment["EXO_LIBP2P_NAMESPACE"] = computeNamespace()
         if !hfToken.isEmpty {
             environment["HF_TOKEN"] = hfToken
         }

dashboard/src/lib/components/ChatSidebar.svelte

@@ -12,8 +12,6 @@
     toggleDebugMode,
     topologyOnlyMode,
     toggleTopologyOnlyMode,
-    getInstanceFirstShard,
-    type Instance,
   } from "$lib/stores/app.svelte";
 
   interface Props {
@@ -188,7 +186,7 @@
   function extractInstanceModelId(instanceWrapped: unknown): string | null {
     const [, instance] = getTaggedValue(instanceWrapped);
     if (!instance || typeof instance !== "object") return null;
-    const inst = instance as Instance;
+    const inst = instance as { shardAssignments?: { modelId?: string } };
     return inst.shardAssignments?.modelId ?? null;
   }
 
@@ -206,7 +204,11 @@
     else if (instanceTag === "MlxJacclInstance") instanceType = "MLX RDMA";
 
     let sharding: string | null = null;
-    const firstShardWrapped = getInstanceFirstShard(instance as Instance);
+    const inst = instance as {
+      shardAssignments?: { runnerToShard?: Record<string, unknown> };
+    };
+    const runnerToShard = inst.shardAssignments?.runnerToShard || {};
+    const firstShardWrapped = Object.values(runnerToShard)[0];
     if (firstShardWrapped) {
       const [shardTag] = getTaggedValue(firstShardWrapped);
       if (shardTag === "PipelineShardMetadata") sharding = "Pipeline";

dashboard/src/lib/components/PrefillDecodeDisaggregation.svelte

@@ -9,7 +9,6 @@
     createInstanceLink,
     updateInstanceLink,
     deleteInstanceLink,
-    getInstanceNodeIds,
     type Instance,
   } from "$lib/stores/app.svelte";
   import { deriveBaseModel, deriveFamily } from "$lib/utils/model_family";
@@ -17,6 +16,7 @@
   type InstanceWrapper = {
     MlxRingInstance?: Instance;
     MlxJacclInstance?: Instance;
+    VllmInstance?: Instance;
   };
 
   let interval: ReturnType<typeof setInterval> | null = null;
@@ -43,9 +43,13 @@
     const ids = nodeIdentities();
     for (const [id, raw] of Object.entries(instances())) {
       const wrapper = raw as InstanceWrapper;
-      const inst = wrapper.MlxRingInstance ?? wrapper.MlxJacclInstance;
+      const inst =
+        wrapper.MlxRingInstance ??
+        wrapper.MlxJacclInstance ??
+        wrapper.VllmInstance;
       const modelId = inst?.shardAssignments?.modelId ?? "";
-      const nodeIds = getInstanceNodeIds(inst);
+      const nodeToRunner = inst?.shardAssignments?.nodeToRunner ?? {};
+      const nodeIds = Object.keys(nodeToRunner);
       const nodeNames = nodeIds
         .map((nodeId) => ids[nodeId]?.friendlyName ?? nodeId.slice(0, 6))
         .filter((name) => !!name);

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

@@ -66,40 +66,12 @@ export interface TopologyData {
   edges: TopologyEdge[];
 }
 
-export type InstanceShard = [nodeId: string, runnerId: string, shard: unknown];
-
-export interface ShardAssignments {
-  modelId: string;
-  shards: InstanceShard[];
-  primaryOutputNode: number;
-}
-
 export interface Instance {
-  shardAssignments: ShardAssignments;
-}
-
-export function getInstanceShards(
-  instance: Instance | null | undefined,
-): InstanceShard[] {
-  return instance?.shardAssignments.shards ?? [];
-}
-
-export function getInstanceRunnerIds(
-  instance: Instance | null | undefined,
-): string[] {
-  return getInstanceShards(instance).map(([, runnerId]) => runnerId);
-}
-
-export function getInstanceNodeIds(
-  instance: Instance | null | undefined,
-): string[] {
-  return [...new Set(getInstanceShards(instance).map(([nodeId]) => nodeId))];
-}
-
-export function getInstanceFirstShard(
-  instance: Instance | null | undefined,
-): unknown {
-  return getInstanceShards(instance)[0]?.[2];
+  shardAssignments?: {
+    modelId?: string;
+    runnerToShard?: Record<string, unknown>;
+    nodeToRunner?: Record<string, string>;
+  };
 }
 
 export interface RawInstanceLink {
@@ -946,7 +918,7 @@ class AppStore {
   private extractInstanceModelId(instanceWrapped: unknown): string | null {
     const [, instance] = this.getTaggedValue(instanceWrapped);
     if (!instance || typeof instance !== "object") return null;
-    const inst = instance as Instance;
+    const inst = instance as { shardAssignments?: { modelId?: string } };
     return inst.shardAssignments?.modelId ?? null;
   }
 
@@ -964,8 +936,11 @@ class AppStore {
     else if (instanceTag === "MlxJacclInstance") instanceType = "MLX RDMA";
 
     let sharding: string | null = null;
-    const inst = instance as Instance;
-    const firstShardWrapped = getInstanceFirstShard(inst);
+    const inst = instance as {
+      shardAssignments?: { runnerToShard?: Record<string, unknown> };
+    };
+    const runnerToShard = inst.shardAssignments?.runnerToShard || {};
+    const firstShardWrapped = Object.values(runnerToShard)[0];
     if (firstShardWrapped) {
       const [shardTag] = this.getTaggedValue(firstShardWrapped);
       if (shardTag === "PipelineShardMetadata") sharding = "Pipeline";
@@ -2287,7 +2262,7 @@ class AppStore {
         if (keys.length === 1) {
           const instance = (instanceWrapper as Record<string, unknown>)[
             keys[0]
-          ] as Instance;
+          ] as { shardAssignments?: { modelId?: string } };
           const instanceModelId = instance?.shardAssignments?.modelId;
 
           // ensure to only return requestedModelId that matches an instance

dashboard/src/routes/+page.svelte

@@ -65,11 +65,6 @@
     nodeThunderboltBridge,
     nodeIdentities,
     isConnected,
-    getInstanceFirstShard,
-    getInstanceNodeIds,
-    getInstanceRunnerIds,
-    getInstanceShards,
-    type Instance,
     type DownloadProgress,
     type PlacementPreview,
   } from "$lib/stores/app.svelte";
@@ -1003,7 +998,11 @@
     if (keys.length !== 1) return new Set();
     const instance = (instanceWrapped as Record<string, unknown>)[keys[0]];
     if (!instance || typeof instance !== "object") return new Set();
-    return new Set(getInstanceNodeIds(instance as Instance));
+    const inst = instance as {
+      shardAssignments?: { nodeToRunner?: Record<string, string> };
+    };
+    if (!inst.shardAssignments?.nodeToRunner) return new Set();
+    return new Set(Object.keys(inst.shardAssignments.nodeToRunner));
   }
 
   function toggleInstanceDownloadDetails(nodeId: string): void {
@@ -1785,7 +1784,13 @@
       };
     }
 
-    const inst = instance as Instance;
+    const inst = instance as {
+      shardAssignments?: {
+        nodeToRunner?: Record<string, string>;
+        runnerToShard?: Record<string, unknown>;
+        modelId?: string;
+      };
+    };
     const instanceModelId = inst.shardAssignments?.modelId;
 
     if (!instanceModelId) {
@@ -1800,7 +1805,16 @@
       };
     }
 
-    const instanceNodeIds = getInstanceNodeIds(inst);
+    // Get node IDs assigned to this instance
+    const nodeToRunner = inst.shardAssignments?.nodeToRunner || {};
+    const runnerToShard = inst.shardAssignments?.runnerToShard || {};
+    const runnerToNode: Record<string, string> = {};
+    for (const [nodeId, runnerId] of Object.entries(nodeToRunner)) {
+      runnerToNode[runnerId] = nodeId;
+    }
+    const instanceNodeIds = Object.keys(runnerToShard)
+      .map((runnerId) => runnerToNode[runnerId])
+      .filter(Boolean);
 
     const result = collectDownloadStatus(instanceModelId, instanceNodeIds);
 
@@ -1844,7 +1858,6 @@
       case "FAILED":
         return "text-red-400";
       case "SHUTDOWN":
-      case "SHUTTING DOWN":
         return "text-gray-400";
       case "DOWNLOADING":
         return "text-blue-400";
@@ -1852,7 +1865,6 @@
       case "WARMING UP":
       case "WAITING":
       case "INITIALIZING":
-      case "CONNECTING":
         return "text-yellow-400";
       case "RUNNING":
         return "text-teal-400";
@@ -1875,7 +1887,10 @@
       return { statusText: "PREPARING", statusClass: "inactive" };
     }
 
-    const runnerIds = getInstanceRunnerIds(instance as Instance);
+    const inst = instance as {
+      shardAssignments?: { runnerToShard?: Record<string, unknown> };
+    };
+    const runnerIds = Object.keys(inst.shardAssignments?.runnerToShard || {});
 
     const statuses = runnerIds
       .map((rid) => {
@@ -1883,15 +1898,14 @@
         if (!r) return null;
         const [kind] = getTagged(r);
         const statusMap: Record<string, string> = {
-          RunnerIdle: "Idle",
-          RunnerConnecting: "Connecting",
-          RunnerConnected: "Connected",
+          RunnerWaitingForInitialization: "WaitingForInitialization",
+          RunnerInitializingBackend: "InitializingBackend",
+          RunnerWaitingForModel: "WaitingForModel",
           RunnerLoading: "Loading",
           RunnerLoaded: "Loaded",
           RunnerWarmingUp: "WarmingUp",
           RunnerReady: "Ready",
           RunnerRunning: "Running",
-          RunnerShuttingDown: "ShuttingDown",
           RunnerShutdown: "Shutdown",
           RunnerFailed: "Failed",
         };
@@ -1945,15 +1959,14 @@
       return { statusText: "RUNNING", statusClass: "running" };
     if (has("Ready")) return { statusText: "READY", statusClass: "loaded" };
     if (has("Loaded")) return { statusText: "LOADED", statusClass: "loaded" };
-    if (has("Connected"))
+    if (has("WaitingForModel"))
+      return { statusText: "WAITING", statusClass: "starting" };
+    if (has("InitializingBackend"))
+      return { statusText: "INITIALIZING", statusClass: "starting" };
+    if (has("WaitingForInitialization"))
       return { statusText: "INITIALIZING", statusClass: "starting" };
-    if (has("Connecting"))
-      return { statusText: "CONNECTING", statusClass: "starting" };
-    if (has("Idle")) return { statusText: "WAITING", statusClass: "starting" };
-    if (has("ShuttingDown"))
-      return { statusText: "SHUTTING DOWN", statusClass: "inactive" };
 
-    return { statusText: "PREPARING", statusClass: "inactive" };
+    return { statusText: "RUNNING", statusClass: "active" };
   }
 
   function getBytes(value: unknown): number {
@@ -2026,7 +2039,7 @@
   function getInstanceModelId(instanceWrapped: unknown): string {
     const [, instance] = getTagged(instanceWrapped);
     if (!instance || typeof instance !== "object") return "Unknown";
-    const inst = instance as Instance;
+    const inst = instance as { shardAssignments?: { modelId?: string } };
     return inst.shardAssignments?.modelId || "Unknown Model";
   }
 
@@ -2054,11 +2067,17 @@
     if (instanceTag === "MlxRingInstance") instanceType = "MLX Ring";
     else if (instanceTag === "MlxJacclInstance") instanceType = "MLX RDMA";
 
-    const inst = instance as Instance;
+    const inst = instance as {
+      shardAssignments?: {
+        nodeToRunner?: Record<string, string>;
+        runnerToShard?: Record<string, unknown>;
+      };
+    };
 
     // Sharding strategy from first shard
     let sharding = "Unknown";
-    const firstShardWrapped = getInstanceFirstShard(inst);
+    const runnerToShard = inst.shardAssignments?.runnerToShard || {};
+    const firstShardWrapped = Object.values(runnerToShard)[0];
     if (firstShardWrapped) {
       const [shardTag] = getTagged(firstShardWrapped);
       if (shardTag === "PipelineShardMetadata") sharding = "Pipeline";
@@ -2068,7 +2087,8 @@
     }
 
     // Node names from topology
-    const nodeIds = getInstanceNodeIds(inst);
+    const nodeToRunner = inst.shardAssignments?.nodeToRunner || {};
+    const nodeIds = Object.keys(nodeToRunner);
     const nodeNames = nodeIds.map((nodeId) => {
       const node = data?.nodes?.[nodeId];
       return node?.friendly_name || nodeId.slice(0, 8);
@@ -2172,19 +2192,35 @@
   }
 
   function getOrderedRunnerNodes(
-    instance: Instance,
+    instance: Record<string, unknown>,
     shardType: "Pipeline" | "Tensor",
   ) {
-    const runnerEntries = getInstanceShards(instance).map(
-      ([nodeId, runnerId, shardWrapped]) => {
+    const runnerToShard =
+      (
+        instance.shardAssignments as
+          | { runnerToShard?: Record<string, unknown> }
+          | undefined
+      )?.runnerToShard || {};
+    const nodeToRunner =
+      (
+        instance.shardAssignments as
+          | { nodeToRunner?: Record<string, string> }
+          | undefined
+      )?.nodeToRunner || {};
+    const runnerEntries = Object.entries(runnerToShard).map(
+      ([runnerId, shardWrapped]) => {
         const [tag, shard] = getTagged(shardWrapped);
         const meta = shard as
           | {
-              deviceRank?: number;
+              modelMeta?: {
+                worldSize?: number;
+                nLayers?: number;
+                deviceRank?: number;
+              };
             }
           | undefined;
-        const deviceRank = meta?.deviceRank ?? 0;
-        return { nodeId, runnerId, tag, deviceRank };
+        const deviceRank = meta?.modelMeta?.deviceRank ?? 0;
+        return { runnerId, tag, deviceRank };
       },
     );
 
@@ -2195,11 +2231,13 @@
           : r.tag === "TensorShardMetadata",
       )
       .sort((a, b) => a.deviceRank - b.deviceRank)
-      .map((r, idx) => ({
-        nodeId: r.nodeId,
-        runnerId: r.runnerId,
-        order: idx,
-      }));
+      .map((r, idx) => {
+        const nodeId = Object.entries(nodeToRunner).find(
+          ([, rid]) => rid === r.runnerId,
+        )?.[0];
+        return { nodeId, runnerId: r.runnerId, order: idx };
+      })
+      .filter((item) => item.nodeId);
 
     return ordered as Array<{
       nodeId: string;
@@ -2243,7 +2281,10 @@
 
     // Jaccl (RDMA) – show RDMA interfaces from ibvDevices
     if (instanceTag === "MlxJacclInstance") {
-      const ordered = getOrderedRunnerNodes(instance as Instance, "Tensor");
+      const ordered = getOrderedRunnerNodes(
+        instance as Record<string, unknown>,
+        "Tensor",
+      );
       const ibvDevices =
         (instance as { ibvDevices?: Array<Array<string | null>> }).ibvDevices ||
         [];
@@ -2275,7 +2316,10 @@
 
     // Ring – derive ring order from pipeline shard ranks and pick host IPs from hostsByNode
     if (instanceTag === "MlxRingInstance") {
-      const ordered = getOrderedRunnerNodes(instance as Instance, "Pipeline");
+      const ordered = getOrderedRunnerNodes(
+        instance as Record<string, unknown>,
+        "Pipeline",
+      );
       const hostsByNode =
         (
           instance as {
@@ -2562,7 +2606,6 @@
         status.statusText === "WARMING UP" ||
         status.statusText === "WAITING" ||
         status.statusText === "INITIALIZING" ||
-        status.statusText === "CONNECTING" ||
         status.statusText === "PREPARING"
       ) {
         chatLaunchState = "launching";
@@ -5065,10 +5108,7 @@
                   {@const isFailed = statusText === "FAILED"}
                   {@const isLoading = statusText === "LOADING"}
                   {@const isWarmingUp =
-                    statusText === "WARMING UP" ||
-                    statusText === "WAITING" ||
-                    statusText === "INITIALIZING" ||
-                    statusText === "CONNECTING"}
+                    statusText === "WARMING UP" || statusText === "WAITING"}
                   {@const isReady =
                     statusText === "READY" || statusText === "LOADED"}
                   {@const isRunning = statusText === "RUNNING"}
@@ -6204,10 +6244,7 @@
                     {@const isFailed = statusText === "FAILED"}
                     {@const isLoading = statusText === "LOADING"}
                     {@const isWarmingUp =
-                      statusText === "WARMING UP" ||
-                      statusText === "WAITING" ||
-                      statusText === "INITIALIZING" ||
-                      statusText === "CONNECTING"}
+                      statusText === "WARMING UP" || statusText === "WAITING"}
                     {@const isReady =
                       statusText === "READY" || statusText === "LOADED"}
                     {@const isRunning = statusText === "RUNNING"}

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

@@ -3,11 +3,7 @@
   import { fade } from "svelte/transition";
   import HeaderNav from "$lib/components/HeaderNav.svelte";
   import IntegrationCard from "$lib/components/IntegrationCard.svelte";
-  import {
-    instances,
-    refreshState,
-    type Instance,
-  } from "$lib/stores/app.svelte";
+  import { instances, refreshState } from "$lib/stores/app.svelte";
   import { onMount } from "svelte";
 
   const apiUrl = browser
@@ -28,7 +24,9 @@
         if (values.length > 0) {
           const instance = values[0];
           if (instance && typeof instance === "object") {
-            const inst = instance as Instance;
+            const inst = instance as {
+              shardAssignments?: { modelId?: string };
+            };
             const modelId = inst.shardAssignments?.modelId;
             if (modelId && !models.includes(modelId)) {
               models.push(modelId);

flake.lock

@@ -47,11 +47,11 @@
         "rust-analyzer-src": "rust-analyzer-src"
       },
       "locked": {
-        "lastModified": 1777708550,
-        "narHash": "sha256-Qif3UXT0l5OQq8H9pRWt4/ia4gF48MWK2oHKL8uVx8U=",
+        "lastModified": 1775807984,
+        "narHash": "sha256-Redoe3D9zGN5I9QPHWL9vfMVQBehY1fKsMiRXQ83X3w=",
         "owner": "nix-community",
         "repo": "fenix",
-        "rev": "74c1591efaff494756b8d35ebe357c6c2bbdca96",
+        "rev": "fcf90c0c4d368b2ca917a7afa6d08e98a397e5fd",
         "type": "github"
       },
       "original": {
@@ -218,11 +218,11 @@
     "rust-analyzer-src": {
       "flake": false,
       "locked": {
-        "lastModified": 1777639980,
-        "narHash": "sha256-6d7Hdurvbjc5uwJuc0YiK7rZBGj6Gs3uzfBFcTs+xCc=",
+        "lastModified": 1775745684,
+        "narHash": "sha256-8MbfLwd60FNa8dRFkjE+G3TT/x21G3Rsplm1bMBQUtU=",
         "owner": "rust-lang",
         "repo": "rust-analyzer",
-        "rev": "64cdaeb06f69b6b769a492edd88b022ae88e8ca2",
+        "rev": "64ddb549bc9a70d011328746fa46a8883f937b6b",
         "type": "github"
       },
       "original": {

flake.nix

@@ -146,7 +146,7 @@
                 config.treefmt.build.wrapper
 
                 # PYTHON
-                #self'.packages.exo.passthru.evenv
+                self'.packages.exo.passthru.evenv
                 uv
 
                 # RUST

pyproject.toml

@@ -85,6 +85,13 @@ mlx = [
   { url = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv/releases/download/mlx_cuda/mlx-0.32.0-cp313-cp313-manylinux_2_35_aarch64.whl", marker = "sys_platform == 'linux' and platform_machine == 'aarch64'" },
   { url = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv/releases/download/mlx_cuda/mlx-0.32.0-cp313-cp313-manylinux_2_35_x86_64.whl", marker = "sys_platform == 'linux' and platform_machine != 'aarch64'" },
 ]
+mlx-lm = { git = "https://github.com/rltakashige/mlx-lm", branch = "leo/deepseek-v4" }
+mflux = { git = "https://github.com/evanev7/mflux", branch = "exo2" }
+torch = [
+  { index = "pytorch-cpu", marker = "sys_platform == 'linux' and extra == 'mlx-cpu' and extra != 'mlx-cuda13' and extra != 'mlx-cuda12'" },
+  { index = "pytorch-cu128", marker = "sys_platform == 'linux' and extra == 'mlx-cuda12' and extra != 'mlx-cuda13' " },
+  { index = "pytorch-cu130", marker = "sys_platform == 'linux' and extra == 'mlx-cuda13'" },
+]
 mlx-cuda-12 = [
   { url = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv/releases/download/mlx_cuda/mlx_cuda_12-0.32.0-py3-none-manylinux_2_35_aarch64.whl", marker = "sys_platform == 'linux' and platform_machine == 'aarch64'" },
   { url = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv/releases/download/mlx_cuda/mlx_cuda_12-0.32.0-py3-none-manylinux_2_35_x86_64.whl", marker = "sys_platform == 'linux' and platform_machine != 'aarch64'" },
@@ -93,13 +100,6 @@ mlx-cuda-13 = [
   { url = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv/releases/download/mlx_cuda/mlx_cuda_13-0.32.0-py3-none-manylinux_2_35_aarch64.whl", marker = "sys_platform == 'linux' and platform_machine == 'aarch64'" },
   { url = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv/releases/download/mlx_cuda/mlx_cuda_13-0.32.0-py3-none-manylinux_2_35_x86_64.whl", marker = "sys_platform == 'linux' and platform_machine != 'aarch64'" },
 ]
-mlx-lm = { git = "https://github.com/rltakashige/mlx-lm", branch = "leo/deepseek-v4" }
-mflux = { git = "https://github.com/evanev7/mflux", branch = "exo2" }
-torch = [
-  { index = "pytorch-cpu", marker = "sys_platform == 'linux' and extra == 'mlx-cpu' and extra != 'mlx-cuda13' and extra != 'mlx-cuda12'" },
-  { index = "pytorch-cu128", marker = "sys_platform == 'linux' and extra == 'mlx-cuda12' and extra != 'mlx-cuda13' " },
-  { index = "pytorch-cu130", marker = "sys_platform == 'linux' and extra == 'mlx-cuda13'" },
-]
 torchvision = [
   { index = "pytorch-cpu", marker = "sys_platform == 'linux' and extra == 'mlx-cpu' and extra != 'mlx-cuda13' and extra != 'mlx-cuda12'" },
   { index = "pytorch-cu128", marker = "sys_platform == 'linux' and extra == 'mlx-cuda12' and extra != 'mlx-cuda13'" },

rust/exo_rs/Cargo.toml

@@ -22,33 +22,48 @@ doc = false
 workspace = true
 
 [dependencies]
-networking.workspace = true
-extend.workspace = true
+networking = { workspace = true }
 
 # interop
-pyo3 = { workspace = true, features = ["experimental-async"] }
-pyo3-stub-gen.workspace = true
-pyo3-async-runtimes = { workspace = true, features = [
+pyo3 = { version = "0.28.3", features = [
+  # "abi3-py313", # tells pyo3 (and maturin) to build using the stable ABI with minimum Python version 3.13
+  # "nightly", # enables better-supported GIL integration
+  "experimental-async", # async support in #[pyfunction] & #[pymethods]
+  #"experimental-inspect", # inspection of generated binary => easier to automate type-hint generation
+  #"py-clone", # adding Clone-ing of `Py<T>` without GIL (may cause panics - remove if panics happen)
+  # "multiple-pymethods", # allows multiple #[pymethods] sections per class
+
+  # integrations with other libraries
+  # "arc_lock", "bigdecimal", "either", "hashbrown", "indexmap", "num-bigint", "num-complex", "num-rational",
+  # "ordered-float", "rust_decimal", "smallvec",
+  # "anyhow", "chrono", "chrono-local", "chrono-tz", "eyre", "jiff-02", "lock_api", "parking-lot", "time",  "serde",
+] }
+pyo3-stub-gen = { version = "0.22.3" }
+pyo3-async-runtimes = { version = "0.28.0", features = [
   "attributes",
   "tokio-runtime",
   "testing",
 ] }
-pyo3-log.workspace = true
+pyo3-log = "0.13.3"
 
 pidfile-rs = { git = "https://github.com/AndreiCravtov/pidfile-rs" }
 
+# macro dependencies
+extend = { workspace = true }
+delegate = { workspace = true }
+thiserror = "2.0"
+
 # async runtime
-tokio = { workspace = true, features = ["full"] }
-futures-lite.workspace = true
-pin-project.workspace = true
+tokio = { workspace = true, features = ["full", "tracing"] }
+futures-lite = { workspace = true }
+
+# utility dependencies
+util = { workspace = true }
 
 # Tracing
-log.workspace = true
-env_logger.workspace = true
+log = { workspace = true }
+env_logger = "0.11"
 
 # Networking
-zenoh.workspace = true
-zenoh-ext = { workspace = true, features = ["unstable"] }
-rand.workspace = true
-serde_json.workspace = true
-parking_lot.workspace = true
+libp2p = { workspace = true, features = ["full"] }
+pin-project = "1.1.10"

rust/exo_rs/exo_rs.pyi

@@ -2,44 +2,83 @@
 # ruff: noqa: E501, F401, F403, F405
 
 import builtins
-import collections.abc
 import os
 import pathlib
 import typing
 __all__ = [
-    "LVAggregator",
-    "LVPublisher",
-    "LVSubscriber",
+    "AllQueuesFullError",
+    "FromSwarm",
+    "Keypair",
+    "MessageTooLargeError",
     "NetworkingHandle",
+    "NoPeersSubscribedToTopicError",
     "Pidfile",
     "PidfileError",
-    "PyFromSwarm",
-    "SessionHandle",
-    "Storage",
-    "StorageGetter",
-    "TaskChunkSender",
-    "TaskRequest",
-    "TaskRequester",
-    "TaskResponder",
-    "TaskStream",
 ]
 
 @typing.final
-class LVAggregator:
-    def dump(self) -> builtins.dict[builtins.str, builtins.str]: ...
+class AllQueuesFullError(builtins.Exception):
+    def __new__(cls, *args: typing.Any) -> AllQueuesFullError: ...
+    def __repr__(self) -> builtins.str: ...
+    def __str__(self) -> builtins.str: ...
+
+class FromSwarm:
+    @typing.final
+    class Connection(FromSwarm):
+        __match_args__ = ("peer_id", "connected",)
+        @property
+        def peer_id(self) -> builtins.str: ...
+        @property
+        def connected(self) -> builtins.bool: ...
+        def __new__(cls, peer_id: builtins.str, connected: builtins.bool) -> FromSwarm.Connection: ...
+    
+    @typing.final
+    class Message(FromSwarm):
+        __match_args__ = ("origin", "topic", "data",)
+        @property
+        def origin(self) -> builtins.str: ...
+        @property
+        def topic(self) -> builtins.str: ...
+        @property
+        def data(self) -> bytes: ...
+        def __new__(cls, origin: builtins.str, topic: builtins.str, data: bytes) -> FromSwarm.Message: ...
+    
+    ...
 
 @typing.final
-class LVPublisher:
-    def put(self, data: builtins.str) -> collections.abc.Awaitable[None]: ...
+class Keypair:
+    r"""
+    Identity keypair of a node.
+    """
+    @staticmethod
+    def generate() -> Keypair:
+        r"""
+        Generate a new Ed25519 keypair.
+        """
+    @staticmethod
+    def from_bytes(bytes: bytes) -> Keypair:
+        r"""
+        Construct an Ed25519 keypair from secret key bytes
+        """
+    def to_bytes(self) -> bytes:
+        r"""
+        Get the secret key bytes underlying the keypair
+        """
+    def to_node_id(self) -> builtins.str:
+        r"""
+        Convert the `Keypair` into the corresponding `PeerId` string, which we use as our `NodeId`.
+        """
 
 @typing.final
-class LVSubscriber:
-    def recv(self) -> collections.abc.Awaitable[tuple[str, str] | None]: ...
+class MessageTooLargeError(builtins.Exception):
+    def __new__(cls, *args: typing.Any) -> MessageTooLargeError: ...
+    def __repr__(self) -> builtins.str: ...
+    def __str__(self) -> builtins.str: ...
 
 @typing.final
 class NetworkingHandle:
-    @staticmethod
-    def new(identity: builtins.str, listen_port: builtins.int, discovery_service_port: builtins.int) -> NetworkingHandle: ...
+    def __new__(cls, identity: Keypair, bootstrap_peers: typing.Sequence[builtins.str], listen_port: builtins.int) -> NetworkingHandle: ...
+    def recv(self) -> typing.Awaitable[FromSwarm]: ...
     async def gossipsub_subscribe(self, topic: builtins.str) -> builtins.bool:
         r"""
         Subscribe to a `GossipSub` topic.
@@ -58,7 +97,12 @@ class NetworkingHandle:
         
         If no peers are found that subscribe to this topic, throws `NoPeersSubscribedToTopicError` exception.
         """
-    async def recv(self) -> PyFromSwarm: ...
+
+@typing.final
+class NoPeersSubscribedToTopicError(builtins.Exception):
+    def __new__(cls, *args: typing.Any) -> NoPeersSubscribedToTopicError: ...
+    def __repr__(self) -> builtins.str: ...
+    def __str__(self) -> builtins.str: ...
 
 @typing.final
 class Pidfile:
@@ -85,7 +129,6 @@ class Pidfile:
     def __new__(cls, path: builtins.str | os.PathLike | pathlib.Path, mode: builtins.int) -> Pidfile:
         r"""
         Creates a new PID file and locks it.
-        Writes the current process ID to the PID file.
         
         If the PID file cannot be locked, returns `PidfileError::AlreadyRunning` with
         a PID of the already running process, or `None` if no PID has been written to
@@ -117,69 +160,3 @@ class PidfileError(builtins.Exception):
     def __repr__(self) -> builtins.str: ...
     def __str__(self) -> builtins.str: ...
 
-class PyFromSwarm:
-    @typing.final
-    class Connection(PyFromSwarm):
-        __match_args__ = ("connected",)
-        @property
-        def connected(self) -> builtins.bool: ...
-        def __new__(cls, connected: builtins.bool) -> PyFromSwarm.Connection: ...
-    
-    @typing.final
-    class Message(PyFromSwarm):
-        __match_args__ = ("topic", "data",)
-        @property
-        def topic(self) -> builtins.str: ...
-        @property
-        def data(self) -> bytes: ...
-        def __new__(cls, topic: builtins.str, data: bytes) -> PyFromSwarm.Message: ...
-    
-    ...
-
-@typing.final
-class SessionHandle:
-    @staticmethod
-    def new(identity: builtins.str, listen_port: builtins.int, discovery_service_port: builtins.int) -> tuple[SessionHandle, NetworkingHandle]: ...
-    def last_value_aggregator(self, prefix: builtins.str) -> LVAggregator: ...
-    def last_value_subscriber(self, kexpr: builtins.str) -> LVSubscriber: ...
-    def last_value_publisher(self, kexpr: builtins.str) -> LVPublisher: ...
-    def storage_interface(self) -> Storage: ...
-    def task_requester(self) -> TaskRequester: ...
-    def task_responder(self, instance_id: builtins.str) -> TaskResponder: ...
-
-@typing.final
-class Storage:
-    def get(self, key: builtins.str) -> collections.abc.Awaitable[str | None]: ...
-    def get_many(self, key: builtins.str) -> StorageGetter: ...
-    def put(self, key: builtins.str, data: builtins.str) -> collections.abc.Awaitable[None]: ...
-    def delete(self, key: builtins.str) -> collections.abc.Awaitable[None]: ...
-    def dump(self, prefix: builtins.str) -> collections.abc.Awaitable[dict[str, str]]: ...
-
-@typing.final
-class StorageGetter:
-    def recv(self) -> collections.abc.Awaitable[tuple[str, str] | None]: ...
-
-@typing.final
-class TaskChunkSender:
-    def send(self, chunk: builtins.str) -> collections.abc.Awaitable[None]: ...
-
-@typing.final
-class TaskRequest:
-    def reply(self, payload: builtins.str) -> None: ...
-    def reply_err(self, payload: builtins.str) -> None: ...
-
-@typing.final
-class TaskRequester:
-    def submit(self, instance_id: builtins.str, command_id: builtins.str, command: builtins.str) -> collections.abc.Awaitable[TaskStream]: ...
-    def interrupt(self, instance_id: builtins.str, command_id: builtins.str, command: builtins.str) -> collections.abc.Awaitable[None]: ...
-
-@typing.final
-class TaskResponder:
-    def assign_task(self, task_id: builtins.str, task: builtins.str) -> collections.abc.Awaitable[None]: ...
-    def unassign_task(self, task_id: builtins.str) -> collections.abc.Awaitable[None]: ...
-    def recv(self) -> collections.abc.Awaitable[tuple[TaskRequest, TaskChunkSender, str | None] | None]: ...
-
-@typing.final
-class TaskStream:
-    def recv(self) -> collections.abc.Awaitable[str | None]: ...
-

rust/exo_rs/pyproject.toml

@@ -4,12 +4,12 @@ build-backend = "maturin"
 
 [project]
 name = "exo_rs"
-version = "0.3.0"
+version = "0.2.16"
 description = "Add your description here"
 readme = "README.md"
 authors = [
-  { name = "Evan Quiney", email = "evanev7@gmail.com" },
   { name = "Andrei Cravtov", email = "the.andrei.cravtov@gmail.com" },
+  { name = "Evan Quiney", email = "evanev7@gmail.com" },
 ]
 requires-python = ">=3.13"
 dependencies = []
@@ -18,6 +18,8 @@ dependencies = []
 dev = ["exo_rs", "pytest>=8.4.0", "pytest-asyncio>=1.0.0"]
 
 [tool.maturin]
+#purelib = true
+#python-source = "python"
 module-name = "exo_rs"
 features = ["pyo3/extension-module", "pyo3/experimental-async"]
 

rust/exo_rs/src/ident.rs

@@ -0,0 +1,47 @@
+use crate::ext::ResultExt as _;
+use libp2p::identity::Keypair;
+use pyo3::types::{PyBytes, PyBytesMethods as _};
+use pyo3::{Bound, PyResult, Python, pyclass, pymethods};
+use pyo3_stub_gen::derive::{gen_stub_pyclass, gen_stub_pymethods};
+
+/// Identity keypair of a node.
+#[gen_stub_pyclass]
+#[pyclass(name = "Keypair", frozen)]
+#[repr(transparent)]
+pub struct PyKeypair(pub Keypair);
+
+#[gen_stub_pymethods]
+#[pymethods]
+#[allow(clippy::needless_pass_by_value)]
+impl PyKeypair {
+    /// Generate a new Ed25519 keypair.
+    #[staticmethod]
+    fn generate() -> Self {
+        Self(Keypair::generate_ed25519())
+    }
+
+    /// Construct an Ed25519 keypair from secret key bytes
+    #[staticmethod]
+    fn from_bytes(bytes: Bound<'_, PyBytes>) -> PyResult<Self> {
+        let mut bytes = Vec::from(bytes.as_bytes());
+        Ok(Self(Keypair::ed25519_from_bytes(&mut bytes).pyerr()?))
+    }
+
+    /// Get the secret key bytes underlying the keypair
+    fn to_bytes<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyBytes>> {
+        let bytes = self
+            .0
+            .clone()
+            .try_into_ed25519()
+            .pyerr()?
+            .secret()
+            .as_ref()
+            .to_vec();
+        Ok(PyBytes::new(py, &bytes))
+    }
+
+    /// Convert the `Keypair` into the corresponding `PeerId` string, which we use as our `NodeId`.
+    fn to_node_id(&self) -> String {
+        self.0.public().to_peer_id().to_base58()
+    }
+}

rust/exo_rs/src/last_value.rs

@@ -1,181 +0,0 @@
-use networking::{Session, liveliness_aggregator::LivelinessAggregator};
-use std::{
-    collections::{HashMap, HashSet},
-    sync::Arc,
-};
-use zenoh::{Result as ZResult, Wait};
-
-use parking_lot::Mutex;
-use pyo3::{
-    exceptions::{PyConnectionError, PyRuntimeError},
-    prelude::*,
-};
-use pyo3_stub_gen::derive::{gen_stub_pyclass, gen_stub_pymethods};
-use zenoh::{
-    handlers::FifoChannelHandler,
-    sample::{Sample, SampleKind},
-};
-use zenoh_ext::{
-    AdvancedPublisher, AdvancedSubscriber, AdvancedSubscriberBuilderExt, HistoryConfig,
-};
-
-#[gen_stub_pyclass]
-#[pyclass]
-pub struct LVAggregator {
-    pub prefix: Arc<str>,
-    pub store: Arc<Mutex<HashMap<String, String>>>,
-    pub current_live: LivelinessAggregator,
-}
-
-pub fn spawn_lv_aggregator_onto(session: &Session, prefix: Arc<str>) -> ZResult<LVAggregator> {
-    // nota bene: config must be kept in line with SessionHandle::last_value_receiver
-    let store = Arc::new(Mutex::new(HashMap::default()));
-    session
-        .z
-        //assuming all LV aggregators are prefix/node_id/atomic_json
-        .declare_subscriber(format!("{prefix}/*/*"))
-        .advanced()
-        .history(
-            HistoryConfig::default()
-                .max_samples(1)
-                .detect_late_publishers(),
-        )
-        .callback({
-            let store = Arc::clone(&store);
-            let prefix = Arc::clone(&prefix);
-            move |sample| {
-                if let Some(s) = sample
-                    .key_expr()
-                    .to_string()
-                    .strip_prefix(&*prefix)
-                    .and_then(|it| it.strip_prefix('/'))
-                {
-                    let s = s.to_string();
-                    match sample.kind() {
-                        SampleKind::Put => {
-                            store.lock().insert(
-                                s,
-                                sample
-                                    .payload()
-                                    .try_to_string()
-                                    .expect("we only use utf8 encoded strings. someone messed up")
-                                    .to_string(),
-                            );
-                        }
-                        SampleKind::Delete => {
-                            store.lock().remove(&s);
-                        }
-                    }
-                };
-            }
-        })
-        .background()
-        .wait()?;
-    Ok(LVAggregator {
-        prefix,
-        store,
-        current_live: session.liveliness_aggregator.clone(),
-    })
-}
-
-#[gen_stub_pymethods]
-#[pymethods]
-impl LVAggregator {
-    pub fn dump(&self) -> HashMap<String, String> {
-        let mut store = self.store.lock();
-        let currently_alive: HashSet<String> = self.current_live.dump();
-        // remove any keys that are no longer live
-        store.retain(|key, _| {
-            currently_alive.iter().any(|node_id| {
-                key.strip_prefix(node_id)
-                    .is_some_and(|rest| rest.starts_with("/"))
-            })
-        });
-        store.clone()
-    }
-}
-
-#[gen_stub_pyclass]
-#[pyclass]
-pub struct LVSubscriber {
-    pub subscriber: AdvancedSubscriber<FifoChannelHandler<Sample>>,
-}
-#[gen_stub_pymethods]
-#[pymethods]
-impl LVSubscriber {
-    #[gen_stub(override_return_type(
-        type_repr="collections.abc.Awaitable[tuple[str, str] | None]",
-        imports=("collections.abc")
-    ))]
-    pub fn recv<'py>(&'py self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
-        pyo3_async_runtimes::tokio::future_into_py(py, {
-            if self.subscriber.receiver_count() != 1 {
-                return Err(PyRuntimeError::new_err(
-                    "tried to receive twice on the same receiver",
-                ));
-            }
-            let subscriber = self.subscriber.clone();
-            async move {
-                loop {
-                    match subscriber.recv_async().await {
-                        Ok(sample) if sample.kind() == SampleKind::Delete => continue,
-                        Err(_) => {
-                            return Ok(None);
-                        }
-                        Ok(sample) => {
-                            return Ok(Some((
-                                sample.key_expr().to_string(),
-                                sample
-                                    .payload()
-                                    .try_to_string()
-                                    .expect("we only use utf8 encoded strings. someone messed up")
-                                    .to_string(),
-                            )));
-                        }
-                    }
-                }
-            }
-        })
-    }
-}
-
-#[gen_stub_pyclass]
-#[pyclass]
-pub struct LVPublisher {
-    pub state: Arc<AdvancedPublisher<'static>>,
-}
-impl LVPublisher {
-    pub fn new(publisher: AdvancedPublisher<'static>) -> Self {
-        Self {
-            state: Arc::new(publisher),
-        }
-    }
-}
-#[gen_stub_pymethods]
-#[pymethods]
-impl LVPublisher {
-    #[gen_stub(override_return_type(
-        type_repr="collections.abc.Awaitable[None]",
-        imports=("collections.abc")
-    ))]
-    pub fn put<'py>(&'py self, py: Python<'py>, data: String) -> PyResult<Bound<'py, PyAny>> {
-        let state = Arc::clone(&self.state);
-        pyo3_async_runtimes::tokio::future_into_py(py, {
-            // clone the data so py can have it back
-            async move {
-                state
-                    .put(data)
-                    .await
-                    .map_err(|e| PyConnectionError::new_err(e.to_string()))?;
-                Ok(())
-            }
-        })
-    }
-}
-
-pub fn lv_submodule(m: &Bound<'_, PyModule>) -> PyResult<()> {
-    m.add_class::<LVPublisher>()?;
-    m.add_class::<LVSubscriber>()?;
-    m.add_class::<LVAggregator>()?;
-    Ok(())
-}

rust/exo_rs/src/lib.rs

@@ -5,24 +5,23 @@
 //!
 
 mod allow_threading;
-mod pidfile;
-// mod ident;
-pub mod last_value;
+mod ident;
 mod networking;
-pub mod session;
-mod storage;
-mod task;
+mod pidfile;
 
-use crate::last_value::lv_submodule;
+use crate::ident::PyKeypair;
 use crate::networking::networking_submodule;
 use crate::pidfile::pidfile_submodule;
-use crate::session::session_submodule;
-use crate::storage::storage_submodule;
-use crate::task::task_submodule;
 use pyo3::prelude::PyModule;
-use pyo3::{Bound, PyResult, pymodule};
+use pyo3::types::PyModuleMethods;
+use pyo3::{Bound, PyResult, pyclass, pymodule};
 use pyo3_stub_gen::define_stub_info_gatherer;
 
+/// Namespace for all the constants used by this crate.
+pub(crate) mod r#const {
+    pub const MPSC_CHANNEL_SIZE: usize = 1024;
+}
+
 /// Namespace for crate-wide extension traits/methods
 pub(crate) mod ext {
     use crate::allow_threading::AllowThreads;
@@ -154,7 +153,7 @@ pub(crate) mod ext {
 /// A Python module implemented in Rust. The name of this function must match
 /// the `lib.name` setting in the `Cargo.toml`, else Python will not be able to
 /// import the module.
-#[pymodule(name = "exo_rs")]
+#[pymodule(name = "exo_rs", gil_used = true)]
 fn main_module(m: &Bound<'_, PyModule>) -> PyResult<()> {
     // install logger
     pyo3_log::init();
@@ -162,13 +161,15 @@ fn main_module(m: &Bound<'_, PyModule>) -> PyResult<()> {
     builder.enable_all();
     pyo3_async_runtimes::tokio::init(builder);
 
-    // TODO: for now this is all NOT a submodule. KISS
-    pidfile_submodule(m)?;
+    // TODO: for now this is all NOT a submodule, but figure out how to make the submodule system
+    //       work with maturin, where the types generate correctly, in the right folder, without
+    //       too many importing issues...
+    m.add_class::<PyKeypair>()?;
     networking_submodule(m)?;
-    lv_submodule(m)?;
-    session_submodule(m)?;
-    storage_submodule(m)?;
-    task_submodule(m)?;
+    pidfile_submodule(m)?;
+
+    // top-level constructs
+    // TODO: ...
 
     Ok(())
 }

rust/exo_rs/src/networking.rs

@@ -1,40 +1,166 @@
 use std::pin::Pin;
 use std::sync::Arc;
 
+use crate::r#const::MPSC_CHANNEL_SIZE;
 use crate::ext::{ByteArrayExt as _, FutureExt, PyErrExt as _};
 use crate::ext::{ResultExt as _, TokioMpscSenderExt as _};
+use crate::ident::PyKeypair;
+use crate::networking::exception::{
+    PyAllQueuesFullError, PyMessageTooLargeError, PyNoPeersSubscribedToTopicError,
+};
+use crate::pyclass;
 use futures_lite::{Stream, StreamExt as _};
-use networking::Session;
-use networking::swarm::{FromSwarm, Swarm, ToSwarm, create_swarm};
-use pyo3::exceptions::{PyRuntimeError, PyValueError};
-use pyo3::prelude::*;
+use libp2p::gossipsub::PublishError;
+use networking::swarm::{FromSwarm, ToSwarm, create_swarm};
+use pyo3::exceptions::PyRuntimeError;
+use pyo3::prelude::{PyModule, PyModuleMethods as _};
 use pyo3::types::PyBytes;
 use pyo3::{Bound, Py, PyAny, PyErr, PyResult, Python, pymethods};
-use pyo3_stub_gen::derive::{
-    gen_methods_from_python, gen_stub_pyclass, gen_stub_pyclass_complex_enum, gen_stub_pymethods,
-};
+use pyo3_stub_gen::derive::{gen_stub_pyclass, gen_stub_pyclass_complex_enum, gen_stub_pymethods};
 use tokio::sync::{Mutex, mpsc, oneshot};
 
+mod exception {
+    use pyo3::types::PyTuple;
+    use pyo3::{exceptions::PyException, prelude::*};
+    use pyo3_stub_gen::derive::*;
+
+    #[gen_stub_pyclass]
+    #[pyclass(frozen, extends=PyException, name="NoPeersSubscribedToTopicError")]
+    pub struct PyNoPeersSubscribedToTopicError {}
+
+    impl PyNoPeersSubscribedToTopicError {
+        const MSG: &'static str = "\
+        No peers are currently subscribed to receive messages on this topic. \
+        Wait for peers to subscribe or check your network connectivity.";
+
+        ///   Creates a new  [ `PyErr` ]  of this type.
+        ///
+        ///   [`PyErr`] :  https://docs.rs/pyo3/latest/pyo3/struct.PyErr.html   "PyErr in pyo3"
+        pub(crate) fn new_err() -> PyErr {
+            PyErr::new::<Self, _>(()) // TODO: check if this needs to be replaced???
+        }
+    }
+
+    #[gen_stub_pymethods]
+    #[pymethods]
+    impl PyNoPeersSubscribedToTopicError {
+        #[new]
+        #[pyo3(signature = (*args))]
+        #[allow(unused_variables)]
+        pub(crate) fn new(args: &Bound<'_, PyTuple>) -> Self {
+            Self {}
+        }
+
+        fn __repr__(&self) -> String {
+            format!("PeerId(\"{}\")", Self::MSG)
+        }
+
+        fn __str__(&self) -> String {
+            Self::MSG.to_string()
+        }
+    }
+
+    #[gen_stub_pyclass]
+    #[pyclass(frozen, extends=PyException, name="AllQueuesFullError")]
+    pub struct PyAllQueuesFullError {}
+
+    impl PyAllQueuesFullError {
+        const MSG: &'static str =
+            "All libp2p peers are unresponsive, resend the message or reconnect.";
+
+        ///   Creates a new  [ `PyErr` ]  of this type.
+        ///
+        ///   [`PyErr`] :  https://docs.rs/pyo3/latest/pyo3/struct.PyErr.html   "PyErr in pyo3"
+        pub(crate) fn new_err() -> PyErr {
+            PyErr::new::<Self, _>(()) // TODO: check if this needs to be replaced???
+        }
+    }
+
+    #[gen_stub_pymethods]
+    #[pymethods]
+    impl PyAllQueuesFullError {
+        #[new]
+        #[pyo3(signature = (*args))]
+        #[allow(unused_variables)]
+        pub(crate) fn new(args: &Bound<'_, PyTuple>) -> Self {
+            Self {}
+        }
+
+        fn __repr__(&self) -> String {
+            format!("PeerId(\"{}\")", Self::MSG)
+        }
+
+        fn __str__(&self) -> String {
+            Self::MSG.to_string()
+        }
+    }
+
+    #[gen_stub_pyclass]
+    #[pyclass(frozen, extends=PyException, name="MessageTooLargeError")]
+    pub struct PyMessageTooLargeError {}
+
+    impl PyMessageTooLargeError {
+        const MSG: &'static str = "Gossipsub message exceeds max_transmit_size. Reduce prompt length or increase the limit.";
+
+        pub(crate) fn new_err() -> PyErr {
+            PyErr::new::<Self, _>(())
+        }
+    }
+
+    #[gen_stub_pymethods]
+    #[pymethods]
+    impl PyMessageTooLargeError {
+        #[new]
+        #[pyo3(signature = (*args))]
+        #[allow(unused_variables)]
+        pub(crate) fn new(args: &Bound<'_, PyTuple>) -> Self {
+            Self {}
+        }
+
+        fn __repr__(&self) -> String {
+            format!("MessageTooLargeError(\"{}\")", Self::MSG)
+        }
+
+        fn __str__(&self) -> String {
+            Self::MSG.to_string()
+        }
+    }
+}
+
 #[gen_stub_pyclass]
 #[pyclass(name = "NetworkingHandle")]
-pub struct PyNetworkingHandle {
+struct PyNetworkingHandle {
     // channels
     pub to_swarm: mpsc::Sender<ToSwarm>,
     pub swarm: Arc<Mutex<Pin<Box<dyn Stream<Item = FromSwarm> + Send>>>>,
 }
 
 #[gen_stub_pyclass_complex_enum]
-#[pyclass]
-pub enum PyFromSwarm {
-    Connection { connected: bool },
-    Message { topic: String, data: Py<PyBytes> },
+#[pyclass(name = "FromSwarm")]
+enum PyFromSwarm {
+    Connection {
+        peer_id: String,
+        connected: bool,
+    },
+    Message {
+        origin: String,
+        topic: String,
+        data: Py<PyBytes>,
+    },
 }
 impl From<FromSwarm> for PyFromSwarm {
     fn from(value: FromSwarm) -> Self {
         match value {
-            FromSwarm::Discovered {} => Self::Connection { connected: true },
-            FromSwarm::Expired {} => Self::Connection { connected: false },
-            FromSwarm::Message { topic, data } => Self::Message {
+            FromSwarm::Discovered { peer_id } => Self::Connection {
+                peer_id: peer_id.to_base58(),
+                connected: true,
+            },
+            FromSwarm::Expired { peer_id } => Self::Connection {
+                peer_id: peer_id.to_base58(),
+                connected: false,
+            },
+            FromSwarm::Message { from, topic, data } => Self::Message {
+                origin: from.to_base58(),
                 topic: topic,
                 data: data.pybytes(),
             },
@@ -42,20 +168,6 @@ impl From<FromSwarm> for PyFromSwarm {
     }
 }
 
-impl PyNetworkingHandle {
-    pub fn from_session(session: Session) -> Self {
-        let (to_swarm, from_client) = mpsc::channel(1024);
-        let swarm = Swarm {
-            from_client,
-            session,
-        };
-        PyNetworkingHandle {
-            swarm: Arc::new(Mutex::new(swarm.into_stream())),
-            to_swarm,
-        }
-    }
-}
-
 #[gen_stub_pymethods]
 #[pymethods]
 impl PyNetworkingHandle {
@@ -65,49 +177,36 @@ impl PyNetworkingHandle {
 
     // ---- Lifecycle management methods ----
 
-    #[staticmethod]
-    pub fn new<'py>(
-        identity: &str,
+    #[new]
+    #[pyo3(signature = (identity, bootstrap_peers, listen_port))]
+    fn py_new(
+        identity: Bound<'_, PyKeypair>,
+        bootstrap_peers: Vec<String>,
         listen_port: u16,
-        discovery_service_port: u16,
-    ) -> PyResult<PyNetworkingHandle> {
-        // todo: zenoh self assigned peers
-        if listen_port == 0 {
-            todo!();
-        }
+    ) -> PyResult<Self> {
         // create communication channels
-        let (to_swarm, from_client) = mpsc::channel(1024);
+        let (to_swarm, from_client) = mpsc::channel(MPSC_CHANNEL_SIZE);
 
         // get identity
-        if !identity
-            .chars()
-            .all(|c| ('0'..='9').contains(&c) || ('a'..='f').contains(&c))
-            || identity.len() > 32
-        {
-            return Err(PyValueError::new_err(format!(
-                "{identity} is not a valid zenoh identity"
-            )));
-        }
+        let identity = identity.borrow().0.clone();
 
         // create networking swarm (within tokio context!! or it crashes)
-        let swarm = pyo3_async_runtimes::tokio::get_runtime()
-            .block_on(create_swarm(
-                identity,
-                from_client,
-                listen_port,
-                discovery_service_port,
-            ))
-            .pyerr()?;
+        let _guard = pyo3_async_runtimes::tokio::get_runtime().enter();
+        let swarm = create_swarm(identity, from_client, bootstrap_peers, listen_port)
+            .pyerr()?
+            .into_stream();
 
-        Ok(PyNetworkingHandle {
-            swarm: Arc::new(Mutex::new(swarm.into_stream())),
+        Ok(Self {
+            swarm: Arc::new(Mutex::new(swarm)),
             to_swarm,
         })
     }
 
-    #[gen_stub(skip)]
-    pub fn recv<'py>(&'py self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
-        let swarm = Arc::clone(&self.swarm);
+    #[gen_stub(override_return_type(
+        type_repr="typing.Awaitable[FromSwarm]", imports=("typing")
+    ))]
+    fn recv<'py>(&'py self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
+        let swarm = self.swarm.clone();
         pyo3_async_runtimes::tokio::future_into_py(py, async move {
             swarm
                 .try_lock()
@@ -124,7 +223,7 @@ impl PyNetworkingHandle {
     /// Subscribe to a `GossipSub` topic.
     ///
     /// Returns `True` if the subscription worked. Returns `False` if we were already subscribed.
-    pub async fn gossipsub_subscribe(&self, topic: String) -> PyResult<bool> {
+    async fn gossipsub_subscribe(&self, topic: String) -> PyResult<bool> {
         let (tx, rx) = oneshot::channel();
 
         // send off request to subscribe
@@ -146,7 +245,7 @@ impl PyNetworkingHandle {
     /// Unsubscribes from a `GossipSub` topic.
     ///
     /// Returns `True` if we were subscribed to this topic. Returns `False` if we were not subscribed.
-    pub async fn gossipsub_unsubscribe(&self, topic: String) -> PyResult<bool> {
+    async fn gossipsub_unsubscribe(&self, topic: String) -> PyResult<bool> {
         let (tx, rx) = oneshot::channel();
 
         // send off request to unsubscribe
@@ -167,7 +266,7 @@ impl PyNetworkingHandle {
     /// Publishes a message with multiple topics to the `GossipSub` network.
     ///
     /// If no peers are found that subscribe to this topic, throws `NoPeersSubscribedToTopicError` exception.
-    pub async fn gossipsub_publish(&self, topic: String, data: Py<PyBytes>) -> PyResult<()> {
+    async fn gossipsub_publish(&self, topic: String, data: Py<PyBytes>) -> PyResult<()> {
         let (tx, rx) = oneshot::channel();
 
         // send off request to subscribe
@@ -186,21 +285,23 @@ impl PyNetworkingHandle {
             .allow_threads_py() // allow-threads-aware async call
             .await
             .map_err(|_| PyErr::receiver_channel_closed())?
-            .map_err(|e| PyRuntimeError::new_err(e.to_string()))?;
+            .map_err(|e| match e {
+                PublishError::AllQueuesFull(_) => PyAllQueuesFullError::new_err(),
+                PublishError::MessageTooLarge => PyMessageTooLargeError::new_err(),
+                PublishError::NoPeersSubscribedToTopic => {
+                    PyNoPeersSubscribedToTopicError::new_err()
+                }
+                e => PyRuntimeError::new_err(e.to_string()),
+            })?;
         Ok(())
     }
 }
 
-pyo3_stub_gen::inventory::submit! {
-    gen_methods_from_python! {
-        r#"
-            class PyNetworkingHandle:
-                async def recv() -> PyFromSwarm: ...
-        "#
-    }
-}
-
 pub fn networking_submodule(m: &Bound<'_, PyModule>) -> PyResult<()> {
+    m.add_class::<exception::PyNoPeersSubscribedToTopicError>()?;
+    m.add_class::<exception::PyAllQueuesFullError>()?;
+    m.add_class::<exception::PyMessageTooLargeError>()?;
+
     m.add_class::<PyNetworkingHandle>()?;
     m.add_class::<PyFromSwarm>()?;
 

rust/exo_rs/src/pidfile.rs

@@ -77,7 +77,6 @@ impl PyPidfile {
 #[pymethods]
 impl PyPidfile {
     /// Creates a new PID file and locks it.
-    /// Writes the current process ID to the PID file.
     ///
     /// If the PID file cannot be locked, returns `PidfileError::AlreadyRunning` with
     /// a PID of the already running process, or `None` if no PID has been written to

rust/exo_rs/src/session.rs

@@ -1,138 +0,0 @@
-use networking::Session;
-use parking_lot::Mutex;
-use pyo3::{
-    exceptions::{PyConnectionError, PyRuntimeError, PyValueError},
-    prelude::*,
-};
-use pyo3_stub_gen::derive::{gen_stub_pyclass, gen_stub_pymethods};
-use std::{collections::HashMap, sync::Arc};
-use zenoh::Wait;
-use zenoh_ext::{
-    AdvancedPublisherBuilderExt, AdvancedSubscriberBuilderExt, CacheConfig, HistoryConfig,
-    MissDetectionConfig,
-};
-
-use crate::{
-    last_value::{LVAggregator, LVPublisher, LVSubscriber, spawn_lv_aggregator_onto},
-    networking::PyNetworkingHandle,
-    storage::Storage,
-    task::{TaskRequester, TaskResponder},
-};
-
-#[gen_stub_pyclass]
-#[pyclass]
-pub struct SessionHandle {
-    pub session: Session,
-}
-
-#[gen_stub_pymethods]
-#[pymethods]
-impl SessionHandle {
-    // NOTE: `async fn`s here that use `.await` will wrap the future in `.allow_threads_py()`
-    //       immediately beforehand to release the interpreter.
-    //       SEE: https://pyo3.rs/v0.26.0/async-await.html#detaching-from-the-interpreter-across-await
-
-    // ---- Lifecycle management methods ----
-
-    #[staticmethod]
-    pub fn new<'py>(
-        identity: &str,
-        listen_port: u16,
-        discovery_service_port: u16,
-    ) -> PyResult<(SessionHandle, PyNetworkingHandle)> {
-        // get identity
-        if !identity
-            .chars()
-            .all(|c| ('0'..='9').contains(&c) || ('a'..='f').contains(&c))
-            || identity.len() > 32
-        {
-            return Err(PyValueError::new_err(format!(
-                "{identity} is not a valid zenoh identity"
-            )));
-        }
-
-        let cfg = networking::cfg(identity, listen_port).map_err(|e| {
-            PyValueError::new_err(format!("failed to write config: {}", e.to_string()))
-        })?;
-        let session = pyo3_async_runtimes::tokio::get_runtime()
-            .block_on(networking::open(cfg, listen_port, discovery_service_port))
-            .map_err(|e| {
-                PyRuntimeError::new_err(format!(
-                    "failed to spawn networking on tokio runtime: {}",
-                    e.to_string()
-                ))
-            })?;
-        let legacy = PyNetworkingHandle::from_session(session.clone());
-        Ok((Self { session }, legacy))
-    }
-
-    pub fn last_value_aggregator(&self, prefix: String) -> PyResult<LVAggregator> {
-        spawn_lv_aggregator_onto(&self.session, prefix.into()).map_err(|e| {
-            PyConnectionError::new_err(format!("failed to spawn liveliness aggregator: {e}"))
-        })
-    }
-
-    pub fn last_value_subscriber(&self, kexpr: &str) -> PyResult<LVSubscriber> {
-        // nota bene: config must be kept in track with the LVAggregator
-        self.session
-            .z
-            .declare_subscriber(kexpr)
-            .advanced()
-            .history(
-                HistoryConfig::default()
-                    .max_samples(1)
-                    .detect_late_publishers(),
-            )
-            .wait()
-            .map_err(|e| PyConnectionError::new_err(format!("failed to declare subscriber: {e}")))
-            .map(|subscriber| LVSubscriber { subscriber })
-    }
-
-    pub fn last_value_publisher(&self, kexpr: String) -> PyResult<LVPublisher> {
-        self.session
-            .z
-            .declare_publisher(kexpr)
-            .advanced()
-            .publisher_detection()
-            .sample_miss_detection(MissDetectionConfig::default())
-            .cache(CacheConfig::default().max_samples(1))
-            .wait()
-            .map_err(|e| PyConnectionError::new_err(format!("failed to declare publisher: {e}")))
-            .map(LVPublisher::new)
-    }
-
-    pub fn storage_interface(&self) -> Storage {
-        Storage {
-            session: self.session.z.clone(),
-        }
-    }
-
-    pub fn task_requester(&self) -> TaskRequester {
-        TaskRequester {
-            session: self.session.z.clone(),
-        }
-    }
-
-    pub fn task_responder(&self, instance_id: String) -> PyResult<TaskResponder> {
-        let queryable = self
-            .session
-            .z
-            .declare_queryable(format!("task/instances/{instance_id}/tasks/*"))
-            .complete(true)
-            .wait()
-            .map_err(|e| {
-                PyConnectionError::new_err(format!("failed to declare task responder: {e}"))
-            })?;
-        Ok(TaskResponder {
-            instance_id,
-            queryable,
-            session: self.session.z.clone(),
-            assignments: Arc::new(Mutex::new(HashMap::new())),
-        })
-    }
-}
-
-pub fn session_submodule(m: &Bound<'_, PyModule>) -> PyResult<()> {
-    m.add_class::<SessionHandle>()?;
-    Ok(())
-}

rust/exo_rs/src/storage.rs

@@ -1,199 +0,0 @@
-use std::{collections::HashMap, time::Duration};
-
-use networking::STORAGE_PREFIX;
-use pyo3::{
-    exceptions::{PyConnectionError, PyRuntimeError, PyValueError},
-    prelude::*,
-};
-use pyo3_stub_gen::derive::{gen_stub_pyclass, gen_stub_pymethods};
-use zenoh::{
-    Session as ZSession, Wait, handlers::FifoChannelHandler, query::Reply, sample::SampleKind,
-};
-
-#[gen_stub_pyclass]
-#[pyclass]
-pub struct Storage {
-    pub session: ZSession,
-}
-
-#[gen_stub_pymethods]
-#[pymethods]
-impl Storage {
-    #[gen_stub(override_return_type(
-        type_repr="collections.abc.Awaitable[str | None]",
-        imports=("collections.abc")
-    ))]
-    pub fn get<'py>(&'py self, py: Python<'py>, key: String) -> PyResult<Bound<'py, PyAny>> {
-        if key.contains('*') {
-            return Err(PyValueError::new_err(format!(
-                "{key} is invalid -- Storage.get only supports fixed keys"
-            )));
-        }
-        let session = self.session.clone();
-        pyo3_async_runtimes::tokio::future_into_py(py, async move {
-            let subscriber = session
-                .get(format!("{STORAGE_PREFIX}/{key}"))
-                //.allowed_destination(Locality::SessionLocal)
-                .await
-                .map_err(|e| PyConnectionError::new_err(format!("failed to query storage: {e}")))?;
-            tokio::select! {
-                _ = tokio::time::sleep(Duration::from_secs(1)) => {
-                    Ok(None)
-                }
-                reply = subscriber.recv_async() => {
-                    Ok(reply.ok()
-                        .and_then(|reply| reply.into_result().ok())
-                        .and_then(|sample| {
-                            if sample.kind() == SampleKind::Put {
-                                Some(sample
-                                    .payload()
-                                    .try_to_string()
-                                    .expect("we only use utf8 encoded strings. someone messed up").to_string()
-                                )
-                            } else { None }
-                        })
-                    )
-                }
-            }
-        })
-    }
-    pub fn get_many(&self, key: String) -> PyResult<StorageGetter> {
-        self.session
-            .get(key)
-            .wait()
-            .map_err(|e| PyConnectionError::new_err(format!("failed to query storage: {e}")))
-            .map(StorageGetter)
-    }
-
-    #[gen_stub(override_return_type(
-        type_repr="collections.abc.Awaitable[None]",
-        imports=("collections.abc")
-    ))]
-    pub fn put<'py>(
-        &'py self,
-        py: Python<'py>,
-        key: String,
-        data: String,
-    ) -> PyResult<Bound<'py, PyAny>> {
-        if key.contains('*') {
-            return Err(PyValueError::new_err(format!(
-                "{key} is invalid -- Storage.put only supports fixed keys"
-            )));
-        }
-        let session = self.session.clone();
-        pyo3_async_runtimes::tokio::future_into_py(py, async move {
-            session
-                .put(format!("{STORAGE_PREFIX}/{key}"), data)
-                .await
-                .map_err(|e| PyConnectionError::new_err(format!("failed to query storage: {e}")))
-        })
-    }
-
-    #[gen_stub(override_return_type(
-        type_repr="collections.abc.Awaitable[None]",
-        imports=("collections.abc")
-    ))]
-    pub fn delete<'py>(&'py self, py: Python<'py>, key: String) -> PyResult<Bound<'py, PyAny>> {
-        if key.contains('*') {
-            return Err(PyValueError::new_err(format!(
-                "{key} is invalid -- Storage.delete only supports fixed keys"
-            )));
-        }
-        let session = self.session.clone();
-        pyo3_async_runtimes::tokio::future_into_py(py, async move {
-            session
-                .delete(format!("{STORAGE_PREFIX}/{key}"))
-                .await
-                .map_err(|e| PyConnectionError::new_err(format!("failed to query storage: {e}")))
-        })
-    }
-
-    #[gen_stub(override_return_type(
-        type_repr="collections.abc.Awaitable[dict[str, str]]",
-        imports=("collections.abc")
-    ))]
-    pub fn dump<'py>(&'py self, py: Python<'py>, prefix: String) -> PyResult<Bound<'py, PyAny>> {
-        pyo3_async_runtimes::tokio::future_into_py(py, async move {
-            Ok(networking::read_raw_memory_storage()
-                .await
-                .into_iter()
-                .filter_map(|(key, value)| {
-                    Some((
-                        key?.as_str().strip_prefix(prefix.as_str())?.to_string(),
-                        value
-                            .payload
-                            .try_to_string()
-                            .expect("we only use utf8 encoded strings. someone messed up")
-                            .to_string(),
-                    ))
-                })
-                .collect::<HashMap<String, String>>())
-        })
-    }
-}
-
-#[gen_stub_pyclass]
-#[pyclass]
-pub struct StorageGetter(FifoChannelHandler<Reply>);
-
-#[gen_stub_pymethods]
-#[pymethods]
-impl StorageGetter {
-    #[gen_stub(override_return_type(
-        type_repr="collections.abc.Awaitable[tuple[str, str] | None]",
-        imports=("collections.abc")
-    ))]
-    fn recv<'py>(&'py self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
-        if self.0.receiver_count() != 1 {
-            return Err(PyRuntimeError::new_err(
-                "Tried to call StorageGetter.recv twice concurrently",
-            ));
-        }
-        let dupe = self.0.clone();
-        pyo3_async_runtimes::tokio::future_into_py(py, async move {
-            let sample = loop {
-                match dupe.recv_async().await {
-                    Err(_) => return Ok(None),
-                    Ok(reply) => match reply.into_result() {
-                        Err(e) => {
-                            log::warn!("Ignoring reply error: {e}");
-                            continue;
-                        }
-                        Ok(sample) => match sample.kind() {
-                            SampleKind::Put => break sample,
-                            SampleKind::Delete => {
-                                log::warn!(
-                                    "Received unexpected DELETE from queryable: {}",
-                                    sample.key_expr()
-                                );
-                                continue;
-                            }
-                        },
-                    },
-                };
-            };
-
-            let key = sample
-                .key_expr()
-                .to_string()
-                .strip_prefix(format!("{STORAGE_PREFIX}/").as_str())
-                .expect("invalid storage format encountered")
-                .to_string();
-
-            Ok(Some((
-                key,
-                sample
-                    .payload()
-                    .try_to_string()
-                    .expect("we only use utf8 encoded strings. someone messed up")
-                    .to_string(),
-            )))
-        })
-    }
-}
-
-pub fn storage_submodule(m: &Bound<'_, PyModule>) -> PyResult<()> {
-    m.add_class::<Storage>()?;
-    m.add_class::<StorageGetter>()?;
-    Ok(())
-}

rust/exo_rs/src/task.rs

@@ -1,367 +0,0 @@
-use std::{collections::HashMap, sync::Arc, time::Duration};
-
-use pyo3::{
-    exceptions::{PyConnectionError, PyRuntimeError, PyTimeoutError, PyValueError},
-    prelude::*,
-};
-use pyo3_stub_gen::derive::{gen_stub_pyclass, gen_stub_pymethods};
-use zenoh::{
-    Session as ZSession, Wait,
-    handlers::FifoChannelHandler,
-    pubsub::{Publisher, Subscriber},
-    qos::CongestionControl,
-    query::{ConsolidationMode, Query, Queryable},
-    sample::{Sample, SampleKind},
-};
-use zenoh_ext::{AdvancedPublisher, AdvancedPublisherBuilderExt, CacheConfig, MissDetectionConfig};
-
-use parking_lot::Mutex;
-
-#[gen_stub_pyclass]
-#[pyclass]
-pub struct TaskRequester {
-    pub session: ZSession,
-}
-
-#[gen_stub_pyclass]
-#[pyclass]
-pub struct TaskResponder {
-    pub instance_id: String,
-    pub queryable: Queryable<FifoChannelHandler<Query>>,
-    pub session: ZSession,
-    pub assignments: Arc<Mutex<HashMap<String, Arc<AdvancedPublisher<'static>>>>>,
-}
-
-#[gen_stub_pyclass]
-#[pyclass]
-pub struct TaskRequest {
-    pub query: Query,
-    pub key: String,
-}
-
-#[gen_stub_pyclass]
-#[pyclass]
-pub struct TaskChunkSender {
-    pub publisher: Arc<Publisher<'static>>,
-}
-
-#[gen_stub_pyclass]
-#[pyclass]
-pub struct TaskStream {
-    pub receiver: Subscriber<FifoChannelHandler<Sample>>,
-}
-
-#[gen_stub_pymethods]
-#[pymethods]
-impl TaskRequester {
-    #[gen_stub(override_return_type(
-        type_repr="collections.abc.Awaitable[TaskStream]",
-        imports=("collections.abc")
-    ))]
-    pub fn submit<'py>(
-        &'py self,
-        py: Python<'py>,
-        instance_id: String,
-        command_id: String,
-        command: String,
-    ) -> PyResult<Bound<'py, PyAny>> {
-        let session = self.session.clone();
-        pyo3_async_runtimes::tokio::future_into_py(py, async move {
-            let receiver = declare_task_stream(&session, &command_id)?;
-
-            request_task_admission(&session, instance_id, command_id, command).await?;
-            Ok(TaskStream { receiver })
-        })
-    }
-
-    #[gen_stub(override_return_type(
-        type_repr="collections.abc.Awaitable[None]",
-        imports=("collections.abc")
-    ))]
-    pub fn interrupt<'py>(
-        &'py self,
-        py: Python<'py>,
-        instance_id: String,
-        command_id: String,
-        command: String,
-    ) -> PyResult<Bound<'py, PyAny>> {
-        let session = self.session.clone();
-        pyo3_async_runtimes::tokio::future_into_py(py, async move {
-            request_task_admission(&session, instance_id, command_id, command).await
-        })
-    }
-}
-
-fn task_key(instance_id: &str, command_id: &str) -> String {
-    format!("task/instances/{instance_id}/tasks/{command_id}")
-}
-
-fn task_chunks_key(command_id: &str) -> String {
-    format!("task/commands/{command_id}/chunks")
-}
-
-fn task_assignment_key(instance_id: &str, task_id: &str) -> String {
-    format!("task_assignments/{instance_id}/{task_id}")
-}
-
-fn declare_task_stream(
-    session: &ZSession,
-    command_id: &str,
-) -> PyResult<Subscriber<FifoChannelHandler<Sample>>> {
-    session
-        .declare_subscriber(task_chunks_key(command_id))
-        .wait()
-        .map_err(|e| PyConnectionError::new_err(format!("failed to declare task stream: {e}")))
-}
-
-async fn request_task_admission(
-    session: &ZSession,
-    instance_id: String,
-    command_id: String,
-    command: String,
-) -> PyResult<()> {
-    let replies = session
-        .get(task_key(&instance_id, &command_id))
-        .payload(command)
-        .congestion_control(CongestionControl::Block)
-        .consolidation(ConsolidationMode::None)
-        .timeout(Duration::from_secs(5))
-        .wait()
-        .map_err(|e| PyConnectionError::new_err(format!("failed to submit task: {e}")))?;
-
-    let reply = replies.recv_async().await.map_err(|e| {
-        PyConnectionError::new_err(format!("task admission stream closed early: {e}"))
-    })?;
-
-    match reply.into_result() {
-        Ok(sample) => {
-            if sample.kind() == SampleKind::Delete {
-                Err(PyConnectionError::new_err(
-                    "task admission replied with delete",
-                ))
-            } else {
-                let _ = sample_to_string(sample);
-                Ok(())
-            }
-        }
-        Err(error) => error.payload().try_to_string().map_or_else(
-            |err| {
-                Err(PyRuntimeError::new_err(format!(
-                    "task admission failed: {err}"
-                )))
-            },
-            |ok| {
-                if ok == "Timeout" {
-                    Err(PyTimeoutError::new_err("task admission timed out"))
-                } else {
-                    Err(PyRuntimeError::new_err(format!(
-                        "task admission rejected: {ok}"
-                    )))
-                }
-            },
-        ),
-    }
-}
-
-#[gen_stub_pymethods]
-#[pymethods]
-impl TaskResponder {
-    #[gen_stub(override_return_type(
-        type_repr="collections.abc.Awaitable[None]",
-        imports=("collections.abc")
-    ))]
-    pub fn assign_task<'py>(
-        &'py self,
-        py: Python<'py>,
-        task_id: String,
-        task: String,
-    ) -> PyResult<Bound<'py, PyAny>> {
-        let session = self.session.clone();
-        let instance_id = self.instance_id.clone();
-        let assignments = Arc::clone(&self.assignments);
-        pyo3_async_runtimes::tokio::future_into_py(py, async move {
-            let publisher = {
-                let mut assignments = assignments.lock();
-                if let Some(publisher) = assignments.get(&task_id) {
-                    publisher.clone()
-                } else {
-                    let publisher = Arc::new(
-                        session
-                            .declare_publisher(task_assignment_key(&instance_id, &task_id))
-                            .advanced()
-                            .publisher_detection()
-                            .sample_miss_detection(MissDetectionConfig::default())
-                            .cache(CacheConfig::default().max_samples(1))
-                            .wait()
-                            .map_err(|e| {
-                                PyConnectionError::new_err(format!(
-                                    "failed to declare task assignment publisher: {e}"
-                                ))
-                            })?,
-                    );
-                    assignments.insert(task_id.clone(), publisher.clone());
-                    publisher
-                }
-            };
-            publisher
-                .put(task)
-                .await
-                .map_err(|e| PyConnectionError::new_err(format!("failed to assign task: {e}")))
-        })
-    }
-
-    #[gen_stub(override_return_type(
-        type_repr="collections.abc.Awaitable[None]",
-        imports=("collections.abc")
-    ))]
-    pub fn unassign_task<'py>(
-        &'py self,
-        py: Python<'py>,
-        task_id: String,
-    ) -> PyResult<Bound<'py, PyAny>> {
-        let session = self.session.clone();
-        let instance_id = self.instance_id.clone();
-        let assignments = Arc::clone(&self.assignments);
-        pyo3_async_runtimes::tokio::future_into_py(py, async move {
-            let publisher = assignments.lock().remove(&task_id);
-            if let Some(publisher) = publisher {
-                publisher.put("").await.map_err(|e| {
-                    PyConnectionError::new_err(format!("failed to unassign task: {e}"))
-                })?;
-            }
-            session
-                .delete(task_assignment_key(&instance_id, &task_id))
-                .await
-                .map_err(|e| PyConnectionError::new_err(format!("failed to unassign task: {e}")))
-        })
-    }
-
-    #[gen_stub(override_return_type(
-        type_repr="collections.abc.Awaitable[tuple[TaskRequest, TaskChunkSender, str | None] | None]",
-        imports=("collections.abc")
-    ))]
-    pub fn recv<'py>(&'py self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
-        if self.queryable.receiver_count() != 1 {
-            return Err(pyo3::exceptions::PyRuntimeError::new_err(
-                "called recv twice concurrently",
-            ));
-        }
-        let queryable = self.queryable.clone();
-        let session = self.session.clone();
-        let key_prefix = format!("task/instances/{}/tasks/", self.instance_id);
-        pyo3_async_runtimes::tokio::future_into_py(py, async move {
-            loop {
-                match queryable.recv_async().await {
-                    Ok(query) => {
-                        let query_key = query.key_expr().as_str();
-                        if !query_key.starts_with(&key_prefix) {
-                            continue;
-                        }
-                        let key = query_key.to_owned();
-                        let command_id = query_key[key_prefix.len()..].to_string();
-                        let payload = query.payload().map(|payload| {
-                            payload
-                                .try_to_string()
-                                .expect("we only use utf8 encoded strings. someone messed up")
-                                .to_string()
-                        });
-                        let publisher = session
-                            .declare_publisher(task_chunks_key(&command_id))
-                            .wait()
-                            .map_err(|e| {
-                                PyConnectionError::new_err(format!(
-                                    "failed to declare task chunk sender: {e}"
-                                ))
-                            })?;
-                        return Ok(Some((
-                            TaskRequest { query, key },
-                            TaskChunkSender {
-                                publisher: Arc::new(publisher),
-                            },
-                            payload,
-                        )));
-                    }
-                    Err(_) => return Ok(None),
-                }
-            }
-        })
-    }
-}
-
-#[gen_stub_pymethods]
-#[pymethods]
-impl TaskChunkSender {
-    #[gen_stub(override_return_type(
-        type_repr="collections.abc.Awaitable[None]",
-        imports=("collections.abc")
-    ))]
-    pub fn send<'py>(&'py self, py: Python<'py>, chunk: String) -> PyResult<Bound<'py, PyAny>> {
-        let publisher = Arc::clone(&self.publisher);
-        pyo3_async_runtimes::tokio::future_into_py(py, async move {
-            publisher
-                .put(chunk)
-                .await
-                .map_err(|e| PyConnectionError::new_err(format!("failed to send task chunk: {e}")))
-        })
-    }
-}
-
-#[gen_stub_pymethods]
-#[pymethods]
-impl TaskRequest {
-    pub fn reply(&self, payload: String) -> PyResult<()> {
-        self.query
-            .reply(&self.key, payload)
-            .wait()
-            .map_err(|e| PyConnectionError::new_err(format!("failed to reply to task query: {e}")))
-    }
-
-    pub fn reply_err(&self, payload: String) -> PyResult<()> {
-        if payload == "Timeout" {
-            return Err(PyValueError::new_err(
-                "Timeout is reserved for zenoh query timeouts",
-            ));
-        }
-        self.query
-            .reply_err(payload)
-            .wait()
-            .map_err(|e| PyConnectionError::new_err(format!("failed to reject task query: {e}")))
-    }
-}
-
-#[gen_stub_pymethods]
-#[pymethods]
-impl TaskStream {
-    #[gen_stub(override_return_type(
-        type_repr="collections.abc.Awaitable[str | None]",
-        imports=("collections.abc")
-    ))]
-    pub fn recv<'py>(&'py self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
-        if self.receiver.receiver_count() != 1 {
-            return Err(pyo3::exceptions::PyRuntimeError::new_err(
-                "called recv twice concurrently",
-            ));
-        }
-        let receiver = self.receiver.clone();
-        pyo3_async_runtimes::tokio::future_into_py(py, async move {
-            Ok(receiver.recv_async().await.ok().map(sample_to_string))
-        })
-    }
-}
-
-fn sample_to_string(sample: Sample) -> String {
-    sample
-        .payload()
-        .try_to_string()
-        .expect("we only use utf8 encoded strings. someone messed up")
-        .to_string()
-}
-
-pub fn task_submodule(m: &Bound<'_, PyModule>) -> PyResult<()> {
-    m.add_class::<TaskChunkSender>()?;
-    m.add_class::<TaskRequester>()?;
-    m.add_class::<TaskRequest>()?;
-    m.add_class::<TaskResponder>()?;
-    m.add_class::<TaskStream>()?;
-    Ok(())
-}

rust/exo_rs/tests/adv_subs.rs

@@ -1,168 +0,0 @@
-use std::time::{Duration, SystemTime, UNIX_EPOCH};
-
-use zenoh_ext::{
-    AdvancedPublisherBuilderExt, AdvancedSubscriber, AdvancedSubscriberBuilderExt, CacheConfig,
-    HistoryConfig, MissDetectionConfig,
-};
-
-use zenoh::{handlers::FifoChannelHandler, sample::Sample};
-
-// Adjust these imports to your crate/module paths.
-use exo_rs::{
-    last_value::{LVPublisher, LVSubscriber},
-    session::SessionHandle,
-};
-async fn expect_two_values(
-    sub: &AdvancedSubscriber<FifoChannelHandler<Sample>>,
-    key_a: &str,
-    val_a: &str,
-    key_b: &str,
-    val_b: &str,
-) {
-    use std::collections::HashMap;
-    use tokio::time::{Duration, Instant, timeout};
-    use zenoh::sample::SampleKind;
-
-    let deadline = Instant::now() + Duration::from_secs(10);
-    let mut seen: HashMap<String, String> = HashMap::new();
-
-    loop {
-        if seen.get(key_a).map(String::as_str) == Some(val_a)
-            && seen.get(key_b).map(String::as_str) == Some(val_b)
-        {
-            return;
-        }
-
-        let remaining = deadline.saturating_duration_since(Instant::now());
-        assert!(
-            !remaining.is_zero(),
-            "timed out waiting for both historical samples; expected {key_a}={val_a}, {key_b}={val_b}; seen = {seen:?}"
-        );
-
-        match timeout(remaining.min(Duration::from_millis(750)), sub.recv_async()).await {
-            Ok(Ok(sample)) => {
-                if sample.kind() == SampleKind::Delete {
-                    continue;
-                }
-
-                let key = sample.key_expr().to_string();
-                let value = sample
-                    .payload()
-                    .try_to_string()
-                    .expect("payload should be UTF-8")
-                    .to_string();
-
-                if key == key_a || key == key_b {
-                    eprintln!("received relevant {key} = {value}");
-                    seen.insert(key, value);
-                } else {
-                    eprintln!("received unrelated {key} = {value}");
-                }
-            }
-            Ok(Err(e)) => panic!("subscriber receive failed: {e}"),
-            Err(_) => {}
-        }
-    }
-}
-
-#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
-async fn lv_subscriber_receives_last_value_from_multiple_publishers() {
-    let cfg =
-        networking::cfg(&format!("{:x}", rand::random::<u128>()), 52414).expect("create config");
-    let n_session = networking::open(cfg, 52414, 52413)
-        .await
-        .expect("open session");
-
-    let session = SessionHandle { session: n_session };
-
-    let run_id = SystemTime::now()
-        .duration_since(UNIX_EPOCH)
-        .unwrap()
-        .as_nanos();
-
-    let base = format!("zenoh_advanced_history_test/{run_id}");
-
-    let key_a = format!("{base}/a");
-    let key_b = format!("{base}/b");
-    let sub_key = format!("{base}/*");
-
-    let pub1: LVPublisher = session
-        .last_value_publisher(key_a.clone())
-        .expect("declare LV publisher a");
-
-    pub1.state.put("aa").await.expect("publish aa");
-
-    let pub2: LVPublisher = session
-        .last_value_publisher(key_b.clone())
-        .expect("declare LV publisher b");
-
-    pub2.state.put("bb").await.expect("publish bb");
-
-    // Let publisher detection / cache metadata settle before the late subscriber joins.
-    tokio::time::sleep(Duration::from_millis(250)).await;
-
-    let sub: LVSubscriber = session
-        .last_value_subscriber(&*sub_key)
-        .expect("declare LV subscriber");
-
-    expect_two_values(&sub.subscriber, &*key_a, "aa", &*key_b, "bb").await
-}
-#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
-async fn wildcard_advanced_subscriber_receives_history_from_both_publishers() {
-    let cfg =
-        networking::cfg(&format!("{:x}", rand::random::<u128>()), 52412).expect("create config");
-    let n_session = networking::open(cfg, 52412, 52411)
-        .await
-        .expect("open session");
-    let session = n_session.z.clone();
-
-    // Unique prefix so the wildcard subscriber cannot accidentally see unrelated traffic.
-    let run_id = SystemTime::now()
-        .duration_since(UNIX_EPOCH)
-        .unwrap()
-        .as_nanos();
-
-    let base = format!("zenoh_advanced_history_test/{run_id}");
-    let key_a = format!("{base}/a");
-    let key_b = format!("{base}/b");
-    let sub_key = format!("{base}/*");
-
-    let pub1 = session
-        .declare_publisher(key_a.clone())
-        .advanced()
-        .publisher_detection()
-        .sample_miss_detection(MissDetectionConfig::default())
-        .cache(CacheConfig::default().max_samples(1))
-        .await
-        .expect("declare advanced publisher a");
-
-    pub1.put("aa").await.expect("publish aa");
-
-    let pub2 = session
-        .declare_publisher(key_b.clone())
-        .advanced()
-        .sample_miss_detection(MissDetectionConfig::default())
-        .publisher_detection()
-        .cache(CacheConfig::default().max_samples(1))
-        .await
-        .expect("declare advanced publisher b");
-
-    pub2.put("bb").await.expect("publish bb");
-
-    // Give liveliness/cache declarations a brief chance to settle before declaring
-    // the late-joining advanced subscriber.
-    tokio::time::sleep(Duration::from_millis(250)).await;
-
-    let sub = session
-        .declare_subscriber(sub_key)
-        .advanced()
-        .history(
-            HistoryConfig::default()
-                .max_samples(1)
-                .detect_late_publishers(),
-        )
-        .await
-        .expect("declare advanced subscriber");
-
-    expect_two_values(&sub, &*key_a, "aa", &*key_b, "bb").await
-}

rust/exo_rs/tests/dummy.rs

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

rust/exo_rs/tests/test_python.py

@@ -1,28 +1,31 @@
 import asyncio
-import os
 
 import pytest
 from _pytest.capture import CaptureFixture
 from exo_rs import (
+    Keypair,
     NetworkingHandle,
+    NoPeersSubscribedToTopicError,
     Pidfile,
-    PyFromSwarm,
+    FromSwarm,
 )
 
 
 @pytest.mark.asyncio
 async def test_sleep_on_multiple_items() -> None:
     print("PYTHON: starting handle")
-    h = NetworkingHandle.new(os.urandom(16).hex().rstrip("0"), 52412, 52411)
-    print("PYTHON: handle started")
+    h = NetworkingHandle(Keypair.generate(), [], 0)
 
     rt = asyncio.create_task(_await_recv(h))
 
     # sleep for 4 ticks
-    for i in range(10):
+    for i in range(4):
         await asyncio.sleep(1)
 
-        await h.gossipsub_publish("topic", b"somehting or other")
+        try:
+            await h.gossipsub_publish("topic", b"somehting or other")
+        except NoPeersSubscribedToTopicError as e:
+            print("caught it", e)
 
 
 def test_pidfile(capsys: CaptureFixture[str]):
@@ -36,15 +39,11 @@ async def _await_recv(h: NetworkingHandle):
     while True:
         event = await h.recv()
         match event:
-            case PyFromSwarm.Connection() as c:
+            case FromSwarm.Connection() as c:
                 print(f"PYTHON: connection update: {c}")
-            case PyFromSwarm.Message() as m:
+            case FromSwarm.Message() as m:
                 print(f"PYTHON: message: {m}")
 
 
 def scoped_lock_file():
     a = Pidfile("/tmp/lock.pid", 0o0600)
-
-
-if __name__ == "__main__":
-    asyncio.run(test_sleep_on_multiple_items())

rust/exo_rs/tests/test_storage_handle.py

@@ -1,69 +0,0 @@
-import os
-import uuid
-
-import pytest
-from exo_rs import SessionHandle
-
-
-ZENOH_PORT = 52414
-DISCOVERY_PORT = 52413
-
-
-@pytest.fixture(scope="module")
-def storage():
-    node_id = os.urandom(16).hex().rstrip("0")
-
-    session_handle, _nh = SessionHandle.new(
-        node_id,
-        ZENOH_PORT,
-        DISCOVERY_PORT,
-    )
-
-    return session_handle.storage_interface()
-
-
-@pytest.mark.asyncio
-async def test_storage_get_missing_key_returns_none(storage):
-    key = f"tests/storage/{uuid.uuid4().hex}/missing"
-
-    value = await storage.get(key)
-    assert value is None
-
-
-@pytest.mark.asyncio
-async def test_storage_put_then_get_returns_value(storage):
-    key = f"tests/storage/{uuid.uuid4().hex}/value"
-    expected = "hello storage"
-
-    await storage.put(key, expected)
-    assert await storage.get(key) == expected
-
-
-@pytest.mark.asyncio
-async def test_storage_put_overwrites_value(storage):
-    key = f"tests/storage/{uuid.uuid4().hex}/overwrite"
-
-    await storage.put(key, "old")
-    await storage.put(key, "new")
-    assert await storage.get(key) == "new"
-
-
-@pytest.mark.asyncio
-async def test_storage_put_overwrites_value(storage):
-    key = f"tests/storage/{uuid.uuid4().hex}/overwrite"
-
-    await storage.put(key, "old")
-    await storage.delete(key)
-    assert await storage.get(key) == None
-
-
-@pytest.mark.asyncio
-async def test_storage_get_rejects_wildcard_key(storage):
-    with pytest.raises(ValueError, match="only supports fixed keys"):
-        await storage.get("tests/storage/*")
-
-
-@pytest.mark.asyncio
-async def test_storage_put_rejects_wildcard_key(storage):
-    with pytest.raises(ValueError, match="only supports fixed keys"):
-        await storage.put("tests/storage/*", "value")

rust/exo_rs/tests/test_task.py

@@ -1,70 +0,0 @@
-import asyncio
-import os
-import uuid
-
-import pytest
-from exo_rs import SessionHandle
-
-ZENOH_PORT = 52416
-DISCOVERY_PORT = 52413
-
-
-@pytest.fixture(scope="module")
-def session_handle():
-    node_id = os.urandom(16).hex().lstrip("0")
-
-    session_handle, _nh = SessionHandle.new(
-        node_id,
-        ZENOH_PORT,
-        DISCOVERY_PORT,
-    )
-
-    return session_handle
-
-
-@pytest.mark.asyncio
-async def test_task_requester_responder_round_trip(session_handle):
-    instance_id = f"tests-task-instance-{uuid.uuid4().hex}"
-    command_id = f"tests-task-command-{uuid.uuid4().hex}"
-    command = '{"kind":"command"}'
-    chunk = '{"kind":"chunk","finish_reason":"stop"}'
-
-    requester = session_handle.task_requester()
-    responder = session_handle.task_responder(instance_id)
-
-    async def respond_to_submission():
-        received = await responder.recv()
-        assert received is not None
-        request, chunk_sender, payload = received
-        assert payload == command
-        request.reply(command_id)
-        await chunk_sender.send(chunk)
-
-    stream, _ = await asyncio.gather(
-        requester.submit(instance_id, command_id, command),
-        respond_to_submission(),
-    )
-
-    assert await stream.recv() == chunk
-
-
-@pytest.mark.asyncio
-async def test_task_requester_interrupt_round_trip(session_handle):
-    instance_id = f"tests-task-instance-{uuid.uuid4().hex}"
-    command_id = f"tests-task-command-{uuid.uuid4().hex}"
-    command = '{"kind":"interrupt"}'
-
-    requester = session_handle.task_requester()
-    responder = session_handle.task_responder(instance_id)
-
-    async def respond_to_interrupt():
-        received = await responder.recv()
-        assert received is not None
-        request, _chunk_sender, payload = received
-        assert payload == command
-        request.reply(command_id)
-
-    await asyncio.gather(
-        requester.interrupt(instance_id, command_id, command),
-        respond_to_interrupt(),
-    )

rust/networking/Cargo.toml

@@ -1,27 +1,42 @@
 [package]
 name = "networking"
-version.workspace = true
-edition.workspace = true
+version = { workspace = true }
+edition = { workspace = true }
+publish = false
 
-[dependencies]
-async-stream.workspace = true
-futures-lite.workspace = true
-netwatcher = { workspace = true, features = ["tokio"] }
-parking_lot.workspace = true
-tokio = { workspace = true, features = ["full"] }
-zenoh = { workspace = true, features = ["internal", "plugins", "unstable"] }
-zenoh-plugin-storage-manager.workspace = true
-zenoh-plugin-trait.workspace = true
-rand.workspace = true
-log.workspace = true
-bytemuck = { workspace = true, features = ["derive"] }
-socket2 = "0.6.4"
-blake3 = "1.8.5"
+[lib]
+doctest = false
+name = "networking"
+path = "src/lib.rs"
 
 [lints]
 workspace = true
 
-[dev-dependencies]
-env_logger.workspace = true
-smol = "2.0.2"
-tracing = "0.1.44"
+[dependencies]
+# datastructures
+either = { workspace = true }
+
+# macro dependencies
+extend = { workspace = true }
+delegate = { workspace = true }
+
+# async
+async-stream = { workspace = true }
+futures-lite = { workspace = true }
+futures-timer = { workspace = true }
+tokio = { workspace = true, features = ["full"] }
+
+# utility dependencies
+util = { workspace = true }
+tracing-subscriber = { version = "0.3.19", features = [
+  "default",
+  "env-filter",
+] }
+keccak-const = { workspace = true }
+
+# tracing/logging
+log = { workspace = true }
+
+# networking
+libp2p = { workspace = true, features = ["full"] }
+pin-project = "1.1.10"

rust/networking/examples/chatroom.rs

@@ -0,0 +1,86 @@
+use futures_lite::StreamExt;
+use libp2p::identity;
+use networking::swarm;
+use networking::swarm::{FromSwarm, ToSwarm};
+use tokio::sync::{mpsc, oneshot};
+use tokio::{io, io::AsyncBufReadExt as _};
+use tracing_subscriber::EnvFilter;
+use tracing_subscriber::filter::LevelFilter;
+
+#[tokio::main]
+async fn main() {
+    let _ = tracing_subscriber::fmt()
+        .with_env_filter(EnvFilter::from_default_env().add_directive(LevelFilter::INFO.into()))
+        .try_init();
+
+    let (to_swarm, from_client) = mpsc::channel(20);
+
+    // Configure swarm
+    let mut swarm = swarm::create_swarm(
+        identity::Keypair::generate_ed25519(),
+        from_client,
+        vec![],
+        0,
+    )
+    .expect("Swarm creation failed")
+    .into_stream();
+
+    // Create a Gossipsub topic & subscribe
+    let (tx, rx) = oneshot::channel();
+    _ = to_swarm
+        .send(ToSwarm::Subscribe {
+            topic: "test-net".to_string(),
+            result_sender: tx,
+        })
+        .await
+        .expect("should send");
+
+    // Read full lines from stdin
+    let mut stdin = io::BufReader::new(io::stdin()).lines();
+    println!("Enter messages via STDIN and they will be sent to connected peers using Gossipsub");
+
+    tokio::task::spawn(async move {
+        rx.await
+            .expect("tx not dropped")
+            .expect("subscribe shouldn't fail");
+        loop {
+            if let Ok(Some(line)) = stdin.next_line().await {
+                let (tx, rx) = oneshot::channel();
+                if let Err(e) = to_swarm
+                    .send(swarm::ToSwarm::Publish {
+                        topic: "test-net".to_string(),
+                        data: line.as_bytes().to_vec(),
+                        result_sender: tx,
+                    })
+                    .await
+                {
+                    println!("Send error: {e:?}");
+                    return;
+                };
+                match rx.await {
+                    Ok(Err(e)) => println!("Publish error: {e:?}"),
+                    Err(e) => println!("Publish error: {e:?}"),
+                    Ok(_) => {}
+                }
+            }
+        }
+    });
+
+    // Kick it off
+    loop {
+        // on gossipsub outgoing
+        match swarm.next().await {
+            // on gossipsub incoming
+            Some(FromSwarm::Discovered { peer_id }) => {
+                println!("\n\nconnected to {peer_id}\n\n")
+            }
+            Some(FromSwarm::Expired { peer_id }) => {
+                println!("\n\ndisconnected from {peer_id}\n\n")
+            }
+            Some(FromSwarm::Message { from, topic, data }) => {
+                println!("{topic}/{from}:\n{}", String::from_utf8_lossy(&data))
+            }
+            None => {}
+        }
+    }
+}

rust/networking/examples/do_nothing.rs

@@ -1,34 +0,0 @@
-use networking;
-use tracing::{info, warn};
-use zenoh::{Result, Wait};
-
-#[tokio::main]
-async fn main() -> Result<()> {
-    zenoh::init_log_from_env_or("info");
-    info!("Opening session...");
-    let cfg = networking::cfg(&format!("{:x}", rand::random::<u128>()), 52414)?;
-    let session = networking::open(cfg, 52414, 52413).await?;
-    let _tok = session
-        .z
-        .liveliness()
-        .declare_token(format!("nodes/{}/live", session.z.zid()))
-        .wait()?;
-    let subs = session
-        .z
-        .liveliness()
-        .declare_subscriber("**")
-        .history(true)
-        .wait()?;
-    loop {
-        tokio::select! {
-            _ = tokio::signal::ctrl_c() => break,
-            s = subs.recv_async() => {
-                match s {
-                    Err(e) => warn!("{e}"),
-                    Ok(s) => info!("{}: {}", s.kind(), s.key_expr().to_string().split("/").nth(1).unwrap()),
-                }
-            }
-        }
-    }
-    Ok(())
-}

rust/networking/examples/heartbeat.rs

@@ -1,32 +0,0 @@
-use env_logger::Env;
-use log::info;
-use networking;
-use zenoh::{Result, Wait};
-
-#[tokio::main]
-async fn main() -> Result<()> {
-    env_logger::try_init_from_env(Env::new().default_filter_or("info")).expect("logger failed");
-    info!("Opening session...");
-    let cfg = networking::cfg(&format!("{:x}", rand::random::<u128>()), 52414)?;
-    let session = networking::open(cfg, 52414, 52413).await?;
-    let _tok = session
-        .z
-        .liveliness()
-        .declare_token(format!("nodes/{}/live", session.z.zid()))
-        .wait()?;
-    session
-        .z
-        .liveliness()
-        .declare_subscriber("**")
-        .history(true)
-        .callback(|tok| info!("{}: {}", tok.kind(), tok.key_expr().to_string()))
-        .background()
-        .wait()?;
-    loop {
-        tokio::select! {
-            _ = tokio::signal::ctrl_c() => break,
-            _ = session.z.put("hello", "world") => {},
-        }
-    }
-    Ok(())
-}

rust/networking/examples/put_string.rs

@@ -1,48 +0,0 @@
-use std::{env, time::Duration};
-
-use env_logger::Env;
-use log::info;
-use networking;
-use zenoh::Result;
-
-#[tokio::main]
-async fn main() -> Result<()> {
-    env_logger::try_init_from_env(Env::new().default_filter_or("info")).expect("logger failed");
-    let n_bytes = env::args()
-        .nth(1)
-        .and_then(|it| it.parse::<usize>().ok())
-        .expect("USAGE: put_string <n> -- pub a string of n bytes into stream/data");
-    info!("Opening session...");
-    let cfg = networking::cfg(&format!("{:x}", rand::random::<u128>()), 52414)?;
-    let session = networking::open(cfg, 52414, 52413).await?;
-    let _tok = session
-        .z
-        .liveliness()
-        .declare_token(format!("nodes/{}/live", session.z.zid()))
-        .await?;
-    let key_expr = "stream/data";
-    let payload = "n".repeat(n_bytes);
-
-    let pubs = session
-        .z
-        .declare_publisher(key_expr)
-        .congestion_control(zenoh::qos::CongestionControl::Block)
-        .await?;
-    let pubs_l = pubs.matching_listener().await?;
-    if !pubs.matching_status().await?.matching() {
-        while !pubs_l.recv_async().await?.matching() {}
-    }
-
-    tokio::time::sleep(Duration::from_secs(1)).await;
-    info!("Putting Data ('{key_expr}': '{}')...", payload.len());
-    for _ in 0..10 {
-        let t = tokio::time::Instant::now();
-        for _ in 0..5000 {
-            pubs.put(payload.clone()).await?;
-        }
-        info!("{:?}", t.elapsed());
-        tokio::time::sleep(Duration::from_secs(1)).await;
-    }
-    tokio::signal::ctrl_c().await?;
-    Ok(())
-}

rust/networking/examples/serve_storage.rs

@@ -1,74 +0,0 @@
-use std::time::Duration;
-
-use env_logger::Env;
-use log::info;
-use networking;
-use zenoh::Result;
-
-#[tokio::main]
-async fn main() -> Result<()> {
-    env_logger::try_init_from_env(Env::new().default_filter_or("info")).expect("logger failed");
-    info!("Opening session...");
-    let cfg = networking::cfg(&format!("{:x}", rand::random::<u128>()), 52414)?;
-    let session = networking::open(cfg, 52414, 52413).await?;
-    let _tok = session
-        .z
-        .liveliness()
-        .declare_token(format!("nodes/{}/live", session.z.zid()))
-        .await?;
-    let _sub = session
-        .z
-        .liveliness()
-        .declare_subscriber("nodes/*/live")
-        .history(true)
-        .callback(|tok| {
-            info!(
-                "{}: {}",
-                tok.kind(),
-                tok.key_expr()
-                    .to_string()
-                    .strip_prefix("nodes/")
-                    .and_then(|it| it.strip_suffix("/live"))
-                    .unwrap()
-            )
-        })
-        .await?;
-
-    let watch = async {
-        for _ in 0..1000 {
-            tokio::time::sleep(Duration::from_secs(1)).await;
-            session
-                .z
-                .get("**")
-                .callback(|reply| {
-                    let sample = reply.into_result().expect("no errs");
-                    info!(
-                        "got {} bytes on {}",
-                        sample.payload().len(),
-                        sample.key_expr()
-                    )
-                })
-                .await?;
-        }
-        Result::<()>::Ok(())
-    };
-    let subs = session.z.declare_subscriber("**").await?;
-
-    let mut i = 0;
-    let _a = async {
-        while let Ok(sample) = subs.recv_async().await {
-            i += 1;
-            info!(
-                "[{i}] received {} bytes on {}",
-                sample.payload().len(),
-                sample.key_expr()
-            )
-        }
-    };
-    tokio::select! {
-        _ = watch => {},
-        _ = _a => {},
-        _ = tokio::signal::ctrl_c() => {},
-    }
-    Ok(())
-}

rust/networking/examples/z_get.rs

@@ -1,41 +0,0 @@
-use std::{borrow::Cow, env};
-
-use env_logger::Env;
-use log::{info, warn};
-use networking;
-use zenoh::{Result, Wait};
-
-#[tokio::main]
-async fn main() -> Result<()> {
-    env_logger::try_init_from_env(Env::new().default_filter_or("info")).expect("logger failed");
-    info!("Opening session...");
-    let cfg = networking::cfg(&format!("{:x}", rand::random::<u128>()), 52414)?;
-    let session = networking::open(cfg, 52414, 52413).await?;
-    let other_live = session
-        .z
-        .liveliness()
-        .declare_subscriber("**")
-        .history(true)
-        .wait()?;
-    _ = other_live.recv_async().await?;
-    let other_live = session.z.liveliness().get("**").wait()?;
-    while let Ok(s) = other_live.recv_async().await {
-        info!("{s:?}");
-    }
-    let query = env::args().nth(1).expect("USAGE: z_get [query]");
-    info!("Querying {query}");
-    let subs = session.z.liveliness().get(query).await?;
-    while let Ok(r) = subs.recv_async().await {
-        match r.into_result() {
-            Ok(s) => info!(
-                "{}: {}",
-                s.key_expr(),
-                s.payload()
-                    .try_to_string()
-                    .unwrap_or_else(|_| Cow::Borrowed("-bytes-"))
-            ),
-            Err(e) => warn!("{e}"),
-        }
-    }
-    Ok(())
-}

rust/networking/src/RESEARCH_NOTES.txt

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

rust/networking/src/discovery.rs

@@ -1,344 +1,390 @@
-use std::{
-    env, io,
-    net::{Ipv6Addr, SocketAddr, SocketAddrV6},
-    sync::Arc,
-    time::Duration,
+use crate::ext::MultiaddrExt;
+use delegate::delegate;
+use either::Either;
+use futures_lite::FutureExt;
+use futures_timer::Delay;
+use libp2p::core::transport::PortUse;
+use libp2p::core::{ConnectedPoint, Endpoint};
+use libp2p::swarm::behaviour::ConnectionEstablished;
+use libp2p::swarm::dial_opts::DialOpts;
+use libp2p::swarm::{
+    CloseConnection, ConnectionClosed, ConnectionDenied, ConnectionHandler,
+    ConnectionHandlerSelect, ConnectionId, FromSwarm, NetworkBehaviour, THandler, THandlerInEvent,
+    THandlerOutEvent, ToSwarm, dummy,
 };
+use libp2p::{Multiaddr, PeerId, identity, mdns};
+use std::collections::{BTreeSet, HashMap};
+use std::convert::Infallible;
+use std::io;
+use std::net::IpAddr;
+use std::task::{Context, Poll};
+use std::time::Duration;
+use util::wakerdeque::WakerDeque;
 
-use bytemuck::{Pod, Zeroable};
-use log::{debug, trace, warn};
-use netwatcher::WatchHandle;
-use parking_lot::Mutex;
-use tokio::{
-    net::UdpSocket,
-    time::{Interval, interval},
-};
-use zenoh::config::ZenohId;
+const RETRY_CONNECT_INTERVAL: Duration = Duration::from_secs(5);
 
-const GROUP: Ipv6Addr = Ipv6Addr::new(0xff12, 0, 0, 0, 0, 0, 0xe0a1, 0xde89);
-const MAGIC: [u8; 3] = *b"EXO";
+mod managed {
+    use libp2p::swarm::NetworkBehaviour;
+    use libp2p::{identity, mdns, ping};
+    use std::io;
+    use std::time::Duration;
 
-pub struct Discovery {
-    sock: Arc<UdpSocket>,
-    ifaces: Arc<Mutex<Vec<SocketAddrV6>>>,
-    namespace: [u8; 8],
-    last_nonce: Mutex<[u8; 8]>,
-    /// the port of the service we are doing discovery for - transmitted to peers
-    listen_port: u16,
-    zid: ZenohId,
-    tick: Interval,
-    _sync: Mutex<WatchHandle>,
-}
+    const MDNS_RECORD_TTL: Duration = Duration::from_secs(2_500);
+    const MDNS_QUERY_INTERVAL: Duration = Duration::from_secs(1_500);
+    const PING_TIMEOUT: Duration = Duration::from_millis(2_500);
+    const PING_INTERVAL: Duration = Duration::from_millis(2_500);
 
-#[derive(Debug, Clone, Copy)]
-pub struct Discovered {
-    pub zid: ZenohId,
-    pub addr: SocketAddrV6,
-}
+    #[derive(NetworkBehaviour)]
+    pub struct Behaviour {
+        mdns: mdns::tokio::Behaviour,
+        ping: ping::Behaviour,
+    }
 
-impl Discovery {
-    pub async fn new(zid: ZenohId, listen_port: u16, discovery_port: u16) -> io::Result<Self> {
-        let namespace: [u8; 8] = {
-            blake3::hash(
-                env::var("EXO_ZENOH_NAMESPACE")
-                    .unwrap_or_else(|_| "exo".to_string())
-                    .as_bytes(),
-            )
-            .as_bytes()[..8]
-                .try_into()
-                .expect("8=8")
-        };
-        let sock = socket2::Socket::new(
-            socket2::Domain::IPV6,
-            socket2::Type::DGRAM,
-            Some(socket2::Protocol::UDP),
-        )?;
-        sock.set_reuse_address(true)?;
-        #[cfg(unix)]
-        sock.set_reuse_port(true)?;
-        sock.bind(&SocketAddrV6::new(Ipv6Addr::UNSPECIFIED, discovery_port, 0, 0).into())?;
-        sock.set_nonblocking(true)?;
-        sock.set_multicast_loop_v6(true)?;
-        let sock = Arc::new(UdpSocket::from_std(sock.into())?);
-        let ifaces: Arc<Mutex<Vec<SocketAddrV6>>> = Default::default();
-        let _sync = Mutex::new(
-            netwatcher::watch_interfaces_with_callback({
-                let sock = sock.clone();
-                let ifaces = ifaces.clone();
-                move |update| {
-                    for (iface_idx, iface) in update.interfaces.iter() {
-                        if iface
-                            .ipv6_ips()
-                            .all(|addr| addr.is_loopback() || addr.is_unspecified())
-                        {
-                            continue;
-                        }
-
-                        match sock.join_multicast_v6(&GROUP, *iface_idx) {
-                            Ok(()) => ifaces.lock().push(SocketAddrV6::new(
-                                GROUP,
-                                discovery_port,
-                                0,
-                                *iface_idx,
-                            )),
-                            Err(e) if e.kind() != io::ErrorKind::AddrInUse => {
-                                // skip AddrInUse - just means we've already joined the mv6
-                                if let Some(iface) = update.interfaces.get(&iface_idx) {
-                                    warn!(
-                                        "failed to join multicast v6 for interface {}: {e}",
-                                        iface.name
-                                    )
-                                }
-                            }
-                            _ => {}
-                        }
-                    }
-                    for iface_idx in update.diff.removed {
-                        ifaces.lock().retain(|addr| addr.scope_id() != iface_idx);
-
-                        if let Err(e) = sock.leave_multicast_v6(&GROUP, iface_idx) {
-                            if let Some(iface) = update.interfaces.get(&iface_idx) {
-                                warn!(
-                                    "failed to leave multicast v6 for interface {}: {e}",
-                                    iface.name
-                                )
-                            }
-                        }
-                    }
-                }
+    impl Behaviour {
+        pub fn new(keypair: &identity::Keypair) -> io::Result<Self> {
+            Ok(Self {
+                mdns: mdns_behaviour(keypair)?,
+                ping: ping_behaviour(),
             })
-            // todo: better error handling here
-            .expect("failed to bind discovery watcher"),
-        );
+        }
+    }
+
+    fn mdns_behaviour(keypair: &identity::Keypair) -> io::Result<mdns::tokio::Behaviour> {
+        use mdns::{Config, tokio};
+
+        // mDNS config => enable IPv6
+        let mdns_config = Config {
+            ttl: MDNS_RECORD_TTL,
+            query_interval: MDNS_QUERY_INTERVAL,
+
+            // enable_ipv6: true, // TODO: for some reason, TCP+mDNS don't work well with ipv6?? figure out how to make work
+            ..Default::default()
+        };
+
+        let mdns_behaviour = tokio::Behaviour::new(mdns_config, keypair.public().to_peer_id());
+        Ok(mdns_behaviour?)
+    }
+
+    fn ping_behaviour() -> ping::Behaviour {
+        ping::Behaviour::new(
+            ping::Config::new()
+                .with_timeout(PING_TIMEOUT)
+                .with_interval(PING_INTERVAL),
+        )
+    }
+}
+
+/// Events for when a listening connection is truly established and truly closed.
+#[derive(Debug, Clone)]
+pub enum Event {
+    ConnectionEstablished {
+        peer_id: PeerId,
+        connection_id: ConnectionId,
+        remote_ip: IpAddr,
+        remote_tcp_port: u16,
+    },
+    ConnectionClosed {
+        peer_id: PeerId,
+        connection_id: ConnectionId,
+        remote_ip: IpAddr,
+        remote_tcp_port: u16,
+    },
+}
+
+/// Discovery behavior that wraps mDNS to produce truly discovered durable peer-connections.
+///
+/// The behaviour operates as such:
+///  1) All true (listening) connections/disconnections are tracked, emitting corresponding events
+///     to the swarm.
+///  1) mDNS discovered/expired peers are tracked; discovered but not connected peers are dialed
+///     immediately, and expired but connected peers are disconnected from immediately.
+///  2) Every fixed interval: discovered but not connected peers are dialed, and expired but
+///     connected peers are disconnected from.
+pub struct Behaviour {
+    // state-tracking for managed behaviors & mDNS-discovered peers
+    managed: managed::Behaviour,
+    mdns_discovered: HashMap<PeerId, BTreeSet<Multiaddr>>,
+    bootstrap_peers: Vec<Multiaddr>,
+
+    retry_delay: Delay, // retry interval
+
+    // pending events to emmit => waker-backed Deque to control polling
+    pending_events: WakerDeque<ToSwarm<Event, Infallible>>,
+}
+
+impl Behaviour {
+    pub fn new(keypair: &identity::Keypair, bootstrap_peers: Vec<Multiaddr>) -> io::Result<Self> {
         Ok(Self {
-            sock,
-            namespace,
-            ifaces,
-            last_nonce: Mutex::new(rand::random()),
-            listen_port,
-            zid,
-            tick: interval(Duration::from_secs(1)),
-            _sync,
+            managed: managed::Behaviour::new(keypair)?,
+            mdns_discovered: HashMap::new(),
+            bootstrap_peers,
+            retry_delay: Delay::new(RETRY_CONNECT_INTERVAL),
+            pending_events: WakerDeque::new(),
         })
     }
 
-    pub async fn next(&mut self) -> io::Result<Discovered> {
-        let mut buf = [0u8; Hello::buf_size() + WhatsUp::buf_size() + 1];
-        loop {
-            tokio::select! {
-                _ = self.tick.tick() => {
-                    self.announce().await?;
+    fn dial(&mut self, peer_id: PeerId, addr: Multiaddr) {
+        self.pending_events.push_back(ToSwarm::Dial {
+            opts: DialOpts::peer_id(peer_id).addresses(vec![addr]).build(),
+        })
+    }
+
+    fn close_connection(&mut self, peer_id: PeerId, connection: ConnectionId) {
+        // push front to make this IMMEDIATE
+        self.pending_events.push_front(ToSwarm::CloseConnection {
+            peer_id,
+            connection: CloseConnection::One(connection),
+        })
+    }
+
+    fn handle_mdns_discovered(&mut self, peers: Vec<(PeerId, Multiaddr)>) {
+        for (p, ma) in peers {
+            self.dial(p, ma.clone()); // always connect
+
+            // get peer's multi-addresses or insert if missing
+            let Some(mas) = self.mdns_discovered.get_mut(&p) else {
+                self.mdns_discovered.insert(p, BTreeSet::from([ma]));
+                continue;
+            };
+
+            // multiaddress should never already be present - else something has gone wrong
+            let is_new_addr = mas.insert(ma);
+            assert!(is_new_addr, "cannot discover a discovered peer");
+        }
+    }
+
+    fn handle_mdns_expired(&mut self, peers: Vec<(PeerId, Multiaddr)>) {
+        for (p, ma) in peers {
+            // at this point, we *must* have the peer
+            let mas = self
+                .mdns_discovered
+                .get_mut(&p)
+                .expect("nonexistent peer cannot expire");
+
+            // at this point, we *must* have the multiaddress
+            let was_present = mas.remove(&ma);
+            assert!(was_present, "nonexistent multiaddress cannot expire");
+
+            // if empty, remove the peer-id entirely
+            if mas.is_empty() {
+                self.mdns_discovered.remove(&p);
+            }
+        }
+    }
+
+    fn on_connection_established(
+        &mut self,
+        peer_id: PeerId,
+        connection_id: ConnectionId,
+        remote_ip: IpAddr,
+        remote_tcp_port: u16,
+    ) {
+        // send out connected event
+        self.pending_events
+            .push_back(ToSwarm::GenerateEvent(Event::ConnectionEstablished {
+                peer_id,
+                connection_id,
+                remote_ip,
+                remote_tcp_port,
+            }));
+    }
+
+    fn on_connection_closed(
+        &mut self,
+        peer_id: PeerId,
+        connection_id: ConnectionId,
+        remote_ip: IpAddr,
+        remote_tcp_port: u16,
+    ) {
+        // send out disconnected event
+        self.pending_events
+            .push_back(ToSwarm::GenerateEvent(Event::ConnectionClosed {
+                peer_id,
+                connection_id,
+                remote_ip,
+                remote_tcp_port,
+            }));
+    }
+}
+
+impl NetworkBehaviour for Behaviour {
+    type ConnectionHandler =
+        ConnectionHandlerSelect<dummy::ConnectionHandler, THandler<managed::Behaviour>>;
+    type ToSwarm = Event;
+
+    // simply delegate to underlying mDNS behaviour
+
+    delegate! {
+        to self.managed {
+            fn handle_pending_inbound_connection(&mut self, connection_id: ConnectionId, local_addr: &Multiaddr, remote_addr: &Multiaddr) -> Result<(), ConnectionDenied>;
+            fn handle_pending_outbound_connection(&mut self, connection_id: ConnectionId, maybe_peer: Option<PeerId>, addresses: &[Multiaddr], effective_role: Endpoint) -> Result<Vec<Multiaddr>, ConnectionDenied>;
+        }
+    }
+
+    fn handle_established_inbound_connection(
+        &mut self,
+        connection_id: ConnectionId,
+        peer: PeerId,
+        local_addr: &Multiaddr,
+        remote_addr: &Multiaddr,
+    ) -> Result<THandler<Self>, ConnectionDenied> {
+        Ok(ConnectionHandler::select(
+            dummy::ConnectionHandler,
+            self.managed.handle_established_inbound_connection(
+                connection_id,
+                peer,
+                local_addr,
+                remote_addr,
+            )?,
+        ))
+    }
+
+    #[allow(clippy::needless_question_mark)]
+    fn handle_established_outbound_connection(
+        &mut self,
+        connection_id: ConnectionId,
+        peer: PeerId,
+        addr: &Multiaddr,
+        role_override: Endpoint,
+        port_use: PortUse,
+    ) -> Result<THandler<Self>, ConnectionDenied> {
+        Ok(ConnectionHandler::select(
+            dummy::ConnectionHandler,
+            self.managed.handle_established_outbound_connection(
+                connection_id,
+                peer,
+                addr,
+                role_override,
+                port_use,
+            )?,
+        ))
+    }
+
+    fn on_connection_handler_event(
+        &mut self,
+        peer_id: PeerId,
+        connection_id: ConnectionId,
+        event: THandlerOutEvent<Self>,
+    ) {
+        match event {
+            Either::Left(ev) => libp2p::core::util::unreachable(ev),
+            Either::Right(ev) => {
+                self.managed
+                    .on_connection_handler_event(peer_id, connection_id, ev)
+            }
+        }
+    }
+
+    // hook into these methods to drive behavior
+
+    fn on_swarm_event(&mut self, event: FromSwarm) {
+        self.managed.on_swarm_event(event); // let mDNS handle swarm events
+
+        // handle swarm events to update internal state:
+        match event {
+            FromSwarm::ConnectionEstablished(ConnectionEstablished {
+                peer_id,
+                connection_id,
+                endpoint,
+                ..
+            }) => {
+                let remote_address = match endpoint {
+                    ConnectedPoint::Dialer { address, .. } => address,
+                    ConnectedPoint::Listener { send_back_addr, .. } => send_back_addr,
+                };
+
+                if let Some((ip, port)) = remote_address.try_to_tcp_addr() {
+                    // handle connection established event which is filtered correctly
+                    self.on_connection_established(peer_id, connection_id, ip, port)
                 }
-                res = self.sock.recv_from(&mut buf) => {
-                    let Ok((bytes_read, addr)) = res else { continue; };
-                    if let Some(discovered) = self.respond(bytes_read, addr, &buf).await? {
-                        return Ok(discovered)
+            }
+            FromSwarm::ConnectionClosed(ConnectionClosed {
+                peer_id,
+                connection_id,
+                endpoint,
+                ..
+            }) => {
+                let remote_address = match endpoint {
+                    ConnectedPoint::Dialer { address, .. } => address,
+                    ConnectedPoint::Listener { send_back_addr, .. } => send_back_addr,
+                };
+
+                if let Some((ip, port)) = remote_address.try_to_tcp_addr() {
+                    // handle connection closed event which is filtered correctly
+                    self.on_connection_closed(peer_id, connection_id, ip, port)
+                }
+            }
+
+            // since we are running TCP/IP transport layer, we are assuming that
+            // no address changes can occur, hence encountering one is a fatal error
+            FromSwarm::AddressChange(a) => {
+                unreachable!("unhandlable: address change encountered: {:?}", a)
+            }
+            _ => {}
+        }
+    }
+
+    fn poll(&mut self, cx: &mut Context) -> Poll<ToSwarm<Self::ToSwarm, THandlerInEvent<Self>>> {
+        // delegate to managed behaviors for any behaviors they need to perform
+        match self.managed.poll(cx) {
+            Poll::Ready(ToSwarm::GenerateEvent(e)) => {
+                match e {
+                    // handle discovered and expired events from mDNS
+                    managed::BehaviourEvent::Mdns(e) => match e.clone() {
+                        mdns::Event::Discovered(peers) => {
+                            self.handle_mdns_discovered(peers);
+                        }
+                        mdns::Event::Expired(peers) => {
+                            self.handle_mdns_expired(peers);
+                        }
+                    },
+
+                    // handle ping events => if error then disconnect
+                    managed::BehaviourEvent::Ping(e) => {
+                        if let Err(_) = e.result {
+                            self.close_connection(e.peer, e.connection.clone())
+                        }
                     }
                 }
+
+                // since we just consumed an event, we should immediately wake just in case
+                // there are more events to come where that came from
+                cx.waker().wake_by_ref();
             }
-        }
-    }
 
-    async fn respond(
-        &self,
-        bytes_read: usize,
-        addr: SocketAddr,
-        buf: &[u8],
-    ) -> io::Result<Option<Discovered>> {
-        trace!(
-            "raw recv: {bytes_read} bytes from {addr}: {:02x?}",
-            &buf[..bytes_read]
-        );
-        if bytes_read < size_of::<Header>() {
-            trace!("dropped: early EOF");
-            return Ok(None);
-        }
-        let header: &Header = bytemuck::from_bytes(&buf[0..size_of::<Header>()]);
-        if header.magic != MAGIC {
-            trace!("dropped: wrong magic");
-            return Ok(None);
-        }
-        let Ok(kind) = header.kind.try_into() else {
-            trace!("dropped: unknown message kind {}", header.kind);
-            return Ok(None);
-        };
-        match kind {
-            Kind::Hello => {
-                let total = Hello::buf_size();
-                if bytes_read != total {
-                    trace!("dropped: hello wrong size");
-                    return Ok(None);
-                }
-                let hello: &Hello = bytemuck::from_bytes(&buf[size_of::<Header>()..total]);
-                if hello.nonce == *self.last_nonce.lock() {
-                    trace!("dropped: local hello nonce");
-                    return Ok(None);
-                }
-                if hello.namespace != self.namespace {
-                    trace!("dropped: different namespace");
-                    return Ok(None);
-                }
-
-                // reply
-                trace!("replying to Hello({:?})", hello.nonce);
-                let reply = WhatsUp {
-                    nonce: hello.nonce,
-                    zid: self.zid.to_le_bytes(),
-                    port_le: self.listen_port.to_le_bytes(),
-                }
-                .alloc();
-
-                for i in 1..6 {
-                    if self
-                        .sock
-                        .send_to(&reply, addr)
-                        .await
-                        .inspect_err(|e| debug!("send to {addr} failed: {e}"))
-                        .is_ok_and(|sent| sent == WhatsUp::buf_size())
-                    {
-                        trace!(
-                            "sent {} bytes to {addr} after {} attempt(s)",
-                            WhatsUp::buf_size(),
-                            i
-                        );
-                        break;
-                    }
-                    tokio::time::sleep(Duration::from_millis(300)).await;
-                }
-                Ok(None)
+            // forward any other mDNS event to the swarm or its connection handler(s)
+            Poll::Ready(e) => {
+                return Poll::Ready(
+                    e.map_out(|_| unreachable!("events returning to swarm already handled"))
+                        .map_in(Either::Right),
+                );
             }
-            Kind::WhatsUp => {
-                let total = WhatsUp::buf_size();
-                if bytes_read != total {
-                    trace!("dropped: whatsup wrong size");
-                    return Ok(None);
-                }
-                let whats_up: &WhatsUp = bytemuck::from_bytes(&buf[size_of::<Header>()..total]);
-                if whats_up.nonce != *self.last_nonce.lock() {
-                    trace!("dropped: stale nonce");
-                    return Ok(None);
-                }
-                let SocketAddr::V6(v6) = addr else {
-                    trace!("dropped: v4 addr used");
-                    return Ok(None);
-                };
-                let Ok(zid) = ZenohId::try_from(&whats_up.zid[..]) else {
-                    trace!("dropped: zenoh conversion failed");
-                    return Ok(None);
-                };
-                if zid == self.zid {
-                    trace!("dropped: self zenoh id");
-                    return Ok(None);
-                }
-                // discovery success!
-                // the incoming port is our listen port;
-                // overwrite it with the whats_up port corresponding to the remote zenoh service
-                let addr = {
-                    let mut x = v6;
-                    x.set_port(u16::from_le_bytes(whats_up.port_le));
-                    x
-                };
-                Ok(Some(Discovered { addr, zid }))
-            }
-        }
-    }
 
-    async fn announce(&self) -> io::Result<()> {
-        let nonce = rand::random();
-        *self.last_nonce.lock() = nonce;
-        let buf = Hello {
-            nonce,
-            namespace: self.namespace,
+            Poll::Pending => {}
         }
-        .alloc();
 
-        let addrs = self.ifaces.lock().clone();
-        debug!("announcing Hello({nonce:?}) to {addrs:?}");
-        // rev so .remove() doesn't break things
-        for (i, addr) in addrs.into_iter().enumerate().rev() {
-            match self.sock.send_to(&buf, addr).await {
-                Ok(bytes) => trace!("sent {bytes} to {addr}"),
-                Err(e) if e.kind() == io::ErrorKind::HostUnreachable => {
-                    debug!("disabling discovery address {addr}: {e}");
-                    _ = self.ifaces.lock().swap_remove(i);
+        // retry connecting to all mDNS peers periodically (fails safely if already connected)
+        if self.retry_delay.poll(cx).is_ready() {
+            for (p, mas) in self.mdns_discovered.clone() {
+                for ma in mas {
+                    self.dial(p, ma)
                 }
-                Err(e) => debug!("failed to reach {addr}: {e}"),
             }
+            // dial bootstrap peers (for environments where mDNS is unavailable)
+            for addr in &self.bootstrap_peers {
+                self.pending_events.push_back(ToSwarm::Dial {
+                    opts: DialOpts::unknown_peer_id().address(addr.clone()).build(),
+                })
+            }
+            self.retry_delay.reset(RETRY_CONNECT_INTERVAL) // reset timeout
         }
-        Ok(())
+
+        // send out any pending events from our own service
+        if let Some(e) = self.pending_events.pop_front(cx) {
+            return Poll::Ready(e.map_in(Either::Left));
+        }
+
+        // wait for pending events
+        Poll::Pending
     }
 }
-
-#[repr(u8)]
-#[derive(Debug, Clone, Copy)]
-// packet & version
-pub enum Kind {
-    Hello = 0,
-    WhatsUp = 1,
-}
-
-pub struct UnknownKind;
-impl TryFrom<u8> for Kind {
-    type Error = UnknownKind;
-    fn try_from(value: u8) -> Result<Self, Self::Error> {
-        match value {
-            0 => Ok(Self::Hello),
-            1 => Ok(Self::WhatsUp),
-            _ => Err(UnknownKind),
-        }
-    }
-}
-
-pub trait Message: Pod {
-    const KIND: Kind;
-}
-// should be part of the Message trait, but const in traits isnt stabilized. this lets alloc :: Self -> [u8; Self::buf_size()]
-macro_rules! impl_alloc {
-    ($a:ident) => {
-        impl $a {
-            const fn buf_size() -> usize {
-                size_of::<Header>() + size_of::<Self>()
-            }
-            pub fn alloc(self) -> [u8; Self::buf_size()] {
-                let mut buf = [0u8; Self::buf_size()];
-                buf[0..size_of::<Header>()].copy_from_slice(bytemuck::bytes_of(&Header {
-                    magic: MAGIC,
-                    kind: Self::KIND as u8,
-                }));
-                buf[size_of::<Header>()..Self::buf_size()]
-                    .copy_from_slice(bytemuck::bytes_of(&self));
-                buf
-            }
-        }
-    };
-}
-
-#[repr(C)]
-#[derive(Debug, Clone, Copy, Pod, Zeroable)]
-pub struct Header {
-    magic: [u8; 3],
-    kind: u8,
-}
-
-#[repr(C)]
-#[derive(Debug, Clone, Copy, Pod, Zeroable)]
-pub struct Hello {
-    pub nonce: [u8; 8],
-    pub namespace: [u8; 8],
-}
-impl Message for Hello {
-    const KIND: Kind = Kind::Hello;
-}
-impl_alloc!(Hello);
-
-#[repr(C)]
-#[derive(Debug, Clone, Copy, Pod, Zeroable)]
-pub struct WhatsUp {
-    pub nonce: [u8; 8],
-    pub zid: [u8; 16],
-    pub port_le: [u8; 2],
-}
-impl Message for WhatsUp {
-    const KIND: Kind = Kind::WhatsUp;
-}
-impl_alloc!(WhatsUp);

rust/networking/src/lib.rs

@@ -1,120 +1,44 @@
-use std::sync::Arc;
-
-use tokio::task::JoinHandle;
-use zenoh::{Result, Session as ZSession, config::Locator};
-use zenoh_plugin_storage_manager::StoragesPlugin;
-use zenoh_plugin_trait::PluginsManager;
-
-pub use zenoh::{Config, config::ZenohId};
-
-use crate::{
-    discovery::Discovery,
-    liveliness_aggregator::{LivelinessAggregator, spawn_liveliness_aggregator},
-};
-
-pub use zenoh_plugin_storage_manager::read_raw_memory_storage;
-pub const STORAGE_PREFIX: &str = "storage/mem1";
-
+//! TODO: crate documentation
+//!
+//! this is here as a placeholder documentation
+//!
+//!
 pub mod discovery;
-pub mod liveliness_aggregator;
 pub mod swarm;
 
-pub fn cfg(identity: &str, listen_port: u16) -> Result<zenoh::Config> {
-    assert!(
-        identity
-            .chars()
-            .all(|c| ('0'..='9').contains(&c) || ('a'..='f').contains(&c))
-    );
-    assert!(identity.len() <= 32);
-    assert!(listen_port != 0, "must used defined listen port");
-    let mut cfg = zenoh::Config::default();
-    // todo: cleanup
-    cfg.insert_json5("id", &format!("\"{identity}\""))?;
-    cfg.insert_json5("mode", "\"router\"")?;
-    cfg.insert_json5("listen/endpoints", &format!("[\"tcp/[::]:{listen_port}\"]"))?;
-    cfg.insert_json5("scouting/multicast/enabled", "false")?;
-    cfg.insert_json5("scouting/multicast/autoconnect", "[]")?;
-    cfg.insert_json5("scouting/gossip/multihop", "true")?;
-    cfg.insert_json5("adminspace/enabled", "true")?;
-    cfg.insert_json5("transport/link/tx/batch_size", "9216")?;
-    cfg.insert_json5("transport/link/rx/buffer_size", "16777216")?;
-    cfg.insert_json5("timestamping/enabled", "true")?;
-    cfg.insert_json5("plugins/storage_manager/__required__", "true")?;
-    cfg.insert_json5(
-        "plugins/storage_manager/storages/mem1/key_expr",
-        &format!("\"{STORAGE_PREFIX}/**\""),
-    )?;
-    cfg.insert_json5(
-        "plugins/storage_manager/storages/mem1/strip_prefix",
-        &format!("\"{STORAGE_PREFIX}\""),
-    )?;
-    cfg.insert_json5("plugins/storage_manager/storages/mem1/volume", "\"memory\"")?;
-    cfg.insert_json5(
-        "plugins/storage_manager/storages/mem1/replication/interval",
-        "2",
-    )?;
-    Ok(cfg)
+/// Namespace for all the type/trait aliases used by this crate.
+pub(crate) mod alias {
+    use std::error::Error;
+
+    pub type AnyError = Box<dyn Error + Send + Sync + 'static>;
+    pub type AnyResult<T> = Result<T, AnyError>;
 }
 
-pub async fn open(
-    cfg: zenoh::Config,
-    listen_port: u16,
-    discovery_service_port: u16,
-) -> Result<Session> {
-    assert!(listen_port != 0, "must used defined listen port");
-    let mut plugins = PluginsManager::static_plugins_only();
-    plugins.declare_static_plugin::<StoragesPlugin, _>("storage_manager", true);
-    let mut runtime = zenoh::internal::runtime::RuntimeBuilder::new(cfg)
-        .plugins_manager(plugins)
-        .build()
-        .await?;
-    let z = zenoh::session::init(runtime.clone().into()).await?;
-    runtime.start().await?;
-    let mut discovery = Discovery::new(z.zid(), listen_port, discovery_service_port).await?;
-    let _jh = Arc::new(AbortOnDrop(tokio::task::spawn(async move {
-        loop {
-            let Ok(discovered) = discovery.next().await.inspect_err(|e| {
-                log::warn!("discovery error {e}");
-            }) else {
-                continue;
+/// Namespace for crate-wide extension traits/methods
+pub(crate) mod ext {
+    use extend::ext;
+    use libp2p::Multiaddr;
+    use libp2p::multiaddr::Protocol;
+    use std::net::IpAddr;
+
+    #[ext(pub, name = MultiaddrExt)]
+    impl Multiaddr {
+        /// If the multiaddress corresponds to a TCP address, extracts it
+        fn try_to_tcp_addr(&self) -> Option<(IpAddr, u16)> {
+            let mut ps = self.into_iter();
+            let ip = if let Some(p) = ps.next() {
+                match p {
+                    Protocol::Ip4(ip) => IpAddr::V4(ip),
+                    Protocol::Ip6(ip) => IpAddr::V6(ip),
+                    _ => return None,
+                }
+            } else {
+                return None;
             };
-
-            if discovered.zid > runtime.zid() {
-                log::debug!("not connecting to peer with greater zid");
-                continue;
-            }
-
-            let Ok(locator) =
-                Locator::new("tcp", discovered.addr.to_string(), "").inspect_err(|e| {
-                    log::warn!("failed to parse locator from addr: {e}");
-                })
-            else {
-                continue;
+            let Some(Protocol::Tcp(port)) = ps.next() else {
+                return None;
             };
-
-            runtime
-                .connect_peer(&discovered.zid.into(), &[locator])
-                .await;
+            Some((ip, port))
         }
-    })));
-    let liveliness_aggregator = spawn_liveliness_aggregator(&z)?;
-    Ok(Session {
-        z,
-        liveliness_aggregator,
-        _jh,
-    })
-}
-
-pub struct AbortOnDrop(pub JoinHandle<()>);
-impl Drop for AbortOnDrop {
-    fn drop(&mut self) {
-        self.0.abort();
     }
 }
-
-#[derive(Clone)]
-pub struct Session {
-    pub z: ZSession,
-    pub liveliness_aggregator: LivelinessAggregator,
-    _jh: Arc<AbortOnDrop>,
-}

rust/networking/src/liveliness_aggregator.rs

@@ -1,44 +0,0 @@
-use std::{collections::HashSet, sync::Arc};
-
-use parking_lot::Mutex;
-use zenoh::{Result, Session, Wait, sample::SampleKind};
-
-pub fn spawn_liveliness_aggregator(session: &Session) -> Result<LivelinessAggregator> {
-    let store = Arc::new(Mutex::new(HashSet::default()));
-    session
-        .liveliness()
-        .declare_subscriber("live/*")
-        .history(true)
-        .callback({
-            let store = Arc::clone(&store);
-            move |sample| {
-                let Some(nid) = sample
-                    .key_expr()
-                    .to_string()
-                    .strip_prefix("live/")
-                    .map(str::to_owned)
-                else {
-                    return;
-                };
-                let mut mg = store.lock();
-                match sample.kind() {
-                    SampleKind::Put => mg.insert(nid),
-                    SampleKind::Delete => mg.remove(&nid),
-                };
-            }
-        })
-        .background()
-        .wait()?;
-    Ok(LivelinessAggregator { store })
-}
-
-#[derive(Clone)]
-pub struct LivelinessAggregator {
-    // need two arcs as the sub owns an arc to the store.
-    store: Arc<Mutex<HashSet<String>>>,
-}
-impl LivelinessAggregator {
-    pub fn dump(&self) -> HashSet<String> {
-        self.store.lock().clone()
-    }
-}

rust/networking/src/swarm.rs

@@ -1,22 +1,24 @@
-//! Compat shim for the old libp2p code
-
-use std::collections::HashMap;
 use std::pin::Pin;
 
-use futures_lite::Stream;
-use tokio::sync::mpsc;
-use tokio::sync::oneshot;
-use zenoh::Result;
-use zenoh::Session;
-use zenoh::handlers::FifoChannelHandler;
-use zenoh::liveliness::LivelinessToken;
-use zenoh::pubsub::Publisher;
-use zenoh::pubsub::Subscriber;
-use zenoh::qos::CongestionControl;
-use zenoh::sample::Sample;
-use zenoh::sample::SampleKind;
+use crate::swarm::transport::tcp_transport;
+use crate::{alias, discovery};
+pub use behaviour::{Behaviour, BehaviourEvent};
+use futures_lite::{Stream, StreamExt};
+use libp2p::{PeerId, SwarmBuilder, gossipsub, identity, swarm::SwarmEvent};
+use tokio::sync::{mpsc, oneshot};
 
-#[derive(Debug)]
+/// The current version of the network: this prevents devices running different versions of the
+/// software from interacting with each other.
+///
+/// TODO: right now this is a hardcoded constant; figure out what the versioning semantics should
+///       even be, and how to inject the right version into this config/initialization. E.g. should
+///       this be passed in as a parameter? What about rapidly changing versions in debug builds?
+///       this is all VERY very hard to figure out and needs to be mulled over as a team.
+pub const NETWORK_VERSION: &[u8] = b"v0.0.1";
+pub const OVERRIDE_VERSION_ENV_VAR: &str = "EXO_LIBP2P_NAMESPACE";
+
+// Uses oneshot senders to emulate function calling apis while avoiding requiring unique ownership
+// of the Swarm.
 pub enum ToSwarm {
     Unsubscribe {
         topic: String,
@@ -24,66 +26,52 @@ pub enum ToSwarm {
     },
     Subscribe {
         topic: String,
-        result_sender: oneshot::Sender<Result<bool>>,
+        result_sender: oneshot::Sender<Result<bool, gossipsub::SubscriptionError>>,
     },
     Publish {
         topic: String,
         data: Vec<u8>,
-        result_sender: oneshot::Sender<Result<()>>,
+        result_sender: oneshot::Sender<Result<gossipsub::MessageId, gossipsub::PublishError>>,
     },
 }
-#[derive(Debug)]
 pub enum FromSwarm {
-    Message { topic: String, data: Vec<u8> },
-    Discovered {},
-    Expired {},
+    Message {
+        from: PeerId,
+        topic: String,
+        data: Vec<u8>,
+    },
+    Discovered {
+        peer_id: PeerId,
+    },
+    Expired {
+        peer_id: PeerId,
+    },
 }
 
-pub type Topics = HashMap<String, (Subscriber<()>, Publisher<'static>)>;
 pub struct Swarm {
-    pub session: crate::Session,
-    pub from_client: mpsc::Receiver<ToSwarm>,
+    swarm: libp2p::Swarm<Behaviour>,
+    from_client: mpsc::Receiver<ToSwarm>,
 }
 
 impl Swarm {
     pub fn into_stream(self) -> Pin<Box<dyn Stream<Item = FromSwarm> + Send>> {
         let Swarm {
-            session,
+            mut swarm,
             mut from_client,
         } = self;
         let stream = async_stream::stream! {
-            let mut session = session;
-            let (mut to_topics, mut from_topics) = mpsc::channel(1024);
-            let mut topics = Topics::new();
-            let Ok((_token, discovery)) = register_liveness(&mut session.z).await else { return; };
             loop {
                 tokio::select! {
                     msg = from_client.recv() => {
                         let Some(msg) = msg else { break };
-                        on_message(&mut session.z, &mut topics, &mut to_topics, msg).await;
+                        on_message(&mut swarm, msg);
                     }
-                    event = from_topics.recv() => {
-                        if let Some(event) = event {
-                            yield event
+                    event = swarm.next() => {
+                        let Some(event) = event else { break };
+                        if let Some(item) = filter_swarm_event(event) {
+                            yield item;
                         }
                     }
-                    token = discovery.recv_async() => {
-                        if let Ok(token) = token {
-                            let key_expr = token.key_expr().as_str().to_owned();
-                            let nid = key_expr.strip_prefix("live/");
-                            yield match token.kind() {
-                                SampleKind::Put => {
-                                    log::info!("discovered: {nid:?}");
-                                    FromSwarm::Discovered {}
-                                }
-                                SampleKind::Delete => {
-                                    log::info!("expired: {nid:?}");
-                                    FromSwarm::Expired {}
-                                }
-                            }
-                        }
-
-                    }
                 }
             }
         };
@@ -91,117 +79,208 @@ impl Swarm {
     }
 }
 
-async fn register_liveness(
-    session: &mut Session,
-) -> Result<(LivelinessToken, Subscriber<FifoChannelHandler<Sample>>)> {
-    let token = session
-        .liveliness()
-        .declare_token(format!("live/{}", session.zid()))
-        .await?;
-    let sub = session
-        .liveliness()
-        .declare_subscriber("live/*")
-        .history(true)
-        .await?;
-    Ok((token, sub))
-}
-
-async fn on_message(
-    session: &mut Session,
-    topics: &mut Topics,
-    to_topics: &mut mpsc::Sender<FromSwarm>,
-    msg: ToSwarm,
-) {
-    match msg {
-        ToSwarm::Publish {
+fn on_message(swarm: &mut libp2p::Swarm<Behaviour>, message: ToSwarm) {
+    match message {
+        ToSwarm::Subscribe {
             topic,
-            data,
             result_sender,
         } => {
-            let res = match topics.get(&topic) {
-                Some(topic) => topic.1.put(data).await,
-                None => {
-                    // TODO: this should be an error but the python FromSwarm is somewhat nondeterministic
-                    Ok(()) //Err("not subscribed to topic!".into()),
-                }
-            };
-            _ = result_sender.send(res);
+            let result = swarm
+                .behaviour_mut()
+                .gossipsub
+                .subscribe(&gossipsub::IdentTopic::new(topic));
+            _ = result_sender.send(result);
         }
         ToSwarm::Unsubscribe {
             topic,
             result_sender,
         } => {
-            let Some((_, (subscriber, publisher))) = topics.remove_entry(&topic) else {
-                _ = result_sender.send(false);
-                return;
-            };
-            _ = publisher.undeclare().await;
-            _ = subscriber.undeclare().await;
-            _ = result_sender.send(true);
+            let result = swarm
+                .behaviour_mut()
+                .gossipsub
+                .unsubscribe(&gossipsub::IdentTopic::new(topic));
+            _ = result_sender.send(result);
         }
-        ToSwarm::Subscribe {
+        ToSwarm::Publish {
             topic,
+            data,
             result_sender,
         } => {
-            assert!(topic.is_ascii());
-            if topics.contains_key(&topic) {
-                _ = result_sender.send(Ok(false));
-                return;
-            }
-
-            let publisher_res = session
-                .declare_publisher(format!("topics/{topic}"))
-                .congestion_control(CongestionControl::Block)
-                .await;
-            let publisher = match publisher_res {
-                Ok(p) => p,
-                Err(e) => {
-                    _ = result_sender.send(Err(e));
-                    return;
-                }
-            };
-
-            let subscriber_res = session
-                .declare_subscriber(format!("topics/{topic}"))
-                .allowed_origin(zenoh::sample::Locality::Remote)
-                .callback({
-                    let sender = to_topics.clone();
-                    let topic = topic.clone();
-                    move |sample| {
-                        if sample.kind() != SampleKind::Put {
-                            return;
-                        }
-                        _ = sender.try_send(FromSwarm::Message {
-                            topic: topic.clone(),
-                            data: sample.payload().to_bytes().to_vec(),
-                        });
-                    }
-                })
-                .await;
-            let subscriber = match subscriber_res {
-                Ok(s) => s,
-                Err(e) => {
-                    _ = result_sender.send(Err(e));
-                    return;
-                }
-            };
-
-            assert!(topics.insert(topic, (subscriber, publisher)).is_none());
-            _ = result_sender.send(Ok(true));
+            let result = swarm
+                .behaviour_mut()
+                .gossipsub
+                .publish(gossipsub::IdentTopic::new(topic), data);
+            _ = result_sender.send(result);
         }
     }
 }
 
-pub async fn create_swarm(
-    identity: &str,
-    from_client: mpsc::Receiver<ToSwarm>,
-    listen_port: u16,
-    discovery_service_port: u16,
-) -> Result<Swarm> {
-    let cfg = crate::cfg(identity, listen_port)?;
-    let session = crate::open(cfg, listen_port, discovery_service_port).await?;
-    Ok(Swarm {
-        session,
-        from_client,
-    })
+fn filter_swarm_event(event: SwarmEvent<BehaviourEvent>) -> Option<FromSwarm> {
+    match event {
+        SwarmEvent::Behaviour(BehaviourEvent::Gossipsub(gossipsub::Event::Message {
+            message:
+                gossipsub::Message {
+                    source: Some(peer_id),
+                    topic,
+                    data,
+                    ..
+                },
+            ..
+        })) => Some(FromSwarm::Message {
+            from: peer_id,
+            topic: topic.into_string(),
+            data,
+        }),
+        SwarmEvent::Behaviour(BehaviourEvent::Discovery(
+            discovery::Event::ConnectionEstablished { peer_id, .. },
+        )) => Some(FromSwarm::Discovered { peer_id }),
+        SwarmEvent::Behaviour(BehaviourEvent::Discovery(discovery::Event::ConnectionClosed {
+            peer_id,
+            ..
+        })) => Some(FromSwarm::Expired { peer_id }),
+        _ => None,
+    }
+}
+
+/// Create and configure a swarm.
+///
+/// - `listen_port`: TCP port to listen on. `0` lets the OS assign one.
+/// - `bootstrap_peers`: multiaddrs to dial for environments without mDNS.
+pub fn create_swarm(
+    keypair: identity::Keypair,
+    from_client: mpsc::Receiver<ToSwarm>,
+    bootstrap_peers: Vec<String>,
+    listen_port: u16,
+) -> alias::AnyResult<Swarm> {
+    let parsed_bootstrap_peers: Vec<libp2p::Multiaddr> = bootstrap_peers
+        .iter()
+        .filter(|s| !s.is_empty())
+        .filter_map(|s| s.parse().ok())
+        .collect();
+
+    let mut swarm = SwarmBuilder::with_existing_identity(keypair)
+        .with_tokio()
+        .with_other_transport(tcp_transport)?
+        .with_behaviour(|keypair| Behaviour::new(keypair, parsed_bootstrap_peers))?
+        .build();
+
+    swarm.listen_on(format!("/ip4/0.0.0.0/tcp/{listen_port}").parse()?)?;
+    Ok(Swarm { swarm, from_client })
+}
+
+mod transport {
+    use crate::alias;
+    use crate::swarm::{NETWORK_VERSION, OVERRIDE_VERSION_ENV_VAR};
+    use futures_lite::{AsyncRead, AsyncWrite};
+    use keccak_const::Sha3_256;
+    use libp2p::core::muxing;
+    use libp2p::core::transport::Boxed;
+    use libp2p::pnet::{PnetError, PnetOutput};
+    use libp2p::{PeerId, Transport, identity, noise, pnet, yamux};
+    use std::{env, sync::LazyLock};
+
+    /// Key used for networking's private network; parametrized on the [`NETWORK_VERSION`].
+    /// See [`pnet_upgrade`] for more.
+    static PNET_PRESHARED_KEY: LazyLock<[u8; 32]> = LazyLock::new(|| {
+        let builder = Sha3_256::new().update(b"exo_discovery_network");
+
+        if let Ok(var) = env::var(OVERRIDE_VERSION_ENV_VAR) {
+            let bytes = var.into_bytes();
+            builder.update(&bytes)
+        } else {
+            builder.update(NETWORK_VERSION)
+        }
+        .finalize()
+    });
+
+    /// Make the Swarm run on a private network, as to not clash with public libp2p nodes and
+    /// also different-versioned instances of this same network.
+    /// This is implemented as an additional "upgrade" ontop of existing [`libp2p::Transport`] layers.
+    async fn pnet_upgrade<TSocket>(
+        socket: TSocket,
+        _: impl Sized,
+    ) -> Result<PnetOutput<TSocket>, PnetError>
+    where
+        TSocket: AsyncRead + AsyncWrite + Send + Unpin + 'static,
+    {
+        use pnet::{PnetConfig, PreSharedKey};
+        PnetConfig::new(PreSharedKey::new(*PNET_PRESHARED_KEY))
+            .handshake(socket)
+            .await
+    }
+
+    /// TCP/IP transport layer configuration.
+    pub fn tcp_transport(
+        keypair: &identity::Keypair,
+    ) -> alias::AnyResult<Boxed<(PeerId, muxing::StreamMuxerBox)>> {
+        use libp2p::{
+            core::upgrade::Version,
+            tcp::{Config, tokio},
+        };
+
+        // `TCP_NODELAY` enabled => avoid latency
+        let tcp_config = Config::default().nodelay(true);
+
+        // V1 + lazy flushing => 0-RTT negotiation
+        let upgrade_version = Version::V1Lazy;
+
+        // Noise is faster than TLS + we don't care much for security
+        let noise_config = noise::Config::new(keypair)?;
+
+        // Use default Yamux config for multiplexing
+        let yamux_config = yamux::Config::default();
+
+        // Create new Tokio-driven TCP/IP transport layer
+        let base_transport = tokio::Transport::new(tcp_config)
+            .and_then(pnet_upgrade)
+            .upgrade(upgrade_version)
+            .authenticate(noise_config)
+            .multiplex(yamux_config);
+
+        // Return boxed transport (to flatten complex type)
+        Ok(base_transport.boxed())
+    }
+}
+
+mod behaviour {
+    use crate::{alias, discovery};
+    use libp2p::swarm::NetworkBehaviour;
+    use libp2p::{gossipsub, identity};
+
+    /// Behavior of the Swarm which composes all desired behaviors:
+    /// Right now its just [`discovery::Behaviour`] and [`gossipsub::Behaviour`].
+    #[derive(NetworkBehaviour)]
+    pub struct Behaviour {
+        pub discovery: discovery::Behaviour,
+        pub gossipsub: gossipsub::Behaviour,
+    }
+
+    impl Behaviour {
+        pub fn new(
+            keypair: &identity::Keypair,
+            bootstrap_peers: Vec<libp2p::Multiaddr>,
+        ) -> alias::AnyResult<Self> {
+            Ok(Self {
+                discovery: discovery::Behaviour::new(keypair, bootstrap_peers)?,
+                gossipsub: gossipsub_behaviour(keypair),
+            })
+        }
+    }
+
+    fn gossipsub_behaviour(keypair: &identity::Keypair) -> gossipsub::Behaviour {
+        use gossipsub::{ConfigBuilder, MessageAuthenticity, ValidationMode};
+
+        // build a gossipsub network behaviour
+        //  => signed message authenticity + strict validation mode means the message-ID is
+        //     automatically provided by gossipsub w/out needing to provide custom message-ID function
+        gossipsub::Behaviour::new(
+            MessageAuthenticity::Signed(keypair.clone()),
+            ConfigBuilder::default()
+                .max_transmit_size(8 * 1024 * 1024)
+                .validation_mode(ValidationMode::Strict)
+                .build()
+                .expect("the configuration should always be valid"),
+        )
+        .expect("creating gossipsub behavior should always work")
+    }
 }

rust/networking/tests/bootstrap_peers.rs

@@ -0,0 +1,107 @@
+use futures_lite::StreamExt;
+use networking::swarm::{FromSwarm, create_swarm};
+use std::time::Duration;
+use tokio::sync::mpsc;
+use tokio::time::timeout;
+
+/// Helper: find a free TCP port.
+fn free_port() -> u16 {
+    let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap();
+    listener.local_addr().unwrap().port()
+}
+
+/// Two nodes connect via bootstrap peers — no mDNS needed.
+///
+/// Node A listens on a fixed port. Node B bootstraps to A's address.
+/// We verify that B emits `FromSwarm::Discovered` for A's peer ID.
+#[tokio::test]
+async fn two_nodes_connect_via_bootstrap_peers() {
+    let port_a = free_port();
+
+    // Node A: listens on a known port, no bootstrap peers
+    let keypair_a = libp2p::identity::Keypair::generate_ed25519();
+    let peer_id_a = keypair_a.public().to_peer_id();
+    let (_tx_a, rx_a) = mpsc::channel(16);
+    let swarm_a = create_swarm(keypair_a, rx_a, vec![], port_a).expect("create swarm A");
+    let mut stream_a = swarm_a.into_stream();
+
+    // Node B: bootstraps to A's address
+    let keypair_b = libp2p::identity::Keypair::generate_ed25519();
+    let (_tx_b, rx_b) = mpsc::channel(16);
+    let swarm_b = create_swarm(
+        keypair_b,
+        rx_b,
+        vec![format!("/ip4/127.0.0.1/tcp/{port_a}")],
+        0,
+    )
+    .expect("create swarm B");
+    let mut stream_b = swarm_b.into_stream();
+
+    // Wait for B to discover A (connection established)
+    let connected = timeout(Duration::from_secs(10), async {
+        loop {
+            tokio::select! {
+                Some(event) = stream_a.next() => {
+                    // A will also see B connect, but we check from B's perspective
+                    let _ = event;
+                }
+                Some(event) = stream_b.next() => {
+                    if let FromSwarm::Discovered { peer_id } = event {
+                        if peer_id == peer_id_a {
+                            return true;
+                        }
+                    }
+                }
+            }
+        }
+    })
+    .await;
+
+    assert!(
+        connected.is_ok() && connected.unwrap(),
+        "Node B should discover Node A via bootstrap peer"
+    );
+}
+
+/// Empty bootstrap peers should work (backward compatible).
+#[tokio::test]
+async fn create_swarm_with_empty_bootstrap_peers() {
+    let keypair = libp2p::identity::Keypair::generate_ed25519();
+    let (_tx, rx) = mpsc::channel(16);
+    let swarm = create_swarm(keypair, rx, vec![], 0);
+    assert!(
+        swarm.is_ok(),
+        "create_swarm with no bootstrap peers should succeed"
+    );
+}
+
+/// Invalid multiaddr strings are silently filtered out.
+#[tokio::test]
+async fn create_swarm_ignores_invalid_bootstrap_addrs() {
+    let keypair = libp2p::identity::Keypair::generate_ed25519();
+    let (_tx, rx) = mpsc::channel(16);
+    let swarm = create_swarm(
+        keypair,
+        rx,
+        vec![
+            "not-a-valid-multiaddr".to_string(),
+            "".to_string(),
+            "/ip4/10.0.0.1/tcp/30000".to_string(), // valid
+        ],
+        0,
+    );
+    assert!(
+        swarm.is_ok(),
+        "create_swarm should succeed even with invalid bootstrap addrs"
+    );
+}
+
+/// Fixed listen port works correctly.
+#[tokio::test]
+async fn create_swarm_with_fixed_port() {
+    let port = free_port();
+    let keypair = libp2p::identity::Keypair::generate_ed25519();
+    let (_tx, rx) = mpsc::channel(16);
+    let swarm = create_swarm(keypair, rx, vec![], port);
+    assert!(swarm.is_ok(), "create_swarm with fixed port should succeed");
+}

rust/networking/tests/dummy.rs

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

rust/networking/tests/runtime_polling.rs

@@ -1,99 +0,0 @@
-use std::time::Duration;
-
-use zenoh::Wait;
-
-fn unique_key(name: &str) -> String {
-    format!("test/zenoh-runtime-polling/{}/{}", std::process::id(), name)
-}
-
-#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
-async fn get_and_recv_work_on_tokio_baseline() {
-    let session = zenoh::open(zenoh::Config::default())
-        .await
-        .expect("open session");
-
-    let key = unique_key("tokio-baseline");
-    let reply_key = key.clone();
-
-    let _queryable = session
-        .declare_queryable(key.clone())
-        .callback(move |query| {
-            query
-                .reply(reply_key.clone(), "hello-from-queryable")
-                .wait()
-                .expect("reply from queryable");
-        })
-        .await
-        .expect("declare queryable");
-
-    let replies = session.get(key).await.expect("get");
-
-    let reply = tokio::time::timeout(Duration::from_secs(5), replies.recv_async())
-        .await
-        .expect("timed out waiting for reply")
-        .expect("reply channel closed");
-
-    let sample = reply.result().expect("reply result was error");
-    let payload = sample
-        .payload()
-        .try_to_string()
-        .expect("payload should be utf8");
-
-    assert_eq!(payload.as_ref(), "hello-from-queryable");
-}
-
-#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
-async fn get_and_recv_work_when_polled_by_smol_without_tokio_context() {
-    let session = zenoh::open(zenoh::Config::default())
-        .await
-        .expect("open session under tokio");
-
-    let key = unique_key("smol-no-tokio-context");
-    let reply_key = key.clone();
-
-    let _queryable = session
-        .declare_queryable(key.clone())
-        .callback(move |query| {
-            query
-                .reply(reply_key.clone(), "hello-from-queryable")
-                .wait()
-                .expect("reply from queryable");
-        })
-        .await
-        .expect("declare queryable under tokio");
-
-    let session_for_smol = session.clone();
-
-    let (tx, rx) = std::sync::mpsc::channel();
-
-    std::thread::spawn(move || {
-        // This thread was not entered by Tokio.
-        // If Zenoh's get/recv path requires an ambient Tokio Handle in the polling thread,
-        // this is where it should panic, hang, or error.
-        let result = {
-            smol::block_on(async move {
-                let replies = session_for_smol.get(key).await.expect("get under smol");
-
-                let reply = replies.recv_async().await.expect("reply channel closed");
-
-                let sample = reply.result().expect("reply result was error");
-                let payload = sample
-                    .payload()
-                    .try_to_string()
-                    .expect("payload should be utf8");
-
-                payload.to_string()
-            })
-        };
-
-        tx.send(result).expect("send test result");
-    });
-
-    let result = rx
-        .recv_timeout(Duration::from_secs(5))
-        .expect("smol thread timed out; likely hung waiting for get/reply");
-
-    let payload = result;
-
-    assert_eq!(payload, "hello-from-queryable");
-}

rust/util/Cargo.toml

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

rust/util/src/lib.rs

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

rust/util/src/wakerdeque.rs

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

src/exo/api/load_balancing.py

@@ -1,200 +0,0 @@
-from pydantic import ValidationError
-
-from exo.master.placement_utils import find_ip_prioritised
-from exo.shared.models.model_cards import ModelId
-from exo.shared.types.commands import ImageEdits, ImageGeneration, TextGeneration
-from exo.shared.types.instance_link import InstanceLink
-from exo.shared.types.state import State
-from exo.shared.types.tasks import (
-    ImageEdits as ImageEditsTask,
-)
-from exo.shared.types.tasks import (
-    ImageGeneration as ImageGenerationTask,
-)
-from exo.shared.types.tasks import (
-    Task,
-    TaskId,
-    TaskStatus,
-)
-from exo.shared.types.tasks import (
-    TextGeneration as TextGenerationTask,
-)
-from exo.shared.types.worker.instances import InstanceId
-
-
-def decode_instance_for_text_generation(
-    state: State, model_id: ModelId, instance_links: list[InstanceLink]
-) -> InstanceId | None:
-    prefill_only: set[InstanceId] = set()
-    for link in instance_links:
-        prefill_only.update(link.prefill_instances)
-    for link in instance_links:
-        prefill_only.difference_update(link.decode_instances)
-
-    instance_task_counts = _instance_task_counts_for_model(state, model_id)
-    for instance_id in prefill_only:
-        instance_task_counts.pop(instance_id, None)
-
-    if not instance_task_counts:
-        return None
-
-    return min(
-        instance_task_counts, key=lambda instance_id: instance_task_counts[instance_id]
-    )
-
-
-def instance_for_generation(state: State, model_id: ModelId) -> InstanceId | None:
-    instance_task_counts = _instance_task_counts_for_model(state, model_id)
-    if not instance_task_counts:
-        return None
-
-    return min(
-        instance_task_counts, key=lambda instance_id: instance_task_counts[instance_id]
-    )
-
-
-def text_generation_task(
-    state: State, command: TextGeneration, instance_links: list[InstanceLink]
-) -> TextGenerationTask:
-    instance_id = decode_instance_for_text_generation(
-        state, command.task_params.model, instance_links
-    )
-    assert instance_id is not None
-    task_params = command.task_params.model_copy(
-        update={
-            "prefill_endpoint": prefill_endpoint_for(
-                state,
-                instance_links,
-                instance_id,
-            ),
-        }
-    )
-    return TextGenerationTask(
-        task_id=TaskId(),
-        command_id=command.command_id,
-        instance_id=instance_id,
-        task_status=TaskStatus.Pending,
-        task_params=task_params,
-    )
-
-
-def image_generation_task(
-    state: State, command: ImageGeneration
-) -> ImageGenerationTask:
-    instance_id = instance_for_generation(state, ModelId(command.task_params.model))
-    assert instance_id is not None
-    return ImageGenerationTask(
-        task_id=TaskId(),
-        command_id=command.command_id,
-        instance_id=instance_id,
-        task_status=TaskStatus.Pending,
-        task_params=command.task_params,
-    )
-
-
-def image_edits_task(state: State, command: ImageEdits) -> ImageEditsTask:
-    instance_id = instance_for_generation(state, ModelId(command.task_params.model))
-    assert instance_id is not None
-    return ImageEditsTask(
-        task_id=TaskId(),
-        command_id=command.command_id,
-        instance_id=instance_id,
-        task_status=TaskStatus.Pending,
-        task_params=command.task_params,
-    )
-
-
-def task_from_command(
-    state: State,
-    command: TextGeneration | ImageGeneration | ImageEdits,
-    instance_links: list[InstanceLink],
-) -> Task:
-    match command:
-        case TextGeneration():
-            return text_generation_task(state, command, instance_links)
-        case ImageGeneration():
-            return image_generation_task(state, command)
-        case ImageEdits():
-            return image_edits_task(state, command)
-
-
-def instance_id_for_command(
-    state: State,
-    command: TextGeneration | ImageGeneration | ImageEdits,
-    instance_links: list[InstanceLink],
-) -> InstanceId | None:
-    match command:
-        case TextGeneration():
-            return decode_instance_for_text_generation(
-                state, command.task_params.model, instance_links
-            )
-        case ImageGeneration() | ImageEdits():
-            return instance_for_generation(state, ModelId(command.task_params.model))
-
-
-def load_instance_links(values: list[str]) -> list[InstanceLink]:
-    instance_links: list[InstanceLink] = []
-    for value in values:
-        try:
-            instance_links.append(InstanceLink.model_validate_json(value))
-        except ValidationError:
-            continue
-    return instance_links
-
-
-def prefill_endpoint_for(
-    state: State, instance_links: list[InstanceLink], decode_instance_id: InstanceId
-) -> str | None:
-    decode = state.instances.get(decode_instance_id)
-    if decode is None:
-        return None
-    decode_node = decode.shard_assignments.shards[
-        decode.shard_assignments.primary_output_node
-    ].node_id
-
-    sources: set[InstanceId] = set()
-    for link in instance_links:
-        if decode_instance_id in link.decode_instances:
-            sources.update(link.prefill_instances)
-    sources.discard(decode_instance_id)
-
-    in_flight = {TaskStatus.Pending, TaskStatus.Running}
-    task_counts: dict[InstanceId, int] = {
-        src_id: sum(
-            1
-            for task in state.tasks.values()
-            if task.instance_id == src_id and task.task_status in in_flight
-        )
-        for src_id in sources
-    }
-    for src_id in sorted(sources, key=lambda sid: task_counts[sid]):
-        instance = state.instances.get(src_id)
-        if instance is None:
-            continue
-        for node_id, runner_id, _ in instance.shard_assignments.shards:
-            port = state.prefill_server_ports.get(runner_id)
-            if port is None:
-                continue
-            ip = find_ip_prioritised(
-                decode_node, node_id, state.topology, state.node_network, ring=True
-            )
-            if ip is None:
-                continue
-            return f"{ip}:{port}"
-    return None
-
-
-def _instance_task_counts_for_model(
-    state: State, model_id: ModelId
-) -> dict[InstanceId, int]:
-    in_flight = {TaskStatus.Pending, TaskStatus.Running}
-    return {
-        instance.instance_id: sum(
-            1
-            for task in state.tasks.values()
-            if task.instance_id == instance.instance_id
-            and task.task_status in in_flight
-        )
-        for instance in state.instances.values()
-        if instance.shard_assignments.model_id == model_id
-    }

src/exo/api/tests/test_cancel_command.py

@@ -7,7 +7,6 @@ from fastapi.testclient import TestClient
 
 from exo.api.main import API
 from exo.shared.types.common import CommandId
-from exo.shared.types.worker.instances import InstanceId
 
 
 def _make_api() -> Any:
@@ -16,9 +15,9 @@ def _make_api() -> Any:
     app = FastAPI()
     api = object.__new__(API)
     api.app = app
-    api._bridge_command_instances = {}  # pyright: ignore[reportPrivateUsage]
-    api.task_requester = MagicMock()
-    api.task_requester.interrupt = AsyncMock()
+    api._text_generation_queues = {}  # pyright: ignore[reportPrivateUsage]
+    api._image_generation_queues = {}  # pyright: ignore[reportPrivateUsage]
+    api._send = AsyncMock()  # pyright: ignore[reportPrivateUsage]
     api._setup_exception_handlers()  # pyright: ignore[reportPrivateUsage]
     app.post("/v1/cancel/{command_id}")(api.cancel_command)
     return api
@@ -39,42 +38,40 @@ def test_cancel_nonexistent_command_returns_404() -> None:
 
 
 def test_cancel_active_text_generation() -> None:
-    """Cancel an active text generation command: returns 200, interrupt sent."""
+    """Cancel an active text generation command: returns 200, sender.close() called."""
     api = _make_api()
     client = TestClient(api.app)
 
     cid = CommandId("text-cmd-123")
-    instance_id = InstanceId("instance-a")
-    api._bridge_command_instances[cid] = instance_id
+    sender = MagicMock()
+    api._text_generation_queues[cid] = sender
 
     response = client.post(f"/v1/cancel/{cid}")
     assert response.status_code == 200
     data: dict[str, Any] = response.json()
     assert data["message"] == "Command cancelled."
     assert data["command_id"] == str(cid)
-    api.task_requester.interrupt.assert_called_once()
-    args = api.task_requester.interrupt.call_args.args
-    assert args[0] == instance_id
-    assert args[1] == cid
-    assert '"cancelled_command_id":"text-cmd-123"' in args[2]
+    sender.close.assert_called_once()
+    api._send.assert_called_once()
+    task_cancelled = api._send.call_args[0][0]
+    assert task_cancelled.cancelled_command_id == cid
 
 
 def test_cancel_active_image_generation() -> None:
-    """Cancel an active image generation command: returns 200, interrupt sent."""
+    """Cancel an active image generation command: returns 200, sender.close() called."""
     api = _make_api()
     client = TestClient(api.app)
 
     cid = CommandId("img-cmd-456")
-    instance_id = InstanceId("instance-b")
-    api._bridge_command_instances[cid] = instance_id
+    sender = MagicMock()
+    api._image_generation_queues[cid] = sender
 
     response = client.post(f"/v1/cancel/{cid}")
     assert response.status_code == 200
     data: dict[str, Any] = response.json()
     assert data["message"] == "Command cancelled."
     assert data["command_id"] == str(cid)
-    api.task_requester.interrupt.assert_called_once()
-    args = api.task_requester.interrupt.call_args.args
-    assert args[0] == instance_id
-    assert args[1] == cid
-    assert '"cancelled_command_id":"img-cmd-456"' in args[2]
+    sender.close.assert_called_once()
+    api._send.assert_called_once()
+    task_cancelled = api._send.call_args[0][0]
+    assert task_cancelled.cancelled_command_id == cid

src/exo/api/types/api.py

@@ -329,6 +329,7 @@ class InstanceLinkBody(BaseModel):
 
 class InstanceLinkResponse(BaseModel):
     message: str
+    command_id: CommandId
 
 
 ImageSize = Literal[
@@ -393,6 +394,7 @@ class ImageEditsTaskParams(BaseModel):
     """Internal task params for image-editing requests."""
 
     image_data: str = ""  # Base64-encoded image (empty when using chunked transfer)
+    total_input_chunks: int = 0
     prompt: str
     model: str
     n: int | None = 1

src/exo/main.py

@@ -10,7 +10,7 @@ from typing import Self
 import anyio
 from anyio.lowlevel import checkpoint as anyio_checkpoint
 from daemon import DaemonContext  # pyright: ignore[reportMissingTypeStubs]
-from exo_rs import Pidfile, PidfileError, SessionHandle
+from exo_rs import Pidfile, PidfileError
 from loguru import logger
 from pydantic import PositiveInt
 
@@ -20,7 +20,7 @@ from exo.download.coordinator import DownloadCoordinator
 from exo.download.impl_shard_downloader import exo_shard_downloader
 from exo.master.main import Master
 from exo.routing.event_router import EventRouter
-from exo.routing.router import Router
+from exo.routing.router import Router, get_node_id_keypair
 from exo.shared.constants import EXO_DEFAULT_MODELS_DIR, EXO_LOG, EXO_PID_FILE
 from exo.shared.election import Election, ElectionResult
 from exo.shared.logging import logger_cleanup, logger_setup
@@ -46,17 +46,18 @@ class Node:
     node_id: NodeId
     offline: bool
     _api_port: int
-    _sh: SessionHandle
     _tg: TaskGroup = field(init=False, default_factory=TaskGroup)
 
     @classmethod
     async def create(cls, args: "Args") -> Self:
-        identity = os.urandom(16).hex().lstrip("0")
-        node_id = NodeId(identity)
+        keypair = get_node_id_keypair()
+        node_id = NodeId(keypair.to_node_id())
         session_id = SessionId(master_node_id=node_id, election_clock=0)
-        session_handle, _nh = SessionHandle.new(identity, args.zenoh_port, 52413)
-        router = Router(_nh)
-
+        router = Router.create(
+            keypair,
+            bootstrap_peers=args.bootstrap_peers,
+            listen_port=args.libp2p_port,
+        )
         await router.register_topic(topics.GLOBAL_EVENTS)
         await router.register_topic(topics.LOCAL_EVENTS)
         await router.register_topic(topics.COMMANDS)
@@ -95,7 +96,6 @@ class Node:
                 command_sender=router.sender(topics.COMMANDS),
                 download_command_sender=router.sender(topics.DOWNLOAD_COMMANDS),
                 election_receiver=router.receiver(topics.ELECTION_MESSAGES),
-                session_handle=session_handle,
             )
         else:
             api = None
@@ -107,7 +107,6 @@ class Node:
                 event_sender=event_router.sender(),
                 command_sender=router.sender(topics.COMMANDS),
                 download_command_sender=router.sender(topics.DOWNLOAD_COMMANDS),
-                session_handle=session_handle,
                 api_port=args.api_port,
             )
         else:
@@ -122,8 +121,6 @@ class Node:
             local_event_receiver=router.receiver(topics.LOCAL_EVENTS),
             command_receiver=router.receiver(topics.COMMANDS),
             download_command_sender=router.sender(topics.DOWNLOAD_COMMANDS),
-            aggregator=session_handle.last_value_aggregator("node_metrics"),
-            storage=session_handle.storage_interface(),
         )
 
         er_send, er_recv = channel[ElectionResult]()
@@ -152,7 +149,6 @@ class Node:
             node_id,
             args.offline,
             args.api_port,
-            session_handle,
         )
 
     async def run(self):
@@ -228,8 +224,6 @@ class Node:
                         download_command_sender=self.router.sender(
                             topics.DOWNLOAD_COMMANDS
                         ),
-                        aggregator=self._sh.last_value_aggregator("node_metrics"),
-                        storage=self._sh.storage_interface(),
                     )
                     self._tg.start_soon(self.master.run)
                 elif (
@@ -269,7 +263,6 @@ class Node:
                             download_command_sender=self.router.sender(
                                 topics.DOWNLOAD_COMMANDS
                             ),
-                            session_handle=self._sh,
                             api_port=self._api_port,
                         )
                         self._tg.start_soon(self.worker.run)
@@ -345,18 +338,16 @@ def main_inner(args: "Args"):
     # TODO: Refactor the current verbosity system
     logger_setup(EXO_LOG, args.verbosity)
 
-    logger.info(f"pid = {os.getpid()}")
-    if os.getenv("EXO_LIBP2P_NAMESPACE"):
-        raise ValueError(
-            "EXO_LIBP2P_NAMESPACE has been removed - use EXO_ZENOH_NAMESPACE instead"
-        )
-    logger.info(f"EXO_ZENOH_NAMESPACE: {os.getenv('EXO_ZENOH_NAMESPACE')}")
+    logger.info(f"{'=' * 40}")
+    logger.info(f"Starting EXO | pid={os.getpid()}")
+    logger.info(f"{'=' * 40}")
+    logger.info(f"EXO_LIBP2P_NAMESPACE: {os.getenv('EXO_LIBP2P_NAMESPACE')}")
 
     if args.offline:
         logger.info("Running in OFFLINE mode — no internet checks, local models only")
 
     if args.bootstrap_peers:
-        raise ValueError("Bootstrap peers has been temporarily removed")
+        logger.info(f"Bootstrap peers: {args.bootstrap_peers}")
 
     if args.no_batch:
         os.environ["EXO_NO_BATCH"] = "1"
@@ -384,18 +375,19 @@ def main_inner(args: "Args"):
 
 
 class Args(FrozenModel):
-    verbosity: int
-    force_master: bool
-    spawn_api: bool
-    api_port: PositiveInt
+    verbosity: int = 0
+    force_master: bool = False
+    spawn_api: bool = False
+    api_port: PositiveInt = 52415
+    tb_only: bool = False
     no_worker: bool = False
     no_downloads: bool = False
-    offline: bool
-    no_batch: bool
+    offline: bool = os.getenv("EXO_OFFLINE", "false").lower() == "true"
+    no_batch: bool = False
     fast_synch: bool | None = None  # None = auto, True = force on, False = force off
-    legacy_daemon: bool
+    legacy_daemon: bool = False
     bootstrap_peers: list[str] = []
-    zenoh_port: int
+    libp2p_port: int
 
     @classmethod
     def parse(cls) -> Self:
@@ -468,11 +460,11 @@ class Args(FrozenModel):
             help="Comma-separated libp2p multiaddrs to dial on startup (env: EXO_BOOTSTRAP_PEERS)",
         )
         parser.add_argument(
-            "--zenoh-port",
+            "--libp2p-port",
             type=int,
-            default=52414,
-            dest="zenoh_port",
-            help="Fixed port for zenoh to listen on.",
+            default=0,
+            dest="libp2p_port",
+            help="Fixed TCP port for libp2p to listen on (0 = OS-assigned).",
         )
         fast_synch_group = parser.add_mutually_exclusive_group()
         fast_synch_group.add_argument(

src/exo/master/main.py

@@ -1,9 +1,7 @@
 from datetime import datetime, timedelta, timezone
 
 import anyio
-from exo_rs import LVAggregator, Storage
 from loguru import logger
-from pydantic import ValidationError
 
 from exo.master.placement import (
     add_instance_to_placements,
@@ -20,14 +18,19 @@ from exo.routing.event_router import (
 from exo.shared.apply import apply
 from exo.shared.constants import EXO_EVENT_LOG_DIR, EXO_TRACING_ENABLED
 from exo.shared.types.commands import (
+    AddCustomModelCard,
     CreateInstance,
+    DeleteCustomModelCard,
     DeleteInstance,
+    DeleteInstanceLink,
     ForwarderCommand,
     ForwarderDownloadCommand,
     ImageEdits,
     ImageGeneration,
     PlaceInstance,
     RequestEventLog,
+    SendInputChunk,
+    SetInstanceLink,
     TaskCancelled,
     TaskFinished,
     TestCommand,
@@ -35,10 +38,15 @@ from exo.shared.types.commands import (
 )
 from exo.shared.types.common import CommandId, NodeId, SessionId, SystemId
 from exo.shared.types.events import (
+    CustomModelCardAdded,
+    CustomModelCardDeleted,
     Event,
     GlobalForwarderEvent,
     IndexedEvent,
+    InputChunkReceived,
     InstanceDeleted,
+    InstanceLinkCreated,
+    InstanceLinkDeleted,
     LocalForwarderEvent,
     NodeGatheredInfo,
     NodeTimedOut,
@@ -71,18 +79,16 @@ from exo.utils.event_buffer import MultiSourceBuffer
 from exo.utils.task_group import TaskGroup
 
 
-def _prefill_endpoint_for(
-    state: State, instance_links: list[InstanceLink], decode_instance_id: InstanceId
-) -> str | None:
+def _prefill_endpoint_for(state: State, decode_instance_id: InstanceId) -> str | None:
     decode = state.instances.get(decode_instance_id)
     if decode is None:
         return None
-    decode_node = decode.shard_assignments.shards[
-        decode.shard_assignments.primary_output_node
-    ].node_id
+    decode_node = next(iter(decode.shard_assignments.node_to_runner.keys()), None)
+    if decode_node is None:
+        return None
 
     sources: set[InstanceId] = set()
-    for link in instance_links:
+    for link in state.instance_links.values():
         if decode_instance_id in link.decode_instances:
             sources.update(link.prefill_instances)
     sources.discard(decode_instance_id)
@@ -100,7 +106,7 @@ def _prefill_endpoint_for(
         instance = state.instances.get(src_id)
         if instance is None:
             continue
-        for node_id, runner_id, _ in instance.shard_assignments.shards:
+        for node_id, runner_id in instance.shard_assignments.node_to_runner.items():
             port = state.prefill_server_ports.get(runner_id)
             if port is None:
                 continue
@@ -124,8 +130,6 @@ class Master:
         local_event_receiver: Receiver[LocalForwarderEvent],
         global_event_sender: Sender[GlobalForwarderEvent],
         download_command_sender: Sender[ForwarderDownloadCommand],
-        aggregator: LVAggregator,
-        storage: Storage,
     ):
         self.node_id = node_id
         self.session_id = session_id
@@ -141,9 +145,7 @@ class Master:
         self._multi_buffer = MultiSourceBuffer[SystemId, Event]()
         self._event_log = DiskEventLog(EXO_EVENT_LOG_DIR / "master")
         self._pending_traces: dict[TaskId, dict[int, list[TraceEventData]]] = {}
-        self._world_sizes: dict[TaskId, int] = {}
-        self.aggregator = aggregator
-        self.storage = storage
+        self._expected_ranks: dict[TaskId, set[int]] = {}
 
     async def run(self):
         logger.info("Starting Master")
@@ -180,21 +182,10 @@ class Master:
                             pass
                         case TextGeneration():
                             # set-difference => prefill-only nodes
-                            instance_links: list[InstanceLink] = []
                             prefill_only: set[InstanceId] = set()
-                            for _, link in (
-                                await self.storage.dump("instance_links/")
-                            ).items():
-                                try:
-                                    instance_links.append(
-                                        InstanceLink.model_validate_json(link)
-                                    )
-                                except ValidationError:
-                                    continue
-
-                            for link in instance_links:
+                            for link in self.state.instance_links.values():
                                 prefill_only.update(link.prefill_instances)
-                            for link in instance_links:
+                            for link in self.state.instance_links.values():
                                 prefill_only.difference_update(link.decode_instances)
 
                             for instance in self.state.instances.values():
@@ -234,9 +225,7 @@ class Master:
                             params = command.task_params.model_copy(
                                 update={
                                     "prefill_endpoint": _prefill_endpoint_for(
-                                        self.state.with_aggregator(self.aggregator),
-                                        instance_links,
-                                        decode_instance_id,
+                                        self.state, decode_instance_id
                                     ),
                                 }
                             )
@@ -304,9 +293,11 @@ class Master:
                                     selected_instance_id
                                 )
                                 if selected_instance:
-                                    self._world_sizes[task_id] = len(
-                                        selected_instance.shard_assignments.shards
+                                    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 (
@@ -358,9 +349,11 @@ class Master:
                                     selected_instance_id
                                 )
                                 if selected_instance:
-                                    self._world_sizes[task_id] = len(
-                                        selected_instance.shard_assignments.shards
+                                    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(
@@ -376,16 +369,15 @@ class Master:
                                 )
                             generated_events.extend(transition_events)
                         case PlaceInstance():
-                            state = self.state.with_aggregator(self.aggregator)
                             placement = place_instance(
                                 command,
-                                state.topology,
-                                state.instances,
-                                state.node_memory,
-                                state.node_network,
-                                state.node_backends,
-                                download_status=state.downloads,
-                                node_rdma_ctl=state.node_rdma_ctl,
+                                self.state.topology,
+                                self.state.instances,
+                                self.state.node_memory,
+                                self.state.node_network,
+                                self.state.node_backends,
+                                download_status=self.state.downloads,
+                                node_rdma_ctl=self.state.node_rdma_ctl,
                             )
                             transition_events = get_transition_events(
                                 self.state.instances, placement, self.state.tasks
@@ -401,6 +393,13 @@ class Master:
                                 self.state.instances, placement, self.state.tasks
                             )
                             generated_events.extend(transition_events)
+                        case SendInputChunk(chunk=chunk):
+                            generated_events.append(
+                                InputChunkReceived(
+                                    command_id=chunk.command_id,
+                                    chunk=chunk,
+                                )
+                            )
                         case TaskCancelled():
                             if (
                                 task_id := self.command_task_mapping.get(
@@ -429,6 +428,29 @@ class Master:
                                     f"Finished command {command.finished_command_id} finished"
                                 )
 
+                        case AddCustomModelCard():
+                            generated_events.append(
+                                CustomModelCardAdded(model_card=command.model_card)
+                            )
+                        case DeleteCustomModelCard():
+                            generated_events.append(
+                                CustomModelCardDeleted(model_id=command.model_id)
+                            )
+                        case SetInstanceLink():
+                            link = InstanceLink(
+                                link_id=command.link_id,
+                                prefill_instances=list(
+                                    dict.fromkeys(command.prefill_instances)
+                                ),
+                                decode_instances=list(
+                                    dict.fromkeys(command.decode_instances)
+                                ),
+                            )
+                            generated_events.append(InstanceLinkCreated(link=link))
+                        case DeleteInstanceLink():
+                            generated_events.append(
+                                InstanceLinkDeleted(link_id=command.link_id)
+                            )
                         case RequestEventLog():
                             # We should just be able to send everything, since other buffers will ignore old messages
                             # rate limit to 1000 at a time
@@ -442,18 +464,16 @@ class Master:
                                 )
                     for event in generated_events:
                         await self.event_sender.send(event)
-                except Exception as e:
+                except ValueError as e:
                     logger.opt(exception=e).warning("Error in command processor")
 
     # These plan loops are the cracks showing in our event sourcing architecture - more things could be commands
     async def _plan(self) -> None:
         while True:
             # kill broken instances
-            connected_node_ids = set(
-                self.state.with_aggregator(self.aggregator).topology.list_nodes()
-            )
+            connected_node_ids = set(self.state.topology.list_nodes())
             for instance_id, instance in self.state.instances.items():
-                for node_id, _, _ in instance.shard_assignments.shards:
+                for node_id in instance.shard_assignments.node_to_runner:
                     if node_id not in connected_node_ids:
                         await self.event_sender.send(
                             InstanceDeleted(instance_id=instance_id)
@@ -461,9 +481,7 @@ class Master:
                         break
 
             # time out dead nodes
-            for node_id, time in self.state.with_aggregator(
-                self.aggregator
-            ).last_seen.items():
+            for node_id, time in self.state.last_seen.items():
                 now = datetime.now(tz=timezone.utc)
                 if now - time > timedelta(seconds=30):
                     logger.info(f"Manually removing node {node_id} due to inactivity")
@@ -522,8 +540,9 @@ class Master:
         self._pending_traces[task_id][event.rank] = event.traces
 
         if (
-            task_id in self._world_sizes
-            and len(self._pending_traces[task_id]) >= self._world_sizes[task_id]
+            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)
 
@@ -537,5 +556,5 @@ class Master:
         )
 
         del self._pending_traces[task_id]
-        if task_id in self._world_sizes:
-            del self._world_sizes[task_id]
+        if task_id in self._expected_ranks:
+            del self._expected_ranks[task_id]

src/exo/master/placement.py

@@ -262,7 +262,20 @@ def place_instance(
 
     match command.instance_meta:
         case InstanceMeta.MlxJaccl:
-            coordinator_node_id = shard_assignments.shards[0].node_id
+            # TODO(evan): shard assignments should contain information about ranks, this is ugly
+            def get_device_rank(node_id: NodeId) -> int:
+                runner_id = shard_assignments.node_to_runner[node_id]
+                shard_metadata = shard_assignments.runner_to_shard.get(runner_id)
+                assert shard_metadata is not None
+                return shard_metadata.device_rank
+
+            zero_node_ids = [
+                node_id
+                for node_id in selected_cycle.node_ids
+                if get_device_rank(node_id) == 0
+            ]
+            assert len(zero_node_ids) == 1
+            coordinator_node_id = zero_node_ids[0]
 
             mlx_jaccl_devices = get_mlx_jaccl_devices_matrix(
                 [node_id for node_id in selected_cycle],
@@ -363,10 +376,10 @@ def cancel_unnecessary_downloads(
     active_models = set(
         (
             node_id,
-            instance.shard_assignments.model_id,
+            instance.shard_assignments.runner_to_shard[runner_id].model_card.model_id,
         )
         for instance in instances.values()
-        for node_id, _, _ in instance.shard_assignments.shards
+        for node_id, runner_id in instance.shard_assignments.node_to_runner.items()
     )
     for pair in currently_downloading:
         if pair not in active_models:

src/exo/master/placement_utils.py

@@ -8,11 +8,12 @@ from exo.shared.types.common import Host, NodeId
 from exo.shared.types.memory import Memory
 from exo.shared.types.profiling import MemoryUsage, NodeNetworkInfo
 from exo.shared.types.topology import Cycle, RDMAConnection, SocketConnection
-from exo.shared.types.worker.runners import RunnerId, ShardAssignments, ShardWithId
+from exo.shared.types.worker.runners import RunnerId, ShardAssignments
 from exo.shared.types.worker.shards import (
     CfgShardMetadata,
     PipelineShardMetadata,
     Sharding,
+    ShardMetadata,
     TensorShardMetadata,
 )
 
@@ -151,27 +152,27 @@ def _get_shard_assignments_for_cfg_parallel(
     _validate_cycle(cycle)
 
     world_size = len(cycle)
-    pipeline_world_size = world_size // 2
+    cfg_world_size = 2
+    pipeline_world_size = world_size // cfg_world_size
 
     # Allocate layers for one pipeline group (both groups run the same layers)
     pipeline_node_ids = cycle.node_ids[:pipeline_world_size]
     pipeline_memory = _compute_total_memory(pipeline_node_ids, node_memory)
-
-    # nb: only validates the forward path...
     layer_allocations = _allocate_and_validate_layers(
         pipeline_node_ids, node_memory, pipeline_memory, model_card
     )
 
     # Ring topology: group 0 ascending [0,1,2,...], group 1 descending [...,2,1,0]
     # This places both last stages as neighbors for CFG exchange.
-    position_to_cfg_pipeline = list(range(pipeline_world_size)) + list(
-        reversed(range(pipeline_world_size))
-    )
+    position_to_cfg_pipeline = [(0, r) for r in range(pipeline_world_size)] + [
+        (1, r) for r in reversed(range(pipeline_world_size))
+    ]
 
-    shards: list[ShardWithId] = []
+    runner_to_shard: dict[RunnerId, ShardMetadata] = {}
+    node_to_runner: dict[NodeId, RunnerId] = {}
 
     for device_rank, node_id in enumerate(cycle.node_ids):
-        pipeline_rank = position_to_cfg_pipeline[device_rank]
+        cfg_rank, pipeline_rank = position_to_cfg_pipeline[device_rank]
         layers_before = sum(layer_allocations[:pipeline_rank])
         node_layers = layer_allocations[pipeline_rank]
 
@@ -182,15 +183,20 @@ def _get_shard_assignments_for_cfg_parallel(
             start_layer=layers_before,
             end_layer=layers_before + node_layers,
             n_layers=model_card.n_layers,
+            cfg_rank=cfg_rank,
+            cfg_world_size=cfg_world_size,
+            pipeline_rank=pipeline_rank,
+            pipeline_world_size=pipeline_world_size,
         )
 
         runner_id = RunnerId()
-        shards.append(ShardWithId(node_id, runner_id, shard))
+        runner_to_shard[runner_id] = shard
+        node_to_runner[node_id] = runner_id
 
     return ShardAssignments(
         model_id=model_card.model_id,
-        shards=shards,
-        primary_output_node=pipeline_world_size - 1,
+        runner_to_shard=runner_to_shard,
+        node_to_runner=node_to_runner,
     )
 
 
@@ -202,13 +208,13 @@ def _get_shard_assignments_for_pure_pipeline(
     """Create shard assignments for pure pipeline execution."""
     _validate_cycle(cycle)
     total_memory = _compute_total_memory(cycle.node_ids, node_memory)
-    world_size = len(cycle)
 
     layer_allocations = _allocate_and_validate_layers(
         cycle.node_ids, node_memory, total_memory, model_card
     )
 
-    shards: list[ShardWithId] = []
+    runner_to_shard: dict[RunnerId, ShardMetadata] = {}
+    node_to_runner: dict[NodeId, RunnerId] = {}
 
     for pipeline_rank, node_id in enumerate(cycle.node_ids):
         layers_before = sum(layer_allocations[:pipeline_rank])
@@ -217,17 +223,20 @@ def _get_shard_assignments_for_pure_pipeline(
         shard = PipelineShardMetadata(
             model_card=model_card,
             device_rank=pipeline_rank,
-            world_size=world_size,
+            world_size=len(cycle),
             start_layer=layers_before,
             end_layer=layers_before + node_layers,
             n_layers=model_card.n_layers,
         )
 
         runner_id = RunnerId()
-        shards.append(ShardWithId(node_id, runner_id, shard))
+        runner_to_shard[runner_id] = shard
+        node_to_runner[node_id] = runner_id
 
     return ShardAssignments(
-        model_id=model_card.model_id, shards=shards, primary_output_node=world_size - 1
+        model_id=model_card.model_id,
+        runner_to_shard=runner_to_shard,
+        node_to_runner=node_to_runner,
     )
 
 
@@ -237,7 +246,8 @@ def get_shard_assignments_for_tensor_parallel(
 ):
     total_layers = model_card.n_layers
     world_size = len(cycle)
-    shards: list[ShardWithId] = []
+    runner_to_shard: dict[RunnerId, ShardMetadata] = {}
+    node_to_runner: dict[NodeId, RunnerId] = {}
 
     for i, node_id in enumerate(cycle):
         shard = TensorShardMetadata(
@@ -250,10 +260,14 @@ def get_shard_assignments_for_tensor_parallel(
         )
 
         runner_id = RunnerId()
-        shards.append(ShardWithId(node_id, runner_id, shard))
+
+        runner_to_shard[runner_id] = shard
+        node_to_runner[node_id] = runner_id
 
     shard_assignments = ShardAssignments(
-        model_id=model_card.model_id, shards=shards, primary_output_node=world_size - 1
+        model_id=model_card.model_id,
+        runner_to_shard=runner_to_shard,
+        node_to_runner=node_to_runner,
     )
 
     return shard_assignments

src/exo/master/tests/test_master.py

@@ -6,6 +6,7 @@ import pytest
 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.types.backends import Backend
 from exo.shared.types.commands import (
@@ -41,34 +42,15 @@ from exo.shared.types.worker.instances import (
     MlxRingInstance,
     ShardAssignments,
 )
-from exo.shared.types.worker.runners import ShardWithId
 from exo.shared.types.worker.shards import PipelineShardMetadata, Sharding
 from exo.utils.channels import channel
 from exo.utils.info_gatherer.info_gatherer import NodeBackends
 
 
-class MockAggregator:
-    def dump(self) -> dict[str, str]:
-        return {}
-
-
-class MockStorage:
-    async def get(self, _: str) -> None:
-        return None
-
-    async def put(self, _1: str, _2: str) -> None:
-        return None
-
-    async def delete(self, _: str) -> None:
-        return None
-
-    async def dump(self, _: str) -> dict[str, str]:
-        return {}
-
-
 @pytest.mark.asyncio
 async def test_master():
-    node_id = NodeId("yoooo")
+    keypair = get_node_id_keypair()
+    node_id = NodeId(keypair.to_node_id())
     session_id = SessionId(master_node_id=node_id, election_clock=0)
 
     ge_sender, global_event_receiver = channel[GlobalForwarderEvent]()
@@ -113,8 +95,6 @@ async def test_master():
         local_event_receiver=le_receiver,
         command_receiver=co_receiver,
         download_command_sender=fcds,
-        aggregator=MockAggregator(),  # pyright: ignore[reportArgumentType]
-        storage=MockStorage(),  # pyright: ignore[reportArgumentType]
     )
     logger.info("run the master")
     async with anyio.create_task_group() as tg:
@@ -226,33 +206,29 @@ async def test_master():
         assert isinstance(events[2].event, InstanceCreated)
         created_instance = events[2].event.instance
         assert isinstance(created_instance, MlxRingInstance)
-        runner_id = created_instance.shard_assignments.shards[0].runner_id
+        runner_id = list(created_instance.shard_assignments.runner_to_shard.keys())[0]
         # Validate the shard assignments
         expected_shard_assignments = ShardAssignments(
             model_id=ModelId("llama-3.2-1b"),
-            shards=[
-                ShardWithId(
-                    node_id,
-                    runner_id,
-                    PipelineShardMetadata(
-                        start_layer=0,
-                        end_layer=16,
+            runner_to_shard={
+                (runner_id): PipelineShardMetadata(
+                    start_layer=0,
+                    end_layer=16,
+                    n_layers=16,
+                    model_card=ModelCard(
+                        model_id=ModelId("llama-3.2-1b"),
                         n_layers=16,
-                        model_card=ModelCard(
-                            model_id=ModelId("llama-3.2-1b"),
-                            n_layers=16,
-                            storage_size=Memory.from_bytes(678948),
-                            hidden_size=7168,
-                            supports_tensor=True,
-                            tasks=[ModelTask.TextGeneration],
-                            backends=[Backend.MlxMetal],
-                        ),
-                        device_rank=0,
-                        world_size=1,
+                        storage_size=Memory.from_bytes(678948),
+                        hidden_size=7168,
+                        supports_tensor=True,
+                        tasks=[ModelTask.TextGeneration],
+                        backends=[Backend.MlxMetal],
                     ),
+                    device_rank=0,
+                    world_size=1,
                 )
-            ],
-            primary_output_node=0,
+            },
+            node_to_runner={node_id: runner_id},
         )
         assert created_instance.shard_assignments == expected_shard_assignments
         # For single-node, hosts_by_node should have one entry with self-binding

src/exo/master/tests/test_placement.py

@@ -49,36 +49,16 @@ from exo.shared.types.worker.instances import (
     MlxJacclInstance,
     MlxRingInstance,
 )
-from exo.shared.types.worker.runners import RunnerId, ShardAssignments, ShardWithId
+from exo.shared.types.worker.runners import ShardAssignments
 from exo.shared.types.worker.shards import PipelineShardMetadata, Sharding
 
 
-class MockShard:
-    def is_primary_output(self) -> bool:
-        return True
-
-
 @pytest.fixture
-def instance(model_card: ModelCard) -> Instance:
+def instance() -> Instance:
     return MlxRingInstance(
         instance_id=InstanceId(),
         shard_assignments=ShardAssignments(
-            model_id=ModelId("test-model"),
-            shards=[
-                ShardWithId(
-                    NodeId(),
-                    RunnerId(),
-                    PipelineShardMetadata(
-                        model_card=model_card,
-                        device_rank=0,
-                        world_size=1,
-                        start_layer=0,
-                        end_layer=model_card.n_layers,
-                        n_layers=model_card.n_layers,
-                    ),
-                )
-            ],
-            primary_output_node=0,
+            model_id=ModelId("test-model"), runner_to_shard={}, node_to_runner={}
         ),
         hosts_by_node={},
         ephemeral_port=50000,
@@ -143,11 +123,6 @@ def test_get_instance_placements_create_instance(
     node_id_a = NodeId()
     node_id_b = NodeId()
     node_id_c = NodeId()
-    node_to_layers = {
-        node_id_a: expected_layers[0],
-        node_id_b: expected_layers[1],
-        node_id_c: expected_layers[2],
-    }
 
     # fully connected (directed) between the 3 nodes
     conn_a_b = Connection(
@@ -200,11 +175,22 @@ def test_get_instance_placements_create_instance(
     instance = placements[instance_id]
     assert instance.shard_assignments.model_id == model_card.model_id
 
-    for nid, _, shard in (shards := instance.shard_assignments.shards):
-        assert shard.end_layer - shard.start_layer == node_to_layers[nid]
+    runner_id_a = instance.shard_assignments.node_to_runner[node_id_a]
+    runner_id_b = instance.shard_assignments.node_to_runner[node_id_b]
+    runner_id_c = instance.shard_assignments.node_to_runner[node_id_c]
 
-    assert shards[0].shard.start_layer == 0
-    assert shards[-1].shard.end_layer == total_layers
+    shard_a = instance.shard_assignments.runner_to_shard[runner_id_a]
+    shard_b = instance.shard_assignments.runner_to_shard[runner_id_b]
+    shard_c = instance.shard_assignments.runner_to_shard[runner_id_c]
+
+    assert shard_a.end_layer - shard_a.start_layer == expected_layers[0]
+    assert shard_b.end_layer - shard_b.start_layer == expected_layers[1]
+    assert shard_c.end_layer - shard_c.start_layer == expected_layers[2]
+
+    shards = [shard_a, shard_b, shard_c]
+    shards_sorted = sorted(shards, key=lambda s: s.start_layer)
+    assert shards_sorted[0].start_layer == 0
+    assert shards_sorted[-1].end_layer == total_layers
 
 
 def test_get_instance_placements_one_node_exact_fit() -> None:
@@ -232,7 +218,9 @@ def test_get_instance_placements_one_node_exact_fit() -> None:
     instance_id = list(placements.keys())[0]
     instance = placements[instance_id]
     assert instance.shard_assignments.model_id == "test-model"
-    assert len(instance.shard_assignments.shards) == 1
+    assert len(instance.shard_assignments.node_to_runner) == 1
+    assert len(instance.shard_assignments.runner_to_shard) == 1
+    assert len(instance.shard_assignments.runner_to_shard) == 1
 
 
 def test_get_instance_placements_one_node_fits_with_extra_memory() -> None:
@@ -260,7 +248,9 @@ def test_get_instance_placements_one_node_fits_with_extra_memory() -> None:
     instance_id = list(placements.keys())[0]
     instance = placements[instance_id]
     assert instance.shard_assignments.model_id == "test-model"
-    assert len(instance.shard_assignments.shards) == 1
+    assert len(instance.shard_assignments.node_to_runner) == 1
+    assert len(instance.shard_assignments.runner_to_shard) == 1
+    assert len(instance.shard_assignments.runner_to_shard) == 1
 
 
 def test_get_instance_placements_one_node_not_fit() -> None:
@@ -391,7 +381,7 @@ def test_placement_selects_leaf_nodes(
     assert len(placements) == 1
     instance = list(placements.values())[0]
 
-    assigned_nodes = set(map(lambda it: it.node_id, instance.shard_assignments.shards))
+    assigned_nodes = set(instance.shard_assignments.node_to_runner.keys())
     assert assigned_nodes == set((node_id_a, node_id_b)) or assigned_nodes == set(
         (
             node_id_c,
@@ -508,8 +498,8 @@ def test_tensor_rdma_backend_connectivity_matrix(
     for i in range(3):
         assert matrix[i][i] is None
 
-    assigned_nodes = list(instance.shard_assignments.shards)
-    node_to_idx = {node_id: idx for idx, (node_id, _, _) in enumerate(assigned_nodes)}
+    assigned_nodes = list(instance.shard_assignments.node_to_runner.keys())
+    node_to_idx = {node_id: idx for idx, node_id in enumerate(assigned_nodes)}
 
     idx_a = node_to_idx[node_a]
     idx_b = node_to_idx[node_b]
@@ -521,7 +511,7 @@ def test_tensor_rdma_backend_connectivity_matrix(
 
     # Verify coordinators are set for all nodes
     assert len(instance.jaccl_coordinators) == 3
-    for node_id, _, _ in assigned_nodes:
+    for node_id in assigned_nodes:
         assert node_id in instance.jaccl_coordinators
         coordinator = instance.jaccl_coordinators[node_id]
         assert ":" in coordinator
@@ -835,7 +825,7 @@ def test_placement_prefers_cycle_with_downloaded_model(
 
     assert len(placements) == 1
     instance = list(placements.values())[0]
-    assigned_nodes = set(map(lambda it: it.node_id, instance.shard_assignments.shards))
+    assigned_nodes = set(instance.shard_assignments.node_to_runner.keys())
     assert assigned_nodes == {node_b}
 
 
@@ -913,7 +903,7 @@ def test_placement_prefers_cycle_with_higher_download_progress(
 
     assert len(placements) == 1
     instance = list(placements.values())[0]
-    assigned_nodes = set(map(lambda it: it.node_id, instance.shard_assignments.shards))
+    assigned_nodes = set(instance.shard_assignments.node_to_runner.keys())
     assert assigned_nodes == {node_b}
 
 
@@ -967,7 +957,7 @@ def test_placement_does_not_prefer_cycle_with_failed_download(
 
     assert len(placements) == 1
     instance = list(placements.values())[0]
-    assigned_nodes = set(map(lambda it: it.node_id, instance.shard_assignments.shards))
+    assigned_nodes = set(instance.shard_assignments.node_to_runner.keys())
     # node_a should win on RAM tiebreaker since failed download scores 0.0
     assert assigned_nodes == {node_a}
 

src/exo/master/tests/test_placement_utils.py

@@ -204,11 +204,6 @@ def test_get_shard_assignments(
     node_a_id = NodeId()
     node_b_id = NodeId()
     node_c_id = NodeId()
-    layers_by_node = {
-        node_a_id: expected_layers[0],
-        node_b_id: expected_layers[1],
-        node_c_id: expected_layers[2],
-    }
 
     # create connections (A -> B -> C -> A forms a 3-cycle, plus B -> A also exists)
     connection1 = Connection(
@@ -263,8 +258,25 @@ def test_get_shard_assignments(
     )
 
     # assert
-    for nid, _, shard in shard_assignments.shards:
-        assert shard.end_layer - shard.start_layer == layers_by_node[nid]
+    runner_id_a = shard_assignments.node_to_runner[node_a_id]
+    runner_id_b = shard_assignments.node_to_runner[node_b_id]
+    runner_id_c = shard_assignments.node_to_runner[node_c_id]
+
+    assert (
+        shard_assignments.runner_to_shard[runner_id_a].end_layer
+        - shard_assignments.runner_to_shard[runner_id_a].start_layer
+        == expected_layers[0]
+    )
+    assert (
+        shard_assignments.runner_to_shard[runner_id_b].end_layer
+        - shard_assignments.runner_to_shard[runner_id_b].start_layer
+        == expected_layers[1]
+    )
+    assert (
+        shard_assignments.runner_to_shard[runner_id_c].end_layer
+        - shard_assignments.runner_to_shard[runner_id_c].start_layer
+        == expected_layers[2]
+    )
 
 
 def test_get_mlx_jaccl_coordinators():
@@ -531,11 +543,11 @@ class TestCfgParallelPlacement:
             model_card, cycle, node_memory
         )
 
-        shards = list(assignments.shards)
+        shards = list(assignments.runner_to_shard.values())
         assert len(shards) == 2
 
         # CFG models should get CfgShardMetadata
-        for _, _, shard in shards:
+        for shard in shards:
             assert isinstance(shard, CfgShardMetadata)
             # Both nodes should have all layers (no pipeline split)
             assert shard.start_layer == 0
@@ -546,7 +558,7 @@ class TestCfgParallelPlacement:
             assert shard.pipeline_rank == 0
 
         cfg_ranks = sorted(
-            s.shard.cfg_rank for s in shards if isinstance(s.shard, CfgShardMetadata)
+            s.cfg_rank for s in shards if isinstance(s, CfgShardMetadata)
         )
         assert cfg_ranks == [0, 1]
 
@@ -575,11 +587,11 @@ class TestCfgParallelPlacement:
             model_card, cycle, node_memory
         )
 
-        shards = assignments.shards
+        shards = list(assignments.runner_to_shard.values())
         assert len(shards) == 4
 
         # CFG models should get CfgShardMetadata
-        for _, _, shard in shards:
+        for shard in shards:
             assert isinstance(shard, CfgShardMetadata)
             assert shard.cfg_world_size == 2
             assert shard.pipeline_world_size == 2
@@ -587,14 +599,10 @@ class TestCfgParallelPlacement:
 
         # Check we have 2 nodes in each CFG group
         cfg_0_shards = [
-            s.shard
-            for s in shards
-            if isinstance(s.shard, CfgShardMetadata) and s.shard.cfg_rank == 0
+            s for s in shards if isinstance(s, CfgShardMetadata) and s.cfg_rank == 0
         ]
         cfg_1_shards = [
-            s.shard
-            for s in shards
-            if isinstance(s.shard, CfgShardMetadata) and s.shard.cfg_rank == 1
+            s for s in shards if isinstance(s, CfgShardMetadata) and s.cfg_rank == 1
         ]
         assert len(cfg_0_shards) == 2
         assert len(cfg_1_shards) == 2
@@ -629,11 +637,11 @@ class TestCfgParallelPlacement:
             model_card, cycle, node_memory
         )
 
-        shards = list(assignments.shards)
+        shards = list(assignments.runner_to_shard.values())
         assert len(shards) == 3
 
         # Odd node count with CFG model falls back to PipelineShardMetadata (sequential CFG)
-        for _, _, shard in shards:
+        for shard in shards:
             assert isinstance(shard, PipelineShardMetadata)
 
     def test_two_nodes_non_cfg_model_uses_pipeline(self):
@@ -665,18 +673,18 @@ class TestCfgParallelPlacement:
             model_card, cycle, node_memory
         )
 
-        shards = list(assignments.shards)
+        shards = list(assignments.runner_to_shard.values())
         assert len(shards) == 2
 
         # Non-CFG models should get PipelineShardMetadata
-        for _, _, shard in shards:
+        for shard in shards:
             assert isinstance(shard, PipelineShardMetadata)
 
         # Should have actual layer sharding (pipeline)
         layer_ranges = sorted(
-            (s.shard.start_layer, s.shard.end_layer)
+            (s.start_layer, s.end_layer)
             for s in shards
-            if isinstance(s.shard, PipelineShardMetadata)
+            if isinstance(s, PipelineShardMetadata)
         )
         # First shard starts at 0, last shard ends at 57
         assert layer_ranges[0][0] == 0

src/exo/routing/connection_message.py

@@ -1,13 +1,15 @@
-from exo_rs import PyFromSwarm
+from exo_rs import FromSwarm
 
+from exo.shared.types.common import NodeId
 from exo.utils.pydantic_ext import FrozenModel
 
 """Serialisable types for Connection Updates/Messages"""
 
 
 class ConnectionMessage(FrozenModel):
+    node_id: NodeId
     connected: bool
 
     @classmethod
-    def from_update(cls, update: PyFromSwarm.Connection) -> "ConnectionMessage":
-        return cls(connected=update.connected)
+    def from_update(cls, update: FromSwarm.Connection) -> "ConnectionMessage":
+        return cls(node_id=NodeId(update.peer_id), connected=update.connected)

src/exo/routing/router.py

@@ -1,7 +1,8 @@
-import os
+from collections.abc import Sequence
 from copy import copy
 from itertools import count
 from math import inf
+from os import PathLike
 from pathlib import Path
 from typing import cast
 
@@ -12,13 +13,17 @@ from anyio import (
     sleep_forever,
 )
 from exo_rs import (
+    AllQueuesFullError,
+    FromSwarm,
+    Keypair,
+    MessageTooLargeError,
     NetworkingHandle,
-    PyFromSwarm,
+    NoPeersSubscribedToTopicError,
 )
+from filelock import FileLock
 from loguru import logger
 
-from exo.shared.constants import EXO_NODE_ZID
-from exo.shared.types.common import NodeId
+from exo.shared.constants import EXO_NODE_ID_KEYPAIR
 from exo.utils.channels import Receiver, Sender, channel
 from exo.utils.pydantic_ext import FrozenModel
 from exo.utils.task_group import TaskGroup
@@ -100,12 +105,12 @@ class Router:
     @classmethod
     def create(
         cls,
-        identity: str,
-        listen_port: int,
-        discovery_service_port: int,
+        identity: Keypair,
+        bootstrap_peers: Sequence[str] = (),
+        listen_port: int = 0,
     ) -> "Router":
         return cls(
-            handle=NetworkingHandle.new(identity, listen_port, discovery_service_port)
+            handle=NetworkingHandle(identity, list(bootstrap_peers), listen_port)
         )
 
     def __init__(self, handle: NetworkingHandle):
@@ -186,8 +191,10 @@ class Router:
                 from_swarm = await self._net.recv()
                 logger.debug(from_swarm)
                 match from_swarm:
-                    case PyFromSwarm.Message(topic, data):
-                        logger.trace(f"Received message on {topic} with payload {data}")
+                    case FromSwarm.Message(origin, topic, data):
+                        logger.trace(
+                            f"Received message on {topic} from {origin} with payload {data}"
+                        )
                         if topic not in self.topic_routers:
                             logger.warning(
                                 f"Received message on unknown or inactive topic {topic}"
@@ -195,7 +202,7 @@ class Router:
                             continue
                         router = self.topic_routers[topic]
                         await router.publish_bytes(data)
-                    case PyFromSwarm.Connection():
+                    case FromSwarm.Connection():
                         message = ConnectionMessage.from_update(from_swarm)
                         logger.trace(
                             f"Received message on connection_messages with payload {message}"
@@ -218,25 +225,33 @@ class Router:
     async def _networking_publish(self):
         with self.networking_receiver as networked_items:
             async for topic, data in networked_items:
-                logger.trace(f"Sending message on {topic} with payload {data}")
-                if len(data) > 1024 * 1024:
+                try:
+                    logger.trace(f"Sending message on {topic} with payload {data}")
+                    if len(data) > 1024 * 1024:
+                        logger.warning(
+                            "Sending overlarge payload, network performance may be temporarily degraded"
+                        )
+                    await self._net.gossipsub_publish(topic, data)
+                except NoPeersSubscribedToTopicError:
+                    pass
+                except AllQueuesFullError:
+                    logger.warning(f"All peer queues full, dropping message on {topic}")
+                except MessageTooLargeError:
                     logger.warning(
-                        "Sending overlarge payload, network performance may be temporarily degraded"
+                        f"Message too large for gossipsub on {topic} ({len(data)} bytes), dropping"
                     )
-                await self._net.gossipsub_publish(topic, data)
 
 
-def get_node_zid(
-    path: Path = EXO_NODE_ZID,
-) -> NodeId:
+def get_node_id_keypair(
+    path: str | bytes | PathLike[str] | PathLike[bytes] = EXO_NODE_ID_KEYPAIR,
+) -> 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 NodeId(os.urandom(16).hex())
+    return Keypair.generate()
 
-    """
     def lock_path(path: str | bytes | PathLike[str] | PathLike[bytes]) -> Path:
         return Path(str(path) + ".lock")
 
@@ -258,4 +273,3 @@ def get_node_zid(
             keypair = Keypair.generate()
             f.write(keypair.to_bytes())
             return keypair
-    """

src/exo/shared/apply.py

@@ -4,13 +4,19 @@ from datetime import datetime
 
 from loguru import logger
 
-from exo.shared.types.common import NodeId
+from exo.shared.models.model_cards import ModelCard
+from exo.shared.types.common import ModelId, NodeId
 from exo.shared.types.events import (
     ChunkGenerated,
+    CustomModelCardAdded,
+    CustomModelCardDeleted,
     Event,
     IndexedEvent,
+    InputChunkReceived,
     InstanceCreated,
     InstanceDeleted,
+    InstanceLinkCreated,
+    InstanceLinkDeleted,
     NodeDownloadProgress,
     NodeGatheredInfo,
     NodeTimedOut,
@@ -26,6 +32,7 @@ from exo.shared.types.events import (
     TracesCollected,
     TracesMerged,
 )
+from exo.shared.types.instance_link import InstanceLink, InstanceLinkId
 from exo.shared.types.profiling import (
     NodeIdentity,
     NodeNetworkInfo,
@@ -35,6 +42,7 @@ from exo.shared.types.profiling import (
 )
 from exo.shared.types.state import State
 from exo.shared.types.tasks import Task, 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.runners import (
@@ -59,6 +67,18 @@ from exo.utils.info_gatherer.info_gatherer import (
 )
 
 
+def _is_rdma_ctl_enabled(
+    node_id: NodeId, node_rdma_ctl: Mapping[NodeId, NodeRdmaCtlStatus]
+) -> bool:
+    """A node is RDMA-capable only if rdma_ctl status has been observed as enabled.
+
+    Missing entries default to ``False`` — if we have not yet observed (or the node
+    cannot run) ``rdma_ctl``, it must not participate in an RDMA-backed instance.
+    """
+    status = node_rdma_ctl.get(node_id)
+    return status is not None and status.enabled
+
+
 def event_apply(event: Event, state: State) -> State:
     """Apply an event to state."""
     match event:
@@ -66,10 +86,15 @@ def event_apply(event: Event, state: State) -> State:
             TestEvent()
             | ChunkGenerated()
             | TaskAcknowledged()
+            | InputChunkReceived()
             | TracesCollected()
             | TracesMerged()
         ):  # Pass-through events that don't modify state
             return state
+        case CustomModelCardAdded():
+            return apply_custom_model_card_added(event, state)
+        case CustomModelCardDeleted():
+            return apply_custom_model_card_deleted(event, state)
         case InstanceCreated():
             return apply_instance_created(event, state)
         case InstanceDeleted():
@@ -94,6 +119,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 InstanceLinkCreated():
+            return apply_instance_link_created(event, state)
+        case InstanceLinkDeleted():
+            return apply_instance_link_deleted(event, state)
 
 
 def apply(state: State, event: IndexedEvent) -> State:
@@ -193,7 +222,38 @@ def apply_instance_deleted(event: InstanceDeleted, state: State) -> State:
     new_instances: Mapping[InstanceId, Instance] = {
         iid: inst for iid, inst in state.instances.items() if iid != event.instance_id
     }
-    return state.model_copy(update={"instances": new_instances})
+    new_links: dict[InstanceLinkId, InstanceLink] = {}
+    for link_id, link in state.instance_links.items():
+        prefill = [i for i in link.prefill_instances if i != event.instance_id]
+        decode = [i for i in link.decode_instances if i != event.instance_id]
+        if not prefill or not decode:
+            continue
+        if prefill == list(link.prefill_instances) and decode == list(
+            link.decode_instances
+        ):
+            new_links[link_id] = link
+        else:
+            new_links[link_id] = link.model_copy(
+                update={"prefill_instances": prefill, "decode_instances": decode}
+            )
+    return state.model_copy(
+        update={"instances": new_instances, "instance_links": new_links}
+    )
+
+
+def apply_instance_link_created(event: InstanceLinkCreated, state: State) -> State:
+    new_links: Mapping[InstanceLinkId, InstanceLink] = {
+        **state.instance_links,
+        event.link.link_id: event.link,
+    }
+    return state.model_copy(update={"instance_links": new_links})
+
+
+def apply_instance_link_deleted(event: InstanceLinkDeleted, state: State) -> State:
+    new_links: Mapping[InstanceLinkId, InstanceLink] = {
+        lid: link for lid, link in state.instance_links.items() if lid != event.link_id
+    }
+    return state.model_copy(update={"instance_links": new_links})
 
 
 def apply_runner_status_updated(event: RunnerStatusUpdated, state: State) -> State:
@@ -348,26 +408,59 @@ def apply_node_gathered_info(event: NodeGatheredInfo, state: State) -> State:
                 event.node_id: NodeThunderboltInfo(interfaces=info.idents),
             }
         case MacThunderboltConnections():
-            update["node_thunderbolt_connections"] = {
-                **state.node_thunderbolt_connections,
-                event.node_id: info,
+            conn_map = {
+                tb_ident.domain_uuid: (nid, tb_ident.rdma_interface)
+                for nid in state.node_thunderbolt
+                for tb_ident in state.node_thunderbolt[nid].interfaces
             }
+            source_is_rdma_enabled = _is_rdma_ctl_enabled(
+                event.node_id, state.node_rdma_ctl
+            )
+            as_rdma_conns = [
+                Connection(
+                    source=event.node_id,
+                    sink=conn_map[tb_conn.sink_uuid][0],
+                    edge=RDMAConnection(
+                        source_rdma_iface=conn_map[tb_conn.source_uuid][1],
+                        sink_rdma_iface=conn_map[tb_conn.sink_uuid][1],
+                    ),
+                )
+                for tb_conn in info.conns
+                if tb_conn.source_uuid in conn_map
+                if tb_conn.sink_uuid in conn_map
+                if source_is_rdma_enabled
+                and _is_rdma_ctl_enabled(
+                    conn_map[tb_conn.sink_uuid][0], state.node_rdma_ctl
+                )
+            ]
+            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
-            update["thunderbolt_bridge_cycles"] = (
-                topology.get_thunderbolt_bridge_cycles(
-                    new_tb_bridge, state.node_network
+            # 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
+                    )
                 )
-            )
         case RdmaCtlStatus():
             update["node_rdma_ctl"] = {
                 **state.node_rdma_ctl,
                 event.node_id: NodeRdmaCtlStatus(enabled=info.enabled),
             }
+            # If RDMA just got disabled on this node, drop any RDMA edges touching it
+            # so placement / topology consumers cannot pick a disabled node for an
+            # RDMA-backed instance. (Edges will repopulate on the next
+            # MacThunderboltConnections poll once both endpoints are enabled again.)
+            if not info.enabled:
+                topology.remove_all_rdma_connections_touching(event.node_id)
         case NodeBackends():
             update["node_backends"] = {
                 **state.node_backends,
@@ -378,45 +471,32 @@ def apply_node_gathered_info(event: NodeGatheredInfo, state: State) -> State:
 
 
 def apply_topology_edge_created(event: TopologyEdgeCreated, state: State) -> State:
-    source_connections = state.node_socket_connections.get(event.conn.source, {})
-    sink_connections = source_connections.get(event.conn.sink, [])
-
-    update = {
-        "node_socket_connections": {
-            **state.node_socket_connections,
-            event.conn.source: {
-                **source_connections,
-                event.conn.sink: sink_connections
-                if event.conn.edge in sink_connections
-                else [*sink_connections, event.conn.edge],
-            },
-        }
-    }
-    return state.model_copy(update=update)
+    topology = copy.deepcopy(state.topology)
+    topology.add_connection(event.conn)
+    return state.model_copy(update={"topology": topology})
 
 
 def apply_topology_edge_deleted(event: TopologyEdgeDeleted, state: State) -> State:
-    inner_update = {
-        sink: final_edges
-        for sink, edges in state.node_socket_connections.get(
-            event.conn.source, {}
-        ).items()
-        if (
-            final_edges := [
-                edge
-                for edge in edges
-                if (edge != event.conn.edge or sink != event.conn.sink)
-            ]
-        )
+    topology = copy.deepcopy(state.topology)
+    topology.remove_connection(event.conn)
+    # TODO: Clean up removing the reverse connection
+    return state.model_copy(update={"topology": topology})
+
+
+def apply_custom_model_card_added(event: CustomModelCardAdded, state: State) -> State:
+    new_cards: Mapping[ModelId, ModelCard] = {
+        **state.custom_model_cards,
+        event.model_card.model_id: event.model_card,
     }
-    update = {
-        "node_socket_connections": {
-            source: maps
-            for source, maps in {
-                **state.node_socket_connections,
-                event.conn.source: inner_update,
-            }.items()
-            if maps
-        }
+    return state.model_copy(update={"custom_model_cards": new_cards})
+
+
+def apply_custom_model_card_deleted(
+    event: CustomModelCardDeleted, state: State
+) -> State:
+    new_cards: Mapping[ModelId, ModelCard] = {
+        model_id: card
+        for model_id, card in state.custom_model_cards.items()
+        if model_id != event.model_id
     }
-    return state.model_copy(update=update)
+    return state.model_copy(update={"custom_model_cards": new_cards})

src/exo/shared/constants.py

@@ -76,7 +76,7 @@ EXO_TEST_LOG = EXO_CACHE_HOME / "exo_test.log"
 EXO_PID_FILE = EXO_CACHE_HOME / "exo.pid"
 
 # Identity (config)
-EXO_NODE_ZID = EXO_CACHE_HOME / "node_zid"
+EXO_NODE_ID_KEYPAIR = EXO_CONFIG_HOME / "node_id.keypair"
 EXO_CONFIG_FILE = EXO_CONFIG_HOME / "config.toml"
 
 # libp2p topics for event forwarding

src/exo/shared/logging.py

@@ -46,8 +46,7 @@ class _InterceptHandler(logging.Handler):
 def logger_setup(log_file: Path | None, verbosity: int = 0):
     """Set up logging for this process - formatting, file handles, verbosity and output"""
 
-    logging.getLogger("exo_rs").setLevel(logging.INFO)
-    logging.getLogger("networking").setLevel(logging.INFO)
+    logging.getLogger("exo_rs").setLevel(logging.WARNING)
     logging.getLogger("httpx").setLevel(logging.WARNING)
     logging.getLogger("httpcore").setLevel(logging.WARNING)
 

src/exo/shared/models/model_cards.py

@@ -55,7 +55,7 @@ class _CardCache:
         except OSError as e:
             logger.warning(f"failed to save custom model card ({e.strerror})")
 
-    async def delete(self, model_id: ModelId) -> "ModelCard | None":
+    async def pop(self, model_id: ModelId) -> "ModelCard | None":
         """Delete a user-added custom model card. Returns True if deleted."""
         card_path = _custom_cards_dir / (ModelId(model_id).normalize() + ".toml")
         try:

src/exo/shared/tests/test_apply/test_apply_custom_model_cards.py

@@ -0,0 +1,46 @@
+from exo.shared.apply import apply
+from exo.shared.models.model_cards import ModelCard, ModelTask
+from exo.shared.types.backends import Backend
+from exo.shared.types.common import ModelId
+from exo.shared.types.events import (
+    CustomModelCardAdded,
+    CustomModelCardDeleted,
+    IndexedEvent,
+)
+from exo.shared.types.memory import Memory
+from exo.shared.types.state import State
+
+
+def _model_card(model_id: ModelId) -> ModelCard:
+    return ModelCard(
+        model_id=model_id,
+        n_layers=1,
+        storage_size=Memory.from_bytes(1),
+        hidden_size=1,
+        supports_tensor=True,
+        tasks=[ModelTask.TextGeneration],
+        backends=[Backend.MlxMetal],
+    )
+
+
+def test_custom_model_card_added_is_reduced_into_state() -> None:
+    card = _model_card(ModelId("custom/model"))
+
+    state = apply(
+        State(),
+        IndexedEvent(idx=0, event=CustomModelCardAdded(model_card=card)),
+    )
+
+    assert state.custom_model_cards == {card.model_id: card}
+
+
+def test_custom_model_card_deleted_removes_card_from_state() -> None:
+    card = _model_card(ModelId("custom/model"))
+    state = State(custom_model_cards={card.model_id: card}, last_event_applied_idx=0)
+
+    state = apply(
+        state,
+        IndexedEvent(idx=1, event=CustomModelCardDeleted(model_id=card.model_id)),
+    )
+
+    assert state.custom_model_cards == {}

src/exo/shared/tests/test_apply/test_apply_instance_link.py

@@ -0,0 +1,72 @@
+from exo.shared.apply import (
+    apply_instance_deleted,
+    apply_instance_link_created,
+    apply_instance_link_deleted,
+)
+from exo.shared.types.events import (
+    InstanceDeleted,
+    InstanceLinkCreated,
+    InstanceLinkDeleted,
+)
+from exo.shared.types.instance_link import InstanceLink, InstanceLinkId
+from exo.shared.types.state import State
+from exo.shared.types.worker.instances import InstanceId
+
+
+def _link(
+    prefill: list[InstanceId],
+    decode: list[InstanceId],
+    link_id: InstanceLinkId | None = None,
+) -> InstanceLink:
+    return InstanceLink(
+        link_id=link_id or InstanceLinkId(),
+        prefill_instances=prefill,
+        decode_instances=decode,
+    )
+
+
+def test_create_link() -> None:
+    state = State()
+    link = _link([InstanceId("a")], [InstanceId("b")])
+    new_state = apply_instance_link_created(InstanceLinkCreated(link=link), state)
+    assert new_state.instance_links == {link.link_id: link}
+
+
+def test_update_replaces_existing_link() -> None:
+    a, b, c = InstanceId("a"), InstanceId("b"), InstanceId("c")
+    link = _link([a], [b])
+    state = State(instance_links={link.link_id: link})
+
+    updated = link.model_copy(update={"decode_instances": [b, c]})
+    new_state = apply_instance_link_created(InstanceLinkCreated(link=updated), state)
+    assert set(new_state.instance_links[link.link_id].decode_instances) == {b, c}
+
+
+def test_delete_link() -> None:
+    link = _link([InstanceId("a")], [InstanceId("b")])
+    state = State(instance_links={link.link_id: link})
+
+    new_state = apply_instance_link_deleted(
+        InstanceLinkDeleted(link_id=link.link_id), state
+    )
+    assert new_state.instance_links == {}
+
+
+def test_instance_deleted_strips_from_links() -> None:
+    a, b, c = InstanceId("a"), InstanceId("b"), InstanceId("c")
+    link = _link([a, c], [b])
+    state = State(instance_links={link.link_id: link})
+
+    new_state = apply_instance_deleted(InstanceDeleted(instance_id=a), state)
+    remaining = new_state.instance_links[link.link_id]
+    assert remaining.prefill_instances == [c]
+    assert remaining.decode_instances == [b]
+
+
+def test_instance_deleted_drops_link_when_role_empties() -> None:
+    a, b = InstanceId("a"), InstanceId("b")
+    link = _link([a], [b])
+    state = State(instance_links={link.link_id: link})
+
+    new_state = apply_instance_deleted(InstanceDeleted(instance_id=a), state)
+    assert link.link_id not in new_state.instance_links

src/exo/shared/tests/test_apply/test_apply_rdma_gating.py

@@ -217,7 +217,7 @@ def test_topology_remove_all_rdma_connections_touching_keeps_socket_edges():
         )
     )
     socket_edge = SocketConnection(
-        sink_multiaddr=Multiaddr(address="/ip4/10.0.0.1/tcp/8000"),
+        sink_multiaddr=Multiaddr(address="/ip4/10.0.0.1/tcp/8000")
     )
     topology.add_connection(Connection(source=node_a, sink=node_b, edge=socket_edge))
 

src/exo/shared/tests/test_election.py

@@ -327,7 +327,7 @@ async def test_connection_message_triggers_new_round_broadcast() -> None:
             tg.start_soon(election.run)
 
             # Send any connection message object; we close quickly to cancel before result creation
-            await cm_tx.send(ConnectionMessage(connected=True))
+            await cm_tx.send(ConnectionMessage(node_id=NodeId(), connected=True))
 
             # Expect a broadcast for the new round at clock=1
             while True:

src/exo/shared/tests/test_node_id_persistence.py

@@ -10,8 +10,8 @@ from multiprocessing.synchronize import Semaphore as SemaphoreT
 from loguru import logger
 from pytest import LogCaptureFixture, mark
 
-from exo.routing.router import get_node_zid
-from exo.shared.constants import EXO_NODE_ZID
+from exo.routing.router import get_node_id_keypair
+from exo.shared.constants import EXO_NODE_ID_KEYPAIR
 
 NUM_CONCURRENT_PROCS = 10
 
@@ -23,7 +23,7 @@ def _get_keypair_concurrent_subprocess_task(
     sem.release()
     # wait to be told to begin simultaneous read
     ev.wait()
-    queue.put(get_node_zid().encode())
+    queue.put(get_node_id_keypair().to_bytes())
 
 
 def _get_keypair_concurrent(num_procs: int) -> bytes:
@@ -79,7 +79,7 @@ def test_node_id_fetching(caplog: LogCaptureFixture):
     reps = 10
 
     # delete current file and write a new one
-    _delete_if_exists(EXO_NODE_ZID)
+    _delete_if_exists(EXO_NODE_ID_KEYPAIR)
     kp = _get_keypair_concurrent(NUM_CONCURRENT_PROCS)
 
     with caplog.at_level(101):  # supress logs
@@ -88,6 +88,6 @@ def test_node_id_fetching(caplog: LogCaptureFixture):
             assert kp == _get_keypair_concurrent(NUM_CONCURRENT_PROCS)
 
         # make sure that after deleting, we are not fetching the same value
-        _delete_if_exists(EXO_NODE_ZID)
+        _delete_if_exists(EXO_NODE_ID_KEYPAIR)
         for _ in range(reps):
             assert kp != _get_keypair_concurrent(NUM_CONCURRENT_PROCS)

src/exo/shared/tests/test_state_serialization.py

@@ -0,0 +1,35 @@
+from exo.shared.types.common import NodeId
+from exo.shared.types.multiaddr import Multiaddr
+from exo.shared.types.state import State
+from exo.shared.types.topology import Connection, SocketConnection
+
+
+def test_state_serialization_roundtrip() -> None:
+    """Verify that State → JSON → State round-trip preserves topology."""
+
+    # --- build a simple state ------------------------------------------------
+    node_a = NodeId("node-a")
+    node_b = NodeId("node-b")
+
+    connection = Connection(
+        source=node_a,
+        sink=node_b,
+        edge=SocketConnection(
+            sink_multiaddr=Multiaddr(address="/ip4/127.0.0.1/tcp/10001"),
+        ),
+    )
+
+    state = State()
+    state.topology.add_connection(connection)
+
+    json_repr = state.model_dump_json()
+    restored_state = State.model_validate_json(json_repr)
+
+    assert (
+        state.topology.to_snapshot().nodes
+        == restored_state.topology.to_snapshot().nodes
+    )
+    assert set(state.topology.to_snapshot().connections) == set(
+        restored_state.topology.to_snapshot().connections
+    )
+    assert restored_state.model_dump_json() == json_repr

src/exo/shared/tests/test_xdg_paths.py

@@ -97,6 +97,13 @@ def test_macos_uses_traditional_paths():
         assert home / ".exo" == constants.EXO_CACHE_HOME
 
 
+def test_node_id_in_config_dir():
+    """Test that node ID keypair is in the config directory."""
+    import exo.shared.constants as constants
+
+    assert constants.EXO_NODE_ID_KEYPAIR.parent == constants.EXO_CONFIG_HOME
+
+
 def test_models_in_data_dir():
     """Test that default models directory is in the data directory."""
     # Clear EXO_MODELS_DIRS to test default behavior

src/exo/shared/topology.py

@@ -1,8 +1,10 @@
+import contextlib
 from collections.abc import Mapping, Sequence
 from dataclasses import dataclass, field
 from typing import Iterable
 
 import rustworkx as rx
+from pydantic import BaseModel, ConfigDict
 
 from exo.shared.types.common import NodeId
 from exo.shared.types.profiling import (
@@ -18,6 +20,15 @@ from exo.shared.types.topology import (
 )
 
 
+class TopologySnapshot(BaseModel):
+    nodes: Sequence[NodeId]
+    connections: Mapping[
+        NodeId, Mapping[NodeId, Sequence[SocketConnection | RDMAConnection]]
+    ]
+
+    model_config = ConfigDict(frozen=True, extra="forbid")
+
+
 @dataclass
 class Topology:
     _graph: rx.PyDiGraph[NodeId, SocketConnection | RDMAConnection] = field(
@@ -25,6 +36,28 @@ class Topology:
     )
     _vertex_indices: dict[NodeId, int] = field(init=False, default_factory=dict)
 
+    def to_snapshot(self) -> TopologySnapshot:
+        return TopologySnapshot(
+            nodes=list(self.list_nodes()), connections=self.map_connections()
+        )
+
+    @classmethod
+    def from_snapshot(cls, snapshot: TopologySnapshot) -> "Topology":
+        topology = cls()
+
+        for node_id in snapshot.nodes:
+            with contextlib.suppress(ValueError):
+                topology.add_node(node_id)
+
+        for source in snapshot.connections:
+            for sink in snapshot.connections[source]:
+                for edge in snapshot.connections[source][sink]:
+                    topology.add_connection(
+                        Connection(source=source, sink=sink, edge=edge)
+                    )
+
+        return topology
+
     def add_node(self, node_id: NodeId) -> None:
         if node_id in self._vertex_indices:
             return

src/exo/shared/types/chunks.py

@@ -13,6 +13,8 @@ from exo.shared.models.model_cards import ModelId
 from exo.utils.pydantic_ext import TaggedModel
 from exo.worker.runner.diagnostics import KnownRunnerDiagnostic
 
+from .common import CommandId
+
 
 class BaseChunk(TaggedModel):
     model: ModelId
@@ -66,6 +68,21 @@ class ImageChunk(BaseChunk):
                 yield name, value
 
 
+class InputImageChunk(BaseChunk):
+    command_id: CommandId
+    data: str
+    chunk_index: int
+    total_chunks: 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]
+            if name == "data" and hasattr(value, "__len__"):  # pyright: ignore[reportAny]
+                yield name, f"<{len(self.data)} chars>"
+            elif name is not None:
+                yield name, value
+
+
 class PrefillProgressChunk(BaseChunk):
     """Data class for prefill progress events during streaming."""
 

src/exo/shared/types/commands.py

@@ -5,7 +5,9 @@ from exo.api.types import (
     ImageGenerationTaskParams,
 )
 from exo.shared.models.model_cards import ModelCard, ModelId
+from exo.shared.types.chunks import InputImageChunk
 from exo.shared.types.common import CommandId, NodeId, SystemId
+from exo.shared.types.instance_link import InstanceLinkId
 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
@@ -55,6 +57,12 @@ class TaskFinished(BaseCommand):
     finished_command_id: CommandId
 
 
+class SendInputChunk(BaseCommand):
+    """Command to send an input image chunk (converted to event by master)."""
+
+    chunk: InputImageChunk
+
+
 class RequestEventLog(BaseCommand):
     since_idx: int
 
@@ -74,6 +82,24 @@ class CancelDownload(BaseCommand):
     model_id: ModelId
 
 
+class AddCustomModelCard(BaseCommand):
+    model_card: ModelCard
+
+
+class DeleteCustomModelCard(BaseCommand):
+    model_id: ModelId
+
+
+class SetInstanceLink(BaseCommand):
+    link_id: InstanceLinkId
+    prefill_instances: list[InstanceId]
+    decode_instances: list[InstanceId]
+
+
+class DeleteInstanceLink(BaseCommand):
+    link_id: InstanceLinkId
+
+
 DownloadCommand = StartDownload | DeleteDownload | CancelDownload
 
 
@@ -88,6 +114,11 @@ Command = (
     | DeleteInstance
     | TaskCancelled
     | TaskFinished
+    | SendInputChunk
+    | AddCustomModelCard
+    | DeleteCustomModelCard
+    | SetInstanceLink
+    | DeleteInstanceLink
 )
 
 

src/exo/shared/types/events.py

@@ -3,9 +3,11 @@ from typing import final
 
 from pydantic import Field
 
+from exo.shared.models.model_cards import ModelCard
 from exo.shared.topology import Connection
-from exo.shared.types.chunks import Chunk
-from exo.shared.types.common import CommandId, Id, NodeId, SessionId, SystemId
+from exo.shared.types.chunks import Chunk, InputImageChunk
+from exo.shared.types.common import CommandId, Id, ModelId, NodeId, SessionId, SystemId
+from exo.shared.types.instance_link import InstanceLink, InstanceLinkId
 from exo.shared.types.tasks import Task, TaskId, TaskStatus
 from exo.shared.types.worker.downloads import DownloadProgress
 from exo.shared.types.worker.instances import Instance, InstanceId
@@ -93,6 +95,11 @@ class ChunkGenerated(BaseEvent):
     chunk: Chunk
 
 
+class InputChunkReceived(BaseEvent):
+    command_id: CommandId
+    chunk: InputImageChunk
+
+
 class TopologyEdgeCreated(BaseEvent):
     conn: Connection
 
@@ -101,6 +108,14 @@ class TopologyEdgeDeleted(BaseEvent):
     conn: Connection
 
 
+class CustomModelCardAdded(BaseEvent):
+    model_card: ModelCard
+
+
+class CustomModelCardDeleted(BaseEvent):
+    model_id: ModelId
+
+
 @final
 class TraceEventData(FrozenModel):
     name: str
@@ -123,6 +138,14 @@ class TracesMerged(BaseEvent):
     traces: list[TraceEventData]
 
 
+class InstanceLinkCreated(BaseEvent):
+    link: InstanceLink
+
+
+class InstanceLinkDeleted(BaseEvent):
+    link_id: InstanceLinkId
+
+
 Event = (
     TestEvent
     | TaskCreated
@@ -137,10 +160,15 @@ Event = (
     | NodeGatheredInfo
     | NodeDownloadProgress
     | ChunkGenerated
+    | InputChunkReceived
     | TopologyEdgeCreated
     | TopologyEdgeDeleted
     | TracesCollected
     | TracesMerged
+    | CustomModelCardAdded
+    | CustomModelCardDeleted
+    | InstanceLinkCreated
+    | InstanceLinkDeleted
 )
 
 

src/exo/shared/types/profiling.py

@@ -16,10 +16,6 @@ class MemoryUsage(FrozenModel):
     swap_total: Memory
     swap_available: Memory
 
-    @classmethod
-    def tag(cls) -> str:
-        return cls.__name__
-
     @classmethod
     def from_bytes(
         cls, *, ram_total: int, ram_available: int, swap_total: int, swap_available: int

src/exo/shared/types/state.py

@@ -1,15 +1,15 @@
 from collections.abc import Mapping, Sequence
 from datetime import datetime
-from typing import Any
+from typing import Any, cast
 
-from exo_rs import LVAggregator
-from pydantic import ConfigDict, Field, model_serializer
+from pydantic import ConfigDict, Field, field_serializer, field_validator
 from pydantic.alias_generators import to_camel
-from pydantic_core.core_schema import SerializerFunctionWrapHandler
 
-from exo.shared.topology import Topology
+from exo.shared.models.model_cards import ModelCard
+from exo.shared.topology import Topology, TopologySnapshot
 from exo.shared.types.backends import Backend
-from exo.shared.types.common import NodeId
+from exo.shared.types.common import ModelId, NodeId
+from exo.shared.types.instance_link import InstanceLink, InstanceLinkId
 from exo.shared.types.profiling import (
     DiskUsage,
     MemoryUsage,
@@ -21,15 +21,9 @@ from exo.shared.types.profiling import (
     ThunderboltBridgeStatus,
 )
 from exo.shared.types.tasks import Task, TaskId
-from exo.shared.types.topology import (
-    Connection,
-    RDMAConnection,
-    SocketConnection,
-)
 from exo.shared.types.worker.downloads import DownloadProgress
 from exo.shared.types.worker.instances import Instance, InstanceId
 from exo.shared.types.worker.runners import RunnerId, RunnerStatus
-from exo.utils.info_gatherer.info_gatherer import MacThunderboltConnections
 from exo.utils.pydantic_ext import FrozenModel
 
 
@@ -45,6 +39,7 @@ class State(FrozenModel):
         alias_generator=to_camel,
         validate_by_name=True,
         extra="forbid",
+        # I want to reenable this ASAP, but it's causing an issue with TaskStatus
         strict=True,
         arbitrary_types_allowed=True,
     )
@@ -53,6 +48,7 @@ class State(FrozenModel):
     downloads: Mapping[NodeId, Sequence[DownloadProgress]] = {}
     tasks: Mapping[TaskId, Task] = {}
     last_seen: Mapping[NodeId, datetime] = {}
+    topology: Topology = Field(default_factory=Topology)
     last_event_applied_idx: int = Field(default=-1, ge=-1)
 
     # Granular node state mappings (update independently at different frequencies)
@@ -65,94 +61,34 @@ class State(FrozenModel):
     node_thunderbolt_bridge: Mapping[NodeId, ThunderboltBridgeStatus] = {}
     node_rdma_ctl: Mapping[NodeId, NodeRdmaCtlStatus] = {}
     node_backends: Mapping[NodeId, list[Backend]] = {}
-    node_socket_connections: Mapping[
-        NodeId, Mapping[NodeId, Sequence[SocketConnection]]
-    ] = {}
-    node_thunderbolt_connections: Mapping[NodeId, MacThunderboltConnections] = {}
 
     # Detected cycles where all nodes have Thunderbolt bridge enabled (>2 nodes)
     thunderbolt_bridge_cycles: Sequence[Sequence[NodeId]] = []
 
+    instance_links: Mapping[InstanceLinkId, InstanceLink] = {}
     prefill_server_ports: Mapping[RunnerId, int] = {}
 
-    @model_serializer(mode="wrap")
-    def _serialize(self, handler: SerializerFunctionWrapHandler) -> dict[str, Any]:
-        data = handler(self)  # pyright: ignore[reportAny]
-        data["topology"] = {
-            "nodes": list(self.node_identities.keys()),
-            "connections": self.topology.map_connections(),
-        }
-        return data  # pyright: ignore[reportAny]
+    # User-added model cards. Workers can reconcile their on-disk custom card cache
+    custom_model_cards: Mapping[ModelId, ModelCard] = {}
 
-    @property
-    def topology(self) -> Topology:
-        topology = Topology()
-        thunderbolt_by_uuid = {
-            ident.domain_uuid: (node_id, ident.rdma_interface)
-            for node_id, info in self.node_thunderbolt.items()
-            for ident in info.interfaces
-        }
-        for node_id in self.node_identities:
-            topology.add_node(node_id)
+    @field_serializer("topology", mode="plain")
+    def _encode_topology(self, value: Topology) -> TopologySnapshot:
+        return value.to_snapshot()
 
-        for source, data in self.node_socket_connections.items():
-            for sink, conns in data.items():
-                for conn in conns:
-                    topology.add_connection(
-                        Connection(source=source, sink=sink, edge=conn)
-                    )
+    @field_validator("topology", mode="before")
+    @classmethod
+    def _deserialize_topology(cls, value: object) -> Topology:  # noqa: D401 – Pydantic validator signature
+        """Convert an incoming *value* into a :class:`Topology` instance.
 
-        for source, connections in self.node_thunderbolt_connections.items():
-            if not self.node_rdma_ctl.get(
-                source, NodeRdmaCtlStatus(enabled=False)
-            ).enabled:
-                continue
-            for connection in connections.conns:
-                if (
-                    source_iface := thunderbolt_by_uuid.get(connection.source_uuid)
-                ) is None or (
-                    sink_iface := thunderbolt_by_uuid.get(connection.sink_uuid)
-                ) is None:
-                    continue
-                if not self.node_rdma_ctl.get(
-                    sink_iface[0], NodeRdmaCtlStatus(enabled=False)
-                ).enabled:
-                    continue
-                assert source_iface[0] == source, "registered invalid source uuid"
-                topology.add_connection(
-                    Connection(
-                        source=source_iface[0],
-                        sink=sink_iface[0],
-                        edge=RDMAConnection(
-                            source_rdma_iface=source_iface[1],
-                            sink_rdma_iface=sink_iface[1],
-                        ),
-                    )
-                )
+        Accepts either an already constructed :class:`Topology` or a mapping
+        representing :class:`~shared.topology.TopologySnapshot`.
+        """
 
-        return topology
+        if isinstance(value, Topology):
+            return value
 
-    def with_aggregator(self, aggregator: LVAggregator) -> "State":
-        from datetime import datetime, timezone
+        if isinstance(value, Mapping):  # likely a snapshot-dict coming from JSON
+            snapshot = TopologySnapshot(**cast(dict[str, Any], value))  # type: ignore[arg-type]
+            return Topology.from_snapshot(snapshot)
 
-        from pydantic import TypeAdapter
-
-        from exo.shared.apply import event_apply
-        from exo.shared.types.events import NodeGatheredInfo
-        from exo.utils.info_gatherer.info_gatherer import GatheredInfo
-
-        state = self.model_copy()
-        for key, value in aggregator.dump().items():
-            try:
-                data = TypeAdapter[GatheredInfo](GatheredInfo).validate_json(value)
-                node_id = NodeId(key.split("/")[0])
-                event = NodeGatheredInfo(
-                    node_id=node_id, when=str(datetime.now(tz=timezone.utc)), info=data
-                )
-                state = event_apply(event, state)
-            except Exception as e:
-                print(
-                    f"\n{'=' * 10}key: {key} with exception {str(e)}\nvalue: {value}{'=' * 10}\n"
-                )
-
-        return state
+        raise TypeError("Invalid representation for Topology field in State")

src/exo/shared/types/text_generation.py

@@ -69,6 +69,10 @@ class Base64Image(TruncatingString):
     truncate_length = 10
 
 
+class Base64ImageHash(TruncatingString):
+    truncate_length = 10
+
+
 def _wrap_chat_value(x: Any) -> Any:  # pyright: ignore[reportAny]
     if isinstance(x, (InputMessageContent, Base64Image)):
         return x
@@ -126,6 +130,7 @@ class TextGenerationTaskParams(BaseModel, frozen=True):
     presence_penalty: float | None = None
     frequency_penalty: float | None = None
     images: list[Base64Image] = Field(default_factory=list)
+    image_hashes: dict[int, Base64ImageHash] = Field(default_factory=dict)
 
     prefill_endpoint: str | None = None
 

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

@@ -1,4 +1,3 @@
-from collections.abc import Iterable
 from enum import Enum
 
 from pydantic import model_validator
@@ -23,14 +22,7 @@ class BaseInstance(TaggedModel):
     shard_assignments: ShardAssignments
 
     def shard(self, runner_id: RunnerId) -> ShardMetadata | None:
-        for _, rid, shard in self.shard_assignments.shards:
-            if rid == runner_id:
-                return shard
-
-    def runners_for(self, node_id: NodeId) -> Iterable[RunnerId]:
-        for nid, rid, _ in self.shard_assignments.shards:
-            if nid == node_id:
-                yield rid
+        return self.shard_assignments.runner_to_shard.get(runner_id, None)
 
 
 class MlxRingInstance(BaseInstance):
@@ -52,12 +44,6 @@ class BoundInstance(FrozenModel):
     bound_runner_id: RunnerId
     bound_node_id: NodeId
 
-    def is_primary_output_node(self) -> bool:
-        return (
-            self.instance.shard_assignments.primary_output_node
-            == self.bound_shard.device_rank
-        )
-
     @property
     def bound_shard(self) -> ShardMetadata:
         shard = self.instance.shard(self.bound_runner_id)
@@ -73,9 +59,8 @@ class BoundInstance(FrozenModel):
 
     @model_validator(mode="after")
     def validate_shard_exists(self) -> "BoundInstance":
-        assert any(
-            rid == self.bound_runner_id
-            for (_, rid, _) in self.instance.shard_assignments.shards
+        assert (
+            self.bound_runner_id in self.instance.shard_assignments.runner_to_shard
         ), (
             "Bound Instance must be constructed with a runner_id that is in the instances assigned shards"
         )

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

@@ -1,5 +1,4 @@
-from collections.abc import Sequence
-from typing import NamedTuple
+from collections.abc import Mapping
 
 from pydantic import model_validator
 
@@ -84,26 +83,16 @@ RunnerStatus = (
 )
 
 
-class ShardWithId(NamedTuple):
-    node_id: NodeId
-    runner_id: RunnerId
-    shard: ShardMetadata
-
-
 class ShardAssignments(FrozenModel):
     model_id: ModelId
-    shards: Sequence[ShardWithId]
-    # this node needs to be connected to the API node for the stream to be considered ready
-    # (this is a device rank)
-    primary_output_node: int
+    runner_to_shard: Mapping[RunnerId, ShardMetadata]
+    node_to_runner: Mapping[NodeId, RunnerId]
 
     @model_validator(mode="after")
     def validate_runners_exist(self) -> "ShardAssignments":
-        for position, shard in enumerate(self.shards):
-            if shard.shard.device_rank != position:
-                raise ValueError("shard position does not correspond to device rank")
-
-        if not self.shards[self.primary_output_node].shard.is_primary_output():
-            raise ValueError("primary output node does not correspond to primary shard")
-
+        for runner_id in self.node_to_runner.values():
+            if runner_id not in self.runner_to_shard:
+                raise ValueError(
+                    f"Runner {runner_id} in node_to_runner does not exist in runner_to_shard"
+                )
         return self

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

@@ -15,14 +15,18 @@ class Sharding(str, Enum):
 class BaseShardMetadata(TaggedModel):
     """
     Defines a specific shard of the model that is ready to be run on a device.
-    Layers are represented as a half-open interval [start_layer, end_layer),
-    where start_layer is inclusive and end_layer is exclusive.
+    Replaces previous `Shard` object.
     """
 
     model_card: ModelCard
     device_rank: int
     world_size: int
 
+    # Error handling; equivalent to monkey-patch, but we can't monkey-patch runner.py
+    # This is kinda annoying because it allocates memory in the ShardMetadata object. Can be rethought after Shanghai.
+    immediate_exception: bool = False
+    should_timeout: float | None = None
+
     start_layer: int = Field(ge=0)
     end_layer: int = Field(ge=0)
     n_layers: int = Field(ge=0)
@@ -47,59 +51,27 @@ class BaseShardMetadata(TaggedModel):
             )
         )
 
-    def is_primary_output(self) -> bool:
-        return self.device_rank == self.world_size - 1
-
-    def is_primary_output_node(self) -> bool:
-        return self.is_primary_output()
-
 
 @final
 class PipelineShardMetadata(BaseShardMetadata):
-    pass
+    """
+    Pipeline parallelism shard meta.
+
+    Layers are represented as a half-open interval [start_layer, end_layer),
+    where start_layer is inclusive and end_layer is exclusive.
+    """
 
 
 @final
 class CfgShardMetadata(BaseShardMetadata):
-    # example
-    # world_size 6
-    # rank prank crank
-    #  0     0     0
-    #  1     1     0
-    #  2     2     0
-    #  3     2     1
-    #  4     1     1
-    #  5     0     1
+    """Shard metadata for CFG-parallel image generation models."""
 
-    @property
-    def cfg_rank(self) -> int:
-        # 0 = positive branch, 1 = negative branch
-        return 0 if self.device_rank < self.world_size // 2 else 1
+    cfg_rank: int  # 0 = positive branch, 1 = negative branch
+    cfg_world_size: int = 2
 
-    @property
-    def cfg_world_size(self) -> int:
-        return 2
-
-    @property
-    def pipeline_rank(self) -> int:
-        return (
-            self.device_rank
-            if self.cfg_rank == 0
-            else (self.world_size - self.device_rank - 1)
-        )
-
-    @property
-    def pipeline_world_size(self) -> int:
-        return self.world_size // 2
-
-    def is_primary_output(self) -> bool:
-        """
-        For CFG models: the last pipeline stage in CFG group 0 (positive prompt).
-        For non-CFG models: the last pipeline stage.
-        """
-        assert self.pipeline_world_size == self.world_size // 2
-        assert self.world_size % 2 == 0
-        return self.device_rank == (self.world_size // 2) - 1
+    # Pipeline-relative coordinates (computed at placement time)
+    pipeline_rank: int  # rank within the pipeline group (0, 1, 2, ...)
+    pipeline_world_size: int  # number of nodes per pipeline group
 
 
 @final

src/exo/utils/__init__.py

@@ -1,4 +1,4 @@
-from typing import Any, Callable, Iterable, Iterator, Type, TypeGuard
+from typing import Any, Type
 
 from .phantom import PhantomData
 
@@ -19,11 +19,3 @@ def todo[T](
     _phantom: PhantomData[T] = None,
 ) -> T:
     raise NotImplementedError(msg)
-
-
-def not_none[T](t: T | None) -> TypeGuard[T]:
-    return t is not None
-
-
-def fmap[T, U](f: Callable[[T], U | None], s: Iterable[T]) -> Iterator[U]:
-    return filter(not_none, map(f, s))

src/exo/utils/banner.py

@@ -38,7 +38,7 @@ def print_startup_banner(port: int) -> None:
 
 ╔═══════════════════════════════════════════════════════════════════════╗
 ║                                                                       ║
-║  Dashboard & API Ready                                                ║
+║  🌐 Dashboard & API Ready                                             ║
 ║                                                                       ║
 ║  {dashboard_url}{" " * (69 - len(dashboard_url))}║
 ║                                                                       ║

src/exo/utils/info_gatherer/info_gatherer.py

@@ -10,13 +10,11 @@ from typing import Self, cast
 import anyio
 from anyio import fail_after, open_process, to_thread
 from anyio.streams.buffered import BufferedByteReceiveStream
-from exo_rs import LVPublisher, SessionHandle
 from loguru import logger
 from pydantic import ValidationError
 
 from exo.shared.constants import EXO_CONFIG_FILE, EXO_DEFAULT_MODELS_DIR
 from exo.shared.types.backends import Backend
-from exo.shared.types.common import NodeId
 from exo.shared.types.memory import Memory
 from exo.shared.types.profiling import (
     DiskUsage,
@@ -29,6 +27,7 @@ from exo.shared.types.thunderbolt import (
     ThunderboltConnectivity,
     ThunderboltIdentifier,
 )
+from exo.utils.channels import Sender
 from exo.utils.pydantic_ext import TaggedModel
 from exo.utils.task_group import TaskGroup
 
@@ -402,42 +401,10 @@ GatheredInfo = (
 
 @dataclass
 class InfoGatherer:
-    session_handle: SessionHandle
-    node_id: NodeId
-    info_senders: dict[str, LVPublisher] = field(init=False, default_factory=dict)
+    info_sender: Sender[GatheredInfo]
     _tg: TaskGroup = field(init=False, default_factory=TaskGroup)
     _psutil_enabled: bool = field(init=False, default=False)
 
-    async def send(self, info: GatheredInfo):
-        if (tag := info.tag()) not in self.info_senders:
-            self.info_senders[tag] = self.session_handle.last_value_publisher(
-                f"node_metrics/{self.node_id}/{tag}"
-            )
-        await self.info_senders[tag].put(info.model_dump_json())
-
-    async def run(self):
-        async with self._tg as tg:
-            if IS_DARWIN:
-                tg.start_soon(self._monitor_macmon, 1)
-                tg.start_soon(self._monitor_system_profiler_thunderbolt_data, 5)
-                tg.start_soon(self._monitor_thunderbolt_bridge_status, 10)
-                tg.start_soon(self._monitor_rdma_ctl_status, 10)
-            if not IS_DARWIN:
-                tg.start_soon(self._monitor_memory_usage, 1)
-            tg.start_soon(self._watch_system_info, 10)
-            tg.start_soon(self._monitor_misc, 60)
-            tg.start_soon(self._monitor_static_info, 60)
-            tg.start_soon(self._monitor_disk_usage, 30)
-
-            nc = await NodeConfig.gather()
-            if nc is not None:
-                await self.send(nc)
-
-            await self.send(await NodeBackends.gather())
-
-    def shutdown(self):
-        self._tg.cancel_tasks()
-
     async def _can_read_macmon_metrics(self, macmon_path: str) -> bool:
         try:
             with fail_after(5):
@@ -474,11 +441,34 @@ class InfoGatherer:
 
         return True
 
+    async def run(self):
+        async with self._tg as tg:
+            if IS_DARWIN:
+                tg.start_soon(self._monitor_macmon, 1)
+                tg.start_soon(self._monitor_system_profiler_thunderbolt_data, 5)
+                tg.start_soon(self._monitor_thunderbolt_bridge_status, 10)
+                tg.start_soon(self._monitor_rdma_ctl_status, 10)
+            if not IS_DARWIN:
+                tg.start_soon(self._monitor_memory_usage, 1)
+            tg.start_soon(self._watch_system_info, 10)
+            tg.start_soon(self._monitor_misc, 60)
+            tg.start_soon(self._monitor_static_info, 60)
+            tg.start_soon(self._monitor_disk_usage, 30)
+
+            nc = await NodeConfig.gather()
+            if nc is not None:
+                await self.info_sender.send(nc)
+
+            await self.info_sender.send(await NodeBackends.gather())
+
+    def shutdown(self):
+        self._tg.cancel_tasks()
+
     async def _monitor_static_info(self, static_info_poll_interval: float):
         while True:
             try:
                 with fail_after(30):
-                    await self.send(await StaticNodeInformation.gather())
+                    await self.info_sender.send(await StaticNodeInformation.gather())
             except Exception as e:
                 logger.opt(exception=e).warning("Error gathering static node info")
             await anyio.sleep(static_info_poll_interval)
@@ -487,7 +477,7 @@ class InfoGatherer:
         while True:
             try:
                 with fail_after(10):
-                    await self.send(await MiscData.gather())
+                    await self.info_sender.send(await MiscData.gather())
             except Exception as e:
                 logger.opt(exception=e).warning("Error gathering misc data")
             await anyio.sleep(misc_poll_interval)
@@ -508,10 +498,12 @@ class InfoGatherer:
                     idents = [
                         it for i in data if (it := i.ident(iface_map)) is not None
                     ]
-                    await self.send(MacThunderboltIdentifiers(idents=idents))
+                    await self.info_sender.send(
+                        MacThunderboltIdentifiers(idents=idents)
+                    )
 
                     conns = [it for i in data if (it := i.conn()) is not None]
-                    await self.send(MacThunderboltConnections(conns=conns))
+                    await self.info_sender.send(MacThunderboltConnections(conns=conns))
             except Exception as e:
                 logger.opt(exception=e).warning("Error gathering Thunderbolt data")
             await anyio.sleep(system_profiler_interval)
@@ -528,7 +520,7 @@ class InfoGatherer:
         )
         while True:
             try:
-                await self.send(
+                await self.info_sender.send(
                     MemoryUsage.from_psutil(override_memory=override_memory)
                 )
             except Exception as e:
@@ -540,7 +532,7 @@ class InfoGatherer:
             try:
                 with fail_after(10):
                     nics = await get_network_interfaces()
-                    await self.send(NodeNetworkInterfaces(ifaces=nics))
+                    await self.info_sender.send(NodeNetworkInterfaces(ifaces=nics))
             except Exception as e:
                 logger.opt(exception=e).warning("Error gathering network interfaces")
             await anyio.sleep(interface_watcher_interval)
@@ -553,7 +545,7 @@ class InfoGatherer:
                 with fail_after(30):
                     curr = await ThunderboltBridgeInfo.gather()
                     if curr is not None:
-                        await self.send(curr)
+                        await self.info_sender.send(curr)
             except Exception as e:
                 logger.opt(exception=e).warning(
                     "Error gathering Thunderbolt Bridge status"
@@ -565,7 +557,7 @@ class InfoGatherer:
             try:
                 curr = await RdmaCtlStatus.gather()
                 if curr is not None:
-                    await self.send(curr)
+                    await self.info_sender.send(curr)
             except Exception as e:
                 logger.opt(exception=e).warning("Error gathering RDMA ctl status")
             await anyio.sleep(rdma_ctl_poll_interval)
@@ -574,7 +566,7 @@ class InfoGatherer:
         while True:
             try:
                 with fail_after(5):
-                    await self.send(await NodeDiskUsage.gather())
+                    await self.info_sender.send(await NodeDiskUsage.gather())
             except Exception as e:
                 logger.opt(exception=e).warning("Error gathering disk usage")
             await anyio.sleep(disk_poll_interval)
@@ -619,7 +611,7 @@ class InfoGatherer:
                             )
                             text = data.decode("utf-8", errors="replace").strip()
                             metrics = MacmonMetrics.from_raw_json(text)
-                        await self.send(metrics)
+                        await self.info_sender.send(metrics)
             except TimeoutError:
                 logger.warning(
                     f"MacMon produced no output for {read_timeout}s, restarting"

src/exo/utils/pydantic_ext.py

@@ -19,22 +19,18 @@ class FrozenModel(BaseModel):
 
 
 class TaggedModel(FrozenModel):
-    @classmethod
-    def tag(cls) -> str:
-        return cls.__name__
-
     @model_serializer(mode="wrap")
     def _serialize(self, handler: SerializerFunctionWrapHandler):
         inner = handler(self)  # pyright: ignore[reportAny]
-        return {self.tag(): inner}
+        return {self.__class__.__name__: inner}
 
     @model_validator(mode="wrap")
     @classmethod
     def _validate(cls, v: Any, handler: ValidatorFunctionWrapHandler) -> Self:  # pyright: ignore[reportAny]
-        if isinstance(v, dict) and len(v) == 1 and cls.tag() in v:  # pyright: ignore[reportUnknownArgumentType]
-            return handler(v[cls.tag()])  # pyright: ignore[reportAny]
+        if isinstance(v, dict) and len(v) == 1 and cls.__name__ in v:  # pyright: ignore[reportUnknownArgumentType]
+            return handler(v[cls.__name__])  # pyright: ignore[reportAny]
 
         return handler(v)  # pyright: ignore[reportAny]
 
     def __str__(self) -> str:
-        return f"{self.tag()}({super().__str__()})"
+        return f"{self.__class__.__name__}({super().__str__()})"

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

@@ -30,6 +30,8 @@ from exo.shared.types.worker.runner_response import (
     ModelLoadingResponse,
 )
 from exo.shared.types.worker.shards import (
+    CfgShardMetadata,
+    PipelineShardMetadata,
     ShardMetadata,
 )
 from exo.utils.channels import MpReceiver, MpSender
@@ -47,6 +49,22 @@ from exo.worker.engines.mlx.utils_mlx import (
 )
 
 
+def _is_primary_output_node(shard_metadata: ShardMetadata) -> bool:
+    """Check if this node is the primary output node for image generation.
+
+    For CFG models: the last pipeline stage in CFG group 0 (positive prompt).
+    For non-CFG models: the last pipeline stage.
+    """
+    if isinstance(shard_metadata, CfgShardMetadata):
+        is_pipeline_last = (
+            shard_metadata.pipeline_rank == shard_metadata.pipeline_world_size - 1
+        )
+        return is_pipeline_last and shard_metadata.cfg_rank == 0
+    elif isinstance(shard_metadata, PipelineShardMetadata):
+        return shard_metadata.device_rank == shard_metadata.world_size - 1
+    return False
+
+
 def _send_traces_if_enabled(
     event_sender: MpSender[Event],
     task_id: TaskId,
@@ -153,7 +171,7 @@ class ImageEngine(Engine):
             resp = next(self.current_gen, None)
         return (
             (resp,)
-            if resp is not None and self.shard_metadata.is_primary_output()
+            if resp is not None and _is_primary_output_node(self.shard_metadata)
             else ()
         )
 
@@ -184,10 +202,10 @@ class ImageEngine(Engine):
                 task=task_params,
                 cancel_checker=cancel_checker,
             ):
-                if self.shard_metadata.is_primary_output():
+                if _is_primary_output_node(self.shard_metadata):
                     yield (task_id, response)
         except Exception as e:
-            if self.shard_metadata.is_primary_output():
+            if _is_primary_output_node(self.shard_metadata):
                 yield (
                     task_id,
                     ErrorChunk(

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

@@ -38,13 +38,11 @@ class MlxBuilder(Builder):
     tokenizer: TokenizerWrapper | None = None
     group: mx.distributed.Group | None = None
     vision_processor: VisionProcessor | None = None
-    is_primary_output_node: bool = False
 
     def connect(self, bound_instance: BoundInstance) -> None:
         self.group = initialize_mlx(bound_instance)
 
     def load(self, bound_instance: BoundInstance) -> Generator[ModelLoadingResponse]:
-        self.is_primary_output_node = bound_instance.is_primary_output_node()
         (
             self.inference_model,
             self.tokenizer,
@@ -84,6 +82,7 @@ class MlxBuilder(Builder):
 
         kv_prefix_cache = KVPrefixCache(self.group)
 
+        device_rank = 0 if self.group is None else self.group.rank()
         if os.environ.get("EXO_NO_BATCH"):
             logger.info("using SequentialGenerator (batching disabled)")
             return SequentialGenerator(
@@ -93,7 +92,7 @@ class MlxBuilder(Builder):
                 tool_parser=tool_parser,
                 kv_prefix_cache=kv_prefix_cache,
                 model_id=self.model_id,
-                is_primary_output_node=self.is_primary_output_node,
+                device_rank=device_rank,
                 cancel_receiver=self.cancel_receiver,
                 event_sender=self.event_sender,
                 vision_processor=vision_processor,
@@ -107,7 +106,7 @@ class MlxBuilder(Builder):
                 tool_parser=tool_parser,
                 kv_prefix_cache=kv_prefix_cache,
                 model_id=self.model_id,
-                is_primary_output_node=self.is_primary_output_node,
+                device_rank=device_rank,
                 cancel_receiver=self.cancel_receiver,
                 event_sender=self.event_sender,
                 vision_processor=vision_processor,

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

@@ -229,6 +229,47 @@ def has_non_kv_caches(cache: KVCacheType) -> bool:
     return any(is_non_trimmable_cache_entry(c) for c in cache)
 
 
+# Max snapshots retained per cache entry. Each CacheSnapshot pins detached GPU
+# copies of every non-trimmable (SSM/ArraysCache, RotatingKVCache) layer, so
+# retaining one per ~4096-token prefill chunk makes snapshot memory grow linearly
+# with context — the dominant residual cost when a single entry is grown to long
+# contexts on hybrid models (~56 MB/snapshot on Qwen3.5-122B, so a full 256K
+# context = 64 snapshots ≈ 3.6 GB). A sliding window of the most-recent N caps
+# this at N×per-snapshot (~0.9 GB here) while preserving the restore points
+# in-place grows actually use (they always extend from the tip).
+_MAX_RETAINED_SNAPSHOTS = 16
+
+
+def _bounded_snapshots(snapshots: list[CacheSnapshot]) -> list[CacheSnapshot]:
+    """Deduplicate snapshots by token position and bound the retained count.
+
+    Returned list is sorted ascending by ``token_count``.
+    """
+    # Deduplicate by position, keeping the most-recently-appended snapshot per
+    # position. Repeated in-place grows re-snapshot positions the kept old
+    # snapshots already cover, which would otherwise grow `_snapshots`
+    # unbounded even at constant context.
+    # TODO: keying on token_count alone is safe only while a position uniquely
+    # identifies the prefix within an entry (grows are strict prefix-extensions).
+    # If edit-and-regenerate, sliding-window/prefix trimming, cross-entry
+    # snapshot sharing, per-request adapter/LoRA swap, or branchy decoding
+    # (beam/parallel/speculative) is added, enrich the key to
+    # (token_count, prefix_hash[, media/adapter id]) — else a stale snapshot
+    # could be restored for a different prefix (silent wrong output).
+    by_position: dict[int, CacheSnapshot] = {}
+    for snapshot in snapshots:
+        by_position[snapshot.token_count] = snapshot
+    deduped = [by_position[pos] for pos in sorted(by_position)]
+
+    # Sliding window: keep only the most-recent N positions. In-place grows
+    # always extend from the tip, so the newest snapshots are the ones future
+    # grows restore from — dropping the oldest is never incorrect: a later hit on
+    # a prefix older than the window finds no snapshot <= target, so get_kv_cache
+    # returns a fresh cache (matched_index=None) and the request takes a full cold
+    # prefill — correct, just slower than a partial-hit reuse for that one request.
+    return deduped[-_MAX_RETAINED_SNAPSHOTS:]
+
+
 class KVPrefixCache:
     def __init__(self, group: mx.distributed.Group | None):
         self.prompts: list[mx.array] = []  # mx array of tokens (ints)
@@ -261,7 +302,9 @@ class KVPrefixCache:
         self._evict_if_needed()
         self.prompts.append(prompt_tokens)
         self.caches.append(deepcopy(cache))
-        self._snapshots.append(ssm_snapshots)
+        self._snapshots.append(
+            _bounded_snapshots(ssm_snapshots) if ssm_snapshots else None
+        )
         self._media_regions.append(media_regions or [])
         self.prefill_tps.append(prefill_tps)
         self._access_counter += 1
@@ -288,7 +331,7 @@ class KVPrefixCache:
 
         self.prompts[index] = prompt_tokens
         self.caches[index] = deepcopy(cache)
-        self._snapshots[index] = merged or None
+        self._snapshots[index] = _bounded_snapshots(merged) or None
         self._media_regions[index] = media_regions or []
         self.prefill_tps[index] = prefill_tps
         self._access_counter += 1

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

@@ -154,7 +154,7 @@ def initialize_mlx(
     # TODO: pass in seed from params
     mx.random.seed(42)
 
-    assert len(bound_instance.instance.shard_assignments.shards) > 1, (
+    assert len(bound_instance.instance.shard_assignments.node_to_runner) > 1, (
         "Tried to initialize mlx for a single node instance"
     )
     return mlx_distributed_init(bound_instance)

src/exo/worker/main.py

@@ -1,12 +1,12 @@
+import hashlib
 from collections import defaultdict
 from datetime import datetime, timezone
 
 import anyio
 from anyio import fail_after, to_thread
-from exo_rs import LVAggregator, SessionHandle
 from loguru import logger
-from pydantic import ValidationError
 
+from exo.api.types import ImageEditsTaskParams
 from exo.download.download_utils import is_read_only_model_dir, resolve_existing_model
 from exo.routing.event_router import (
     EventRouterBrokenResourceError,
@@ -14,18 +14,19 @@ from exo.routing.event_router import (
 )
 from exo.shared.apply import apply
 from exo.shared.constants import EXO_MAX_INSTANCE_RETRIES
-from exo.shared.models import model_cards
-from exo.shared.models.model_cards import ModelCard, ModelId
+from exo.shared.models.model_cards import ModelId, card_cache
+from exo.shared.types.chunks import InputImageChunk
 from exo.shared.types.commands import (
     DeleteInstance,
     ForwarderCommand,
     ForwarderDownloadCommand,
     StartDownload,
 )
-from exo.shared.types.common import NodeId, SystemId
+from exo.shared.types.common import CommandId, NodeId, SystemId
 from exo.shared.types.events import (
     Event,
     IndexedEvent,
+    InputChunkReceived,
     InstanceDeleted,
     NodeDownloadProgress,
     NodeGatheredInfo,
@@ -40,16 +41,19 @@ from exo.shared.types.tasks import (
     CancelTask,
     CreateRunner,
     DownloadModel,
+    ImageEdits,
     LoadModel,
     Shutdown,
     Task,
     TaskStatus,
+    TextGeneration,
 )
+from exo.shared.types.text_generation import Base64Image, Base64ImageHash
 from exo.shared.types.topology import Connection, SocketConnection
 from exo.shared.types.worker.downloads import DownloadCompleted
 from exo.shared.types.worker.instances import InstanceId
 from exo.shared.types.worker.runners import RunnerId
-from exo.utils.channels import Receiver, Sender
+from exo.utils.channels import Receiver, Sender, channel
 from exo.utils.info_gatherer.info_gatherer import GatheredInfo, InfoGatherer
 from exo.utils.info_gatherer.net_profile import check_reachable
 from exo.utils.keyed_backoff import KeyedBackoff
@@ -69,7 +73,6 @@ class Worker:
         # but I think it's the correct way to be thinking about commands
         command_sender: Sender[ForwarderCommand],
         download_command_sender: Sender[ForwarderDownloadCommand],
-        session_handle: SessionHandle,
         api_port: int,
     ):
         self.node_id: NodeId = node_id
@@ -85,25 +88,27 @@ class Worker:
 
         self._system_id = SystemId()
 
+        # Buffer for input image chunks (for image editing)
+        self.input_chunk_buffer: dict[CommandId, dict[int, InputImageChunk]] = {}
+        self.input_chunk_counts: dict[CommandId, int] = {}
+        self.image_cache: dict[Base64ImageHash, Base64Image] = {}
 
         self._download_backoff: KeyedBackoff[ModelId] = KeyedBackoff(base=0.5, cap=10.0)
         self._instance_backoff: KeyedBackoff[InstanceId] = KeyedBackoff(
             base=0.5, cap=10.0
         )
         self._stopped: anyio.Event = anyio.Event()
-        self._sh: SessionHandle = session_handle
-        self.aggregator: LVAggregator = session_handle.last_value_aggregator(
-            "node_metrics"
-        )
 
     async def run(self):
         logger.info("Starting Worker")
 
-        info_gatherer: InfoGatherer = InfoGatherer(self._sh, self.node_id)
+        info_send, info_recv = channel[GatheredInfo]()
+        info_gatherer: InfoGatherer = InfoGatherer(info_send)
 
         try:
             async with self._tg as tg:
                 tg.start_soon(info_gatherer.run)
+                tg.start_soon(self._forward_info, info_recv)
                 tg.start_soon(self.plan_step)
                 tg.start_soon(self._event_applier)
                 tg.start_soon(self._poll_connection_updates)
@@ -142,26 +147,48 @@ class Worker:
                 if isinstance(event, InstanceDeleted):
                     self._instance_backoff.reset(event.instance_id)
 
+                # Buffer input image chunks for image editing
+                if isinstance(event, InputChunkReceived):
+                    cmd_id = event.command_id
+                    if cmd_id not in self.input_chunk_buffer:
+                        self.input_chunk_buffer[cmd_id] = {}
+                        self.input_chunk_counts[cmd_id] = event.chunk.total_chunks
+
+                    self.input_chunk_buffer[cmd_id][event.chunk.chunk_index] = (
+                        event.chunk
+                    )
+                    if (
+                        len(self.input_chunk_buffer[cmd_id])
+                        == self.input_chunk_counts[cmd_id]
+                    ):
+                        per_image: defaultdict[int, list[InputImageChunk]] = (
+                            defaultdict(list)
+                        )
+                        for chunk in self.input_chunk_buffer[cmd_id].values():
+                            per_image[chunk.image_index].append(chunk)
+                        for chunks_for_image in per_image.values():
+                            sorted_chunks = sorted(
+                                chunks_for_image, key=lambda c: c.chunk_index
+                            )
+                            img = Base64Image("".join(c.data for c in sorted_chunks))
+                            self.image_cache[
+                                Base64ImageHash(
+                                    hashlib.sha256(img.encode("ascii")).hexdigest()
+                                )
+                            ] = img
 
     async def _reconcile_custom_cards(self) -> None:
-        storage = self._sh.storage_interface()
         while True:
-            await anyio.sleep(10)
-            target: list[ModelId] = []
-            for _, value in (await storage.dump("custom_model_cards/")).items():
-                try:
-                    card = ModelCard.model_validate_json(value)
-                except ValidationError:
+            await anyio.sleep(1)
+            target = dict(self.state.custom_model_cards)
+            for model_id, card in target.items():
+                if card_cache.get(model_id) == card:
                     continue
-                target.append(card.model_id)
-                if model_cards.card_cache.get(card.model_id) == card:
-                    continue
-                logger.info(f"Registered new custom model card for {card.model_id}")
-                await model_cards.card_cache.save(card)
+                await card_cache.save(card)
 
-            for card in await model_cards.card_cache.list_all():
+            for card in await card_cache.list_all():
                 if card.model_id not in target:
-                    await model_cards.card_cache.delete(card.model_id)
+                    await card_cache.pop(card.model_id)
 
     async def plan_step(self):
         while True:
@@ -173,6 +200,8 @@ class Worker:
                 self.state.instances,
                 self.state.runners,
                 self.state.tasks,
+                self.input_chunk_buffer,
+                self.image_cache,
                 self._instance_backoff,
                 self._download_backoff,
             )
@@ -271,6 +300,63 @@ class Worker:
                             task_id=task.task_id, task_status=TaskStatus.Complete
                         )
                     )
+                case ImageEdits() if task.task_params.total_input_chunks > 0:
+                    # Assemble image from chunks and inject into task
+                    cmd_id = task.command_id
+                    chunks = self.input_chunk_buffer.get(cmd_id, {})
+                    assembled = "".join(chunks[i].data for i in range(len(chunks)))
+                    logger.info(
+                        f"Assembled input image from {len(chunks)} chunks, "
+                        f"total size: {len(assembled)} bytes"
+                    )
+                    # Create modified task with assembled image data
+                    modified_task = ImageEdits(
+                        task_id=task.task_id,
+                        command_id=task.command_id,
+                        instance_id=task.instance_id,
+                        task_status=task.task_status,
+                        task_params=ImageEditsTaskParams(
+                            image_data=assembled,
+                            total_input_chunks=task.task_params.total_input_chunks,
+                            prompt=task.task_params.prompt,
+                            model=task.task_params.model,
+                            n=task.task_params.n,
+                            quality=task.task_params.quality,
+                            output_format=task.task_params.output_format,
+                            response_format=task.task_params.response_format,
+                            size=task.task_params.size,
+                            image_strength=task.task_params.image_strength,
+                            bench=task.task_params.bench,
+                            stream=task.task_params.stream,
+                            partial_images=task.task_params.partial_images,
+                            advanced_params=task.task_params.advanced_params,
+                        ),
+                    )
+                    # Cleanup buffers
+                    if cmd_id in self.input_chunk_buffer:
+                        del self.input_chunk_buffer[cmd_id]
+                    if cmd_id in self.input_chunk_counts:
+                        del self.input_chunk_counts[cmd_id]
+                    await self._start_runner_task(modified_task)
+
+                case TextGeneration() if task.task_params.image_hashes:
+                    cmd_id = task.command_id
+                    resolved_images = [
+                        self.image_cache[h]
+                        for _, h in sorted(task.task_params.image_hashes.items())
+                    ]
+                    modified_task = task.model_copy(
+                        update={
+                            "task_params": task.task_params.model_copy(
+                                update={"images": resolved_images}
+                            )
+                        }
+                    )
+                    if cmd_id in self.input_chunk_buffer:
+                        del self.input_chunk_buffer[cmd_id]
+                    if cmd_id in self.input_chunk_counts:
+                        del self.input_chunk_counts[cmd_id]
+                    await self._start_runner_task(modified_task)
                 case LoadModel(instance_id=instance_id):
                     if (instance := self.state.instances.get(instance_id)) is not None:
                         model_id = instance.shard_assignments.model_id
@@ -286,23 +372,15 @@ class Worker:
 
     async def _start_runner_task(self, task: Task):
         if (instance := self.state.instances.get(task.instance_id)) is not None:
-            for rid in instance.runners_for(self.node_id):
-                await self.runners[rid].start_task(task)
+            await self.runners[
+                instance.shard_assignments.node_to_runner[self.node_id]
+            ].start_task(task)
 
     async def _create_supervisor(self, task: CreateRunner) -> RunnerSupervisor:
         """Creates and stores a new AssignedRunner with initial downloading status."""
-        task_responder = (
-            self._sh.task_responder(task.instance_id)
-            if task.bound_instance.is_primary_output_node()
-            else None
-        )
         runner = await RunnerSupervisor.create(
             bound_instance=task.bound_instance,
             event_sender=self.event_sender.clone(),
-            task_assignment_subscriber=self._sh.last_value_subscriber(
-                f"task_assignments/{task.instance_id}/*"
-            ),
-            task_responder=task_responder,
         )
         self.runners[task.bound_instance.bound_runner_id] = runner
         self._tg.start_soon(runner.run)
@@ -310,13 +388,14 @@ class Worker:
 
     async def _poll_connection_updates(self):
         while True:
-            state = self.state.with_aggregator(self.aggregator)
-            edges = set(conn.edge for conn in state.topology.out_edges(self.node_id))
+            edges = set(
+                conn.edge for conn in self.state.topology.out_edges(self.node_id)
+            )
             conns: defaultdict[NodeId, set[str]] = defaultdict(set)
             async for ip, nid in check_reachable(
-                state.topology,
+                self.state.topology,
                 self.node_id,
-                state.node_network,
+                self.state.node_network,
                 api_port=self.api_port,
             ):
                 if ip in conns[nid]:
@@ -337,7 +416,7 @@ class Worker:
                         )
                     )
 
-            for conn in state.topology.out_edges(self.node_id):
+            for conn in self.state.topology.out_edges(self.node_id):
                 if not isinstance(conn.edge, SocketConnection):
                     continue
                 # ignore mDNS discovered connections

src/exo/worker/plan.py

@@ -2,19 +2,24 @@
 
 from collections.abc import Mapping, Sequence
 
-from exo.shared.types.common import ModelId, NodeId
+from exo.shared.types.chunks import InputImageChunk
+from exo.shared.types.common import CommandId, ModelId, NodeId
 from exo.shared.types.tasks import (
     CancelTask,
     ConnectToGroup,
     CreateRunner,
     DownloadModel,
+    ImageEdits,
+    ImageGeneration,
     LoadModel,
     Shutdown,
     StartWarmup,
     Task,
     TaskId,
     TaskStatus,
+    TextGeneration,
 )
+from exo.shared.types.text_generation import Base64Image, Base64ImageHash
 from exo.shared.types.worker.downloads import (
     DownloadCompleted,
     DownloadFailed,
@@ -30,10 +35,11 @@ from exo.shared.types.worker.runners import (
     RunnerIdle,
     RunnerLoaded,
     RunnerLoading,
+    RunnerReady,
+    RunnerRunning,
     RunnerStatus,
     RunnerWarmingUp,
 )
-from exo.utils import fmap
 from exo.utils.keyed_backoff import KeyedBackoff
 from exo.worker.runner.supervisor import RunnerSupervisor
 
@@ -46,6 +52,8 @@ def plan(
     instances: Mapping[InstanceId, Instance],
     all_runners: Mapping[RunnerId, RunnerStatus],  # all global
     tasks: Mapping[TaskId, Task],
+    input_chunk_buffer: Mapping[CommandId, Mapping[int, InputImageChunk]],
+    image_cache: Mapping[Base64ImageHash, Base64Image],
     instance_backoff: KeyedBackoff[InstanceId],
     download_backoff: KeyedBackoff[ModelId],
 ) -> Task | None:
@@ -60,6 +68,7 @@ def plan(
         or _init_distributed_backend(runners, all_runners)
         or _load_model(runners, all_runners, global_download_status)
         or _ready_to_warmup(runners, all_runners)
+        or _pending_tasks(runners, tasks, all_runners, input_chunk_buffer, image_cache)
     )
 
 
@@ -79,10 +88,8 @@ def _kill_runner(
             )
 
         for (
-            _,
-            global_runner_id,
-            _,
-        ) in runner.bound_instance.instance.shard_assignments.shards:
+            global_runner_id
+        ) in runner.bound_instance.instance.shard_assignments.node_to_runner.values():
             if runner_id == global_runner_id:
                 continue
 
@@ -101,13 +108,7 @@ def _create_runner(
     instance_backoff: KeyedBackoff[InstanceId],
 ) -> CreateRunner | None:
     for instance in instances.values():
-        runner_id = next(
-            fmap(
-                lambda it: it.runner_id if it.node_id == node_id else None,
-                instance.shard_assignments.shards,
-            ),
-            None,
-        )
+        runner_id = instance.shard_assignments.node_to_runner.get(node_id, None)
         if runner_id is None:
             continue
 
@@ -117,7 +118,7 @@ def _create_runner(
         # don't create runners if any other nodes have runners that have failed - wait for them to fix themselves first.
         instance_has_failed_runner = any(
             isinstance(all_runners.get(remote_runner_id), RunnerFailed)
-            for (_, remote_runner_id, _) in instance.shard_assignments.shards
+            for remote_runner_id in instance.shard_assignments.node_to_runner.values()
             if remote_runner_id != runner_id
         )
         we_have_failed_before = isinstance(all_runners.get(runner_id), RunnerFailed)
@@ -174,7 +175,7 @@ def _init_distributed_backend(
         instance = runner.bound_instance.instance
         shard_assignments = instance.shard_assignments
 
-        is_single_node_instance = len(shard_assignments.shards) == 1
+        is_single_node_instance = len(shard_assignments.runner_to_shard) == 1
         if is_single_node_instance:
             continue
 
@@ -184,7 +185,7 @@ def _init_distributed_backend(
                 all_runners.get(global_runner_id),
                 (RunnerConnecting, RunnerIdle),
             )
-            for (_, global_runner_id, _) in shard_assignments.shards
+            for global_runner_id in shard_assignments.runner_to_shard
         )
 
         if not (runner_is_idle and all_runners_connecting):
@@ -204,7 +205,7 @@ def _init_distributed_backend(
         # Rank = n-1
         connecting_rank_ready = device_rank == world_size - 1 and all(
             isinstance(all_runners.get(global_runner_id, None), RunnerConnecting)
-            for (_, global_runner_id, _) in shard_assignments.shards
+            for global_runner_id in shard_assignments.runner_to_shard
             if global_runner_id != runner_id
         )
 
@@ -232,12 +233,12 @@ def _load_model(
                 and dp.shard_metadata.model_card.model_id == shard_assignments.model_id
                 for dp in global_download_status[nid]
             )
-            for (nid, _, _) in shard_assignments.shards
+            for nid in shard_assignments.node_to_runner
         )
         if not all_local_downloads_complete:
             continue
 
-        is_single_node_instance = len(instance.shard_assignments.shards) == 1
+        is_single_node_instance = len(instance.shard_assignments.runner_to_shard) == 1
         if is_single_node_instance and isinstance(runner.status, RunnerIdle):
             return LoadModel(instance_id=instance.instance_id)
 
@@ -248,7 +249,7 @@ def _load_model(
                 all_runners.get(global_runner_id, None),
                 (RunnerConnected, RunnerLoading, RunnerLoaded),
             )
-            for (_, global_runner_id, _) in shard_assignments.shards
+            for global_runner_id in shard_assignments.runner_to_shard
         )
 
         if is_runner_waiting and all_ready_for_model:
@@ -280,13 +281,13 @@ def _ready_to_warmup(
                 all_runners.get(global_runner_id, None),
                 (RunnerLoaded, RunnerWarmingUp),
             )
-            for (_, global_runner_id, _) in shard_assignments.shards
+            for global_runner_id in shard_assignments.runner_to_shard
         )
 
         # Rank = 0
         connecting_rank_ready = device_rank == 0 and all(
             isinstance(all_runners.get(global_runner_id, None), RunnerWarmingUp)
-            for (_, global_runner_id, _) in shard_assignments.shards
+            for global_runner_id in shard_assignments.runner_to_shard
             if global_runner_id != runner_id
         )
 
@@ -296,6 +297,52 @@ def _ready_to_warmup(
     return None
 
 
+def _pending_tasks(
+    runners: Mapping[RunnerId, RunnerSupervisor],
+    tasks: Mapping[TaskId, Task],
+    all_runners: Mapping[RunnerId, RunnerStatus],
+    input_chunk_buffer: Mapping[CommandId, Mapping[int, InputImageChunk]],
+    image_cache: Mapping[Base64ImageHash, Base64Image],
+) -> Task | None:
+    for task in tasks.values():
+        # for now, just forward chat completions
+        # TODO(ciaran): do this better!
+        if not isinstance(task, (TextGeneration, ImageGeneration, ImageEdits)):
+            continue
+        if task.task_status not in (TaskStatus.Pending, TaskStatus.Running):
+            continue
+
+        if isinstance(task, ImageEdits) and task.task_params.total_input_chunks > 0:
+            received = len(input_chunk_buffer.get(task.command_id, {}))
+            if received < task.task_params.total_input_chunks:
+                continue  # Wait for all chunks to arrive
+
+        if (
+            isinstance(task, TextGeneration)
+            and task.task_params.image_hashes
+            and not all(
+                h in image_cache for h in task.task_params.image_hashes.values()
+            )
+        ):
+            continue  # Wait for all images to be assembled into the cache
+
+        for runner in runners.values():
+            if task.instance_id != runner.bound_instance.instance.instance_id:
+                continue
+
+            # the task status _should_ be set to completed by the LAST runner
+            # it is currently set by the first
+            # this is definitely a hack
+            if task.task_id in runner.completed or task.task_id in runner.in_progress:
+                continue
+
+            if isinstance(runner.status, (RunnerReady, RunnerRunning)) and all(
+                isinstance(all_runners[global_runner_id], (RunnerReady, RunnerRunning))
+                for global_runner_id in runner.bound_instance.instance.shard_assignments.runner_to_shard
+            ):
+                return task
+
+
 def _cancel_tasks(
     runners: Mapping[RunnerId, RunnerSupervisor],
     tasks: Mapping[TaskId, Task],

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

@@ -94,7 +94,7 @@ class SequentialGenerator(Engine):
     kv_prefix_cache: KVPrefixCache | None
     tool_parser: ToolParser | None
     model_id: ModelId
-    is_primary_output_node: bool
+    device_rank: int
     cancel_receiver: MpReceiver[TaskId]
     event_sender: MpSender[Event]
     vision_processor: VisionProcessor | None = None
@@ -201,8 +201,7 @@ class SequentialGenerator(Engine):
 
         return filter(
             lambda chunk: (
-                self.is_primary_output_node
-                or isinstance(chunk[1], (CancelledResponse, FinishedResponse))
+                not isinstance(chunk[1], GenerationChunk) or self.device_rank == 0
             ),
             itertools.chain(
                 output,
@@ -236,7 +235,7 @@ class SequentialGenerator(Engine):
         self._active = (task, gen, queue, output_generator)
 
     def _send_error(self, task: TextGeneration, e: Exception) -> None:
-        if self.is_primary_output_node:
+        if self.device_rank == 0:
             self.event_sender.send(
                 ChunkGenerated(
                     command_id=task.command_id,
@@ -253,7 +252,7 @@ class SequentialGenerator(Engine):
         prompt = apply_chat_template(self.tokenizer, task.task_params)
 
         def on_prefill_progress(processed: int, total: int) -> None:
-            if self.is_primary_output_node:
+            if self.device_rank == 0:
                 self.event_sender.send(
                     ChunkGenerated(
                         command_id=task.command_id,
@@ -326,7 +325,7 @@ class BatchGenerator(Engine):
     kv_prefix_cache: KVPrefixCache | None
     tool_parser: ToolParser | None
     model_id: ModelId
-    is_primary_output_node: bool
+    device_rank: int
     cancel_receiver: MpReceiver[TaskId]
     event_sender: MpSender[Event]
     check_for_cancel_every: int = 50
@@ -461,8 +460,7 @@ class BatchGenerator(Engine):
 
         return filter(
             lambda chunk: (
-                self.is_primary_output_node
-                or isinstance(chunk[1], (CancelledResponse, FinishedResponse))
+                not isinstance(chunk[1], GenerationChunk) or self.device_rank == 0
             ),
             itertools.chain(output, self._apply_cancellations()),
         )
@@ -496,7 +494,7 @@ class BatchGenerator(Engine):
         return iter(results)
 
     def _send_error(self, task: TextGeneration, e: Exception) -> None:
-        if self.is_primary_output_node:
+        if self.device_rank == 0:
             self.event_sender.send(
                 ChunkGenerated(
                     command_id=task.command_id,
@@ -513,7 +511,7 @@ class BatchGenerator(Engine):
         prompt = apply_chat_template(self.tokenizer, task.task_params)
 
         def on_prefill_progress(processed: int, total: int) -> None:
-            if self.is_primary_output_node:
+            if self.device_rank == 0:
                 self.event_sender.send(
                     ChunkGenerated(
                         command_id=task.command_id,

src/exo/worker/runner/runner.py

@@ -102,6 +102,10 @@ class Runner:
         self.device_rank = self.shard_metadata.device_rank
 
         logger.info("hello from the runner")
+        if getattr(self.shard_metadata, "immediate_exception", False):
+            raise Exception("Fake exception - runner failed to spin up.")
+        if timeout := getattr(self.shard_metadata, "should_timeout", 0):
+            time.sleep(timeout)
 
         self.setup_start_time = time.time()
 

src/exo/worker/runner/supervisor.py

@@ -12,34 +12,15 @@ from anyio import (
     CancelScope,
     ClosedResourceError,
 )
-from exo_rs import LVSubscriber, TaskChunkSender, TaskRequest, TaskResponder
 from loguru import logger
-from pydantic import TypeAdapter, ValidationError
 
 from exo.shared.constants import EXO_RUNNER_STDERR_LOG, EXO_RUNNER_STDOUT_LOG
-from exo.shared.types.chunks import ErrorChunk, PrefillProgressChunk
-from exo.shared.types.commands import (
-    Command,
-    TaskCancelled,
-    TaskFinished,
-)
-from exo.shared.types.commands import (
-    ImageEdits as ImageEditsCommand,
-)
-from exo.shared.types.commands import (
-    ImageGeneration as ImageGenerationCommand,
-)
-from exo.shared.types.commands import (
-    TextGeneration as TextGenerationCommand,
-)
-from exo.shared.types.common import CommandId
+from exo.shared.types.chunks import ErrorChunk
 from exo.shared.types.events import (
     ChunkGenerated,
     Event,
     RunnerStatusUpdated,
     TaskAcknowledged,
-    TaskCreated,
-    TaskDeleted,
     TaskStatusUpdated,
 )
 from exo.shared.types.tasks import (
@@ -57,7 +38,6 @@ from exo.shared.types.worker.runners import (
     RunnerFailed,
     RunnerIdle,
     RunnerLoading,
-    RunnerReady,
     RunnerRunning,
     RunnerShuttingDown,
     RunnerStatus,
@@ -76,22 +56,6 @@ from exo.worker.runner.diagnostics import (
 
 PREFILL_TIMEOUT_SECONDS = 60
 DECODE_TIMEOUT_SECONDS = 5
-type BridgeTask = TextGeneration | ImageGeneration | ImageEdits
-_BRIDGE_COMMAND_ADAPTER: TypeAdapter[Command] = TypeAdapter(Command)
-_BRIDGE_TASK_ADAPTER: TypeAdapter[BridgeTask] = TypeAdapter(BridgeTask)
-
-
-def _task_assignment_ids(key: str) -> tuple[str, TaskId] | None:
-    prefix = "task_assignments/"
-    if not key.startswith(prefix):
-        return None
-
-    suffix = key.removeprefix(prefix)
-    parts = suffix.split("/", maxsplit=1)
-    if len(parts) != 2 or not parts[0] or not parts[1] or "/" in parts[1]:
-        return None
-
-    return parts[0], TaskId(parts[1])
 
 
 @dataclass(eq=False)
@@ -225,22 +189,12 @@ class RunnerSupervisor:
     _task_sender: MpSender[Task]
     _event_sender: Sender[Event]
     _cancel_sender: MpSender[TaskId]
-    _task_responder: TaskResponder | None = None
-    _task_assignment_subscriber: LVSubscriber | None = None
-    _assigned_tasks: dict[TaskId, BridgeTask] = field(default_factory=dict, init=False)
     _tg: TaskGroup = field(default_factory=TaskGroup, init=False)
     status: RunnerStatus = field(default_factory=RunnerIdle, init=False)
     pending: dict[TaskId, anyio.Event] = field(default_factory=dict, init=False)
     in_progress: dict[TaskId, Task] = field(default_factory=dict, init=False)
     completed: set[TaskId] = field(default_factory=set, init=False)
     cancelled: set[TaskId] = field(default_factory=set, init=False)
-    bridge_tasks: dict[TaskId, BridgeTask] = field(default_factory=dict, init=False)
-    bridge_command_tasks: dict[CommandId, TaskId] = field(
-        default_factory=dict, init=False
-    )
-    bridge_chunk_senders: dict[CommandId, TaskChunkSender] = field(
-        default_factory=dict, init=False
-    )
     _cancel_watch_runner: anyio.CancelScope = field(
         default_factory=anyio.CancelScope, init=False
     )
@@ -251,8 +205,6 @@ class RunnerSupervisor:
         *,
         bound_instance: BoundInstance,
         event_sender: Sender[Event],
-        task_assignment_subscriber: LVSubscriber | None = None,
-        task_responder: TaskResponder | None = None,
         initialize_timeout: float = 400,
     ) -> Self:
         ev_send, ev_recv = mp_channel[Event | RunnerTerminationError]()
@@ -286,8 +238,6 @@ class RunnerSupervisor:
             _task_sender=task_sender,
             _cancel_sender=cancel_sender,
             _event_sender=event_sender,
-            _task_responder=task_responder,
-            _task_assignment_subscriber=task_assignment_subscriber,
         )
 
         return self
@@ -301,13 +251,6 @@ class RunnerSupervisor:
 
                 tg.start_soon(self._watch_runner)
                 tg.start_soon(self._forward_events)
-                if self._task_responder is not None:
-                    tg.start_soon(self._run_task_responder, self._task_responder)
-                if self._task_assignment_subscriber is not None:
-                    tg.start_soon(
-                        self._run_task_assignment_subscriber,
-                        self._task_assignment_subscriber,
-                    )
         finally:
             logger.info("Runner supervisor shutting down")
             if not self._cancel_watch_runner.cancel_called:
@@ -332,44 +275,6 @@ class RunnerSupervisor:
     def shutdown(self):
         self._tg.cancel_tasks()
 
-    async def _run_task_assignment_subscriber(
-        self, subscriber: LVSubscriber
-    ) -> None:
-        instance_id = self.bound_instance.instance.instance_id
-        while (received := await subscriber.recv()) is not None:
-            key, payload = received
-            if (ids := _task_assignment_ids(key)) is None:
-                continue
-            assigned_instance_id, assigned_task_id = ids
-            if assigned_instance_id != instance_id:
-                continue
-
-            if payload == "":
-                self._assigned_tasks.pop(assigned_task_id, None)
-                continue
-
-            try:
-                task = _BRIDGE_TASK_ADAPTER.validate_json(payload)
-            except ValidationError:
-                logger.warning(f"Ignoring invalid task assignment from {key}")
-                continue
-
-            if task.instance_id != instance_id or task.task_id != assigned_task_id:
-                logger.warning(f"Ignoring mismatched task assignment from {key}")
-                continue
-
-            self._assigned_tasks[task.task_id] = task
-            await self._reconcile_assigned_tasks()
-
-    async def _reconcile_assigned_tasks(self) -> None:
-        if not isinstance(self.status, (RunnerReady, RunnerRunning)):
-            return
-
-        for task in list(self._assigned_tasks.values()):
-            if task.task_id in self.in_progress or task.task_id in self.completed:
-                continue
-            await self.start_task(task)
-
     async def start_task(self, task: Task):
         if task.task_id in self.pending:
             logger.warning(
@@ -415,58 +320,33 @@ class RunnerSupervisor:
         try:
             with self._ev_recv as events:
                 async for event in events:
-                    match event:
-                        case RunnerTerminationError():
-                            await self._check_runner(event)
-                            break
-
-                        case RunnerStatusUpdated(runner_status=runner_status):
-                            self.status = runner_status
-                            await self._event_sender.send(event)
-                            await self._reconcile_assigned_tasks()
-
-                        case TaskAcknowledged(task_id=task_id):
-                            self.pending.pop(task_id).set()
-
-                        case TaskStatusUpdated(
-                            task_id=task_id, task_status=TaskStatus.Complete
-                        ):
-                            # If a task has just been completed, we should be working on it.
-                            assert isinstance(
-                                self.status,
-                                (
-                                    RunnerRunning,
-                                    RunnerWarmingUp,
-                                    RunnerLoading,
-                                    RunnerConnecting,
-                                    RunnerShuttingDown,
-                                ),
-                            )
-                            self.in_progress.pop(task_id, None)
-                            self.completed.add(task_id)
-                            self._assigned_tasks.pop(task_id, None)
-                            await self._event_sender.send(event)
-                            if task_id in self.bridge_tasks:
-                                await self._finish_bridge_task_id(task_id)
-
-                        case ChunkGenerated(command_id=command_id, chunk=chunk):
-                            task_id = self.bridge_command_tasks.get(command_id)
-                            chunk_sender = self.bridge_chunk_senders.get(command_id)
-                            if task_id is None or chunk_sender is None:
-                                logger.debug(
-                                    f"Dropping bridge chunk for inactive command {command_id}"
-                                )
-                                continue
-
-                            await chunk_sender.send(chunk.model_dump_json())
-                            if (
-                                not isinstance(chunk, PrefillProgressChunk)
-                                and chunk.finish_reason is not None
-                            ):
-                                self.bridge_chunk_senders.pop(command_id, None)
-
-                        case _:
-                            await self._event_sender.send(event)
+                    if isinstance(event, RunnerTerminationError):
+                        # try to get exception if possible
+                        await self._check_runner(event)
+                        break
+                    if isinstance(event, RunnerStatusUpdated):
+                        self.status = event.runner_status
+                    if isinstance(event, TaskAcknowledged):
+                        self.pending.pop(event.task_id).set()
+                        continue
+                    if (
+                        isinstance(event, TaskStatusUpdated)
+                        and event.task_status == TaskStatus.Complete
+                    ):
+                        # If a task has just been completed, we should be working on it.
+                        assert isinstance(
+                            self.status,
+                            (
+                                RunnerRunning,
+                                RunnerWarmingUp,
+                                RunnerLoading,
+                                RunnerConnecting,
+                                RunnerShuttingDown,
+                            ),
+                        )
+                        self.in_progress.pop(event.task_id, None)
+                        self.completed.add(event.task_id)
+                    await self._event_sender.send(event)
         except (ClosedResourceError, BrokenResourceError):
             # this is the happy path shutdown - we don't need to spam log with it
             await self._check_runner()
@@ -474,109 +354,6 @@ class RunnerSupervisor:
             for tid in self.pending:
                 self.pending[tid].set()
 
-    async def _run_task_responder(self, responder: TaskResponder) -> None:
-        while True:
-            received = await responder.recv()
-            if received is None:
-                return
-            request, chunk_sender, payload = received
-            if payload is None:
-                request.reply_err("Task command query did not include a payload")
-                continue
-
-            try:
-                command = _BRIDGE_COMMAND_ADAPTER.validate_json(payload)
-                match command:
-                    case TextGenerationCommand():
-                        await self._submit_bridge_task(request, chunk_sender, command)
-                    case ImageGenerationCommand():
-                        await self._submit_bridge_task(request, chunk_sender, command)
-                    case ImageEditsCommand():
-                        await self._submit_bridge_task(request, chunk_sender, command)
-                    case TaskCancelled(cancelled_command_id=command_id):
-                        await self._cancel_bridge_task(command_id)
-                        request.reply(command_id)
-                    case TaskFinished(finished_command_id=command_id):
-                        await self._finish_bridge_task(command_id)
-                        request.reply(command_id)
-                    case _:
-                        request.reply_err(f"Unsupported bridge command: {command}")
-            except Exception as exception:
-                logger.opt(exception=exception).warning(
-                    "Failed to admit bridge command"
-                )
-                request.reply_err(str(exception))
-
-    async def _submit_bridge_task(
-        self,
-        request: TaskRequest,
-        chunk_sender: TaskChunkSender,
-        command: TextGenerationCommand | ImageGenerationCommand | ImageEditsCommand,
-    ) -> None:
-        task_id = TaskId()
-        match command:
-            case TextGenerationCommand():
-                task = TextGeneration(
-                    task_id=task_id,
-                    command_id=command.command_id,
-                    instance_id=self.bound_instance.instance.instance_id,
-                    task_status=TaskStatus.Pending,
-                    task_params=command.task_params,
-                )
-            case ImageGenerationCommand():
-                task = ImageGeneration(
-                    task_id=task_id,
-                    command_id=command.command_id,
-                    instance_id=self.bound_instance.instance.instance_id,
-                    task_status=TaskStatus.Pending,
-                    task_params=command.task_params,
-                )
-            case ImageEditsCommand():
-                task = ImageEdits(
-                    task_id=task_id,
-                    command_id=command.command_id,
-                    instance_id=self.bound_instance.instance.instance_id,
-                    task_status=TaskStatus.Pending,
-                    task_params=command.task_params,
-                )
-
-        self.bridge_tasks[task.task_id] = task
-        self.bridge_command_tasks[command.command_id] = task.task_id
-        self.bridge_chunk_senders[command.command_id] = chunk_sender
-        await self._event_sender.send(TaskCreated(task_id=task.task_id, task=task))
-        assert self._task_responder is not None
-        await self._task_responder.assign_task(task.task_id, task.model_dump_json())
-        request.reply(command.command_id)
-
-    async def _cancel_bridge_task(self, command_id: CommandId) -> None:
-        task_id = self.bridge_command_tasks.get(command_id)
-        if task_id is None:
-            logger.warning(f"Unable to cancel unknown bridge command {command_id}")
-            return
-        await self.cancel_task(task_id)
-        await self._event_sender.send(
-            TaskStatusUpdated(task_id=task_id, task_status=TaskStatus.Cancelled)
-        )
-        await self._finish_bridge_task_id(task_id)
-
-    async def _finish_bridge_task(self, command_id: CommandId) -> None:
-        task_id = self.bridge_command_tasks.get(command_id)
-        if task_id is None:
-            logger.warning(f"Unable to finish unknown bridge command {command_id}")
-            return
-        await self._finish_bridge_task_id(task_id)
-
-    async def _finish_bridge_task_id(self, task_id: TaskId) -> None:
-        task = self.bridge_tasks.pop(task_id, None)
-        if task is None:
-            logger.warning(f"Unable to finish unknown bridge task {task_id}")
-            return
-        self.bridge_command_tasks.pop(task.command_id, None)
-        self.bridge_chunk_senders.pop(task.command_id, None)
-        await self._event_sender.send(TaskDeleted(task_id=task_id))
-        if self._task_responder is not None:
-            await self._task_responder.unassign_task(task_id)
-
     async def _watch_runner(self) -> None:
         with self._cancel_watch_runner:
             while True:

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

@@ -11,12 +11,7 @@ from exo.shared.types.worker.instances import (
     InstanceId,
     MlxRingInstance,
 )
-from exo.shared.types.worker.runners import (
-    RunnerId,
-    RunnerStatus,
-    ShardAssignments,
-    ShardWithId,
-)
+from exo.shared.types.worker.runners import RunnerId, RunnerStatus, ShardAssignments
 from exo.shared.types.worker.shards import PipelineShardMetadata, ShardMetadata
 
 
@@ -57,22 +52,16 @@ def get_pipeline_shard_metadata(
     )
 
 
-# todo: clean up legacy formatted shards
 def get_shard_assignments(
     model_id: ModelId,
     node_to_runner: dict[NodeId, RunnerId],
     runner_to_shard: dict[RunnerId, ShardMetadata],
 ) -> ShardAssignments:
-    pon = 0
-    shards = [
-        ShardWithId(nid, rid := node_to_runner[nid], runner_to_shard[rid])
-        for nid in node_to_runner
-    ]
-    for i, (_, _, shard) in enumerate(shards):
-        if shard.is_primary_output():
-            pon = i
-
-    return ShardAssignments(model_id=model_id, shards=shards, primary_output_node=pon)
+    return ShardAssignments(
+        model_id=model_id,
+        node_to_runner=node_to_runner,
+        runner_to_shard=runner_to_shard,
+    )
 
 
 def get_mlx_ring_instance(

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

@@ -11,6 +11,7 @@ from mlx_lm.sample_utils import make_sampler
 from exo.shared.types.common import ModelId
 from exo.shared.types.text_generation import InputMessage, TextGenerationTaskParams
 from exo.worker.engines.mlx.cache import (
+    CacheSnapshot,
     KVPrefixCache,
     cache_length,
     encode_prompt,
@@ -77,6 +78,74 @@ class TestGetPrefixLength:
         assert get_prefix_length(a, b) == 0
 
 
+class TestSnapshotAccumulation:
+    """Locks in the fix for the actual per-grow Metal leak on hybrid (SSM)
+    models: `update_kv_cache` must not let `_snapshots` grow without bound when
+    the same entry is grown in place many times."""
+
+    def test_repeated_update_does_not_accumulate_snapshots(self):
+        with patch(
+            "exo.worker.engines.mlx.cache.get_memory_used_percentage",
+            return_value=0.0,
+        ):
+            kv_prefix_cache = KVPrefixCache(None)
+            initial = [
+                CacheSnapshot(states=[None], token_count=4096),
+                CacheSnapshot(states=[None], token_count=8192),
+            ]
+            kv_prefix_cache.add_kv_cache(
+                mx.arange(10000), [KVCache()], ssm_snapshots=initial
+            )
+
+            # Each in-place grow re-prefills from restore_pos and produces a
+            # fresh snapshot at a position the retained old snapshots already
+            # cover. Pre-fix this appended one snapshot per grow forever.
+            for _ in range(50):
+                fresh = [CacheSnapshot(states=[None], token_count=8192)]
+                kv_prefix_cache.update_kv_cache(
+                    0, mx.arange(10000), [KVCache()], fresh, restore_pos=8192
+                )
+
+            stored = kv_prefix_cache._snapshots[0]
+            assert stored is not None
+            # Bounded by the number of distinct snapshot positions (here 2),
+            # not by the 50 grows.
+            assert len(stored) == 2
+            assert sorted(s.token_count for s in stored) == [4096, 8192]
+            # The kept 8192 snapshot must be the most recently supplied one.
+            assert stored[1] is fresh[0]
+
+    def test_extension_caps_snapshots_to_sliding_window(self):
+        """Extending a single entry to a long context (one snapshot per ~4096
+        tokens) must cap retained snapshots to a sliding window of the most-recent
+        N, not keep all of them — that linear-in-context retention was the
+        residual OOM cause."""
+        from exo.worker.engines.mlx.cache import _MAX_RETAINED_SNAPSHOTS
+
+        with patch(
+            "exo.worker.engines.mlx.cache.get_memory_used_percentage",
+            return_value=0.0,
+        ):
+            kv_prefix_cache = KVPrefixCache(None)
+            # 64 distinct positions = a 262144-token context at 4096/chunk.
+            num_positions = 64
+            snaps = [
+                CacheSnapshot(states=[None], token_count=4096 * (i + 1))
+                for i in range(num_positions)
+            ]
+            kv_prefix_cache.add_kv_cache(
+                mx.arange(10), [KVCache()], ssm_snapshots=snaps
+            )
+
+            stored = kv_prefix_cache._snapshots[0]
+            assert stored is not None
+            # Capped at the window; the most-recent N positions are retained
+            # (in-place grows extend from the tip, so these are what get used).
+            assert len(stored) == _MAX_RETAINED_SNAPSHOTS
+            assert stored == snaps[-_MAX_RETAINED_SNAPSHOTS:]
+            assert stored[-1] is snaps[-1]  # tip always kept
+
+
 class TestKVPrefix:
     @pytest.fixture
     def mock_tokenizer(self):

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

@@ -53,6 +53,8 @@ def test_plan_requests_download_when_waiting_and_shard_not_downloaded():
         instances=instances,
         all_runners=all_runners,
         tasks={},
+        input_chunk_buffer={},
+        image_cache={},
         instance_backoff=KeyedBackoff(),
         download_backoff=KeyedBackoff(),
     )
@@ -107,6 +109,8 @@ def test_plan_loads_model_when_all_shards_downloaded_and_waiting():
         instances=instances,
         all_runners=all_runners,
         tasks={},
+        input_chunk_buffer={},
+        image_cache={},
         instance_backoff=KeyedBackoff(),
         download_backoff=KeyedBackoff(),
     )
@@ -151,6 +155,8 @@ def test_plan_does_not_request_download_when_shard_already_downloaded():
         instances=instances,
         all_runners=all_runners,
         tasks={},
+        input_chunk_buffer={},
+        image_cache={},
         instance_backoff=KeyedBackoff(),
         download_backoff=KeyedBackoff(),
     )
@@ -200,6 +206,8 @@ def test_plan_does_not_load_model_until_all_shards_downloaded_globally():
         instances=instances,
         all_runners=all_runners,
         tasks={},
+        input_chunk_buffer={},
+        image_cache={},
         instance_backoff=KeyedBackoff(),
         download_backoff=KeyedBackoff(),
     )
@@ -222,6 +230,8 @@ def test_plan_does_not_load_model_until_all_shards_downloaded_globally():
         instances=instances,
         all_runners=all_runners,
         tasks={},
+        input_chunk_buffer={},
+        image_cache={},
         instance_backoff=KeyedBackoff(),
         download_backoff=KeyedBackoff(),
     )

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

@@ -53,6 +53,8 @@ def test_plan_kills_runner_when_instance_missing():
         instances=instances,
         all_runners=all_runners,
         tasks={},
+        input_chunk_buffer={},
+        image_cache={},
         instance_backoff=KeyedBackoff(),
         download_backoff=KeyedBackoff(),
     )
@@ -94,6 +96,8 @@ def test_plan_kills_runner_when_sibling_failed():
         instances=instances,
         all_runners=all_runners,
         tasks={},
+        input_chunk_buffer={},
+        image_cache={},
         instance_backoff=KeyedBackoff(),
         download_backoff=KeyedBackoff(),
     )
@@ -127,6 +131,8 @@ def test_plan_creates_runner_when_missing_for_node():
         instances=instances,
         all_runners=all_runners,
         tasks={},
+        input_chunk_buffer={},
+        image_cache={},
         instance_backoff=KeyedBackoff(),
         download_backoff=KeyedBackoff(),
     )
@@ -167,6 +173,8 @@ def test_plan_does_not_create_runner_when_supervisor_already_present():
         instances=instances,
         all_runners=all_runners,
         tasks={},
+        input_chunk_buffer={},
+        image_cache={},
         instance_backoff=KeyedBackoff(),
         download_backoff=KeyedBackoff(),
     )
@@ -198,6 +206,8 @@ def test_plan_does_not_create_runner_for_unassigned_node():
         instances=instances,
         all_runners=all_runners,
         tasks={},
+        input_chunk_buffer={},
+        image_cache={},
         instance_backoff=KeyedBackoff(),
         download_backoff=KeyedBackoff(),
     )

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

@@ -0,0 +1,303 @@
+from typing import cast
+
+import exo.worker.plan as plan_mod
+from exo.shared.types.tasks import Task, TaskId, TaskStatus, TextGeneration
+from exo.shared.types.text_generation import (
+    InputMessage,
+    InputMessageContent,
+    TextGenerationTaskParams,
+)
+from exo.shared.types.worker.instances import BoundInstance, InstanceId
+from exo.shared.types.worker.runners import (
+    RunnerIdle,
+    RunnerReady,
+    RunnerRunning,
+)
+from exo.utils.keyed_backoff import KeyedBackoff
+from exo.worker.tests.constants import (
+    COMMAND_1_ID,
+    INSTANCE_1_ID,
+    MODEL_A_ID,
+    NODE_A,
+    NODE_B,
+    RUNNER_1_ID,
+    RUNNER_2_ID,
+    TASK_1_ID,
+)
+from exo.worker.tests.unittests.conftest import (
+    FakeRunnerSupervisor,
+    OtherTask,
+    get_mlx_ring_instance,
+    get_pipeline_shard_metadata,
+)
+
+
+def test_plan_forwards_pending_chat_completion_when_runner_ready():
+    """
+    When there is a pending TextGeneration for the local instance and all
+    runners are Ready/Running, plan() should forward that task.
+    """
+    shard0 = get_pipeline_shard_metadata(MODEL_A_ID, device_rank=0, world_size=2)
+    shard1 = get_pipeline_shard_metadata(MODEL_A_ID, device_rank=1, world_size=2)
+    instance = get_mlx_ring_instance(
+        instance_id=INSTANCE_1_ID,
+        model_id=MODEL_A_ID,
+        node_to_runner={NODE_A: RUNNER_1_ID, NODE_B: RUNNER_2_ID},
+        runner_to_shard={RUNNER_1_ID: shard0, RUNNER_2_ID: shard1},
+    )
+    bound_instance = BoundInstance(
+        instance=instance, bound_runner_id=RUNNER_1_ID, bound_node_id=NODE_A
+    )
+    local_runner = FakeRunnerSupervisor(
+        bound_instance=bound_instance, status=RunnerReady()
+    )
+
+    runners = {RUNNER_1_ID: local_runner}
+    instances = {INSTANCE_1_ID: instance}
+    all_runners = {
+        RUNNER_1_ID: RunnerReady(),
+        RUNNER_2_ID: RunnerReady(),
+    }
+
+    task = TextGeneration(
+        task_id=TASK_1_ID,
+        instance_id=INSTANCE_1_ID,
+        task_status=TaskStatus.Pending,
+        command_id=COMMAND_1_ID,
+        task_params=TextGenerationTaskParams(
+            model=MODEL_A_ID,
+            input=[InputMessage(role="user", content=InputMessageContent(""))],
+        ),
+    )
+
+    result = plan_mod.plan(
+        node_id=NODE_A,
+        runners=runners,  # type: ignore
+        global_download_status={NODE_A: []},
+        instances=instances,
+        all_runners=all_runners,
+        tasks={TASK_1_ID: task},
+        input_chunk_buffer={},
+        image_cache={},
+        instance_backoff=KeyedBackoff(),
+        download_backoff=KeyedBackoff(),
+    )
+
+    assert result is task
+
+
+def test_plan_does_not_forward_chat_completion_if_any_runner_not_ready():
+    """
+    Even with a pending TextGeneration, plan() should not forward it unless
+    all runners for the instance are Ready/Running.
+    """
+    shard1 = get_pipeline_shard_metadata(MODEL_A_ID, device_rank=0, world_size=2)
+    shard2 = get_pipeline_shard_metadata(MODEL_A_ID, device_rank=1, world_size=2)
+    instance = get_mlx_ring_instance(
+        instance_id=INSTANCE_1_ID,
+        model_id=MODEL_A_ID,
+        node_to_runner={NODE_A: RUNNER_1_ID, NODE_B: RUNNER_2_ID},
+        runner_to_shard={RUNNER_1_ID: shard1, RUNNER_2_ID: shard2},
+    )
+    bound_instance = BoundInstance(
+        instance=instance, bound_runner_id=RUNNER_1_ID, bound_node_id=NODE_A
+    )
+    local_runner = FakeRunnerSupervisor(
+        bound_instance=bound_instance, status=RunnerReady()
+    )
+
+    runners = {RUNNER_1_ID: local_runner}
+    instances = {INSTANCE_1_ID: instance}
+    all_runners = {
+        RUNNER_1_ID: RunnerReady(),
+        RUNNER_2_ID: RunnerIdle(),
+    }
+
+    task = TextGeneration(
+        task_id=TASK_1_ID,
+        instance_id=INSTANCE_1_ID,
+        task_status=TaskStatus.Pending,
+        command_id=COMMAND_1_ID,
+        task_params=TextGenerationTaskParams(
+            model=MODEL_A_ID,
+            input=[InputMessage(role="user", content=InputMessageContent(""))],
+        ),
+    )
+
+    result = plan_mod.plan(
+        node_id=NODE_A,
+        runners=runners,  # type: ignore
+        global_download_status={NODE_A: [], NODE_B: []},
+        instances=instances,
+        all_runners=all_runners,
+        tasks={TASK_1_ID: task},
+        input_chunk_buffer={},
+        image_cache={},
+        instance_backoff=KeyedBackoff(),
+        download_backoff=KeyedBackoff(),
+    )
+
+    assert result is None
+
+
+def test_plan_does_not_forward_tasks_for_other_instances():
+    """
+    plan() should ignore pending TextGeneration tasks whose instance_id does
+    not match the local instance.
+    """
+    shard = get_pipeline_shard_metadata(model_id=MODEL_A_ID, device_rank=0)
+    local_instance = get_mlx_ring_instance(
+        instance_id=INSTANCE_1_ID,
+        model_id=MODEL_A_ID,
+        node_to_runner={NODE_A: RUNNER_1_ID},
+        runner_to_shard={RUNNER_1_ID: shard},
+    )
+    bound_instance = BoundInstance(
+        instance=local_instance, bound_runner_id=RUNNER_1_ID, bound_node_id=NODE_A
+    )
+    local_runner = FakeRunnerSupervisor(
+        bound_instance=bound_instance, status=RunnerReady()
+    )
+
+    runners = {RUNNER_1_ID: local_runner}
+    instances = {INSTANCE_1_ID: local_instance}
+    all_runners = {RUNNER_1_ID: RunnerReady()}
+
+    other_instance_id = InstanceId("instance-2")
+    foreign_task = TextGeneration(
+        task_id=TaskId("other-task"),
+        instance_id=other_instance_id,
+        task_status=TaskStatus.Pending,
+        command_id=COMMAND_1_ID,
+        task_params=TextGenerationTaskParams(
+            model=MODEL_A_ID,
+            input=[InputMessage(role="user", content=InputMessageContent(""))],
+        ),
+    )
+
+    result = plan_mod.plan(
+        node_id=NODE_A,
+        runners=runners,  # type: ignore
+        global_download_status={NODE_A: []},
+        instances=instances,
+        all_runners=all_runners,
+        tasks={foreign_task.task_id: foreign_task},
+        input_chunk_buffer={},
+        image_cache={},
+        instance_backoff=KeyedBackoff(),
+        download_backoff=KeyedBackoff(),
+    )
+
+    assert result is None
+
+
+def test_plan_ignores_non_pending_or_non_chat_tasks():
+    """
+    _pending_tasks should not forward tasks that are either not TextGeneration
+    or not in Pending/Running states.
+    """
+    shard0 = get_pipeline_shard_metadata(MODEL_A_ID, device_rank=0, world_size=2)
+    shard1 = get_pipeline_shard_metadata(MODEL_A_ID, device_rank=1, world_size=2)
+    instance = get_mlx_ring_instance(
+        instance_id=INSTANCE_1_ID,
+        model_id=MODEL_A_ID,
+        node_to_runner={NODE_A: RUNNER_1_ID, NODE_B: RUNNER_2_ID},
+        runner_to_shard={RUNNER_1_ID: shard0, RUNNER_2_ID: shard1},
+    )
+    bound_instance = BoundInstance(
+        instance=instance, bound_runner_id=RUNNER_1_ID, bound_node_id=NODE_A
+    )
+
+    local_runner = FakeRunnerSupervisor(
+        bound_instance=bound_instance, status=RunnerReady()
+    )
+
+    runners = {RUNNER_1_ID: local_runner}
+    instances = {INSTANCE_1_ID: instance}
+    all_runners = {
+        RUNNER_1_ID: RunnerReady(),
+        RUNNER_2_ID: RunnerReady(),
+    }
+
+    completed_task = TextGeneration(
+        task_id=TASK_1_ID,
+        instance_id=INSTANCE_1_ID,
+        task_status=TaskStatus.Complete,
+        command_id=COMMAND_1_ID,
+        task_params=TextGenerationTaskParams(
+            model=MODEL_A_ID,
+            input=[InputMessage(role="user", content=InputMessageContent(""))],
+        ),
+    )
+
+    other_task_id = TaskId("other-task")
+
+    other_task = cast(
+        Task,
+        cast(
+            object,
+            OtherTask(
+                task_id=other_task_id,
+                instance_id=INSTANCE_1_ID,
+                task_status=TaskStatus.Pending,
+            ),
+        ),
+    )
+
+    result = plan_mod.plan(
+        node_id=NODE_A,
+        runners=runners,  # type: ignore
+        global_download_status={NODE_A: [], NODE_B: []},
+        instances=instances,
+        all_runners=all_runners,
+        tasks={TASK_1_ID: completed_task, other_task_id: other_task},
+        input_chunk_buffer={},
+        image_cache={},
+        instance_backoff=KeyedBackoff(),
+        download_backoff=KeyedBackoff(),
+    )
+
+    assert result is None
+
+
+def test_plan_returns_none_when_nothing_to_do():
+    """
+    If there are healthy runners, no downloads needed, and no pending tasks,
+    plan() should return None (steady state).
+    """
+    shard0 = get_pipeline_shard_metadata(MODEL_A_ID, device_rank=0, world_size=2)
+    shard1 = get_pipeline_shard_metadata(MODEL_A_ID, device_rank=1, world_size=2)
+    instance = get_mlx_ring_instance(
+        instance_id=INSTANCE_1_ID,
+        model_id=MODEL_A_ID,
+        node_to_runner={NODE_A: RUNNER_1_ID, NODE_B: RUNNER_2_ID},
+        runner_to_shard={RUNNER_1_ID: shard0, RUNNER_2_ID: shard1},
+    )
+    bound_instance = BoundInstance(
+        instance=instance, bound_runner_id=RUNNER_1_ID, bound_node_id=NODE_A
+    )
+    local_runner = FakeRunnerSupervisor(
+        bound_instance=bound_instance, status=RunnerRunning()
+    )
+
+    runners = {RUNNER_1_ID: local_runner}
+    instances = {INSTANCE_1_ID: instance}
+    all_runners = {
+        RUNNER_1_ID: RunnerRunning(),
+        RUNNER_2_ID: RunnerRunning(),
+    }
+
+    result = plan_mod.plan(
+        node_id=NODE_A,
+        runners=runners,  # type: ignore
+        global_download_status={NODE_A: [], NODE_B: []},
+        instances=instances,
+        all_runners=all_runners,
+        tasks={},
+        input_chunk_buffer={},
+        image_cache={},
+        instance_backoff=KeyedBackoff(),
+        download_backoff=KeyedBackoff(),
+    )
+
+    assert result is None

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

@@ -62,6 +62,8 @@ def test_plan_starts_warmup_for_accepting_rank_when_all_loaded_or_warming():
         instances=instances,
         all_runners=all_runners,
         tasks={},
+        input_chunk_buffer={},
+        image_cache={},
         instance_backoff=KeyedBackoff(),
         download_backoff=KeyedBackoff(),
     )
@@ -105,6 +107,8 @@ def test_plan_starts_warmup_for_rank_zero_after_others_warming():
         instances=instances,
         all_runners=all_runners,
         tasks={},
+        input_chunk_buffer={},
+        image_cache={},
         instance_backoff=KeyedBackoff(),
         download_backoff=KeyedBackoff(),
     )
@@ -147,6 +151,8 @@ def test_plan_does_not_start_warmup_for_non_zero_rank_until_all_loaded_or_warmin
         instances=instances,
         all_runners=all_runners,
         tasks={},
+        input_chunk_buffer={},
+        image_cache={},
         instance_backoff=KeyedBackoff(),
         download_backoff=KeyedBackoff(),
     )
@@ -193,6 +199,8 @@ def test_plan_does_not_start_warmup_for_rank_zero_until_others_warming():
         instances=instances,
         all_runners=all_runners,
         tasks={},
+        input_chunk_buffer={},
+        image_cache={},
         instance_backoff=KeyedBackoff(),
         download_backoff=KeyedBackoff(),
     )
@@ -211,6 +219,8 @@ def test_plan_does_not_start_warmup_for_rank_zero_until_others_warming():
         instances=instances,
         all_runners=all_runners,
         tasks={},
+        input_chunk_buffer={},
+        image_cache={},
         instance_backoff=KeyedBackoff(),
         download_backoff=KeyedBackoff(),
     )
@@ -256,6 +266,8 @@ def test_plan_starts_warmup_for_connecting_rank_after_others_warming():
         instances=instances,
         all_runners=all_runners,
         tasks={},
+        input_chunk_buffer={},
+        image_cache={},
         instance_backoff=KeyedBackoff(),
         download_backoff=KeyedBackoff(),
     )
@@ -300,6 +312,8 @@ def test_plan_does_not_start_warmup_for_accepting_rank_until_all_loaded_or_warmi
         instances=instances,
         all_runners=all_runners,
         tasks={},
+        input_chunk_buffer={},
+        image_cache={},
         instance_backoff=KeyedBackoff(),
         download_backoff=KeyedBackoff(),
     )
@@ -343,6 +357,8 @@ def test_plan_does_not_start_warmup_for_connecting_rank_until_others_warming():
         instances=instances,
         all_runners=all_runners,
         tasks={},
+        input_chunk_buffer={},
+        image_cache={},
         instance_backoff=KeyedBackoff(),
         download_backoff=KeyedBackoff(),
     )