ljos_cli/
tool_bound.rs

1//! The bound on one MCP tool call.
2//!
3//! A pack socket, a git command, a habitat child and an advisory lock stop
4//! when this runs out, and the line names the one that did. The default is
5//! 30 seconds. `LJOS_MCP_TIMEOUT_MS` is milliseconds. `LJOS_MCP_TIMEOUT` is
6//! whole seconds and is used when the millisecond form is unset.
7
8use std::sync::{Arc, Mutex};
9use std::time::{Duration, Instant};
10
11/// The bound when neither environment variable is set.
12pub const DEFAULT: Duration = Duration::from_secs(30);
13
14struct State {
15    deadline: Instant,
16    stage: Mutex<String>,
17    failed: Mutex<Option<String>>,
18}
19
20tokio::task_local! {
21    static CALL: Arc<State>;
22}
23
24thread_local! {
25    static THREAD: std::cell::RefCell<Option<Arc<State>>> =
26        const { std::cell::RefCell::new(None) };
27}
28
29/// Restores the bound the thread had before [`Ticket::bind`].
30pub struct ThreadGuard {
31    previous: Option<Arc<State>>,
32}
33
34impl Drop for ThreadGuard {
35    fn drop(&mut self) {
36        let previous = self.previous.take();
37        THREAD.with(|slot| *slot.borrow_mut() = previous);
38    }
39}
40
41/// The bound captured on the calling task, for a `spawn_blocking` thread.
42pub struct Ticket(Arc<State>);
43
44impl Ticket {
45    /// Make this thread see the call's bound until the guard drops.
46    #[must_use]
47    pub fn bind(self) -> ThreadGuard {
48        bind_thread(self.0)
49    }
50}
51
52fn bind_thread(state: Arc<State>) -> ThreadGuard {
53    let previous = THREAD.with(|slot| slot.replace(Some(state)));
54    ThreadGuard { previous }
55}
56
57fn current() -> Option<Arc<State>> {
58    if let Some(state) = THREAD.with(|slot| slot.borrow().clone()) {
59        return Some(state);
60    }
61    CALL.try_with(Arc::clone).ok()
62}
63
64fn locked<T>(mutex: &Mutex<T>) -> std::sync::MutexGuard<'_, T> {
65    mutex
66        .lock()
67        .unwrap_or_else(std::sync::PoisonError::into_inner)
68}
69
70/// How long a tool call may run.
71#[must_use]
72pub fn budget() -> Duration {
73    if let Some(ms) = positive("LJOS_MCP_TIMEOUT_MS") {
74        return Duration::from_millis(ms);
75    }
76    if let Some(secs) = positive("LJOS_MCP_TIMEOUT") {
77        return Duration::from_secs(secs);
78    }
79    DEFAULT
80}
81
82fn positive(key: &str) -> Option<u64> {
83    std::env::var(key)
84        .ok()
85        .and_then(|value| value.trim().parse().ok())
86        .filter(|n| *n > 0)
87}
88
89/// Whether this task or thread is inside a tool call's bound.
90#[must_use]
91pub fn armed() -> bool {
92    current().is_some()
93}
94
95/// When the call must be done. `None` outside a tool call.
96#[must_use]
97pub fn deadline_instant() -> Option<Instant> {
98    current().map(|state| state.deadline)
99}
100
101/// Time left in the call. `None` outside a tool call.
102#[must_use]
103pub fn remaining() -> Option<Duration> {
104    deadline_instant().map(|deadline| deadline.saturating_duration_since(Instant::now()))
105}
106
107/// The bound on this task, so a blocking thread can take it with it.
108#[must_use]
109pub fn ticket() -> Option<Ticket> {
110    current().map(Ticket)
111}
112
113/// Name the wait in progress. A timeout line uses the latest name.
114pub fn note(what: &str) {
115    if let Some(state) = current() {
116        *locked(&state.stage) = what.to_string();
117    }
118}
119
120/// Name `what` until the guard drops, unless the bound has already run out.
121pub fn stage(what: &str) -> StageGuard {
122    let previous = current().map(|state| locked(&state.stage).clone());
123    note(what);
124    StageGuard { previous }
125}
126
127/// Puts the previous stage back. A spent bound keeps the name it timed out on.
128pub struct StageGuard {
129    previous: Option<String>,
130}
131
132impl Drop for StageGuard {
133    fn drop(&mut self) {
134        if remaining().is_some_and(|left| left.is_zero()) {
135            return;
136        }
137        if let Some(previous) = self.previous.take() {
138            note(&previous);
139        }
140    }
141}
142
143/// Remember the first timeout line for this call.
144pub fn mark_failed(line: impl Into<String>) {
145    let line = line.into();
146    if let Some(state) = current() {
147        let mut slot = locked(&state.failed);
148        if slot.is_none() {
149            *slot = Some(line);
150        }
151    }
152}
153
154/// The timeout line a wait already recorded, if one has.
155#[must_use]
156pub fn failure() -> Option<String> {
157    current().and_then(|state| locked(&state.failed).clone())
158}
159
160/// One line naming the wait that was in progress.
161#[must_use]
162pub fn timeout_line() -> String {
163    let what = current()
164        .map(|state| locked(&state.stage).clone())
165        .filter(|what| !what.is_empty())
166        .unwrap_or_else(|| "tool".to_string());
167    format!("{what} timed out")
168}
169
170/// Run `fut` inside a fresh bound named `tool`.
171pub async fn scope<F: std::future::Future>(tool: &str, fut: F) -> F::Output {
172    let state = Arc::new(State {
173        deadline: Instant::now() + budget(),
174        stage: Mutex::new(tool.to_string()),
175        failed: Mutex::new(None),
176    });
177    CALL.scope(state, fut).await
178}
179
180/// `spawn_blocking` that keeps this call's bound on the worker.
181pub async fn blocking<T: Send + 'static>(
182    f: impl FnOnce() -> T + Send + 'static,
183) -> Result<T, tokio::task::JoinError> {
184    let ticket = ticket();
185    tokio::task::spawn_blocking(move || {
186        let _guard = ticket.map(Ticket::bind);
187        f()
188    })
189    .await
190}
191
192/// Arm a bound on this thread. Tests use this. The guard clears it.
193#[cfg(test)]
194#[must_use]
195pub(crate) fn arm(total: Duration) -> ThreadGuard {
196    bind_thread(Arc::new(State {
197        deadline: Instant::now() + total,
198        stage: Mutex::new(String::new()),
199        failed: Mutex::new(None),
200    }))
201}
202
203#[cfg(test)]
204mod tests {
205    use super::*;
206
207    #[test]
208    fn the_tool_budget_is_thirty_seconds_unless_configured() {
209        let _g = crate::tests::env_guard();
210        unsafe {
211            std::env::remove_var("LJOS_MCP_TIMEOUT");
212            std::env::remove_var("LJOS_MCP_TIMEOUT_MS");
213        }
214        assert_eq!(budget(), Duration::from_secs(30));
215        unsafe { std::env::set_var("LJOS_MCP_TIMEOUT", "12") };
216        assert_eq!(budget(), Duration::from_secs(12));
217        unsafe { std::env::set_var("LJOS_MCP_TIMEOUT_MS", "250") };
218        assert_eq!(budget(), Duration::from_millis(250));
219    }
220
221    /// A writer that accepts and then says nothing used to hold the call.
222    #[test]
223    fn a_silent_packsetd_stops_inside_the_tool_budget() {
224        let _g = crate::tests::env_guard();
225        let home = tempfile::tempdir().unwrap();
226        let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap();
227        let addr = listener.local_addr().unwrap();
228        std::thread::spawn(move || {
229            for stream in listener.incoming() {
230                let Ok(mut stream) = stream else { break };
231                let mut buf = [0u8; 1024];
232                let _ = std::io::Read::read(&mut stream, &mut buf);
233                std::thread::sleep(Duration::from_secs(30));
234            }
235        });
236        unsafe {
237            std::env::set_var("HOME", home.path());
238            std::env::set_var("PACKSET_URL", format!("http://{addr}"));
239            std::env::remove_var("PACKSET_TOKEN");
240            std::env::remove_var("PACKSET_TOKEN_FILE");
241            std::env::remove_var("PACKSET_TIMEOUT_MS");
242            std::env::remove_var("LJOS_IN_HOOK");
243        }
244        let _bound = arm(Duration::from_millis(400));
245        let started = Instant::now();
246        let err = crate::pack_link::PacksetClient::from_env()
247            .expect("url")
248            .status(None)
249            .expect_err("a silent writer answers");
250        assert!(
251            started.elapsed() < Duration::from_secs(2),
252            "waited {:?}",
253            started.elapsed()
254        );
255        let text = err.to_string();
256        assert!(text.contains("packsetd timed out"), "{text}");
257    }
258
259    /// A habitat child that never exits, and a git-shaped child, both stop.
260    #[test]
261    fn a_hung_subprocess_stops_inside_the_tool_budget() {
262        let _g = crate::tests::env_guard();
263        let _bound = arm(Duration::from_millis(400));
264        let started = Instant::now();
265        let err = crate::run_captured("sh", &["-c", "exec sleep 30"]).unwrap_err();
266        assert!(
267            started.elapsed() < Duration::from_secs(2),
268            "habitat child waited {:?}",
269            started.elapsed()
270        );
271        assert!(err.to_string().contains("sh timed out"), "{err}");
272
273        let _bound = arm(Duration::from_millis(400));
274        let started = Instant::now();
275        let mut cmd = std::process::Command::new("sh");
276        cmd.args(["-c", "exec sleep 30"]);
277        let err = crate::bounded_output(&mut cmd).unwrap_err();
278        assert!(
279            started.elapsed() < Duration::from_secs(2),
280            "git-shaped child waited {:?}",
281            started.elapsed()
282        );
283        assert!(err.to_string().contains("sh timed out"), "{err}");
284    }
285
286    /// An exclusive lock another holder keeps is given up at the bound.
287    #[test]
288    fn a_held_lock_stops_inside_the_tool_budget() {
289        let _g = crate::tests::env_guard();
290        let dir = tempfile::tempdir().unwrap();
291        let path = dir.path().join("held.lock");
292        let held = std::fs::OpenOptions::new()
293            .create(true)
294            .append(true)
295            .open(&path)
296            .unwrap();
297        {
298            use std::os::unix::io::AsRawFd;
299            assert_eq!(unsafe { libc::flock(held.as_raw_fd(), libc::LOCK_EX) }, 0);
300        }
301        let _bound = arm(Duration::from_millis(300));
302        let started = Instant::now();
303        let _lock = crate::CommitLock::acquire(&path);
304        assert!(
305            started.elapsed() < Duration::from_secs(2),
306            "lock waited {:?}",
307            started.elapsed()
308        );
309        assert_eq!(failure().as_deref(), Some("lock timed out"));
310    }
311}
ljos_cli/tool_bound.rs

ljos_cli/
tool_bound.rs
