fix(shell): replace global static output buffer with per-instance TempDir (#7632)

Signed-off-by: Adrian Cole <adrian@tetrate.io>
This commit is contained in:
Adrian Cole
2026-03-05 07:11:53 +08:00
committed by GitHub
parent b6a3ab5edb
commit 42fc5152bf
2 changed files with 77 additions and 48 deletions
@@ -70,7 +70,7 @@ impl DeveloperClient {
Ok(Self { Ok(Self {
info, info,
shell_tool: Arc::new(ShellTool::new()), shell_tool: Arc::new(ShellTool::new()?),
edit_tools: Arc::new(EditTools::new()), edit_tools: Arc::new(EditTools::new()),
tree_tool: Arc::new(TreeTool::new()), tree_tool: Arc::new(TreeTool::new()),
}) })
@@ -1,7 +1,7 @@
use std::collections::HashMap;
use std::path::PathBuf; use std::path::PathBuf;
use std::process::Stdio; use std::process::Stdio;
use std::sync::{Mutex, OnceLock}; use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::OnceLock;
use std::time::Duration; use std::time::Duration;
use rmcp::model::{CallToolResult, Content}; use rmcp::model::{CallToolResult, Content};
@@ -16,6 +16,8 @@ const OUTPUT_LIMIT_LINES: usize = 2000;
const OUTPUT_LIMIT_BYTES: usize = 50_000; const OUTPUT_LIMIT_BYTES: usize = 50_000;
const OUTPUT_PREVIEW_LINES: usize = 50; const OUTPUT_PREVIEW_LINES: usize = 50;
const OUTPUT_SLOTS: usize = 8;
#[derive(Debug, Deserialize, JsonSchema)] #[derive(Debug, Deserialize, JsonSchema)]
pub struct ShellParams { pub struct ShellParams {
pub command: String, pub command: String,
@@ -78,11 +80,17 @@ fn user_login_path() -> Option<&'static str> {
CACHED.get_or_init(resolve_login_shell_path).as_deref() CACHED.get_or_init(resolve_login_shell_path).as_deref()
} }
pub struct ShellTool; pub struct ShellTool {
output_dir: tempfile::TempDir,
call_index: AtomicUsize,
}
impl ShellTool { impl ShellTool {
pub fn new() -> Self { pub fn new() -> std::io::Result<Self> {
Self Ok(Self {
output_dir: tempfile::tempdir()?,
call_index: AtomicUsize::new(0),
})
} }
pub async fn shell(&self, params: ShellParams) -> CallToolResult { pub async fn shell(&self, params: ShellParams) -> CallToolResult {
@@ -106,10 +114,12 @@ impl ShellTool {
// Derive stdout, stderr, and interleaved display from the single tagged-line buffer // Derive stdout, stderr, and interleaved display from the single tagged-line buffer
let (raw_stdout, raw_stderr, interleaved) = split_lines(&execution.lines); let (raw_stdout, raw_stderr, interleaved) = split_lines(&execution.lines);
let output_dir = self.output_dir.path();
let slot = self.call_index.fetch_add(1, Ordering::Relaxed) % OUTPUT_SLOTS;
let truncated_stdout = if raw_stdout.is_empty() { let truncated_stdout = if raw_stdout.is_empty() {
String::new() String::new()
} else { } else {
match truncate_output(&raw_stdout, "stdout") { match truncate_output(&raw_stdout, &format!("stdout-{slot}"), output_dir) {
Ok(t) => t, Ok(t) => t,
Err(error) => return Self::error_result(&error, None), Err(error) => return Self::error_result(&error, None),
} }
@@ -117,7 +127,7 @@ impl ShellTool {
let truncated_stderr = if raw_stderr.is_empty() { let truncated_stderr = if raw_stderr.is_empty() {
String::new() String::new()
} else { } else {
match truncate_output(&raw_stderr, "stderr") { match truncate_output(&raw_stderr, &format!("stderr-{slot}"), output_dir) {
Ok(t) => t, Ok(t) => t,
Err(error) => return Self::error_result(&error, None), Err(error) => return Self::error_result(&error, None),
} }
@@ -130,7 +140,8 @@ impl ShellTool {
timed_out: execution.timed_out, timed_out: execution.timed_out,
}; };
let structured_content = serde_json::to_value(&shell_output).ok(); let structured_content = serde_json::to_value(&shell_output).ok();
let mut rendered = match render_output(&interleaved, "output") { let mut rendered = match render_output(&interleaved, &format!("output-{slot}"), output_dir)
{
Ok(rendered) => rendered, Ok(rendered) => rendered,
Err(error) => return Self::error_result(&error, None), Err(error) => return Self::error_result(&error, None),
}; };
@@ -177,12 +188,6 @@ impl ShellTool {
} }
} }
impl Default for ShellTool {
fn default() -> Self {
Self::new()
}
}
struct ExecutionOutput { struct ExecutionOutput {
/// Lines in arrival order, tagged by source: (is_stderr, text) /// Lines in arrival order, tagged by source: (is_stderr, text)
lines: Vec<(bool, String)>, lines: Vec<(bool, String)>,
@@ -325,14 +330,22 @@ async fn collect_tagged_lines(
Ok(lines) Ok(lines)
} }
fn render_output(full_output: &str, label: &str) -> Result<String, String> { fn render_output(
full_output: &str,
label: &str,
output_dir: &std::path::Path,
) -> Result<String, String> {
if full_output.is_empty() { if full_output.is_empty() {
return Ok("(no output)".to_string()); return Ok("(no output)".to_string());
} }
truncate_output(full_output, label) truncate_output(full_output, label, output_dir)
} }
fn truncate_output(full_output: &str, label: &str) -> Result<String, String> { fn truncate_output(
full_output: &str,
label: &str,
output_dir: &std::path::Path,
) -> Result<String, String> {
let lines: Vec<&str> = full_output.split('\n').collect(); let lines: Vec<&str> = full_output.split('\n').collect();
let total_lines = lines.len(); let total_lines = lines.len();
let total_bytes = full_output.len(); let total_bytes = full_output.len();
@@ -344,7 +357,7 @@ fn truncate_output(full_output: &str, label: &str) -> Result<String, String> {
return Ok(full_output.to_string()); return Ok(full_output.to_string());
} }
let output_path = save_full_output(full_output, label)?; let output_path = save_full_output(full_output, label, output_dir)?;
let preview_start = total_lines.saturating_sub(OUTPUT_PREVIEW_LINES); let preview_start = total_lines.saturating_sub(OUTPUT_PREVIEW_LINES);
let preview = lines[preview_start..].join("\n"); let preview = lines[preview_start..].join("\n");
@@ -366,24 +379,12 @@ fn truncate_output(full_output: &str, label: &str) -> Result<String, String> {
)) ))
} }
fn output_buffer_path(label: &str) -> Result<PathBuf, String> { fn save_full_output(
static PATHS: Mutex<Option<HashMap<String, PathBuf>>> = Mutex::new(None); output: &str,
let mut guard = PATHS.lock().map_err(|e| format!("Lock poisoned: {e}"))?; label: &str,
let map = guard.get_or_insert_with(HashMap::new); output_dir: &std::path::Path,
if let Some(path) = map.get(label) { ) -> Result<PathBuf, String> {
return Ok(path.clone()); let path = output_dir.join(label);
}
let temp_file =
tempfile::NamedTempFile::new().map_err(|e| format!("Failed to create temp file: {e}"))?;
let (_, path) = temp_file
.keep()
.map_err(|e| format!("Failed to persist temp file: {}", e.error))?;
map.insert(label.to_string(), path.clone());
Ok(path)
}
fn save_full_output(output: &str, label: &str) -> Result<PathBuf, String> {
let path = output_buffer_path(label)?;
std::fs::write(&path, output).map_err(|e| format!("Failed to write output buffer: {e}"))?; std::fs::write(&path, output).map_err(|e| format!("Failed to write output buffer: {e}"))?;
Ok(path) Ok(path)
} }
@@ -402,7 +403,7 @@ mod tests {
#[tokio::test] #[tokio::test]
async fn shell_executes_command() { async fn shell_executes_command() {
let tool = ShellTool::new(); let tool = ShellTool::new().unwrap();
let result = tool let result = tool
.shell(ShellParams { .shell(ShellParams {
command: "echo hello".to_string(), command: "echo hello".to_string(),
@@ -417,7 +418,7 @@ mod tests {
#[cfg(not(windows))] #[cfg(not(windows))]
#[tokio::test] #[tokio::test]
async fn shell_returns_error_for_non_zero_exit() { async fn shell_returns_error_for_non_zero_exit() {
let tool = ShellTool::new(); let tool = ShellTool::new().unwrap();
let result = tool let result = tool
.shell(ShellParams { .shell(ShellParams {
command: "echo fail && exit 7".to_string(), command: "echo fail && exit 7".to_string(),
@@ -433,7 +434,7 @@ mod tests {
#[tokio::test] #[tokio::test]
async fn shell_uses_working_dir_for_relative_execution() { async fn shell_uses_working_dir_for_relative_execution() {
let dir = tempfile::tempdir().unwrap(); let dir = tempfile::tempdir().unwrap();
let tool = ShellTool::new(); let tool = ShellTool::new().unwrap();
let result = tool let result = tool
.shell_with_cwd( .shell_with_cwd(
ShellParams { ShellParams {
@@ -452,29 +453,32 @@ mod tests {
#[test] #[test]
fn render_output_returns_full_output_when_under_limit() { fn render_output_returns_full_output_when_under_limit() {
let dir = tempfile::tempdir().unwrap();
let input = (0..100) let input = (0..100)
.map(|i| format!("line {}", i)) .map(|i| format!("line {}", i))
.collect::<Vec<_>>() .collect::<Vec<_>>()
.join("\n"); .join("\n");
let rendered = render_output(&input, "test").unwrap(); let rendered = render_output(&input, "test", dir.path()).unwrap();
assert_eq!(rendered, input); assert_eq!(rendered, input);
} }
#[test] #[test]
fn render_output_shows_empty_message() { fn render_output_shows_empty_message() {
let rendered = render_output("", "test").unwrap(); let dir = tempfile::tempdir().unwrap();
let rendered = render_output("", "test", dir.path()).unwrap();
assert_eq!(rendered, "(no output)"); assert_eq!(rendered, "(no output)");
} }
#[test] #[test]
fn render_output_truncates_when_lines_exceeded() { fn render_output_truncates_when_lines_exceeded() {
let dir = tempfile::tempdir().unwrap();
let input = (0..2500) let input = (0..2500)
.map(|i| format!("line {}", i)) .map(|i| format!("line {}", i))
.collect::<Vec<_>>() .collect::<Vec<_>>()
.join("\n"); .join("\n");
let rendered = render_output(&input, "test_lines").unwrap(); let rendered = render_output(&input, "test_lines", dir.path()).unwrap();
let (preview, metadata) = rendered.split_once("\n\n[").unwrap(); let (preview, metadata) = rendered.split_once("\n\n[").unwrap();
assert_eq!(preview.lines().count(), OUTPUT_PREVIEW_LINES); assert_eq!(preview.lines().count(), OUTPUT_PREVIEW_LINES);
@@ -489,6 +493,7 @@ mod tests {
#[test] #[test]
fn render_output_truncates_when_bytes_exceeded() { fn render_output_truncates_when_bytes_exceeded() {
let dir = tempfile::tempdir().unwrap();
let long_line = "x".repeat(1000); let long_line = "x".repeat(1000);
let input = (0..100) let input = (0..100)
.map(|_| long_line.clone()) .map(|_| long_line.clone())
@@ -497,7 +502,7 @@ mod tests {
assert!(input.len() > OUTPUT_LIMIT_BYTES); assert!(input.len() > OUTPUT_LIMIT_BYTES);
assert!(input.lines().count() <= OUTPUT_LIMIT_LINES); assert!(input.lines().count() <= OUTPUT_LIMIT_LINES);
let rendered = render_output(&input, "test_bytes").unwrap(); let rendered = render_output(&input, "test_bytes", dir.path()).unwrap();
let (_preview, metadata) = rendered.split_once("\n\n[").unwrap(); let (_preview, metadata) = rendered.split_once("\n\n[").unwrap();
assert!(metadata.contains("byte limit")); assert!(metadata.contains("byte limit"));
@@ -507,18 +512,42 @@ mod tests {
#[test] #[test]
fn save_full_output_reuses_same_path() { fn save_full_output_reuses_same_path() {
let path1 = save_full_output("first", "test_reuse").unwrap(); let dir = tempfile::tempdir().unwrap();
let path2 = save_full_output("second", "test_reuse").unwrap(); let path1 = save_full_output("first", "test_reuse", dir.path()).unwrap();
let path2 = save_full_output("second", "test_reuse", dir.path()).unwrap();
assert_eq!(path1, path2); assert_eq!(path1, path2);
assert_eq!(std::fs::read_to_string(&path2).unwrap(), "second"); assert_eq!(std::fs::read_to_string(&path2).unwrap(), "second");
} }
#[test] #[test]
fn save_full_output_uses_separate_files_per_label() { fn save_full_output_uses_separate_files_per_label() {
let path_a = save_full_output("aaa", "label_a").unwrap(); let dir = tempfile::tempdir().unwrap();
let path_b = save_full_output("bbb", "label_b").unwrap(); let path_a = save_full_output("aaa", "label_a", dir.path()).unwrap();
let path_b = save_full_output("bbb", "label_b", dir.path()).unwrap();
assert_ne!(path_a, path_b); assert_ne!(path_a, path_b);
assert_eq!(std::fs::read_to_string(&path_a).unwrap(), "aaa"); assert_eq!(std::fs::read_to_string(&path_a).unwrap(), "aaa");
assert_eq!(std::fs::read_to_string(&path_b).unwrap(), "bbb"); assert_eq!(std::fs::read_to_string(&path_b).unwrap(), "bbb");
} }
#[test]
fn call_index_cycles_through_slots() {
let tool = ShellTool::new().unwrap();
for _cycle in 0..3 {
for expected in 0..OUTPUT_SLOTS {
let slot = tool.call_index.fetch_add(1, Ordering::Relaxed) % OUTPUT_SLOTS;
assert_eq!(slot, expected);
}
}
}
#[test]
fn concurrent_calls_get_distinct_slots() {
let tool = ShellTool::new().unwrap();
let mut slots: Vec<usize> = (0..OUTPUT_SLOTS)
.map(|_| tool.call_index.fetch_add(1, Ordering::Relaxed) % OUTPUT_SLOTS)
.collect();
slots.sort();
let expected: Vec<usize> = (0..OUTPUT_SLOTS).collect();
assert_eq!(slots, expected);
}
} }