Move mx_all_gather_tasks into utils_mlx

what da ya think!

---------

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

Cargo.toml

@@ -1,9 +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"
@@ -23,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

@@ -5,6 +5,7 @@ edition = { workspace = true }
 publish = false
 
 [lib]
+doctest = false
 path = "src/lib.rs"
 name = "exo_pyo3_bindings"
 
@@ -47,11 +48,19 @@ pyo3-log = "0.13.2"
 
 # macro dependencies
 extend = { workspace = true }
+delegate = { workspace = true }
 
 # async runtime
 tokio = { workspace = true, features = ["full", "tracing"] }
 futures-lite = { workspace = true }
 
+# utility dependencies
+util = { workspace = true }
+
 # Tracing
 log = { workspace = true }
 env_logger = "0.11"
+
+# Networking
+libp2p = { workspace = true, features = ["full"] }
+pin-project = "1.1.10"

rust/exo_pyo3_bindings/README.md

@@ -0,0 +1 @@
+TODO: do something here....

rust/exo_pyo3_bindings/src/allow_threading.rs

@@ -1,22 +1,37 @@
-//! See: <https://pyo3.rs/v0.27.2/async-await.html#detaching-from-the-interpreter-across-await>
+//! SEE: https://pyo3.rs/v0.26.0/async-await.html#detaching-from-the-interpreter-across-await
+//!
+
+use pin_project::pin_project;
 use pyo3::prelude::*;
 use std::{
     future::Future,
-    pin::{Pin, pin},
+    pin::Pin,
     task::{Context, Poll},
 };
 
-pub struct AllowThreads<F>(pub(crate) F);
+/// SEE: https://pyo3.rs/v0.26.0/async-await.html#detaching-from-the-interpreter-across-await
+#[pin_project]
+#[repr(transparent)]
+pub(crate) struct AllowThreads<F>(#[pin] F);
+
+impl<F> AllowThreads<F>
+where
+    Self: Future,
+{
+    pub fn new(f: F) -> Self {
+        Self(f)
+    }
+}
 
 impl<F> Future for AllowThreads<F>
 where
-    F: Future + Unpin + Send,
+    F: Future + Send,
     F::Output: Send,
 {
     type Output = F::Output;
 
-    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
+    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
         let waker = cx.waker();
-        Python::attach(|py| py.detach(|| pin!(&mut self.0).poll(&mut Context::from_waker(waker))))
+        Python::attach(|py| py.detach(|| self.project().0.poll(&mut Context::from_waker(waker))))
     }
 }

rust/exo_pyo3_bindings/src/lib.rs

@@ -1,5 +1,174 @@
-use pyo3_stub_gen::define_stub_info_gatherer;
+//! TODO: crate documentation
+//!
+//! this is here as a placeholder documentation
+//!
+//!
 
 mod allow_threading;
