lean_ctx/tools/registered/
ctx_multi_read.rs

1use rmcp::model::Tool;
2use rmcp::ErrorData;
3use serde_json::{json, Map, Value};
4
5use crate::server::tool_trait::{
6    get_bool, get_str, get_str_array, McpTool, ToolContext, ToolOutput,
7};
8use crate::tool_defs::tool_def;
9
10pub struct CtxMultiReadTool;
11
12impl McpTool for CtxMultiReadTool {
13    fn name(&self) -> &'static str {
14        "ctx_multi_read"
15    }
16
17    fn tool_def(&self) -> Tool {
18        tool_def(
19            "ctx_multi_read",
20            "Batch read files in one call. Same modes as ctx_read.",
21            json!({
22                "type": "object",
23                "properties": {
24                    "paths": {
25                        "type": "array",
26                        "items": { "type": "string" },
27                        "description": "Absolute file paths to read, in order"
28                    },
29                    "mode": {
30                        "type": "string",
31                        "description": "Compression mode (default: full). Same modes as ctx_read (auto, full, raw, map, signatures, diff, aggressive, entropy, task, reference, lines:N-M). Use 'raw' for zero-overhead output."
32                    },
33                    "fresh": {
34                        "type": "boolean",
35                        "description": "Bypass cache and force a full re-read for all paths. Use when running as a subagent that may not have the parent's context."
36                    }
37                },
38                "required": ["paths"]
39            }),
40        )
41    }
42
43    fn handle(
44        &self,
45        args: &Map<String, Value>,
46        ctx: &ToolContext,
47    ) -> Result<ToolOutput, ErrorData> {
48        // Panic guard (mirrors ctx_read): a panic in tree-sitter / compression must
49        // never unwind through the dispatch `block_in_place` and kill the MCP server.
50        match std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| self.handle_inner(args, ctx)))
51        {
52            Ok(result) => result,
53            Err(_) => Err(ErrorData::internal_error(
54                "ctx_multi_read panicked while processing the batch. This is a bug — please report it.",
55                None,
56            )),
57        }
58    }
59}
60
61impl CtxMultiReadTool {
62    #[allow(clippy::unused_self)]
63    fn handle_inner(
64        &self,
65        args: &Map<String, Value>,
66        ctx: &ToolContext,
67    ) -> Result<ToolOutput, ErrorData> {
68        let raw_paths = get_str_array(args, "paths")
69            .ok_or_else(|| ErrorData::invalid_params("paths array is required", None))?;
70
71        let session_lock = ctx
72            .session
73            .as_ref()
74            .ok_or_else(|| ErrorData::internal_error("session not available", None))?;
75        let cache_lock = ctx
76            .cache
77            .as_ref()
78            .ok_or_else(|| ErrorData::internal_error("cache not available", None))?;
79
80        let cap = crate::core::limits::max_read_bytes() as u64;
81
82        // Resolve + filter paths and capture the active task under one short read lock.
83        // `bounded_lock` uses `Handle::block_on` directly — NOT a nested
84        // `block_in_place` — because the dispatch layer already wraps this handler in
85        // `block_in_place`. The previous nested `block_in_place` calls could exhaust the
86        // 32-thread blocking pool under concurrent reads and freeze the server (#271).
87        let (paths, current_task) = {
88            let Some(session) =
89                crate::server::bounded_lock::read(session_lock, "ctx_multi_read:session")
90            else {
91                return Err(ErrorData::internal_error(
92                    "session read-lock timeout in ctx_multi_read — another tool may be holding it. Retry in a moment.",
93                    None,
94                ));
95            };
96            let mut paths = Vec::with_capacity(raw_paths.len());
97            for p in &raw_paths {
98                let resolved = super::resolve_path_sync(&session, p)
99                    .map_err(|e| ErrorData::invalid_params(e, None))?;
100                if crate::core::binary_detect::is_binary_file(&resolved) {
101                    continue;
102                }
103                if let Ok(meta) = std::fs::metadata(&resolved) {
104                    if meta.len() > cap {
105                        continue;
106                    }
107                }
108                paths.push(resolved);
109            }
110            let current_task = session.task.as_ref().map(|t| t.description.clone());
111            (paths, current_task)
112        };
113
114        if paths.is_empty() {
115            return Err(ErrorData::invalid_params(
116                "all paths are binary or exceed the size limit",
117                None,
118            ));
119        }
120
121        let mode = get_str(args, "mode").unwrap_or_else(|| {
122            let p = crate::core::profiles::active_profile();
123            let dm = p.read.default_mode_effective();
124            if dm == "auto" {
125                "full".to_string()
126            } else {
127                dm.to_string()
128            }
129        });
130        let fresh = get_bool(args, "fresh").unwrap_or(false);
131
132        // Batch read under one bounded write lock. `bounded_lock` guarantees we never
133        // block the runtime indefinitely and degrade gracefully on contention instead
134        // of hanging; ctx_read's own fast/slow path tolerates this lock being held.
135        let Some(mut cache) =
136            crate::server::bounded_lock::write(cache_lock, "ctx_multi_read:cache")
137        else {
138            return Err(ErrorData::internal_error(
139                "cache write-lock timeout in ctx_multi_read — another tool may be holding it. Retry in a moment.",
140                None,
141            ));
142        };
143        let output = crate::tools::ctx_multi_read::handle_with_task_fresh(
144            &mut cache,
145            &paths,
146            &mode,
147            fresh,
148            ctx.crp_mode,
149            current_task.as_deref(),
150        );
151        let mut total_original: usize = 0;
152        for path in &paths {
153            total_original =
154                total_original.saturating_add(cache.get(path).map_or(0, |e| e.original_tokens));
155        }
156        let tokens = crate::core::tokens::count_tokens(&output);
157        drop(cache);
158
159        Ok(ToolOutput {
160            text: output,
161            original_tokens: total_original,
162            saved_tokens: total_original.saturating_sub(tokens),
163            mode: Some(mode),
164            path: None,
165            changed: false,
166            shell_outcome: None,
167        })
168    }
169}
170
171#[cfg(test)]
172mod tests {
173    use super::*;
174    use std::sync::Arc;
175    use std::time::Duration;
176    use tokio::sync::RwLock;
177
178    use crate::core::cache::SessionCache;
179    use crate::core::session::SessionState;
180    use crate::tools::CrpMode;
181
182    fn ctx_with(
183        cache: Arc<RwLock<SessionCache>>,
184        session: Arc<RwLock<SessionState>>,
185        project_root: &str,
186    ) -> ToolContext {
187        ToolContext {
188            project_root: project_root.to_string(),
189            extra_roots: Vec::new(),
190            minimal: false,
191            resolved_paths: std::collections::HashMap::new(),
192            crp_mode: CrpMode::Off,
193            cache: Some(cache),
194            session: Some(session),
195            tool_calls: None,
196            agent_id: None,
197            workflow: None,
198            ledger: None,
199            client_name: None,
200            pipeline_stats: None,
201            call_count: None,
202            autonomy: None,
203            pressure_snapshot: None,
204            path_errors: std::collections::HashMap::new(),
205            bm25_cache: None,
206            progress_sender: None,
207        }
208    }
209
210    /// Regression for #271 (crash vector 11): under concurrent load,
211    /// `ctx_multi_read` must not hang. The handler runs inside the dispatch
212    /// layer's `block_in_place`, so it must acquire its session/cache locks
213    /// via `Handle::block_on` WITHOUT nesting another `block_in_place` —
214    /// nesting consumes extra blocking-pool threads and, under load, exhausts
215    /// the pool, hanging the call (no JSON-RPC response → client "invoke"
216    /// error).
217    ///
218    /// With only 2 worker threads and 8 concurrent batch reads, a nested
219    /// `block_in_place` regression would deadlock the pool and trip the 20s
220    /// timeout below.
221    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
222    async fn concurrent_multi_read_does_not_hang() {
223        let dir = tempfile::tempdir().unwrap();
224        let mut paths = Vec::new();
225        for i in 0..6 {
226            let p = dir.path().join(format!("file_{i}.rs"));
227            std::fs::write(&p, format!("fn f{i}() {{ let _ = {i}; }}\n")).unwrap();
228            paths.push(p.to_string_lossy().to_string());
229        }
230        let root = dir.path().to_string_lossy().to_string();
231
232        let cache: Arc<RwLock<SessionCache>> = Arc::new(RwLock::new(SessionCache::new()));
233        let session = {
234            let mut s = SessionState::new();
235            s.project_root = Some(root.clone());
236            Arc::new(RwLock::new(s))
237        };
238
239        let mut handles = Vec::new();
240        for _ in 0..8 {
241            let cache = cache.clone();
242            let session = session.clone();
243            let paths = paths.clone();
244            let root = root.clone();
245            handles.push(tokio::spawn(async move {
246                let ctx = ctx_with(cache, session, &root);
247                let args = json!({ "paths": paths, "mode": "full" })
248                    .as_object()
249                    .unwrap()
250                    .clone();
251                tokio::task::block_in_place(|| CtxMultiReadTool.handle(&args, &ctx))
252            }));
253        }
254
255        for h in handles {
256            let joined = tokio::time::timeout(Duration::from_secs(20), h)
257                .await
258                .expect("ctx_multi_read hung (>20s) — nested block_in_place regression?")
259                .expect("spawned task panicked");
260            let out = joined.expect("ctx_multi_read returned an error");
261            assert!(
262                out.text.contains("Read 6 files"),
263                "unexpected output: {}",
264                out.text
265            );
266        }
267    }
268}
lean_ctx/tools/registered/ctx_multi_read.rs

lean_ctx/tools/registered/
ctx_multi_read.rs
