lean_ctx/core/mcp_catalog/
runtime.rs

1// SPDX-License-Identifier: Apache-2.0
2//! One process-lifetime Tokio runtime for all gateway I/O.
3//!
4//! Pooled downstream sessions ([`super::pool`]) spawn their transport tasks on
5//! whichever runtime opened them. Gateway entry points used to build a fresh
6//! runtime per call and drop it afterwards, which silently stopped those tasks:
7//! the session still looked open to the pool's liveness check, was handed out
8//! again, and the next request failed with `Transport closed` (#2040). Every
9//! gateway caller — `ctx_tools` and the compression adapters — now drives its
10//! futures on this one runtime, so a pooled session lives exactly as long as the
11//! pool keeps it.
12
13use std::future::Future;
14use std::sync::{Mutex, OnceLock};
15
16use tokio::runtime::Runtime;
17
18/// Two workers: downstream I/O is light, and the pool keeps at most one
19/// session per wiring.
20const WORKERS: usize = 2;
21
22static RUNTIME: OnceLock<Runtime> = OnceLock::new();
23static INIT: Mutex<()> = Mutex::new(());
24
25/// The gateway runtime, started on first use. A failed start is reported and
26/// retried by the next call rather than remembered for the process lifetime.
27pub fn runtime() -> Result<&'static Runtime, String> {
28    if let Some(rt) = RUNTIME.get() {
29        return Ok(rt);
30    }
31    // Serialise the first start so no losing runtime is ever built and dropped
32    // (dropping a runtime inside async code panics).
33    let _guard = INIT
34        .lock()
35        .unwrap_or_else(std::sync::PoisonError::into_inner);
36    if let Some(rt) = RUNTIME.get() {
37        return Ok(rt);
38    }
39    let rt = tokio::runtime::Builder::new_multi_thread()
40        .worker_threads(WORKERS)
41        .thread_name("lean-ctx-gateway")
42        .enable_all()
43        .build()
44        .map_err(|e| format!("failed to start the gateway runtime: {e}"))?;
45    Ok(RUNTIME.get_or_init(|| rt))
46}
47
48/// Drive `fut` to completion on the gateway runtime from synchronous code,
49/// whatever the caller's context. A thread that has entered another Tokio
50/// runtime (an async task, `block_in_place`, `spawn_blocking`) must not call
51/// `block_on`, so the wait moves to a short scoped thread there; the future
52/// itself and every task it spawns run on the gateway runtime either way.
53pub fn block_on<F>(fut: F) -> Result<F::Output, String>
54where
55    F: Future + Send,
56    F::Output: Send,
57{
58    let rt = runtime()?;
59    if tokio::runtime::Handle::try_current().is_err() {
60        return Ok(rt.block_on(fut));
61    }
62    std::thread::scope(|scope| {
63        std::thread::Builder::new()
64            .name("lean-ctx-gateway-wait".into())
65            .spawn_scoped(scope, || rt.block_on(fut))
66            .map_err(|e| format!("failed to start a gateway wait thread: {e}"))?
67            .join()
68            .map_err(|_| "a gateway call panicked".to_string())
69    })
70}
71
72#[cfg(test)]
73mod tests {
74    use super::*;
75
76    /// The #2040 failure in miniature: a task spawned during one gateway call
77    /// (like a pooled session's transport task) must still run in the next.
78    #[test]
79    fn a_task_spawned_in_one_call_survives_into_the_next() {
80        let (tx, rx) = tokio::sync::oneshot::channel::<u8>();
81        // Wrapped so the call yields the handle itself, not a future to await.
82        let task = block_on(async { Some(tokio::spawn(async move { rx.await.ok() })) })
83            .unwrap()
84            .unwrap();
85        let value = block_on(async move {
86            tx.send(7).expect("the task's receiver is still alive");
87            task.await.expect("the task was not cancelled")
88        })
89        .unwrap();
90        assert_eq!(value, Some(7));
91    }
92
93    /// Called from inside a running async task, `block_on` must neither panic
94    /// ("cannot start a runtime from within a runtime") nor deadlock.
95    #[tokio::test(flavor = "current_thread")]
96    async fn usable_from_inside_an_async_task() {
97        assert_eq!(block_on(async { 41 + 1 }).unwrap(), 42);
98    }
99
100    #[test]
101    fn every_call_shares_one_runtime() {
102        assert!(std::ptr::eq(runtime().unwrap(), runtime().unwrap()));
103    }
104
105    /// A panic inside a gateway future is reported, not propagated into the
106    /// caller's thread.
107    #[tokio::test(flavor = "current_thread")]
108    async fn a_panicking_call_is_an_error_not_a_crash() {
109        let out = block_on(async {
110            assert!(!std::hint::black_box(true), "boom");
111            0_u8
112        });
113        assert_eq!(out.unwrap_err(), "a gateway call panicked");
114    }
115}
lean_ctx/core/mcp_catalog/runtime.rs

lean_ctx/core/mcp_catalog/
runtime.rs