+mod ident;
+mod networking;
+
+use crate::ident::PyKeypair;
+use crate::networking::networking_submodule;
+use pyo3::prelude::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;
+    use extend::ext;
+    use pyo3::exceptions::{PyConnectionError, PyRuntimeError};
+    use pyo3::types::PyBytes;
+    use pyo3::{Py, PyErr, PyResult, Python};
+    use tokio::runtime::Runtime;
+    use tokio::sync::mpsc;
+    use tokio::sync::mpsc::error::TryRecvError;
+    use tokio::task::JoinHandle;
+
+    #[ext(pub, name = ByteArrayExt)]
+    impl [u8] {
+        fn pybytes(&self) -> Py<PyBytes> {
+            Python::attach(|py| PyBytes::new(py, self).unbind())
+        }
+    }
+
+    #[ext(pub, name = ResultExt)]
+    impl<T, E> Result<T, E>
+    where
+        E: ToString,
+    {
+        fn pyerr(self) -> PyResult<T> {
+            self.map_err(|e| PyRuntimeError::new_err(e.to_string()))
+        }
+    }
+
+    pub trait FutureExt: Future + Sized {
+        /// SEE: https://pyo3.rs/v0.26.0/async-await.html#detaching-from-the-interpreter-across-await
+        fn allow_threads_py(self) -> AllowThreads<Self>
+        where
+            AllowThreads<Self>: Future,
+        {
+            AllowThreads::new(self)
+        }
+    }
+
+    impl<T: Future> FutureExt for T {}
+
+    #[ext(pub, name = PyErrExt)]
+    impl PyErr {
+        fn receiver_channel_closed() -> Self {
+            PyConnectionError::new_err("Receiver channel closed unexpectedly")
+        }
+    }
+
+    #[ext(pub, name = PyResultExt)]
+    impl<T> PyResult<T> {
+        fn write_unraisable(self) -> Option<T> {
+            Python::attach(|py| self.write_unraisable_with(py))
+        }
+
+        fn write_unraisable_with(self, py: Python<'_>) -> Option<T> {
+            match self {
+                Ok(v) => Some(v),
+                Err(e) => {
+                    // write error back to python
+                    e.write_unraisable(py, None);
+                    None
+                }
+            }
+        }
+    }
+
+    #[ext(pub, name = TokioRuntimeExt)]
+    impl Runtime {
+        fn spawn_with_scope<F>(&self, py: Python<'_>, future: F) -> PyResult<JoinHandle<F::Output>>
+        where
+            F: Future + Send + 'static,
+            F::Output: Send + 'static,
+        {
+            let locals = pyo3_async_runtimes::tokio::get_current_locals(py)?;
+            Ok(self.spawn(pyo3_async_runtimes::tokio::scope(locals, future)))
+        }
+    }
+
+    #[ext(pub, name = TokioMpscSenderExt)]
+    impl<T> mpsc::Sender<T> {
+        /// Sends a value, waiting until there is capacity.
+        ///
+        /// A successful send occurs when it is determined that the other end of the
+        /// channel has not hung up already. An unsuccessful send would be one where
+        /// the corresponding receiver has already been closed.
+        async fn send_py(&self, value: T) -> PyResult<()> {
+            self.send(value)
+                .await
+                .map_err(|_| PyErr::receiver_channel_closed())
+        }
+    }
+
+    #[ext(pub, name = TokioMpscReceiverExt)]
+    impl<T> mpsc::Receiver<T> {
+        /// Receives the next value for this receiver.
+        async fn recv_py(&mut self) -> PyResult<T> {
+            self.recv().await.ok_or_else(PyErr::receiver_channel_closed)
+        }
+
+        /// Receives at most `limit` values for this receiver and returns them.
+        ///
+        /// For `limit = 0`, an empty collection of messages will be returned immediately.
+        /// For `limit > 0`, if there are no messages in the channel's queue this method
+        /// will sleep until a message is sent.
+        async fn recv_many_py(&mut self, limit: usize) -> PyResult<Vec<T>> {
+            // get updates from receiver channel
+            let mut updates = Vec::with_capacity(limit);
+            let received = self.recv_many(&mut updates, limit).await;
+
+            // if we received zero items, then the channel was unexpectedly closed
+            if limit != 0 && received == 0 {
+                return Err(PyErr::receiver_channel_closed());
+            }
+
+            Ok(updates)
+        }
+
+        /// Tries to receive the next value for this receiver.
+        fn try_recv_py(&mut self) -> PyResult<Option<T>> {
+            match self.try_recv() {
+                Ok(v) => Ok(Some(v)),
+                Err(TryRecvError::Empty) => Ok(None),
+                Err(TryRecvError::Disconnected) => Err(PyErr::receiver_channel_closed()),
+            }
+        }
+    }
+}
+
+/// 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_pyo3_bindings")]
+fn main_module(m: &Bound<'_, PyModule>) -> PyResult<()> {
+    // install logger
+    pyo3_log::init();
+    let mut builder = tokio::runtime::Builder::new_multi_thread();
+    builder.enable_all();
+    pyo3_async_runtimes::tokio::init(builder);
+
+    // 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)?;
+
+    // top-level constructs
+    // TODO: ...
+
+    Ok(())
+}
 
 define_stub_info_gatherer!(stub_info);

rust/exo_pyo3_bindings/src/networking.rs

