lean_ctx/core/
phase_timing.rs

1// SPDX-License-Identifier: Apache-2.0
2//! Named phases of a long tool call (#2063).
3//!
4//! `LEAN_CTX_PHASE_TIMING=1` prints each phase's start and duration to stderr,
5//! so a slow call can be attributed without a debugger — also a call that
6//! never finishes, whose phases would otherwise print nothing. Independently,
7//! the phases that are still running are kept, so `lean-ctx call` can say
8//! where it is stuck instead of going silent.
9//!
10//! Each guard owns its own entry, so concurrent calls (the MCP server runs many
11//! at once) cannot overwrite or resurrect each other's phase when their guards
12//! drop out of order.
13
14use std::collections::BTreeMap;
15use std::sync::Mutex;
16use std::sync::atomic::{AtomicU64, Ordering};
17use std::time::{Duration, Instant};
18
19static NEXT_ID: AtomicU64 = AtomicU64::new(0);
20static RUNNING: Mutex<BTreeMap<u64, (&'static str, Instant)>> = Mutex::new(BTreeMap::new());
21
22/// Whether `LEAN_CTX_PHASE_TIMING=1` asks for phase lines on stderr.
23pub(crate) fn enabled() -> bool {
24    std::env::var_os("LEAN_CTX_PHASE_TIMING").is_some_and(|v| v == "1")
25}
26
27/// Marks `name` as running until the guard drops.
28#[must_use = "the phase ends when the guard drops"]
29pub fn enter(name: &'static str) -> PhaseGuard {
30    let id = NEXT_ID.fetch_add(1, Ordering::Relaxed);
31    let started = Instant::now();
32    if let Ok(mut running) = RUNNING.lock() {
33        running.insert(id, (name, started));
34    }
35    if enabled() {
36        eprintln!("[phase] {name} started");
37    }
38    PhaseGuard { id, name, started }
39}
40
41/// Reports one slow step inside `phase` (one source of a scan, say) when
42/// phase timing is on. `detail` names the step; it goes to stderr only.
43pub fn note_slow(phase: &'static str, detail: &str, spent: Duration) {
44    if enabled() {
45        eprintln!("[phase] {phase} slow step {spent:?}: {detail}");
46    }
47}
48
49/// The most recently entered phase still running and how long it has run.
50pub fn current() -> Option<(&'static str, Duration)> {
51    RUNNING.lock().ok().and_then(|running| {
52        running
53            .last_key_value()
54            .map(|(_, (name, started))| (*name, started.elapsed()))
55    })
56}
57
58pub struct PhaseGuard {
59    id: u64,
60    name: &'static str,
61    started: Instant,
62}
63
64impl Drop for PhaseGuard {
65    fn drop(&mut self) {
66        if let Ok(mut running) = RUNNING.lock() {
67            running.remove(&self.id);
68        }
69        if enabled() {
70            eprintln!("[phase] {} {:?}", self.name, self.started.elapsed());
71        }
72    }
73}
74
75#[cfg(test)]
76mod tests {
77    use super::*;
78
79    fn running_names() -> Vec<&'static str> {
80        RUNNING
81            .lock()
82            .map(|running| running.values().map(|(name, _)| *name).collect())
83            .unwrap_or_default()
84    }
85
86    #[test]
87    fn out_of_order_drops_never_resurrect_a_finished_phase() {
88        let first = enter("phase-test: first call");
89        let second = enter("phase-test: second call");
90        drop(first);
91        assert!(running_names().contains(&"phase-test: second call"));
92        assert!(!running_names().contains(&"phase-test: first call"));
93        drop(second);
94        let names = running_names();
95        assert!(!names.contains(&"phase-test: first call"));
96        assert!(!names.contains(&"phase-test: second call"));
97    }
98}
lean_ctx/core/phase_timing.rs

lean_ctx/core/
phase_timing.rs
