Move mx_all_gather_tasks into utils_mlx

what da ya think!

---------

Co-authored-by: Ryuichi Leo Takashige <leo@exolabs.net>

Cargo.lock

@@ -922,6 +922,7 @@ dependencies = [
  "pyo3-log",
  "pyo3-stub-gen",
  "tokio",
+ "util",
 ]
 
 [[package]]
@@ -2792,6 +2793,7 @@ dependencies = [
  "pin-project",
  "tokio",
  "tracing-subscriber",
+ "util",
 ]
 
 [[package]]
@@ -4621,6 +4623,10 @@ version = "0.2.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "06abde3611657adf66d383f00b093d7faecc7fa57071cce2578660c9f1010821"
 
+[[package]]
+name = "util"
+version = "0.0.1"
+
 [[package]]
 name = "uuid"
 version = "1.19.0"

Cargo.toml

@@ -1,6 +1,6 @@
 [workspace]
 resolver = "3"
-members = ["rust/networking", "rust/exo_pyo3_bindings"]
+members = ["rust/networking", "rust/exo_pyo3_bindings", "rust/util"]
 
 [workspace.package]
 version = "0.0.1"
@@ -20,6 +20,7 @@ opt-level = 3
 [workspace.dependencies]
 ## Crate members as common dependencies
 networking = { path = "rust/networking" }
+util = { path = "rust/util" }
 
 # Macro dependecies
 extend = "1.2"

rust/exo_pyo3_bindings/Cargo.toml

@@ -54,6 +54,9 @@ delegate = { workspace = true }
 tokio = { workspace = true, features = ["full", "tracing"] }
 futures-lite = { workspace = true }
 
+# utility dependencies
+util = { workspace = true }
+
 # Tracing
 log = { workspace = true }
 env_logger = "0.11"

rust/exo_pyo3_bindings/src/networking.rs