@@ -0,0 +1,311 @@
+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 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 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")]
+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]
+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 { 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(),
+            },
+        }
+    }
+}
+
+#[gen_stub_pymethods]
+#[pymethods]
+impl PyNetworkingHandle {
+    // 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 ----
+
+    #[new]
+    fn py_new(identity: Bound<'_, PyKeypair>) -> PyResult<Self> {
+        // create communication channels
+        let (to_swarm, from_client) = mpsc::channel(MPSC_CHANNEL_SIZE);
+
+        // get identity
+        let identity = identity.borrow().0.clone();
+
+        // create networking swarm (within tokio context!! or it crashes)
+        let _guard = pyo3_async_runtimes::tokio::get_runtime().enter();
+        let swarm = create_swarm(identity, from_client).pyerr()?.into_stream();
+
+        Ok(Self {
+            swarm: Arc::new(Mutex::new(swarm)),
+            to_swarm,
+        })
+    }
+
+    #[gen_stub(skip)]
+    fn recv<'py>(&'py self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
+        let swarm = Arc::clone(&self.swarm);
+        pyo3_async_runtimes::tokio::future_into_py(py, async move {
+            swarm
+                .try_lock()
+                .map_err(|_| PyRuntimeError::new_err("called recv twice concurrently"))?
+                .next()
+                .await
+                .ok_or(PyErr::receiver_channel_closed())
+                .map(PyFromSwarm::from)
+        })
+    }
+
+    // ---- Gossipsub management methods ----
+
+    /// Subscribe to a `GossipSub` topic.
+    ///
+    /// Returns `True` if the subscription worked. Returns `False` if we were already subscribed.
+    async fn gossipsub_subscribe(&self, topic: String) -> PyResult<bool> {
+        let (tx, rx) = oneshot::channel();
+
+        // send off request to subscribe
+        self.to_swarm
+            .send_py(ToSwarm::Subscribe {
+                topic,
+                result_sender: tx,
+            })
+            .allow_threads_py() // allow-threads-aware async call
+            .await?;
+
+        // wait for response & return any errors
+        rx.allow_threads_py() // allow-threads-aware async call
+            .await
+            .map_err(|_| PyErr::receiver_channel_closed())?
+            .pyerr()
+    }
+
+    /// Unsubscribes from a `GossipSub` topic.
+    ///
+    /// Returns `True` if we were subscribed to this topic. Returns `False` if we were not subscribed.
+    async fn gossipsub_unsubscribe(&self, topic: String) -> PyResult<bool> {
+        let (tx, rx) = oneshot::channel();
+
+        // send off request to unsubscribe
+        self.to_swarm
+            .send_py(ToSwarm::Unsubscribe {
+                topic,
+                result_sender: tx,
+            })
+            .allow_threads_py() // allow-threads-aware async call
+            .await?;
+
+        // wait for response & convert any errors
+        rx.allow_threads_py() // allow-threads-aware async call
+            .await
+            .map_err(|_| PyErr::receiver_channel_closed())
+    }
+
+    /// Publishes a message with multiple topics to the `GossipSub` network.
+    ///
+    /// If no peers are found that subscribe to this topic, throws `NoPeersSubscribedToTopicError` exception.
+    async fn gossipsub_publish(&self, topic: String, data: Py<PyBytes>) -> PyResult<()> {
+        let (tx, rx) = oneshot::channel();
+
+        // send off request to subscribe
+        let data = Python::attach(|py| Vec::from(data.as_bytes(py)));
+        self.to_swarm
+            .send_py(ToSwarm::Publish {
+                topic,
+                data,
+                result_sender: tx,
+            })
+            .allow_threads_py() // allow-threads-aware async call
+            .await?;
+
+        // wait for response & return any errors => ignore messageID for now!!!
+        let _ = rx
+            .allow_threads_py() // allow-threads-aware async call
+            .await
+            .map_err(|_| PyErr::receiver_channel_closed())?
+            .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>()?;
+
+    Ok(())
+}

rust/exo_pyo3_bindings/tests/dummy.rs

@@ -0,0 +1,54 @@
+#[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;
+
+        let (tx, mut rx) = mpsc::channel::<Ping>(10);
+
+        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;
+                            }
+                        }
+                    }
+                    _ = tokio::time::sleep(Duration::from_secs_f32(0.1)) => {
+                        println!("TASK: heartbeat");
+                    }
+                }
+            }
+
+            println!("TASK: exited");
+        });
+
+        let tx2 = tx.clone();
+
+        tokio::time::sleep(Duration::from_secs_f32(0.11)).await;
+
+        tx.send(Ping).await.expect("Should not fail");
+        drop(tx);
+
+        tokio::time::sleep(Duration::from_secs_f32(0.11)).await;
+
+        tx2.send(Ping).await.expect("Should not fail");
+        drop(tx2);
+
+        tokio::time::sleep(Duration::from_secs_f32(0.11)).await;
+    }
+}

rust/networking/Cargo.toml

@@ -5,6 +5,7 @@ edition = { workspace = true }
 publish = false
 
 [lib]
+doctest = false
 name = "networking"
 path = "src/lib.rs"
 
@@ -12,22 +13,30 @@ path = "src/lib.rs"
 workspace = true
 
 [dependencies]
+# datastructures
+either = { workspace = true }
+
 # macro dependencies
 extend = { workspace = true }
+delegate = { workspace = true }
 
 # async
-tokio = { workspace = true, features = ["full"] }
+async-stream = { workspace = true }
 futures-lite = { workspace = true }
+futures-timer = { workspace = true }
+tokio = { workspace = true, features = ["full"] }
 
 # utility dependencies
-tracing-subscriber = { version = "0.3.19", features = ["default", "env-filter"] }
-pin-project = "1.1.10"
+util = { workspace = true }
+tracing-subscriber = { version = "0.3.19", features = [
+  "default",
+  "env-filter",
+] }
+keccak-const = { workspace = true }
 
 # tracing/logging
 log = { workspace = true }
-iroh = "0.96.1"
-iroh-gossip = "0.96.0"
-iroh-blobs = "0.98.0"
-iroh-docs = "0.96.0"
 
 # networking
+libp2p = { workspace = true, features = ["full"] }
+pin-project = "1.1.10"

rust/networking/examples/chatroom.rs

@@ -1,3 +1,9 @@
+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;
 
@@ -6,4 +12,70 @@ 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)
+        .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/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

@@ -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

@@ -0,0 +1,44 @@
+//! TODO: crate documentation
+//!
+//! 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 {
+    use std::error::Error;
+
+    pub type AnyError = Box<dyn Error + Send + Sync + 'static>;
+    pub type AnyResult<T> = Result<T, AnyError>;
+}
+
+/// 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;
+            };
+            let Some(Protocol::Tcp(port)) = ps.next() else {
+                return None;
+            };
+            Some((ip, port))
+        }
+    }
+}

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/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/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: