diff --git a/src/main.rs b/src/main.rs index 91d9f98..ed6b12a 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1022,40 +1022,21 @@ mod tests { use super::*; use std::fs; - #[test] - fn mcp_proxy_stdio_accepts_hyphen_prefixed_child_args() { - let cli = Cli::try_parse_from([ - "agentk", - "mcp-proxy-stdio", - "--command", - "sh", - "--arg", - "-c", - "--arg", - "printf ok", - ]) - .expect("hyphen-prefixed child args should parse"); - - let Some(Command::McpProxyStdio { args, .. }) = cli.command else { - panic!("expected mcp-proxy-stdio command"); - }; - assert_eq!(args, vec!["-c".to_string(), "printf ok".to_string()]); - } - #[cfg(unix)] - #[test] - fn mcp_proxy_stdio_trace_out_writes_verifiable_events() { - let trace_path = env::temp_dir().join(format!( - "agentk-mcp-proxy-stdio-trace-{}-{}.jsonl", + fn mcp_proxy_trace_out_test_path(label: &str) -> PathBuf { + env::temp_dir().join(format!( + "agentk-mcp-proxy-stdio-{label}-{}-{}.jsonl", std::process::id(), std::time::SystemTime::now() .duration_since(std::time::UNIX_EPOCH) .expect("system clock should be after epoch") .as_nanos() - )); - let _ = fs::remove_file(&trace_path); + )) + } - let server = r#" + #[cfg(unix)] + fn mcp_proxy_trace_out_probe_server() -> String { + r#" while IFS= read -r line; do case "$line" in *'"method":"initialize"'*) @@ -1071,14 +1052,79 @@ while IFS= read -r line; do ;; esac done -"#; - let input = r#" +"# + .to_string() + } + + #[cfg(unix)] + fn mcp_proxy_trace_out_probe_input() -> &'static str { + r#" {"jsonrpc":"2.0","id":1,"method":"initialize","params":{"protocolVersion":"2025-11-25"}} {"jsonrpc":"2.0","method":"notifications/initialized","params":{}} {"jsonrpc":"2.0","id":2,"method":"tools/list","params":{}} -"#; +"# + } + + #[cfg(unix)] + #[derive(Default)] + struct FailOnSecondNewlineWriter { + bytes: Vec, + newline_count: usize, + } + + #[cfg(unix)] + impl std::io::Write for FailOnSecondNewlineWriter { + fn write(&mut self, buf: &[u8]) -> std::io::Result { + for byte in buf { + if *byte == b'\n' { + self.newline_count += 1; + if self.newline_count == 2 { + return Err(std::io::Error::new( + std::io::ErrorKind::BrokenPipe, + "test writer failure after mediated event", + )); + } + } + } + self.bytes.extend_from_slice(buf); + Ok(buf.len()) + } + + fn flush(&mut self) -> std::io::Result<()> { + Ok(()) + } + } + + #[test] + fn mcp_proxy_stdio_accepts_hyphen_prefixed_child_args() { + let cli = Cli::try_parse_from([ + "agentk", + "mcp-proxy-stdio", + "--command", + "sh", + "--arg", + "-c", + "--arg", + "printf ok", + ]) + .expect("hyphen-prefixed child args should parse"); + + let Some(Command::McpProxyStdio { args, .. }) = cli.command else { + panic!("expected mcp-proxy-stdio command"); + }; + assert_eq!(args, vec!["-c".to_string(), "printf ok".to_string()]); + } + + #[cfg(unix)] + #[test] + fn mcp_proxy_stdio_trace_out_writes_verifiable_events() { + let trace_path = mcp_proxy_trace_out_test_path("trace"); + let _ = fs::remove_file(&trace_path); + + let server = mcp_proxy_trace_out_probe_server(); + let input = mcp_proxy_trace_out_probe_input(); let config = McpSubprocessProxyConfig::new("agent://test", "trace-out-probe", "sh") - .with_args(["-c".to_string(), server.to_string()]); + .with_args(["-c".to_string(), server]); let mut output = Vec::new(); mcp_proxy_stdio_with_io( @@ -1097,6 +1143,38 @@ done let _ = fs::remove_file(trace_path); } + #[cfg(unix)] + #[test] + fn mcp_proxy_stdio_trace_out_survives_writer_failure_after_event() { + let trace_path = mcp_proxy_trace_out_test_path("writer-failure"); + let _ = fs::remove_file(&trace_path); + + let server = mcp_proxy_trace_out_probe_server(); + let input = mcp_proxy_trace_out_probe_input(); + let config = McpSubprocessProxyConfig::new("agent://test", "trace-out-probe", "sh") + .with_args(["-c".to_string(), server]); + let mut output = FailOnSecondNewlineWriter::default(); + + let error = mcp_proxy_stdio_with_io( + config, + Some(trace_path.clone()), + BufReader::new(input.as_bytes()), + &mut output, + ) + .expect_err("client writer failure should surface"); + + assert!( + error + .to_string() + .contains("test writer failure after mediated event") + ); + assert_eq!(output.newline_count, 2); + let verify = verify_jsonl(&trace_path).expect("trace-out should survive writer failure"); + assert_eq!(verify.events_checked, 1); + + let _ = fs::remove_file(trace_path); + } + #[test] fn mcp_proxy_allow_env_collects_explicit_parent_values() { let names = vec!["AGENTK_PROXY_DEMO".to_string()];