1use std::path::{Path, PathBuf};
7use std::sync::Arc;
8use std::sync::atomic::{AtomicU64, Ordering};
9use std::time::Duration;
10use std::time::{SystemTime, UNIX_EPOCH};
11
12use rmcp::{
13 ErrorData as McpError,
14 handler::server::ServerHandler,
15 model::*,
16 service::{RequestContext, RoleServer},
17};
18use tokio::sync::Mutex;
19use tokio::task::JoinHandle;
20use tokio_util::sync::CancellationToken;
21use tracing::{debug, error, info, warn};
22
23use crate::background::job::NewRunningJob;
24use crate::background::transfer::{SharedTransferJob, TransferJobRegistry};
25use crate::background::{JobRegistry, JobState, LocalLogSpooler, SharedJobState};
26use crate::config::Config;
27use crate::error::{Result, SshMcpError};
28#[cfg(unix)]
29use crate::platform::O_NOFOLLOW_FLAG;
30use crate::server::handlers::file_edit_common::{FileEditFaultInjection, FileEditPrivilege};
31#[cfg(test)]
32use crate::server::validation::validate_background_log_path;
33use crate::ssh::{
34 CommandOutput, SshConfig, SshConnectionManager, sanitize_command, wrap_sudo_command,
35};
36use crate::tools::ApplyPatchParams;
37use crate::transfer::{
38 TransferEngine, TransferEventSink, TransferParams, TransferResponse, TransferRunContext,
39 TransferSshOptions,
40};
41
42mod args;
43mod exec;
44mod handlers;
45mod testing;
46mod tools;
47mod validation;
48
49const BACKGROUND_START_TIMEOUT: Duration = Duration::from_secs(20);
50
51const JOB_COMPLETED_RETENTION: Duration = Duration::from_secs(60 * 60);
52
53static JOB_COUNTER: AtomicU64 = AtomicU64::new(0);
54
55fn make_job_id() -> String {
56 let counter = JOB_COUNTER.fetch_add(1, Ordering::Relaxed);
57 let epoch_ms = SystemTime::now()
58 .duration_since(UNIX_EPOCH)
59 .map(|d| d.as_millis())
60 .unwrap_or(0);
61 format!("{}-{}", epoch_ms, counter)
62}
63
64#[derive(Clone)]
69pub struct SshMcpServer {
70 config: Config,
72
73 connection: Arc<SshConnectionManager>,
75
76 timeout: Duration,
78
79 max_chars: Option<usize>,
81
82 spooler: Arc<LocalLogSpooler>,
83 job_registry: Arc<JobRegistry>,
84
85 transfer_job_registry: Arc<TransferJobRegistry>,
86 transfer_tasks: Arc<Mutex<Vec<JoinHandle<()>>>>,
87 transfer_shutdown: CancellationToken,
88
89 transfer: TransferEngine,
90}
91
92struct TransferTerminalGuard {
93 job: SharedTransferJob,
94 fallback: Option<TransferResponse>,
95}
96
97impl TransferTerminalGuard {
98 fn new(job: SharedTransferJob, fallback: TransferResponse) -> Self {
99 Self {
100 job,
101 fallback: Some(fallback),
102 }
103 }
104
105 fn finish(mut self, response: &TransferResponse) {
106 if let Ok(mut job) = self.job.lock() {
107 job.finish(response);
108 }
109 self.fallback = None;
110 }
111}
112
113impl Drop for TransferTerminalGuard {
114 fn drop(&mut self) {
115 let Some(fallback) = self.fallback.take() else {
116 return;
117 };
118 if let Ok(mut job) = self.job.lock()
119 && !job.is_terminal()
120 {
121 job.finish(&fallback);
122 }
123 }
124}
125
126impl SshMcpServer {
127 pub async fn new(config: Config) -> Result<Self> {
132 Self::new_with_spool_dir(config, None).await
133 }
134
135 pub async fn new_with_spool_dir(config: Config, spool_dir: Option<PathBuf>) -> Result<Self> {
137 let local_root = std::env::current_dir()?;
138
139 let spooler = Arc::new(resolve_local_spooler(spool_dir)?);
140 spooler.ensure_dir().await.map_err(|e| {
141 SshMcpError::Config(format!(
142 "failed to initialize local log spool dir {}: {e}",
143 spooler.base_dir().display()
144 ))
145 })?;
146 let job_registry = Arc::new(JobRegistry::new(JOB_COMPLETED_RETENTION));
147 let transfer_job_registry = Arc::new(TransferJobRegistry::new(JOB_COMPLETED_RETENTION));
148
149 let mut ssh_config = SshConfig::new(&config.host, &config.user).with_port(config.port);
151
152 if let Some(ref password) = config.password {
154 ssh_config = ssh_config.with_password(password);
155 }
156
157 if let Some(ref key_path) = config.key {
158 let key_content = tokio::fs::read_to_string(key_path)
160 .await
161 .map_err(SshMcpError::Io)?;
162 ssh_config = ssh_config.with_private_key(&key_content);
163 }
164
165 if let Some(ref su_password) = config.su_password {
167 ssh_config = ssh_config.with_su_password(su_password);
168 }
169
170 if let Some(ref sudo_password) = config.sudo_password {
171 ssh_config = ssh_config.with_sudo_password(sudo_password);
172 }
173
174 ssh_config = ssh_config
176 .with_keepalive_interval(config.keepalive_interval)
177 .with_keepalive_max(config.keepalive_max);
178
179 ssh_config = ssh_config
181 .with_reconnect_retries(config.reconnect_retries)
182 .with_reconnect_backoff_ms(config.reconnect_backoff_ms)
183 .with_health_probe_timeout_ms(config.health_probe_timeout_ms);
184
185 ssh_config = ssh_config
187 .with_host_key_checking(config.strict_host_key_checking)
188 .with_known_hosts(config.known_hosts.clone());
189
190 ssh_config = ssh_config.with_max_output_tokens(config.max_output_tokens);
192
193 let connection = Arc::new(SshConnectionManager::new(ssh_config).await);
195
196 let timeout = Duration::from_millis(config.timeout_ms);
197 let max_chars = config.max_chars;
198
199 Ok(Self {
200 config,
201 connection,
202 timeout,
203 max_chars,
204 spooler,
205 job_registry,
206 transfer_job_registry,
207 transfer_tasks: Arc::new(Mutex::new(Vec::new())),
208 transfer_shutdown: CancellationToken::new(),
209 transfer: TransferEngine::new(local_root),
210 })
211 }
212
213 fn connection_id(&self) -> String {
214 format!(
215 "{}@{}:{}",
216 self.config.user, self.config.host, self.config.port
217 )
218 }
219
220 fn default_local_log_path(
221 &self,
222 job_id: &str,
223 ) -> std::result::Result<(PathBuf, String), String> {
224 let path = self
225 .spooler
226 .log_path_for(job_id)
227 .map_err(|e| format!("failed to generate local log path for job_id='{job_id}': {e}"))?;
228 let path_str = path.to_string_lossy().to_string();
229 Ok((path, path_str))
230 }
231
232 async fn ensure_local_log_file(&self, log_path: &Path) -> std::result::Result<(), SshMcpError> {
233 self.spooler.ensure_dir().await.map_err(|e| {
234 SshMcpError::Config(format!(
235 "failed to ensure local log spool dir {}: {e}",
236 self.spooler.base_dir().display()
237 ))
238 })?;
239
240 if log_path.parent() != Some(self.spooler.base_dir()) {
241 return Err(SshMcpError::InvalidParams(format!(
242 "log_path must be directly under {}",
243 self.spooler.base_dir().display()
244 )));
245 }
246
247 match tokio::fs::symlink_metadata(log_path).await {
248 Ok(meta) => {
249 let ft = meta.file_type();
250 if ft.is_symlink() {
251 return Err(SshMcpError::invalid_params(
252 "log_path is a symlink (refusing to follow it)",
253 ));
254 }
255 if !ft.is_file() {
256 return Err(SshMcpError::invalid_params(
257 "log_path exists but is not a regular file",
258 ));
259 }
260 }
261 Err(e) if e.kind() == std::io::ErrorKind::NotFound => {}
262 Err(e) => return Err(SshMcpError::Io(e)),
263 }
264
265 let mut opts = tokio::fs::OpenOptions::new();
266 opts.write(true).create(true).truncate(true);
267
268 #[cfg(unix)]
269 {
270 opts.custom_flags(O_NOFOLLOW_FLAG);
271 }
272
273 let file = match opts.open(log_path).await {
274 Ok(f) => f,
275 Err(e) => {
276 if let Ok(meta) = tokio::fs::symlink_metadata(log_path).await
277 && meta.file_type().is_symlink()
278 {
279 return Err(SshMcpError::invalid_params(
280 "log_path is a symlink (refusing to follow it)",
281 ));
282 }
283 return Err(SshMcpError::Io(e));
284 }
285 };
286
287 file.sync_all().await.map_err(SshMcpError::Io)
288 }
289
290 async fn register_running_job(
291 &self,
292 job_id: &str,
293 pid: u32,
294 log_path: PathBuf,
295 command: &str,
296 ) -> SharedJobState {
297 let job = Arc::new(Mutex::new(JobState::new_running(NewRunningJob {
298 job_id: job_id.to_string(),
299 pid,
300 log_path,
301 command: command.to_string(),
302 connection_id: self.connection_id(),
303 })));
304
305 self.job_registry
306 .insert(job_id.to_string(), Arc::clone(&job))
307 .await;
308
309 let persisted = {
310 let guard = job.lock().await;
311 guard.clone()
312 };
313 if let Err(e) = self.spooler.persist_job_state(&persisted).await {
314 warn!(job_id = ?job_id, error = ?e, "failed to persist running job state");
315 }
316
317 job
318 }
319
320 pub fn connection(&self) -> &Arc<SshConnectionManager> {
322 &self.connection
323 }
324
325 pub async fn shutdown(&self) {
327 info!("Shutting down SSH MCP Server...");
328 self.transfer_shutdown.cancel();
329 let tasks = {
330 let mut tasks = self.transfer_tasks.lock().await;
331 std::mem::take(&mut *tasks)
332 };
333 let deadline = tokio::time::Instant::now() + Duration::from_secs(10);
334 for mut task in tasks {
335 if tokio::time::timeout_at(deadline, &mut task).await.is_err() {
336 task.abort();
337 let _ = task.await;
338 }
339 }
340 self.connection.close().await;
341 }
342
343 async fn execute_command_with_timeout(
345 &self,
346 command: &str,
347 timeout: Duration,
348 ) -> std::result::Result<CallToolResult, McpError> {
349 debug!(
350 "shell tool called: cmd_len={}, background=false, sudo=false, timeout_ms={}",
351 command.len(),
352 timeout.as_millis()
353 );
354
355 let sanitized = match self.sanitize_or_tool_error(command) {
357 Ok(cmd) => cmd,
358 Err(result) => return Ok(result),
359 };
360
361 let requires_elevation = self.connection.get_su_password().is_some();
367 if requires_elevation {
368 if let Err(e) = self.connection.ensure_connected().await {
369 error!(error = ?e, "Failed to ensure SSH connection");
370 return Ok(CallToolResult::error(vec![ContentBlock::text(
371 e.to_string(),
372 )]));
373 }
374
375 if let Err(e) = self.connection.ensure_elevated().await {
376 debug!(error = ?e, "Elevation failed, will run as normal user");
377 }
378 }
379
380 if !requires_elevation && let Err(e) = self.connection.ensure_connected().await {
382 error!(error = ?e, "Failed to ensure SSH connection");
383 return Ok(CallToolResult::error(vec![ContentBlock::text(
384 e.to_string(),
385 )]));
386 }
387
388 self.execute_detachable_foreground_impl(&sanitized, &sanitized, timeout)
389 .await
390 }
391
392 async fn execute_command(
393 &self,
394 command: &str,
395 ) -> std::result::Result<CallToolResult, McpError> {
396 self.execute_command_with_timeout(command, self.timeout)
397 .await
398 }
399
400 async fn execute_background_command(
401 &self,
402 command: &str,
403 log_path: Option<&str>,
404 ) -> std::result::Result<CallToolResult, McpError> {
405 self.execute_background_impl(command, log_path, exec::BackgroundPrivilege::Normal)
406 .await
407 }
408
409 async fn execute_sudo_command_with_timeout(
411 &self,
412 command: &str,
413 timeout: Duration,
414 ) -> std::result::Result<CallToolResult, McpError> {
415 debug!(
416 "sudo_shell tool called: cmd_len={}, background=false, sudo=true, timeout_ms={}",
417 command.len(),
418 timeout.as_millis()
419 );
420
421 let sanitized = match self.sanitize_or_tool_error(command) {
423 Ok(cmd) => cmd,
424 Err(result) => return Ok(result),
425 };
426
427 let sudo_password = self.connection.get_sudo_password();
429 let wrapped_command = wrap_sudo_command(&sanitized, sudo_password);
430 debug!(
431 "Wrapped sudo command (password hidden): sudo -n sh -c '...' or printf '...' | sudo ..."
432 );
433
434 if let Err(e) = self.connection.ensure_connected().await {
435 error!(error = ?e, "Failed to ensure SSH connection");
436 return Ok(CallToolResult::error(vec![ContentBlock::text(
437 e.to_string(),
438 )]));
439 }
440
441 self.execute_detachable_foreground_impl(
442 &wrapped_command,
443 &format!("sudo {sanitized}"),
444 timeout,
445 )
446 .await
447 }
448
449 async fn execute_sudo_command(
450 &self,
451 command: &str,
452 ) -> std::result::Result<CallToolResult, McpError> {
453 self.execute_sudo_command_with_timeout(command, self.timeout)
454 .await
455 }
456
457 async fn execute_background_sudo_command(
458 &self,
459 command: &str,
460 log_path: Option<&str>,
461 ) -> std::result::Result<CallToolResult, McpError> {
462 let sudo_password = self.connection.get_sudo_password();
463 self.execute_background_impl(
464 command,
465 log_path,
466 exec::BackgroundPrivilege::Sudo {
467 password: sudo_password,
468 },
469 )
470 .await
471 }
472
473 fn sanitize_or_tool_error(&self, command: &str) -> std::result::Result<String, CallToolResult> {
474 sanitize_command(command, self.max_chars).map_err(|e| {
475 error!(error = ?e, "Command sanitization failed");
476 CallToolResult::error(vec![ContentBlock::text(format!("Error: {}", e))])
477 })
478 }
479
480 fn calltool_from_command_output(output: CommandOutput) -> CallToolResult {
481 let mut result_text = output.stdout;
483 if !output.stderr.is_empty() {
484 if !result_text.is_empty() {
485 result_text.push_str("\n--- stderr ---\n");
486 }
487 result_text.push_str(&output.stderr);
488 }
489
490 if output.exit_code.map(|code| code != 0).unwrap_or(true) {
494 CallToolResult::error(vec![ContentBlock::text(result_text)])
495 } else {
496 CallToolResult::success(vec![ContentBlock::text(result_text)])
497 }
498 }
499
500 fn shell_tool() -> Tool {
502 tools::shell_tool()
503 }
504
505 fn sudo_shell_tool() -> Tool {
507 tools::sudo_shell_tool()
508 }
509
510 fn transfer_tool() -> Tool {
512 tools::transfer_tool()
513 }
514
515 fn check_process_tool() -> Tool {
517 tools::check_process_tool()
518 }
519
520 fn apply_patch_tool() -> Tool {
522 tools::apply_patch_tool()
523 }
524
525 fn sudo_apply_patch_tool() -> Tool {
527 tools::sudo_apply_patch_tool()
528 }
529
530 fn resolve_timeout(&self, timeout_ms: Option<u64>) -> Duration {
532 timeout_ms
533 .map(Duration::from_millis)
534 .unwrap_or(self.timeout)
535 }
536
537 fn parse_tool_params<T: serde::de::DeserializeOwned>(
539 &self,
540 args: serde_json::Map<String, serde_json::Value>,
541 tool_name: &str,
542 ) -> std::result::Result<T, McpError> {
543 serde_json::from_value(serde_json::Value::Object(args))
544 .map_err(|e| McpError::invalid_params(format!("invalid {tool_name} params: {e}"), None))
545 }
546
547 async fn run_transfer_response(
548 &self,
549 params: TransferParams,
550 cancellation: CancellationToken,
551 progress: Option<TransferEventSink>,
552 ) -> TransferResponse {
553 let timeout = self.resolve_timeout(params.timeout_ms);
554 let key_path = self.config.key.clone();
555
556 self.transfer
557 .run_controlled(
558 &self.connection,
559 params,
560 TransferRunContext {
561 timeout,
562 ssh: TransferSshOptions {
563 host: self.config.host.clone(),
564 port: self.config.port,
565 user: self.config.user.clone(),
566 key_path,
567 host_key_checking: self.config.strict_host_key_checking,
568 known_hosts: self.config.known_hosts.clone(),
569 },
570 },
571 cancellation,
572 progress,
573 )
574 .await
575 }
576
577 async fn execute_transfer(
579 &self,
580 params: TransferParams,
581 verbose: bool,
582 cancellation: CancellationToken,
583 ) -> std::result::Result<CallToolResult, McpError> {
584 let resp = self.run_transfer_response(params, cancellation, None).await;
585 let body = resp
586 .to_json(verbose)
587 .unwrap_or_else(|_| "{\"ok\":false,\"error\":\"serialization_error\"}".to_string());
588 Ok(CallToolResult::success(vec![ContentBlock::text(body)]))
589 }
590
591 async fn execute_background_transfer(
592 &self,
593 params: TransferParams,
594 ) -> std::result::Result<CallToolResult, McpError> {
595 let job_id = format!("transfer-{}", make_job_id());
596 let job = self.transfer_job_registry.register(job_id.clone());
597 let progress_job = Arc::clone(&job);
598 let progress = TransferEventSink::new(move |event| {
599 if let Ok(mut job) = progress_job.lock() {
600 job.apply_event(event);
601 }
602 });
603
604 let server = self.clone();
605 let task_params = params.clone();
606 let fallback = TransferResponse::error(
607 params,
608 self.transfer.local_root(),
609 "background transfer stopped unexpectedly",
610 );
611 let body = serde_json::json!({
612 "ok": true,
613 "background": true,
614 "job_id": job_id,
615 "job_type": "transfer",
616 "state": "running",
617 "phase": "queued",
618 });
619
620 let mut tasks = self.transfer_tasks.lock().await;
621 tasks.retain(|task| !task.is_finished());
622 let task = tokio::spawn(async move {
623 tokio::task::yield_now().await;
624 let terminal = TransferTerminalGuard::new(job, fallback);
625 let response = server
626 .run_transfer_response(
627 task_params,
628 server.transfer_shutdown.child_token(),
629 Some(progress),
630 )
631 .await;
632 terminal.finish(&response);
633 });
634 tasks.push(task);
635 drop(tasks);
636
637 Ok(CallToolResult::success(vec![ContentBlock::text(
638 body.to_string(),
639 )]))
640 }
641}
642
643fn resolve_local_spooler(spool_dir: Option<PathBuf>) -> Result<LocalLogSpooler> {
644 match spool_dir {
645 Some(path) if !path.is_absolute() => Err(SshMcpError::Config(format!(
646 "spool directory must be absolute: {}",
647 path.display()
648 ))),
649 Some(path) => Ok(LocalLogSpooler::new(path)),
650 None => Ok(LocalLogSpooler::new_default()),
651 }
652}
653
654impl ServerHandler for SshMcpServer {
655 fn get_info(&self) -> ServerInfo {
657 ServerInfo::new(ServerCapabilities::builder().enable_tools().build())
658 .with_protocol_version(ProtocolVersion::LATEST)
659 .with_server_info(Implementation::from_build_env())
660 .with_instructions(format!(
661 "SSH MCP Server v{} - Execute commands on {}@{}:{}",
662 env!("CARGO_PKG_VERSION"),
663 self.config.user,
664 self.config.host,
665 self.config.port,
666 ))
667 }
668
669 async fn list_tools(
671 &self,
672 _request: Option<PaginatedRequestParams>,
673 _context: RequestContext<RoleServer>,
674 ) -> std::result::Result<ListToolsResult, McpError> {
675 debug!("list_tools called");
676
677 let mut tools = vec![Self::shell_tool()];
678
679 if !self.config.disable_sudo {
681 tools.push(Self::sudo_shell_tool());
682 tools.push(Self::sudo_apply_patch_tool());
683 }
684 tools.push(Self::check_process_tool());
685 tools.push(Self::transfer_tool());
686 tools.push(Self::apply_patch_tool());
687
688 Ok(ListToolsResult {
689 tools,
690 next_cursor: None,
691 meta: Default::default(),
692 })
693 }
694
695 async fn call_tool(
697 &self,
698 request: CallToolRequestParams,
699 context: RequestContext<RoleServer>,
700 ) -> std::result::Result<CallToolResult, McpError> {
701 let tool_name: &str = request.name.as_ref();
702 debug!("call_tool called: {:?}", tool_name);
703
704 let args = request.arguments.unwrap_or_default();
705
706 match tool_name {
708 "shell" => {
709 let parsed = self.parse_common_tool_args(&args)?;
710 let timeout = self.resolve_timeout(parsed.timeout_ms);
711
712 if parsed.background {
713 self.execute_background_command(&parsed.command, parsed.log_path.as_deref())
714 .await
715 } else {
716 self.execute_command_with_timeout(&parsed.command, timeout)
717 .await
718 }
719 }
720 "sudo_shell" => {
721 if self.config.disable_sudo {
722 return Err(McpError::invalid_params(
723 "sudo_shell tool is disabled",
724 None,
725 ));
726 }
727
728 let parsed = self.parse_common_tool_args(&args)?;
729 let timeout = self.resolve_timeout(parsed.timeout_ms);
730
731 if parsed.background {
732 self.execute_background_sudo_command(
733 &parsed.command,
734 parsed.log_path.as_deref(),
735 )
736 .await
737 } else {
738 self.execute_sudo_command_with_timeout(&parsed.command, timeout)
739 .await
740 }
741 }
742 "transfer" => {
743 let params: TransferParams = self.parse_tool_params(args, "transfer")?;
744 let verbose = params.verbose;
745 if params.background {
746 self.execute_background_transfer(params).await
747 } else {
748 self.execute_transfer(params, verbose, context.ct.clone())
749 .await
750 }
751 }
752 "check_process" => {
753 let params: args::CheckProcessToolArgs =
754 self.parse_tool_params(args, "check_process")?;
755 self.execute_check_process(params.check, params.wait_for, context.ct.cancelled())
756 .await
757 }
758 "apply_patch" => {
759 let params: ApplyPatchParams = self.parse_tool_params(args, "apply_patch")?;
760 self.execute_apply_patch(
761 params,
762 FileEditFaultInjection::None,
763 FileEditPrivilege::User,
764 )
765 .await
766 }
767 "sudo_apply_patch" => {
768 if self.config.disable_sudo {
769 return Err(McpError::invalid_params(
770 "sudo_apply_patch tool is disabled",
771 None,
772 ));
773 }
774
775 let params: ApplyPatchParams = self.parse_tool_params(args, "sudo_apply_patch")?;
776 self.execute_apply_patch(
777 params,
778 FileEditFaultInjection::None,
779 FileEditPrivilege::Sudo,
780 )
781 .await
782 }
783 _ => Err(McpError::invalid_params(
784 format!("Unknown tool: {}", tool_name),
785 None,
786 )),
787 }
788 }
789}
790
791#[cfg(test)]
792mod tests {
793 use super::*;
794 use crate::background::response::{
795 BACKGROUND_JSON_SNIPPET_LIMIT_CHARS, background_json_err, background_json_timeout,
796 };
797 use crate::background::wrapper::{build_background_wrapper_script, remote_job_log_path};
798
799 fn extract_text_from_result(result: &CallToolResult) -> String {
800 result
801 .content
802 .iter()
803 .filter_map(|c| c.as_text().map(|text| text.text.clone()))
804 .collect::<Vec<_>>()
805 .join("\n")
806 }
807
808 #[test]
809 fn test_server_info() {
810 assert!(!env!("CARGO_PKG_VERSION").is_empty());
812 }
813
814 #[test]
815 fn test_resolve_local_spooler_rejects_relative_override() {
816 let error = resolve_local_spooler(Some(PathBuf::from("relative/spool")))
817 .expect_err("relative spool directory must be rejected");
818
819 assert!(matches!(
820 error,
821 SshMcpError::Config(message) if message.contains("must be absolute")
822 ));
823 }
824
825 #[test]
826 fn test_shell_tool_definition() {
827 let tool = SshMcpServer::shell_tool();
828 assert_eq!(tool.name.as_ref(), "shell");
829 assert!(tool.description.is_some());
830 }
831
832 #[test]
833 fn test_sudo_shell_tool_definition() {
834 let tool = SshMcpServer::sudo_shell_tool();
835 assert_eq!(tool.name.as_ref(), "sudo_shell");
836 assert!(tool.description.is_some());
837 }
838
839 #[test]
840 fn test_apply_patch_tool_definition() {
841 let tool = SshMcpServer::apply_patch_tool();
842 assert_eq!(tool.name.as_ref(), "apply_patch");
843 assert!(tool.description.is_some());
844 }
845
846 #[test]
847 fn test_sudo_apply_patch_tool_definition() {
848 let tool = SshMcpServer::sudo_apply_patch_tool();
849 assert_eq!(tool.name.as_ref(), "sudo_apply_patch");
850 assert!(tool.description.is_some());
851 }
852
853 #[test]
854 fn test_build_background_wrapper_escapes_single_quotes_in_user_command() {
855 let remote_log = remote_job_log_path("job-1");
856 let script = build_background_wrapper_script("job-1", "echo 'hello world'", &remote_log);
857 assert!(script.contains("exec sh -c 'set +m; echo '\"'\"'hello world'\"'\"''"));
858 }
859
860 #[test]
861 fn test_build_background_wrapper_is_busybox_friendly() {
862 let remote_log = remote_job_log_path("job-1");
863 let script = build_background_wrapper_script("job-1", "echo test", &remote_log);
864 assert!(!script.contains("dirname --"));
865 assert!(!script.contains("mkdir -p --"));
866 assert!(!script.contains("sh -lc"));
867 assert!(script.contains("exec sh -c"));
868 assert!(!script.contains("nohup"));
869 }
870
871 #[test]
872 fn test_background_wrapper_emits_markers_and_exec() {
873 let remote_log = remote_job_log_path("job-1");
874 let script = build_background_wrapper_script("job-1", "echo test", &remote_log);
875 assert!(script.contains("__SSH_MCP_JOB_ID=job-1"));
876 assert!(script.contains("__SSH_MCP_PID=$$"));
877 assert!(script.contains("__SSH_MCP_LOG=$LOG"));
878 assert!(script.contains("exec sh -c"));
879 }
880
881 #[test]
882 fn test_background_wrapper_does_not_redirect_remote_output() {
883 let remote_log = remote_job_log_path("job-1");
884 let script = build_background_wrapper_script("job-1", "echo test", &remote_log);
885 assert!(!script.contains(">$LOG"));
886 assert!(!script.contains("2>&1"));
887 assert!(!script.contains("$EXIT"));
888 assert!(!script.contains("nohup"));
889 }
890
891 #[test]
892 fn test_validate_background_log_path_rejects_leading_dash() {
893 let err =
894 validate_background_log_path(Path::new("/tmp/ssh-mcp"), "-not-a-path").unwrap_err();
895 assert!(err.contains("start with '-'") || err.contains("start with"));
896 }
897
898 #[test]
899 fn test_validate_background_log_path_rejects_newlines() {
900 assert!(
901 validate_background_log_path(Path::new("/tmp/ssh-mcp"), "/tmp/x\nrm -rf /").is_err()
902 );
903 assert!(
904 validate_background_log_path(Path::new("/tmp/ssh-mcp"), "/tmp/x\rrm -rf /").is_err()
905 );
906 }
907
908 #[test]
909 fn test_background_json_err_omits_unregistered_job_fields() {
910 let long_error = "e".repeat(BACKGROUND_JSON_SNIPPET_LIMIT_CHARS + 10);
911 let long_stderr = "s".repeat(BACKGROUND_JSON_SNIPPET_LIMIT_CHARS + 10);
912
913 let result = background_json_err(&long_error, &long_stderr);
914 let text = extract_text_from_result(&result);
915
916 let value: serde_json::Value =
917 serde_json::from_str(text.trim()).expect("background_json_err should return JSON");
918
919 assert_eq!(value.get("ok").and_then(|v| v.as_bool()), Some(false));
920 assert_eq!(
921 value.get("background").and_then(|v| v.as_bool()),
922 Some(true)
923 );
924 assert_eq!(value.get("truncated").and_then(|v| v.as_bool()), Some(true));
925 assert!(value.get("job_id").is_none());
926 assert!(value.get("log_path").is_none());
927 assert!(value.get("hint").is_none());
928
929 let fields = value
930 .get("truncated_fields")
931 .expect("expected truncated_fields");
932 assert_eq!(fields.get("error").and_then(|v| v.as_bool()), Some(true));
933 assert_eq!(fields.get("stderr").and_then(|v| v.as_bool()), Some(true));
934
935 let error_snippet = value
936 .get("error")
937 .and_then(|v| v.as_str())
938 .expect("expected error field");
939 assert_eq!(
940 error_snippet.chars().count(),
941 BACKGROUND_JSON_SNIPPET_LIMIT_CHARS
942 );
943 let stderr_snippet = value
944 .get("stderr")
945 .and_then(|v| v.as_str())
946 .expect("expected stderr field");
947 assert_eq!(
948 stderr_snippet.chars().count(),
949 BACKGROUND_JSON_SNIPPET_LIMIT_CHARS
950 );
951 }
952
953 #[test]
954 fn test_background_json_timeout_hint_contains_pid_and_check_process_tool() {
955 let result = background_json_timeout(
956 "job-42",
957 4242,
958 "/tmp/ssh-mcp/local.log",
959 &crate::background::response::BackgroundTimeoutSnapshot {
960 state: "running",
961 still_running: true,
962 exit_code: None,
963 state_reason: None,
964 elapsed_time: "00:01",
965 log_exists: true,
966 log_tail: "tail line",
967 tail_lines_used: 50,
968 },
969 );
970 let text = extract_text_from_result(&result);
971
972 let value: serde_json::Value =
973 serde_json::from_str(text.trim()).expect("background_json_timeout should return JSON");
974
975 assert_eq!(value.get("ok").and_then(|v| v.as_bool()), Some(false));
976 assert_eq!(value.get("timeout").and_then(|v| v.as_bool()), Some(true));
977 assert_eq!(
978 value.get("background").and_then(|v| v.as_bool()),
979 Some(true)
980 );
981 assert_eq!(
982 value.get("still_running").and_then(|v| v.as_bool()),
983 Some(true)
984 );
985 assert_eq!(value.get("state").and_then(|v| v.as_str()), Some("running"));
986 assert_eq!(
987 value.get("tail_lines_used").and_then(|v| v.as_u64()),
988 Some(50)
989 );
990 assert_eq!(
991 value.get("elapsed_time").and_then(|v| v.as_str()),
992 Some("00:01")
993 );
994 assert_eq!(
995 value.get("log_tail").and_then(|v| v.as_str()),
996 Some("tail line")
997 );
998
999 let hint = value
1000 .get("hint")
1001 .and_then(|v| v.as_str())
1002 .expect("expected hint field");
1003
1004 assert!(
1006 hint.contains("job_id=job-42"),
1007 "hint should contain the actual job_id value; got: '{hint}'"
1008 );
1009 assert!(
1011 hint.contains("check_process"),
1012 "hint should mention check_process tool; got: '{hint}'"
1013 );
1014 assert!(
1016 hint.contains("DO NOT restart"),
1017 "hint should warn against restarting; got: '{hint}'"
1018 );
1019 assert!(
1021 hint.contains("TIMEOUT_RECOVERY"),
1022 "hint should start with TIMEOUT_RECOVERY; got: '{hint}'"
1023 );
1024 assert!(
1025 hint.contains("MCP client deadlines may be shorter than timeout_ms"),
1026 "hint should distinguish the client deadline from timeout_ms; got: '{hint}'"
1027 );
1028 assert!(
1029 hint.contains("background=true"),
1030 "hint should recommend explicit background mode; got: '{hint}'"
1031 );
1032 assert!(
1034 !hint.contains("<pid>"),
1035 "hint should not contain <pid> placeholder; got: '{hint}'"
1036 );
1037 assert!(
1038 !hint.contains("<log_path>"),
1039 "hint should not contain <log_path> placeholder; got: '{hint}'"
1040 );
1041 }
1042}