@@ -11,7 +11,7 @@ use crate::networking::exception::{
 use crate::pyclass;
 use futures_lite::{Stream, StreamExt as _};
 use libp2p::gossipsub::PublishError;
-use networking::{FromSwarm, ToSwarm, create_swarm};
+use networking::swarm::{FromSwarm, ToSwarm, create_swarm};
 use pyo3::exceptions::PyRuntimeError;
 use pyo3::prelude::{PyModule, PyModuleMethods as _};
 use pyo3::types::PyBytes;

rust/networking/Cargo.toml

@@ -27,6 +27,7 @@ futures-timer = { workspace = true }
 tokio = { workspace = true, features = ["full"] }
 
 # utility dependencies
+util = { workspace = true }
 tracing-subscriber = { version = "0.3.19", features = [
   "default",
   "env-filter",

rust/networking/examples/chatroom.rs

@@ -1,6 +1,7 @@
-use futures_lite::StreamExt as _;
+use futures_lite::StreamExt;
 use libp2p::identity;
-use networking::{FromSwarm, ToSwarm};
+use networking::swarm;
+use networking::swarm::{FromSwarm, ToSwarm};
 use tokio::sync::{mpsc, oneshot};
 use tokio::{io, io::AsyncBufReadExt as _};
 use tracing_subscriber::EnvFilter;
@@ -8,23 +9,23 @@ use tracing_subscriber::filter::LevelFilter;
 
 #[tokio::main]
 async fn main() {
-    tracing_subscriber::fmt()
+    let _ = tracing_subscriber::fmt()
         .with_env_filter(EnvFilter::from_default_env().add_directive(LevelFilter::INFO.into()))
-        .init();
+        .try_init();
 
     let (to_swarm, from_client) = mpsc::channel(20);
 
     // Configure swarm
-    let mut swarm = networking::create_swarm(identity::Keypair::generate_ed25519(), from_client)
+    let mut swarm = swarm::create_swarm(identity::Keypair::generate_ed25519(), from_client)
         .expect("Swarm creation failed")
         .into_stream();
 
     // Create a Gossipsub topic & subscribe
-    let (sub_tx, sub_rx) = oneshot::channel();
-    to_swarm
+    let (tx, rx) = oneshot::channel();
+    _ = to_swarm
         .send(ToSwarm::Subscribe {
-            topic: "test-net".to_owned(),
-            result_sender: sub_tx,
+            topic: "test-net".to_string(),
+            result_sender: tx,
         })
         .await
         .expect("should send");
@@ -34,16 +35,15 @@ async fn main() {
     println!("Enter messages via STDIN and they will be sent to connected peers using Gossipsub");
 
     tokio::task::spawn(async move {
-        sub_rx
-            .await
+        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(ToSwarm::Publish {
-                        topic: "test-net".to_owned(),
+                    .send(swarm::ToSwarm::Publish {
+                        topic: "test-net".to_string(),
                         data: line.as_bytes().to_vec(),
                         result_sender: tx,
                     })
@@ -51,7 +51,7 @@ async fn main() {
                 {
                     println!("Send error: {e:?}");
                     return;
-                }
+                };
                 match rx.await {
                     Ok(Err(e)) => println!("Publish error: {e:?}"),
                     Err(e) => println!("Publish error: {e:?}"),

rust/networking/src/discovery.rs

@@ -0,0 +1,382 @@
+use crate::ext::MultiaddrExt;
+use delegate::delegate;
+use either::Either;
+use futures_lite::FutureExt;
+use futures_timer::Delay;
+use libp2p::core::transport::PortUse;
+use libp2p::core::{ConnectedPoint, Endpoint};
+use libp2p::swarm::behaviour::ConnectionEstablished;
+use libp2p::swarm::dial_opts::DialOpts;
+use libp2p::swarm::{
+    CloseConnection, ConnectionClosed, ConnectionDenied, ConnectionHandler,
+    ConnectionHandlerSelect, ConnectionId, FromSwarm, NetworkBehaviour, THandler, THandlerInEvent,
+    THandlerOutEvent, ToSwarm, dummy,
+};
+use libp2p::{Multiaddr, PeerId, identity, mdns};
+use std::collections::{BTreeSet, HashMap};
+use std::convert::Infallible;
+use std::io;
+use std::net::IpAddr;
+use std::task::{Context, Poll};
+use std::time::Duration;
+use util::wakerdeque::WakerDeque;
+
+const RETRY_CONNECT_INTERVAL: Duration = Duration::from_secs(5);
+
+mod managed {
+    use libp2p::swarm::NetworkBehaviour;
+    use libp2p::{identity, mdns, ping};
+    use std::io;
+    use std::time::Duration;
+
+    const MDNS_RECORD_TTL: Duration = Duration::from_secs(2_500);
+    const MDNS_QUERY_INTERVAL: Duration = Duration::from_secs(1_500);
+    const PING_TIMEOUT: Duration = Duration::from_millis(2_500);
+    const PING_INTERVAL: Duration = Duration::from_millis(2_500);
+
+    #[derive(NetworkBehaviour)]
+    pub struct Behaviour {
+        mdns: mdns::tokio::Behaviour,
+        ping: ping::Behaviour,
+    }
+
+    impl Behaviour {
+        pub fn new(keypair: &identity::Keypair) -> io::Result<Self> {
+            Ok(Self {
+                mdns: mdns_behaviour(keypair)?,
+                ping: ping_behaviour(),
+            })
+        }
+    }
+
+    fn mdns_behaviour(keypair: &identity::Keypair) -> io::Result<mdns::tokio::Behaviour> {
+        use mdns::{Config, tokio};
+
+        // mDNS config => enable IPv6
+        let mdns_config = Config {
+            ttl: MDNS_RECORD_TTL,
+            query_interval: MDNS_QUERY_INTERVAL,
+
+            // enable_ipv6: true, // TODO: for some reason, TCP+mDNS don't work well with ipv6?? figure out how to make work
+            ..Default::default()
+        };
+
+        let mdns_behaviour = tokio::Behaviour::new(mdns_config, keypair.public().to_peer_id());
+        Ok(mdns_behaviour?)
+    }
+
+    fn ping_behaviour() -> ping::Behaviour {
+        ping::Behaviour::new(
+            ping::Config::new()
+                .with_timeout(PING_TIMEOUT)
+                .with_interval(PING_INTERVAL),
+        )
+    }
+}
+
+/// Events for when a listening connection is truly established and truly closed.
+#[derive(Debug, Clone)]
+pub enum Event {
+    ConnectionEstablished {
+        peer_id: PeerId,
+        connection_id: ConnectionId,
+        remote_ip: IpAddr,
+        remote_tcp_port: u16,
+    },
+    ConnectionClosed {
+        peer_id: PeerId,
+        connection_id: ConnectionId,
+        remote_ip: IpAddr,
+        remote_tcp_port: u16,
+    },
+}
+
+/// Discovery behavior that wraps mDNS to produce truly discovered durable peer-connections.
+///
+/// The behaviour operates as such:
+///  1) All true (listening) connections/disconnections are tracked, emitting corresponding events
+///     to the swarm.
+///  1) mDNS discovered/expired peers are tracked; discovered but not connected peers are dialed
+///     immediately, and expired but connected peers are disconnected from immediately.
+///  2) Every fixed interval: discovered but not connected peers are dialed, and expired but
+///     connected peers are disconnected from.
+pub struct Behaviour {
+    // state-tracking for managed behaviors & mDNS-discovered peers
+    managed: managed::Behaviour,
+    mdns_discovered: HashMap<PeerId, BTreeSet<Multiaddr>>,
+
+    retry_delay: Delay, // retry interval
+
+    // pending events to emmit => waker-backed Deque to control polling
+    pending_events: WakerDeque<ToSwarm<Event, Infallible>>,
+}
+
+impl Behaviour {
+    pub fn new(keypair: &identity::Keypair) -> io::Result<Self> {
+        Ok(Self {
+            managed: managed::Behaviour::new(keypair)?,
+            mdns_discovered: HashMap::new(),
+            retry_delay: Delay::new(RETRY_CONNECT_INTERVAL),
+            pending_events: WakerDeque::new(),
+        })
+    }
+
+    fn dial(&mut self, peer_id: PeerId, addr: Multiaddr) {
+        self.pending_events.push_back(ToSwarm::Dial {
+            opts: DialOpts::peer_id(peer_id).addresses(vec![addr]).build(),
+        })
+    }
+
+    fn close_connection(&mut self, peer_id: PeerId, connection: ConnectionId) {
+        // push front to make this IMMEDIATE
+        self.pending_events.push_front(ToSwarm::CloseConnection {
+            peer_id,
+            connection: CloseConnection::One(connection),
+        })
+    }
+
+    fn handle_mdns_discovered(&mut self, peers: Vec<(PeerId, Multiaddr)>) {
+        for (p, ma) in peers {
+            self.dial(p, ma.clone()); // always connect
+
+            // get peer's multi-addresses or insert if missing
+            let Some(mas) = self.mdns_discovered.get_mut(&p) else {
+                self.mdns_discovered.insert(p, BTreeSet::from([ma]));
+                continue;
+            };
+
+            // multiaddress should never already be present - else something has gone wrong
+            let is_new_addr = mas.insert(ma);
+            assert!(is_new_addr, "cannot discover a discovered peer");
+        }
+    }
+
+    fn handle_mdns_expired(&mut self, peers: Vec<(PeerId, Multiaddr)>) {
+        for (p, ma) in peers {
+            // at this point, we *must* have the peer
+            let mas = self
+                .mdns_discovered
+                .get_mut(&p)
+                .expect("nonexistent peer cannot expire");
+
+            // at this point, we *must* have the multiaddress
+            let was_present = mas.remove(&ma);
+            assert!(was_present, "nonexistent multiaddress cannot expire");
+
+            // if empty, remove the peer-id entirely
+            if mas.is_empty() {
+                self.mdns_discovered.remove(&p);
+            }
+        }
+    }
+
+    fn on_connection_established(
+        &mut self,
+        peer_id: PeerId,
+        connection_id: ConnectionId,
+        remote_ip: IpAddr,
+        remote_tcp_port: u16,
+    ) {
+        // send out connected event
+        self.pending_events
+            .push_back(ToSwarm::GenerateEvent(Event::ConnectionEstablished {
+                peer_id,
+                connection_id,
+                remote_ip,
+                remote_tcp_port,
+            }));
+    }
+
+    fn on_connection_closed(
+        &mut self,
+        peer_id: PeerId,
+        connection_id: ConnectionId,
+        remote_ip: IpAddr,
+        remote_tcp_port: u16,
+    ) {
+        // send out disconnected event
+        self.pending_events
+            .push_back(ToSwarm::GenerateEvent(Event::ConnectionClosed {
+                peer_id,
+                connection_id,
+                remote_ip,
+                remote_tcp_port,
+            }));
+    }
+}
+
+impl NetworkBehaviour for Behaviour {
+    type ConnectionHandler =
+        ConnectionHandlerSelect<dummy::ConnectionHandler, THandler<managed::Behaviour>>;
+    type ToSwarm = Event;
+
+    // simply delegate to underlying mDNS behaviour
+
+    delegate! {
+        to self.managed {
+            fn handle_pending_inbound_connection(&mut self, connection_id: ConnectionId, local_addr: &Multiaddr, remote_addr: &Multiaddr) -> Result<(), ConnectionDenied>;
+            fn handle_pending_outbound_connection(&mut self, connection_id: ConnectionId, maybe_peer: Option<PeerId>, addresses: &[Multiaddr], effective_role: Endpoint) -> Result<Vec<Multiaddr>, ConnectionDenied>;
+        }
+    }
+
+    fn handle_established_inbound_connection(
+        &mut self,
+        connection_id: ConnectionId,
+        peer: PeerId,
+        local_addr: &Multiaddr,
+        remote_addr: &Multiaddr,
+    ) -> Result<THandler<Self>, ConnectionDenied> {
+        Ok(ConnectionHandler::select(
+            dummy::ConnectionHandler,
+            self.managed.handle_established_inbound_connection(
+                connection_id,
+                peer,
+                local_addr,
+                remote_addr,
+            )?,
+        ))
+    }
+
+    #[allow(clippy::needless_question_mark)]
+    fn handle_established_outbound_connection(
+        &mut self,
+        connection_id: ConnectionId,
+        peer: PeerId,
+        addr: &Multiaddr,
+        role_override: Endpoint,
+        port_use: PortUse,
+    ) -> Result<THandler<Self>, ConnectionDenied> {
+        Ok(ConnectionHandler::select(
+            dummy::ConnectionHandler,
+            self.managed.handle_established_outbound_connection(
+                connection_id,
+                peer,
+                addr,
+                role_override,
+                port_use,
+            )?,
+        ))
+    }
+
+    fn on_connection_handler_event(
+        &mut self,
+        peer_id: PeerId,
+        connection_id: ConnectionId,
+        event: THandlerOutEvent<Self>,
+    ) {
+        match event {
+            Either::Left(ev) => libp2p::core::util::unreachable(ev),
+            Either::Right(ev) => {
+                self.managed
+                    .on_connection_handler_event(peer_id, connection_id, ev)
+            }
+        }
+    }
+
+    // hook into these methods to drive behavior
+
+    fn on_swarm_event(&mut self, event: FromSwarm) {
+        self.managed.on_swarm_event(event); // let mDNS handle swarm events
+
+        // handle swarm events to update internal state:
+        match event {
+            FromSwarm::ConnectionEstablished(ConnectionEstablished {
+                peer_id,
+                connection_id,
+                endpoint,
+                ..
+            }) => {
+                let remote_address = match endpoint {
+                    ConnectedPoint::Dialer { address, .. } => address,
+                    ConnectedPoint::Listener { send_back_addr, .. } => send_back_addr,
+                };
+
+                if let Some((ip, port)) = remote_address.try_to_tcp_addr() {
+                    // handle connection established event which is filtered correctly
+                    self.on_connection_established(peer_id, connection_id, ip, port)
+                }
+            }
+            FromSwarm::ConnectionClosed(ConnectionClosed {
+                peer_id,
+                connection_id,
+                endpoint,
+                ..
+            }) => {
+                let remote_address = match endpoint {
+                    ConnectedPoint::Dialer { address, .. } => address,
+                    ConnectedPoint::Listener { send_back_addr, .. } => send_back_addr,
+                };
+
+                if let Some((ip, port)) = remote_address.try_to_tcp_addr() {
+                    // handle connection closed event which is filtered correctly
+                    self.on_connection_closed(peer_id, connection_id, ip, port)
+                }
+            }
+
+            // since we are running TCP/IP transport layer, we are assuming that
+            // no address changes can occur, hence encountering one is a fatal error
+            FromSwarm::AddressChange(a) => {
+                unreachable!("unhandlable: address change encountered: {:?}", a)
+            }
+            _ => {}
+        }
+    }
+
+    fn poll(&mut self, cx: &mut Context) -> Poll<ToSwarm<Self::ToSwarm, THandlerInEvent<Self>>> {
+        // delegate to managed behaviors for any behaviors they need to perform
+        match self.managed.poll(cx) {
+            Poll::Ready(ToSwarm::GenerateEvent(e)) => {
+                match e {
+                    // handle discovered and expired events from mDNS
+                    managed::BehaviourEvent::Mdns(e) => match e.clone() {
+                        mdns::Event::Discovered(peers) => {
+                            self.handle_mdns_discovered(peers);
+                        }
+                        mdns::Event::Expired(peers) => {
+                            self.handle_mdns_expired(peers);
+                        }
+                    },
+
+                    // handle ping events => if error then disconnect
+                    managed::BehaviourEvent::Ping(e) => {
+                        if let Err(_) = e.result {
+                            self.close_connection(e.peer, e.connection.clone())
+                        }
+                    }
+                }
+
+                // since we just consumed an event, we should immediately wake just in case
+                // there are more events to come where that came from
+                cx.waker().wake_by_ref();
+            }
+
+            // forward any other mDNS event to the swarm or its connection handler(s)
+            Poll::Ready(e) => {
+                return Poll::Ready(
+                    e.map_out(|_| unreachable!("events returning to swarm already handled"))
+                        .map_in(Either::Right),
+                );
+            }
+
+            Poll::Pending => {}
+        }
+
+        // retry connecting to all mDNS peers periodically (fails safely if already connected)
+        if self.retry_delay.poll(cx).is_ready() {
+            for (p, mas) in self.mdns_discovered.clone() {
+                for ma in mas {
+                    self.dial(p, ma)
+                }
+            }
+            self.retry_delay.reset(RETRY_CONNECT_INTERVAL) // reset timeout
+        }
+
+        // send out any pending events from our own service
+        if let Some(e) = self.pending_events.pop_front(cx) {
+            return Poll::Ready(e.map_in(Either::Left));
+        }
+
+        // wait for pending events
+        Poll::Pending
+    }
+}

rust/networking/src/lib.rs

@@ -3,6 +3,8 @@
 //! this is here as a placeholder documentation
 //!
 //!
+pub mod discovery;
+pub mod swarm;
 
 /// Namespace for all the type/trait aliases used by this crate.
 pub(crate) mod alias {
@@ -11,302 +13,32 @@ pub(crate) mod alias {
     pub type AnyError = Box<dyn Error + Send + Sync + 'static>;
     pub type AnyResult<T> = Result<T, AnyError>;
 }
-use std::collections::{HashMap, HashSet};
-use std::pin::Pin;
 
-pub use behaviour::{Behaviour, BehaviourEvent};
-use futures_lite::{Stream, StreamExt};
-use libp2p::swarm::dial_opts::DialOpts;
-use libp2p::{Multiaddr, mdns};
-use libp2p::{PeerId, SwarmBuilder, gossipsub, identity, swarm::SwarmEvent};
-use tokio::sync::{mpsc, oneshot};
-use transport::tcp_transport;
+/// 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;
 
-/// 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,
-        result_sender: oneshot::Sender<bool>,
-    },
-    Subscribe {
-        topic: String,
-        result_sender: oneshot::Sender<Result<bool, gossipsub::SubscriptionError>>,
-    },
-    Publish {
-        topic: String,
-        data: Vec<u8>,
-        result_sender: oneshot::Sender<Result<gossipsub::MessageId, gossipsub::PublishError>>,
-    },
-}
-pub enum FromSwarm {
-    Message {
-        from: PeerId,
-        topic: String,
-        data: Vec<u8>,
-    },
-    Discovered {
-        peer_id: PeerId,
-    },
-    Expired {
-        peer_id: PeerId,
-    },
-}
-
-pub struct Swarm {
-    swarm: libp2p::Swarm<Behaviour>,
-    from_client: mpsc::Receiver<ToSwarm>,
-}
-
-impl Swarm {
-    #[inline]
-    pub fn into_stream(self) -> Pin<Box<dyn Stream<Item = FromSwarm> + Send>> {
-        let Self {
-            mut swarm,
-            mut from_client,
-        } = self;
-        Box::pin(async_stream::stream! {
-            loop {
-                tokio::select! {
-                    msg = from_client.recv() => {
-                        let Some(msg) = msg else { break };
-                        on_message(&mut swarm, msg);
-                    }
-                    event = swarm.next() => {
-                        let Some(event) = event else { break };
-                        for item in on_swarm_event(&mut swarm, event) {
-                            yield item;
-                        }
-                    }
+    #[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,
                 }
-            }
-        })
-    }
-}
-
-fn on_swarm_event(
-    swarm: &mut libp2p::Swarm<Behaviour>,
-    event: SwarmEvent<BehaviourEvent>,
-) -> Vec<FromSwarm> {
-    match event {
-        SwarmEvent::Behaviour(BehaviourEvent::Gossipsub(gossipsub::Event::Message {
-            message:
-                gossipsub::Message {
-                    source: Some(peer_id),
-                    topic,
-                    data,
-                    ..
-                },
-            ..
-        })) => vec![FromSwarm::Message {
-            from: peer_id,
-            topic: topic.into_string(),
-            data,
-        }],
-        SwarmEvent::Behaviour(BehaviourEvent::Discovery(mdns::Event::Discovered(peer_ids))) => {
-            let mut base = HashMap::<PeerId, Vec<Multiaddr>>::default();
-            for (pid, addr) in peer_ids {
-                base.entry(pid).or_default().push(addr)
-            }
-            let ret = base
-                .keys()
-                .map(|&peer_id| FromSwarm::Discovered { peer_id })
-                .collect::<Vec<_>>();
-            for (pid, addrs) in base {
-                _ = swarm.dial(DialOpts::peer_id(pid).addresses(addrs).build());
-            }
-            ret
-        }
-        SwarmEvent::Behaviour(BehaviourEvent::Discovery(mdns::Event::Expired(peers))) => peers
-            .into_iter()
-            .map(|(peer_id, _)| peer_id)
-            .collect::<HashSet<PeerId>>()
-            .into_iter()
-            .map(|peer_id| FromSwarm::Expired { peer_id })
-            .collect(),
-        _ => vec![],
-    }
-}
-
-fn on_message(swarm: &mut libp2p::Swarm<Behaviour>, message: ToSwarm) {
-    match message {
-        ToSwarm::Subscribe {
-            topic,
-            result_sender,
-        } => {
-            let result = swarm
-                .behaviour_mut()
-                .gossipsub
-                .subscribe(&gossipsub::IdentTopic::new(topic));
-            _ = result_sender.send(result);
-        }
-        ToSwarm::Unsubscribe {
-            topic,
-            result_sender,
-        } => {
-            let result = swarm
-                .behaviour_mut()
-                .gossipsub
-                .unsubscribe(&gossipsub::IdentTopic::new(topic));
-            _ = result_sender.send(result);
-        }
-        ToSwarm::Publish {
-            topic,
-            data,
-            result_sender,
-        } => {
-            let result = swarm
-                .behaviour_mut()
-                .gossipsub
-                .publish(gossipsub::IdentTopic::new(topic), data);
-            _ = result_sender.send(result);
+            } else {
+                return None;
+            };
+            let Some(Protocol::Tcp(port)) = ps.next() else {
+                return None;
+            };
+            Some((ip, port))
         }
     }
 }
-
-/// Create and configure a swarm which listens to all ports on OS
-pub fn create_swarm(
-    keypair: identity::Keypair,
-    from_client: mpsc::Receiver<ToSwarm>,
-) -> alias::AnyResult<Swarm> {
-    let mut swarm = SwarmBuilder::with_existing_identity(keypair)
-        .with_tokio()
-        .with_other_transport(tcp_transport)?
-        .with_behaviour(Behaviour::new)?
-        .build();
-
-    // Listen on all interfaces and whatever port the OS assigns
-    swarm.listen_on("/ip4/0.0.0.0/tcp/0".parse()?)?;
-    Ok(Swarm { swarm, from_client })
-}
-
-mod transport {
-    use crate::alias;
-    use crate::{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;
-    use libp2p::swarm::NetworkBehaviour;
-    use libp2p::{gossipsub, identity, mdns};
-
-    /// Behavior of the Swarm which composes all desired behaviors:
-    /// Right now its just [`mdns::tokio::Behaviour`] and [`gossipsub::Behaviour`].
-    #[derive(NetworkBehaviour)]
-    pub struct Behaviour {
-        pub discovery: mdns::tokio::Behaviour,
-        pub gossipsub: gossipsub::Behaviour,
-    }
-
-    impl Behaviour {
-        pub fn new(keypair: &identity::Keypair) -> alias::AnyResult<Self> {
-            Ok(Self {
-                discovery: mdns_behaviour(keypair)?,
-                gossipsub: gossipsub_behaviour(keypair),
-            })
-        }
-    }
-
-    fn mdns_behaviour(keypair: &identity::Keypair) -> alias::AnyResult<mdns::tokio::Behaviour> {
-        Ok(mdns::tokio::Behaviour::new(
-            mdns::Config::default(),
-            keypair.public().to_peer_id(),
-        )?)
-    }
-
-    fn gossipsub_behaviour(keypair: &identity::Keypair) -> gossipsub::Behaviour {
-        use gossipsub::{ConfigBuilder, MessageAuthenticity, ValidationMode};
-
-        // build a gossipsub network behaviour
-        //  => signed message authenticity + strict validation mode means the message-ID is
-        //     automatically provided by gossipsub w/out needing to provide custom message-ID function
-        gossipsub::Behaviour::new(
-            MessageAuthenticity::Signed(keypair.clone()),
-            ConfigBuilder::default()
-                .max_transmit_size(1024 * 1024)
-                .validation_mode(ValidationMode::Strict)
-                .build()
-                .expect("the configuration should always be valid"),
-        )
-        .expect("creating gossipsub behavior should always work")
-    }
-}

rust/networking/src/swarm.rs

@@ -0,0 +1,273 @@
+use std::pin::Pin;
+
+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};
+
+/// 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,
+        result_sender: oneshot::Sender<bool>,
+    },
+    Subscribe {
+        topic: String,
+        result_sender: oneshot::Sender<Result<bool, gossipsub::SubscriptionError>>,
+    },
+    Publish {
+        topic: String,
+        data: Vec<u8>,
+        result_sender: oneshot::Sender<Result<gossipsub::MessageId, gossipsub::PublishError>>,
+    },
+}
+pub enum FromSwarm {
+    Message {
+        from: PeerId,
+        topic: String,
+        data: Vec<u8>,
+    },
+    Discovered {
+        peer_id: PeerId,
+    },
+    Expired {
+        peer_id: PeerId,
+    },
+}
+
+pub struct Swarm {
+    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 {
+            mut swarm,
+            mut from_client,
+        } = self;
+        let stream = async_stream::stream! {
+            loop {
+                tokio::select! {
+                    msg = from_client.recv() => {
+                        let Some(msg) = msg else { break };
+                        on_message(&mut swarm, msg);
+                    }
+                    event = swarm.next() => {
+                        let Some(event) = event else { break };
+                        if let Some(item) = filter_swarm_event(event) {
+                            yield item;
+                        }
+                    }
+                }
+            }
+        };
+        Box::pin(stream)
+    }
+}
+
+fn on_message(swarm: &mut libp2p::Swarm<Behaviour>, message: ToSwarm) {
+    match message {
+        ToSwarm::Subscribe {
+            topic,
+            result_sender,
+        } => {
+            let result = swarm
+                .behaviour_mut()
+                .gossipsub
+                .subscribe(&gossipsub::IdentTopic::new(topic));
+            _ = result_sender.send(result);
+        }
+        ToSwarm::Unsubscribe {
+            topic,
+            result_sender,
+        } => {
+            let result = swarm
+                .behaviour_mut()
+                .gossipsub
+                .unsubscribe(&gossipsub::IdentTopic::new(topic));
+            _ = result_sender.send(result);
+        }
+        ToSwarm::Publish {
+            topic,
+            data,
+            result_sender,
+        } => {
+            let result = swarm
+                .behaviour_mut()
+                .gossipsub
+                .publish(gossipsub::IdentTopic::new(topic), data);
+            _ = result_sender.send(result);
+        }
+    }
+}
+
+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 which listens to all ports on OS
+pub fn create_swarm(
+    keypair: identity::Keypair,
+    from_client: mpsc::Receiver<ToSwarm>,
+) -> alias::AnyResult<Swarm> {
+    let mut swarm = SwarmBuilder::with_existing_identity(keypair)
+        .with_tokio()
+        .with_other_transport(tcp_transport)?
+        .with_behaviour(Behaviour::new)?
+        .build();
+
+    // Listen on all interfaces and whatever port the OS assigns
+    swarm.listen_on("/ip4/0.0.0.0/tcp/0".parse()?)?;
+    Ok(Swarm { 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) -> alias::AnyResult<Self> {
+            Ok(Self {
+                discovery: discovery::Behaviour::new(keypair)?,
+                gossipsub: gossipsub_behaviour(keypair),
+            })
+        }
+    }
+
+    fn gossipsub_behaviour(keypair: &identity::Keypair) -> gossipsub::Behaviour {
+        use gossipsub::{ConfigBuilder, MessageAuthenticity, ValidationMode};
+
+        // build a gossipsub network behaviour
+        //  => signed message authenticity + strict validation mode means the message-ID is
+        //     automatically provided by gossipsub w/out needing to provide custom message-ID function
+        gossipsub::Behaviour::new(
+            MessageAuthenticity::Signed(keypair.clone()),
+            ConfigBuilder::default()
+                .max_transmit_size(1024 * 1024)
+                .validation_mode(ValidationMode::Strict)
+                .build()
+                .expect("the configuration should always be valid"),
+        )
+        .expect("creating gossipsub behavior should always work")
+    }
+}

rust/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/master/adapters/chat_completions.py

@@ -26,11 +26,7 @@ from exo.shared.types.chunks import (
     ToolCallChunk,
 )
 from exo.shared.types.common import CommandId
-from exo.shared.types.text_generation import (
-    InputMessage,
-    TextGenerationTaskParams,
-    resolve_reasoning_params,
-)
+from exo.shared.types.text_generation import InputMessage, TextGenerationTaskParams
 
 
 def chat_request_to_text_generation(
@@ -79,10 +75,6 @@ def chat_request_to_text_generation(
             dumped: dict[str, Any] = msg_copy.model_dump(exclude_none=True)
             chat_template_messages.append(dumped)
 
-    resolved_effort, resolved_thinking = resolve_reasoning_params(
-        request.reasoning_effort, request.enable_thinking
-    )
-
     return TextGenerationTaskParams(
         model=request.model,
         input=input_messages
@@ -97,8 +89,7 @@ def chat_request_to_text_generation(
         seed=request.seed,
         stream=request.stream,
         tools=request.tools,
-        reasoning_effort=resolved_effort,
-        enable_thinking=resolved_thinking,
+        enable_thinking=request.enable_thinking,
         chat_template_messages=chat_template_messages
         if chat_template_messages
         else None,

src/exo/master/adapters/responses.py

@@ -42,11 +42,7 @@ from exo.shared.types.openai_responses import (
     ResponseTextDoneEvent,
     ResponseUsage,
 )
-from exo.shared.types.text_generation import (
-    InputMessage,
-    TextGenerationTaskParams,
-    resolve_reasoning_params,
-)
+from exo.shared.types.text_generation import InputMessage, TextGenerationTaskParams
 
 
 def _format_sse(event: ResponsesStreamEvent) -> str:
@@ -123,11 +119,6 @@ def responses_request_to_text_generation(
         )
         built_chat_template = chat_template_messages if chat_template_messages else None
 
-    effort_from_reasoning = request.reasoning.effort if request.reasoning else None
-    resolved_effort, resolved_thinking = resolve_reasoning_params(
-        effort_from_reasoning, request.enable_thinking
-    )
-
     return TextGenerationTaskParams(
         model=request.model,
         input=input_value,
@@ -141,8 +132,6 @@ def responses_request_to_text_generation(
         stop=request.stop,
         seed=request.seed,
         chat_template_messages=built_chat_template or request.chat_template_messages,
-        reasoning_effort=resolved_effort,
-        enable_thinking=resolved_thinking,
     )
 
 

src/exo/master/main.py

@@ -328,22 +328,17 @@ class Master:
                                         task_id=task_id,
                                     )
                                 )
-                            else:
-                                logger.warning(
-                                    f"Nonexistent command {command.cancelled_command_id} cancelled"
-                                )
                         case TaskFinished():
-                            if (
-                                task_id := self.command_task_mapping.pop(
-                                    command.finished_command_id, None
+                            generated_events.append(
+                                TaskDeleted(
+                                    task_id=self.command_task_mapping[
+                                        command.finished_command_id
+                                    ]
                                 )
-                            ) is not None:
-                                generated_events.append(TaskDeleted(task_id=task_id))
-                            else:
-                                logger.warning(
-                                    f"Finished command {command.finished_command_id} finished"
-                                )
-
+                            )
+                            self.command_task_mapping.pop(
+                                command.finished_command_id, None
+                            )
                         case RequestEventLog():
                             # We should just be able to send everything, since other buffers will ignore old messages
                             # rate limit to 1000 at a time

src/exo/shared/tests/test_resolve_reasoning_params.py

@@ -1,30 +0,0 @@
-import pytest
-
-from exo.shared.types.text_generation import resolve_reasoning_params
-
-
-def test_both_none_returns_none_none() -> None:
-    assert resolve_reasoning_params(None, None) == (None, None)
-
-
-def test_both_set_passes_through_unchanged() -> None:
-    assert resolve_reasoning_params("high", True) == ("high", True)
-    assert resolve_reasoning_params("none", True) == ("none", True)
-    assert resolve_reasoning_params("low", False) == ("low", False)
-
-
-def test_enable_thinking_true_derives_medium() -> None:
-    assert resolve_reasoning_params(None, True) == ("medium", True)
-
-
-def test_enable_thinking_false_derives_none() -> None:
-    assert resolve_reasoning_params(None, False) == ("none", False)
-
-
-def test_reasoning_effort_none_derives_thinking_false() -> None:
-    assert resolve_reasoning_params("none", None) == ("none", False)
-
-
-@pytest.mark.parametrize("effort", ["minimal", "low", "medium", "high", "xhigh"])
-def test_non_none_effort_derives_thinking_true(effort: str) -> None:
-    assert resolve_reasoning_params(effort, None) == (effort, True)  # pyright: ignore[reportArgumentType]

src/exo/shared/types/api.py

@@ -8,7 +8,6 @@ from pydantic import BaseModel, Field, field_validator
 from exo.shared.models.model_cards import ModelCard, ModelId
 from exo.shared.types.common import CommandId, NodeId
 from exo.shared.types.memory import Memory
-from exo.shared.types.text_generation import ReasoningEffort
 from exo.shared.types.worker.instances import Instance, InstanceId, InstanceMeta
 from exo.shared.types.worker.shards import Sharding, ShardMetadata
 from exo.utils.pydantic_ext import CamelCaseModel
@@ -199,7 +198,6 @@ class ChatCompletionRequest(BaseModel):
     top_p: float | None = None
     top_k: int | None = None
     tools: list[dict[str, Any]] | None = None
-    reasoning_effort: ReasoningEffort | None = None
     enable_thinking: bool | None = None
     tool_choice: str | dict[str, Any] | None = None
     parallel_tool_calls: bool | None = None

src/exo/shared/types/openai_responses.py

@@ -12,7 +12,6 @@ from typing import Any, Literal
 from pydantic import BaseModel, Field
 
 from exo.shared.types.common import ModelId
-from exo.shared.types.text_generation import ReasoningEffort
 
 # Type aliases
 ResponseStatus = Literal["completed", "failed", "in_progress", "incomplete"]
@@ -72,13 +71,6 @@ ResponseInputItem = (
 )
 
 
-class Reasoning(BaseModel, frozen=True):
-    """Reasoning configuration for OpenAI Responses API."""
-
-    effort: ReasoningEffort | None = None
-    summary: Literal["auto", "concise", "detailed"] | None = None
-
-
 class ResponsesRequest(BaseModel, frozen=True):
     """Request body for OpenAI Responses API.
 
@@ -97,15 +89,8 @@ class ResponsesRequest(BaseModel, frozen=True):
     stream: bool = False
     tools: list[dict[str, Any]] | None = None
     metadata: dict[str, str] | None = None
-    reasoning: Reasoning | None = None
 
     # --- exo extensions (not in OpenAI Responses API spec) ---
-    enable_thinking: bool | None = Field(
-        default=None,
-        description="[exo extension] Boolean thinking toggle. Not part of the OpenAI Responses API.",
-        json_schema_extra={"x-exo-extension": True},
-    )
-
     top_k: int | None = Field(
         default=None,
         description="[exo extension] Top-k sampling parameter. Not part of the OpenAI Responses API.",

src/exo/shared/types/text_generation.py

@@ -11,29 +11,6 @@ from pydantic import BaseModel
 from exo.shared.types.common import ModelId
 
 MessageRole = Literal["user", "assistant", "system", "developer"]
-ReasoningEffort = Literal["none", "minimal", "low", "medium", "high", "xhigh"]
-
-
-def resolve_reasoning_params(
-    reasoning_effort: ReasoningEffort | None,
-    enable_thinking: bool | None,
-) -> tuple[ReasoningEffort | None, bool | None]:
-    """
-    enable_thinking=True  -> reasoning_effort="medium"
-    enable_thinking=False -> reasoning_effort="none"
-    reasoning_effort="none" -> enable_thinking=False
-    reasoning_effort=<anything else> -> enable_thinking=True
-    """
-    resolved_effort: ReasoningEffort | None = reasoning_effort
-    resolved_thinking: bool | None = enable_thinking
-
-    if reasoning_effort is None and enable_thinking is not None:
-        resolved_effort = "medium" if enable_thinking else "none"
-
-    if enable_thinking is None and reasoning_effort is not None:
-        resolved_thinking = reasoning_effort != "none"
-
-    return resolved_effort, resolved_thinking
 
 
 class InputMessage(BaseModel, frozen=True):
@@ -63,7 +40,6 @@ class TextGenerationTaskParams(BaseModel, frozen=True):
     stop: str | list[str] | None = None
     seed: int | None = None
     chat_template_messages: list[dict[str, Any]] | None = None
-    reasoning_effort: ReasoningEffort | None = None
     enable_thinking: bool | None = None
     logprobs: bool = False
     top_logprobs: int | None = None

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

@@ -41,6 +41,7 @@ from exo.download.download_utils import build_model_path
 from exo.shared.types.common import Host
 from exo.shared.types.memory import Memory
 from exo.shared.types.mlx import Model
+from exo.shared.types.tasks import TaskId, TextGeneration
 from exo.shared.types.text_generation import TextGenerationTaskParams
 from exo.shared.types.worker.instances import (
     BoundInstance,
@@ -554,8 +555,6 @@ def apply_chat_template(
         # Jinja ignores unknown variables, so passing both is safe.
         extra_kwargs["enable_thinking"] = task_params.enable_thinking
         extra_kwargs["thinking"] = task_params.enable_thinking
-    if task_params.reasoning_effort is not None:
-        extra_kwargs["reasoning_effort"] = task_params.reasoning_effort
 
     patched_template: str | None = None
     if task_params.tools:
@@ -750,3 +749,55 @@ def _parse_kimi_tool_calls(text: str):
         return [_parse_single_tool(match) for match in tool_matches]  # pyright: ignore[reportAny]
     else:
         return [_parse_single_tool(text)]
+
+
+def mx_all_gather_tasks(
+    tasks: list[TextGeneration],
+    group: mx.distributed.Group | None,
+) -> tuple[list[TextGeneration], list[TextGeneration]]:
+    def encode_task_id(task_id: TaskId) -> list[int]:
+        utf8_task_id = task_id.encode()
+        return [
+            int.from_bytes(utf8_task_id[i : i + 1]) for i in range(len(utf8_task_id))
+        ]
+
+    def decode_task_id(encoded_task_id: list[int]) -> TaskId:
+        return TaskId(
+            bytes.decode(b"".join((x).to_bytes(length=1) for x in encoded_task_id))
+        )
+
+    uuid_byte_length = 36
+
+    n_tasks = len(tasks)
+    all_counts = cast(
+        list[int],
+        mx.distributed.all_gather(mx.array([n_tasks]), group=group).tolist(),
+    )
+    max_tasks = max(all_counts)
+    world_size: int = 1 if group is None else group.size()
+
+    if max_tasks == 0:
+        return [], []
+
+    padded = [encode_task_id(task.task_id) for task in tasks] + [
+        [0] * uuid_byte_length
+    ] * (max_tasks - n_tasks)
+    gathered = cast(
+        list[list[list[int]]],
+        mx.distributed.all_gather(mx.array(padded), group=group)
+        .reshape(world_size, max_tasks, -1)
+        .tolist(),
+    )
+    all_task_ids: list[list[TaskId]] = [
+        [decode_task_id(encoded_task_id) for encoded_task_id in rank_tasks[:count]]
+        for rank_tasks, count in zip(gathered, all_counts, strict=True)
+    ]
+
+    agreed_ids: set[TaskId] = set(all_task_ids[0])
+    for rank_tasks in all_task_ids[1:]:
+        agreed_ids &= set(rank_tasks)
+
+    local_tasks = {task.task_id: task for task in tasks}
+    agreed = [local_tasks[tid] for tid in sorted(agreed_ids)]
+    different = [task for task in tasks if task.task_id not in agreed_ids]
+    return agreed, different

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

@@ -5,7 +5,6 @@ from abc import ABC, abstractmethod
 from collections import deque
 from collections.abc import Generator, Iterable
 from dataclasses import dataclass, field
-from typing import cast
 
 import mlx.core as mx
 from mlx_lm.tokenizer_utils import TokenizerWrapper
@@ -26,6 +25,7 @@ from exo.worker.engines.mlx.generator.generate import (
 )
 from exo.worker.engines.mlx.utils_mlx import (
     apply_chat_template,
+    mx_all_gather_tasks,
     mx_any,
 )
 from exo.worker.runner.bootstrap import logger
@@ -218,7 +218,6 @@ class SequentialGenerator(InferenceGenerator):
             self.tokenizer,
             type(self.model),
             self.model_id,
-            task.task_params.tools,
         )
         self._active = (task, mlx_gen, queue, output_generator)
 
@@ -292,55 +291,3 @@ class SequentialGenerator(InferenceGenerator):
 
     def close(self) -> None:
         del self.model, self.tokenizer, self.group
-
-
-def mx_all_gather_tasks(
-    tasks: list[TextGeneration],
-    group: mx.distributed.Group | None,
-) -> tuple[list[TextGeneration], list[TextGeneration]]:
-    def encode_task_id(task_id: TaskId) -> list[int]:
-        utf8_task_id = task_id.encode()
-        return [
-            int.from_bytes(utf8_task_id[i : i + 1]) for i in range(len(utf8_task_id))
-        ]
-
-    def decode_task_id(encoded_task_id: list[int]) -> TaskId:
-        return TaskId(
-            bytes.decode(b"".join((x).to_bytes(length=1) for x in encoded_task_id))
-        )
-
-    uuid_byte_length = 36
-
-    n_tasks = len(tasks)
-    all_counts = cast(
-        list[int],
-        mx.distributed.all_gather(mx.array([n_tasks]), group=group).tolist(),
-    )
-    max_tasks = max(all_counts)
-    world_size: int = 1 if group is None else group.size()
-
-    if max_tasks == 0:
-        return [], []
-
-    padded = [encode_task_id(task.task_id) for task in tasks] + [
-        [0] * uuid_byte_length
-    ] * (max_tasks - n_tasks)
-    gathered = cast(
-        list[list[list[int]]],
-        mx.distributed.all_gather(mx.array(padded), group=group)
-        .reshape(world_size, max_tasks, -1)
-        .tolist(),
-    )
-    all_task_ids: list[list[TaskId]] = [
-        [decode_task_id(encoded_task_id) for encoded_task_id in rank_tasks[:count]]
-        for rank_tasks, count in zip(gathered, all_counts, strict=True)
-    ]
-
-    agreed_ids: set[TaskId] = set(all_task_ids[0])
-    for rank_tasks in all_task_ids[1:]:
-        agreed_ids &= set(rank_tasks)
-
-    local_tasks = {task.task_id: task for task in tasks}
-    agreed = [local_tasks[tid] for tid in sorted(agreed_ids)]
-    different = [task for task in tasks if task.task_id not in agreed_ids]
-    return agreed, different

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

@@ -1,6 +1,5 @@
 from collections.abc import Generator
 from functools import cache
-from typing import Any
 
 from mlx_lm.models.deepseek_v32 import Model as DeepseekV32Model
 from mlx_lm.models.gpt_oss import Model as GptOssModel
@@ -38,7 +37,6 @@ def apply_all_parsers(
     tokenizer: TokenizerWrapper,
     model_type: type[Model],
     model_id: ModelId,
-    tools: list[dict[str, Any]] | None,
 ) -> Generator[GenerationResponse | ToolCallResponse | None]:
     mlx_generator = receiver
 
@@ -57,7 +55,7 @@ def apply_all_parsers(
     ):
         mlx_generator = parse_deepseek_v32(mlx_generator)
     elif tool_parser:
-        mlx_generator = parse_tool_calls(mlx_generator, tool_parser, tools)
+        mlx_generator = parse_tool_calls(mlx_generator, tool_parser)
 
     return mlx_generator
 
@@ -327,9 +325,7 @@ def parse_thinking_models(
 
 
 def parse_tool_calls(
-    responses: Generator[GenerationResponse | None],
-    tool_parser: ToolParser,
-    tools: list[dict[str, Any]] | None,
+    responses: Generator[GenerationResponse | None], tool_parser: ToolParser
 ) -> Generator[GenerationResponse | ToolCallResponse | None]:
     in_tool_call = False
     tool_call_text_parts: list[str] = []
@@ -344,7 +340,9 @@ def parse_tool_calls(
             tool_call_text_parts.append(response.text)
             if response.text.endswith(tool_parser.end_parsing):
                 # parse the actual tool calls from the tool call text
-                parsed = tool_parser.parse("".join(tool_call_text_parts).strip(), tools)
+                parsed = tool_parser.parse_tool_calls(
+                    "".join(tool_call_text_parts).strip()
+                )
                 logger.info(f"parsed {tool_call_text_parts=} into {parsed=}")
                 if parsed is not None:
                     yield ToolCallResponse(

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

@@ -1,5 +1,4 @@
 import json
-import math
 from dataclasses import dataclass
 from typing import Any, Callable
 
@@ -10,177 +9,7 @@ from exo.shared.types.api import ToolCallItem
 class ToolParser:
     start_parsing: str
     end_parsing: str
-    _inner_parser: Callable[[str], list[ToolCallItem] | None]
-
-    def parse(
-        self, text: str, tools: list[dict[str, Any]] | None
-    ) -> list[ToolCallItem] | None:
-        parsed = self._inner_parser(text)
-        if parsed is None:
-            return None
-        if tools is not None:
-            parsed = _coerce_tool_calls_to_schema(parsed, tools)
-        return parsed
-
-
-def _json_type_matches(value: Any, expected_type: str) -> bool:  # pyright: ignore[reportAny]
-    if expected_type == "object":
-        return isinstance(value, dict)
-    if expected_type == "array":
-        return isinstance(value, list)
-    if expected_type == "string":
-        return isinstance(value, str)
-    if expected_type == "integer":
-        return isinstance(value, int) and not isinstance(value, bool)
-    if expected_type == "number":
-        return (isinstance(value, int) and not isinstance(value, bool)) or isinstance(
-            value, float
-        )
-    if expected_type == "boolean":
-        return isinstance(value, bool)
-    if expected_type == "null":
-        return value is None
-    return False
-
-
-def _coerce_tool_arg_with_schema(value: Any, schema: dict[str, Any]) -> Any:  # pyright: ignore[reportAny]
-    schema_type = schema.get("type")
-
-    if isinstance(schema_type, list):
-        for candidate in schema_type:  # pyright: ignore[reportUnknownVariableType]
-            if not isinstance(candidate, str):
-                continue
-            if candidate == "null" and value is None:
-                return None
-            candidate_schema = {**schema, "type": candidate}
-            coerced = _coerce_tool_arg_with_schema(value, candidate_schema)  # pyright: ignore[reportAny]
-            if _json_type_matches(coerced, candidate):
-                return coerced  # pyright: ignore[reportAny]
-        return value  # pyright: ignore[reportAny]
-
-    if not isinstance(schema_type, str):
-        return value  # pyright: ignore[reportAny]
-
-    if schema_type == "object":
-        parsed = value  # pyright: ignore[reportAny]
-        if isinstance(parsed, str):
-            try:
-                parsed = json.loads(parsed)  # pyright: ignore[reportAny]
-            except Exception:
-                return value  # pyright: ignore[reportAny]
-        if not isinstance(parsed, dict):
-            return value  # pyright: ignore[reportAny]
-        properties = schema.get("properties")
-        if not isinstance(properties, dict):
-            return parsed  # pyright: ignore[reportUnknownVariableType]
-        return {
-            key: (
-                _coerce_tool_arg_with_schema(prop_value, prop_schema)  # pyright: ignore[reportUnknownArgumentType]
-                if isinstance(prop_schema, dict)
-                else prop_value
-            )
-            for key, prop_value in parsed.items()  # pyright: ignore[reportUnknownVariableType]
-            for prop_schema in [properties.get(key)]  # type: ignore
-        }
-
-    if schema_type == "array":
-        parsed = value  # pyright: ignore[reportAny]
-        if isinstance(parsed, str):
-            try:
-                parsed = json.loads(parsed)  # pyright: ignore[reportAny]
-            except Exception:
-                return value  # pyright: ignore[reportAny]
-        if not isinstance(parsed, list):
-            return value  # pyright: ignore[reportAny]
-        item_schema = schema.get("items")
-        if not isinstance(item_schema, dict):
-            return parsed  # pyright: ignore[reportUnknownVariableType]
-        return [_coerce_tool_arg_with_schema(item, item_schema) for item in parsed]  # type: ignore
-
-    if schema_type == "integer":
-        if isinstance(value, bool):
-            return value
-        if isinstance(value, int):
-            return value
-        if isinstance(value, float) and value.is_integer():
-            return int(value)
-        if isinstance(value, str):
-            try:
-                return int(value.strip())
-            except ValueError:
-                return value
-        return value
-
-    if schema_type == "number":
-        if isinstance(value, bool):
-            return value
-        if isinstance(value, (int, float)):
-            return value
-        if isinstance(value, str):
-            try:
-                num = float(value.strip())
-                if math.isfinite(num):
-                    return num
-            except ValueError:
-                return value
-        return value
-
-    if schema_type == "boolean":
-        if isinstance(value, bool):
-            return value
-        if isinstance(value, str):
-            lowered = value.strip().lower()
-            if lowered == "true":
-                return True
-            if lowered == "false":
-                return False
-        return value
-
-    return value  # pyright: ignore[reportAny]
-
-
-def _coerce_tool_calls_to_schema(
-    tool_calls: list[ToolCallItem], tools: list[dict[str, Any]]
-) -> list[ToolCallItem]:
-    schema_by_name: dict[str, dict[str, Any]] = {}
-    for tool in tools:
-        function = tool.get("function")
-        if not isinstance(function, dict):
-            continue
-        name = function.get("name")  # type: ignore
-        parameters = function.get("parameters")  # type: ignore
-        if isinstance(name, str) and isinstance(parameters, dict):
-            schema_by_name[name] = parameters
-
-    if not schema_by_name:
-        return tool_calls
-
-    coerced_calls: list[ToolCallItem] = []
-    for tool_call in tool_calls:
-        schema = schema_by_name.get(tool_call.name)
-        if schema is None:
-            coerced_calls.append(tool_call)
-            continue
-
-        try:
-            parsed_args = json.loads(tool_call.arguments)  # pyright: ignore[reportAny]
-        except Exception:
-            coerced_calls.append(tool_call)
-            continue
-
-        if not isinstance(parsed_args, dict):
-            coerced_calls.append(tool_call)
-            continue
-
-        coerced_args = _coerce_tool_arg_with_schema(parsed_args, schema)  # pyright: ignore[reportAny]
-        if not isinstance(coerced_args, dict):
-            coerced_calls.append(tool_call)
-            continue
-
-        coerced_calls.append(
-            tool_call.model_copy(update={"arguments": json.dumps(coerced_args)})
-        )
-    return coerced_calls
+    parse_tool_calls: Callable[[str], list[ToolCallItem] | None]
 
 
 def make_mlx_parser(
@@ -204,7 +33,7 @@ def make_mlx_parser(
     return ToolParser(
         start_parsing=tool_call_start,
         end_parsing=tool_call_end,
-        _inner_parser=parse_tool_calls,
+        parse_tool_calls=parse_tool_calls,
     )
 
 
@@ -233,7 +62,7 @@ def make_json_parser() -> ToolParser:
     return ToolParser(
         start_parsing="<tool_call>",
         end_parsing="</tool_call>",
-        _inner_parser=_parse_json_calls,
+        parse_tool_calls=_parse_json_calls,
     )
 
 

src/exo/worker/runner/runner_supervisor.py

@@ -12,22 +12,13 @@ from anyio import (
 )
 from loguru import logger
 
-from exo.shared.types.chunks import ErrorChunk
 from exo.shared.types.events import (
-    ChunkGenerated,
     Event,
     RunnerStatusUpdated,
     TaskAcknowledged,
     TaskStatusUpdated,
 )
-from exo.shared.types.tasks import (
-    ImageEdits,
-    ImageGeneration,
-    Task,
-    TaskId,
-    TaskStatus,
-    TextGeneration,
-)
+from exo.shared.types.tasks import Task, TaskId, TaskStatus
 from exo.shared.types.worker.instances import BoundInstance
 from exo.shared.types.worker.runners import (
     RunnerConnecting,
@@ -61,7 +52,7 @@ class RunnerSupervisor:
     _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)
+    in_progress: set[TaskId] = field(default_factory=set, init=False)
     completed: set[TaskId] = field(default_factory=set, init=False)
     cancelled: set[TaskId] = field(default_factory=set, init=False)
     _cancel_watch_runner: anyio.CancelScope = field(
@@ -157,11 +148,9 @@ class RunnerSupervisor:
         logger.info(f"Starting task {task}")
         event = anyio.Event()
         self.pending[task.task_id] = event
-        self.in_progress[task.task_id] = task
         try:
             await self._task_sender.send_async(task)
         except ClosedResourceError:
-            self.in_progress.pop(task.task_id, None)
             logger.warning(f"Task {task} dropped, runner closed communication.")
             return
         await event.wait()
@@ -173,13 +162,7 @@ class RunnerSupervisor:
             return
         self.cancelled.add(task_id)
         with anyio.move_on_after(0.5) as scope:
-            try:
-                await self._cancel_sender.send_async(task_id)
-            except ClosedResourceError:
-                # typically occurs when trying to shut down a failed instance
-                logger.warning(
-                    f"Cancelling task {task_id} failed, runner closed communication"
-                )
+            await self._cancel_sender.send_async(task_id)
         if scope.cancel_called:
             logger.error("RunnerSupervisor cancel pipe blocked")
             await self._check_runner(TimeoutError("cancel pipe blocked"))
@@ -192,6 +175,7 @@ class RunnerSupervisor:
                         self.status = event.runner_status
                     if isinstance(event, TaskAcknowledged):
                         self.pending.pop(event.task_id).set()
+                        self.in_progress.add(event.task_id)
                         continue
                     if (
                         isinstance(event, TaskStatusUpdated)
@@ -208,7 +192,7 @@ class RunnerSupervisor:
                                 RunnerShuttingDown,
                             ),
                         )
-                        self.in_progress.pop(event.task_id, None)
+                        self.in_progress.discard(event.task_id)
                         self.completed.add(event.task_id)
                     await self._event_sender.send(event)
         except (ClosedResourceError, BrokenResourceError) as e:
@@ -253,22 +237,6 @@ class RunnerSupervisor:
 
         logger.opt(exception=e).error(f"Runner terminated with {cause}")
 
-        for task in self.in_progress.values():
-            if isinstance(task, (TextGeneration, ImageGeneration, ImageEdits)):
-                with anyio.CancelScope(shield=True):
-                    await self._event_sender.send(
-                        ChunkGenerated(
-                            command_id=task.command_id,
-                            chunk=ErrorChunk(
-                                model=self.shard_metadata.model_card.model_id,
-                                error_message=(
-                                    "Runner shutdown before completing command "
-                                    f"({cause})"
-                                ),
-                            ),
-                        )
-                    )
-
         try:
             self.status = RunnerFailed(error_message=f"Terminated ({cause})")
             with anyio.CancelScope(shield=True):

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

@@ -1,6 +1,5 @@
 """Tests for parse_tool_calls generator, especially unclosed tool call handling."""
 
-import json
 from collections.abc import Generator
 from typing import Any
 
@@ -41,7 +40,6 @@ class TestParseToolCalls:
             parse_tool_calls(
                 _make_responses(texts, finish_on_last=False),
                 _dummy_parser,
-                tools=None,
             )
         )
 
@@ -55,7 +53,6 @@ class TestParseToolCalls:
             parse_tool_calls(
                 _make_responses(texts),
                 _dummy_parser,
-                tools=None,
             )
         )
 
@@ -80,101 +77,9 @@ class TestParseToolCalls:
             parse_tool_calls(
                 _make_responses(texts, finish_on_last=False),
                 make_mlx_parser("<tool_call>", "</tool_call>", _failing_parser),
-                tools=None,
             )
         )
 
         assert len(results) == 1
         assert isinstance(results[0], GenerationResponse)
         assert results[0].text == "<tool_call>bad content</tool_call>"
-
-    def test_tool_schema_coerces_string_arguments_to_expected_types(self):
-        """Tool argument values should be coerced using provided JSON schema."""
-
-        def _parser_with_string_args(_text: str) -> dict[str, Any]:
-            return {
-                "name": "process",
-                "arguments": {
-                    "action": "output",
-                    "id": "0",
-                    "verbose": "true",
-                    "temperature": "0.75",
-                },
-            }
-
-        tools = [
-            {
-                "type": "function",
-                "function": {
-                    "name": "process",
-                    "description": "Manage background processes",
-                    "parameters": {
-                        "type": "object",
-                        "properties": {
-                            "action": {"type": "string"},
-                            "id": {"type": "integer"},
-                            "verbose": {"type": "boolean"},
-                            "temperature": {"type": "number"},
-                        },
-                        "required": ["action"],
-                    },
-                },
-            }
-        ]
-
-        results = list(
-            parse_tool_calls(
-                _make_responses(["<tool_call>", "process", "</tool_call>"]),
-                make_mlx_parser(
-                    "<tool_call>", "</tool_call>", _parser_with_string_args
-                ),
-                tools,
-            )
-        )
-
-        assert len(results) == 1
-        assert isinstance(results[0], ToolCallResponse)
-
-        args = json.loads(results[0].tool_calls[0].arguments)  # pyright: ignore[reportAny]
-        assert args == {
-            "action": "output",
-            "id": 0,
-            "verbose": True,
-            "temperature": 0.75,
-        }
-
-    def test_schema_coercion_skips_unknown_tools(self):
-        """If no matching tool schema exists, arguments should remain unchanged."""
-
-        def _parser_with_string_id(_text: str) -> dict[str, Any]:
-            return {
-                "name": "process",
-                "arguments": {"action": "output", "id": "0"},
-            }
-
-        tools = [
-            {
-                "type": "function",
-                "function": {
-                    "name": "different_tool",
-                    "parameters": {
-                        "type": "object",
-                        "properties": {"id": {"type": "integer"}},
-                    },
-                },
-            }
-        ]
-
-        results = list(
-            parse_tool_calls(
-                _make_responses(["<tool_call>", "process", "</tool_call>"]),
-                make_mlx_parser("<tool_call>", "</tool_call>", _parser_with_string_id),
-                tools,
-            )
-        )
-
-        assert len(results) == 1
-        assert isinstance(results[0], ToolCallResponse)
-
-        args = json.loads(results[0].tool_calls[0].arguments)  # pyright: ignore[reportAny]
-        assert args == {"action": "output", "id": "0"}

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

@@ -1,93 +1 @@
-import multiprocessing as mp
-from typing import cast
-
-import anyio
-import pytest
-
-from exo.shared.models.model_cards import ModelId
-from exo.shared.types.chunks import ErrorChunk
-from exo.shared.types.common import CommandId, NodeId
-from exo.shared.types.events import ChunkGenerated, Event, RunnerStatusUpdated
-from exo.shared.types.tasks import Task, TaskId, TextGeneration
-from exo.shared.types.text_generation import InputMessage, TextGenerationTaskParams
-from exo.shared.types.worker.instances import BoundInstance, InstanceId
-from exo.shared.types.worker.runners import RunnerFailed, RunnerId
-from exo.utils.channels import channel, mp_channel
-from exo.worker.runner.runner_supervisor import RunnerSupervisor
-from exo.worker.tests.unittests.conftest import get_bound_mlx_ring_instance
-
-
-class _DeadProcess:
-    exitcode = -6
-
-    def start(self) -> None:
-        return None
-
-    def is_alive(self) -> bool:
-        return False
-
-    def join(self, _timeout: float | None = None) -> None:
-        return None
-
-    def terminate(self) -> None:
-        return None
-
-    def kill(self) -> None:
-        return None
-
-
-@pytest.mark.asyncio
-async def test_check_runner_emits_error_chunk_for_inflight_text_generation() -> None:
-    event_sender, event_receiver = channel[Event]()
-    task_sender, _ = mp_channel[Task]()
-    cancel_sender, _ = mp_channel[TaskId]()
-    _, ev_recv = mp_channel[Event]()
-
-    bound_instance: BoundInstance = get_bound_mlx_ring_instance(
-        instance_id=InstanceId("instance-a"),
-        model_id=ModelId("mlx-community/Llama-3.2-1B-Instruct-4bit"),
-        runner_id=RunnerId("runner-a"),
-        node_id=NodeId("node-a"),
-    )
-
-    supervisor = RunnerSupervisor(
-        shard_metadata=bound_instance.bound_shard,
-        bound_instance=bound_instance,
-        runner_process=cast("mp.Process", cast(object, _DeadProcess())),
-        initialize_timeout=400,
-        _ev_recv=ev_recv,
-        _task_sender=task_sender,
-        _event_sender=event_sender,
-        _cancel_sender=cancel_sender,
-    )
-
-    command_id = CommandId("cmd-a")
-    task = TextGeneration(
-        task_id=TaskId("task-a"),
-        instance_id=bound_instance.instance.instance_id,
-        command_id=command_id,
-        task_params=TextGenerationTaskParams(
-            model=bound_instance.bound_shard.model_card.model_id,
-            input=[InputMessage(role="user", content="hi")],
-            stream=True,
-        ),
-    )
-    supervisor.in_progress[task.task_id] = task
-    supervisor.shutdown = lambda: None
-
-    await supervisor._check_runner(RuntimeError("boom"))  # pyright: ignore[reportPrivateUsage]
-
-    got_chunk = await event_receiver.receive()
-    got_status = await event_receiver.receive()
-
-    assert isinstance(got_chunk, ChunkGenerated)
-    assert got_chunk.command_id == command_id
-    assert isinstance(got_chunk.chunk, ErrorChunk)
-    assert "Runner shutdown before completing command" in got_chunk.chunk.error_message
-
-    assert isinstance(got_status, RunnerStatusUpdated)
-    assert isinstance(got_status.runner_status, RunnerFailed)
-
-    event_sender.close()
-    with anyio.move_on_after(0.1):
-        await event_receiver.aclose()
+# TODO: