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 tracing::{debug, error, info, warn};
20
21use crate::background::job::NewRunningJob;
22use crate::background::{JobRegistry, JobState, LocalLogSpooler, SharedJobState};
23use crate::config::Config;
24use crate::error::{Result, SshMcpError};
25#[cfg(unix)]
26use crate::platform::O_NOFOLLOW_FLAG;
27use crate::server::handlers::file_edit_common::{FileEditFaultInjection, FileEditPrivilege};
28#[cfg(test)]
29use crate::server::validation::read_file::{
30 READ_FILE_BYTES_PER_TOKEN, READ_FILE_DEFAULT_PREVIEW_LINES, READ_FILE_HARD_MAX_BYTES,
31 READ_FILE_MAX_LINE_WINDOW,
32};
33#[cfg(test)]
34use crate::server::validation::read_file::{
35 estimate_tokens_from_bytes, resolve_read_file_line_limit, resolve_read_file_max_bytes,
36};
37#[cfg(test)]
38use crate::server::validation::validate_background_log_path;
39use crate::ssh::{
40 CommandOutput, SshConfig, SshConnectionManager, sanitize_command, wrap_sudo_command,
41};
42use crate::tools::{ApplyPatchParams, CheckProcessParams, ReadFileMode, ReadFileParams};
43use crate::transfer::{TransferEngine, TransferParams, TransferRunContext, TransferSshOptions};
44
45mod args;
46mod exec;
47mod handlers;
48mod testing;
49mod tools;
50mod validation;
51
52const BACKGROUND_START_TIMEOUT: Duration = Duration::from_secs(20);
53const READ_FILE_ERROR_MARKER: &str = "__SSH_MCP_READ_FILE_ERR__";
54
55const JOB_COMPLETED_RETENTION: Duration = Duration::from_secs(60 * 60);
56
57static JOB_COUNTER: AtomicU64 = AtomicU64::new(0);
58
59fn make_job_id() -> String {
60 let counter = JOB_COUNTER.fetch_add(1, Ordering::Relaxed);
61 let epoch_ms = SystemTime::now()
62 .duration_since(UNIX_EPOCH)
63 .map(|d| d.as_millis())
64 .unwrap_or(0);
65 format!("{}-{}", epoch_ms, counter)
66}
67
68#[derive(Clone)]
73pub struct SshMcpServer {
74 config: Config,
76
77 connection: Arc<SshConnectionManager>,
79
80 timeout: Duration,
82
83 max_chars: Option<usize>,
85
86 spooler: Arc<LocalLogSpooler>,
87 job_registry: Arc<JobRegistry>,
88
89 transfer: TransferEngine,
90}
91
92impl SshMcpServer {
93 pub async fn new(config: Config) -> Result<Self> {
98 let local_root = std::env::current_dir()?;
99
100 let spooler = Arc::new(LocalLogSpooler::new_default());
101 spooler.ensure_dir().await.map_err(|e| {
102 SshMcpError::Config(format!(
103 "failed to initialize local log spool dir {}: {e}",
104 spooler.base_dir().display()
105 ))
106 })?;
107 let job_registry = Arc::new(JobRegistry::new(JOB_COMPLETED_RETENTION));
108
109 let mut ssh_config = SshConfig::new(&config.host, &config.user).with_port(config.port);
111
112 if let Some(ref password) = config.password {
114 ssh_config = ssh_config.with_password(password);
115 }
116
117 if let Some(ref key_path) = config.key {
118 let key_content = tokio::fs::read_to_string(key_path)
120 .await
121 .map_err(SshMcpError::Io)?;
122 ssh_config = ssh_config.with_private_key(&key_content);
123 }
124
125 if let Some(ref su_password) = config.su_password {
127 ssh_config = ssh_config.with_su_password(su_password);
128 }
129
130 if let Some(ref sudo_password) = config.sudo_password {
131 ssh_config = ssh_config.with_sudo_password(sudo_password);
132 }
133
134 ssh_config = ssh_config
136 .with_keepalive_interval(config.keepalive_interval)
137 .with_keepalive_max(config.keepalive_max);
138
139 ssh_config = ssh_config
141 .with_reconnect_retries(config.reconnect_retries)
142 .with_reconnect_backoff_ms(config.reconnect_backoff_ms)
143 .with_health_probe_timeout_ms(config.health_probe_timeout_ms);
144
145 ssh_config = ssh_config
147 .with_host_key_checking(config.strict_host_key_checking)
148 .with_known_hosts(config.known_hosts.clone());
149
150 ssh_config = ssh_config.with_max_output_tokens(config.max_output_tokens);
152
153 let connection = Arc::new(SshConnectionManager::new(ssh_config).await);
155
156 let timeout = Duration::from_millis(config.timeout_ms);
157 let max_chars = config.max_chars;
158
159 Ok(Self {
160 config,
161 connection,
162 timeout,
163 max_chars,
164 spooler,
165 job_registry,
166 transfer: TransferEngine::new(local_root),
167 })
168 }
169
170 fn connection_id(&self) -> String {
171 format!(
172 "{}@{}:{}",
173 self.config.user, self.config.host, self.config.port
174 )
175 }
176
177 fn default_local_log_path(
178 &self,
179 job_id: &str,
180 ) -> std::result::Result<(PathBuf, String), String> {
181 let path = self
182 .spooler
183 .log_path_for(job_id)
184 .map_err(|e| format!("failed to generate local log path for job_id='{job_id}': {e}"))?;
185 let path_str = path.to_string_lossy().to_string();
186 Ok((path, path_str))
187 }
188
189 async fn ensure_local_log_file(&self, log_path: &Path) -> std::result::Result<(), SshMcpError> {
190 self.spooler.ensure_dir().await.map_err(|e| {
191 SshMcpError::Config(format!(
192 "failed to ensure local log spool dir {}: {e}",
193 self.spooler.base_dir().display()
194 ))
195 })?;
196
197 if log_path.parent() != Some(self.spooler.base_dir()) {
198 return Err(SshMcpError::InvalidParams(format!(
199 "log_path must be directly under {}",
200 self.spooler.base_dir().display()
201 )));
202 }
203
204 match tokio::fs::symlink_metadata(log_path).await {
205 Ok(meta) => {
206 let ft = meta.file_type();
207 if ft.is_symlink() {
208 return Err(SshMcpError::invalid_params(
209 "log_path is a symlink (refusing to follow it)",
210 ));
211 }
212 if !ft.is_file() {
213 return Err(SshMcpError::invalid_params(
214 "log_path exists but is not a regular file",
215 ));
216 }
217 }
218 Err(e) if e.kind() == std::io::ErrorKind::NotFound => {}
219 Err(e) => return Err(SshMcpError::Io(e)),
220 }
221
222 let mut opts = tokio::fs::OpenOptions::new();
223 opts.write(true).create(true).truncate(true);
224
225 #[cfg(unix)]
226 {
227 opts.custom_flags(O_NOFOLLOW_FLAG);
228 }
229
230 let file = match opts.open(log_path).await {
231 Ok(f) => f,
232 Err(e) => {
233 if let Ok(meta) = tokio::fs::symlink_metadata(log_path).await
234 && meta.file_type().is_symlink()
235 {
236 return Err(SshMcpError::invalid_params(
237 "log_path is a symlink (refusing to follow it)",
238 ));
239 }
240 return Err(SshMcpError::Io(e));
241 }
242 };
243
244 file.sync_all().await.map_err(SshMcpError::Io)
245 }
246
247 async fn register_running_job(
248 &self,
249 job_id: &str,
250 pid: u32,
251 log_path: PathBuf,
252 command: &str,
253 ) -> SharedJobState {
254 let job = Arc::new(Mutex::new(JobState::new_running(NewRunningJob {
255 job_id: job_id.to_string(),
256 pid,
257 log_path,
258 command: command.to_string(),
259 connection_id: self.connection_id(),
260 })));
261
262 self.job_registry
263 .insert(job_id.to_string(), Arc::clone(&job))
264 .await;
265
266 let persisted = {
267 let guard = job.lock().await;
268 guard.clone()
269 };
270 if let Err(e) = self.spooler.persist_job_state(&persisted).await {
271 warn!(job_id = ?job_id, error = ?e, "failed to persist running job state");
272 }
273
274 job
275 }
276
277 pub fn connection(&self) -> &Arc<SshConnectionManager> {
279 &self.connection
280 }
281
282 pub async fn shutdown(&self) {
284 info!("Shutting down SSH MCP Server...");
285 self.connection.close().await;
286 }
287
288 async fn execute_command_with_timeout(
290 &self,
291 command: &str,
292 timeout: Duration,
293 ) -> std::result::Result<CallToolResult, McpError> {
294 debug!(
295 "shell tool called: cmd_len={}, background=false, sudo=false, timeout_ms={}",
296 command.len(),
297 timeout.as_millis()
298 );
299
300 let sanitized = match self.sanitize_or_tool_error(command) {
302 Ok(cmd) => cmd,
303 Err(result) => return Ok(result),
304 };
305
306 let requires_elevation = self.connection.get_su_password().is_some();
312 if requires_elevation {
313 if let Err(e) = self.connection.ensure_connected().await {
314 error!(error = ?e, "Failed to ensure SSH connection");
315 return Ok(CallToolResult::error(vec![Content::text(e.to_string())]));
316 }
317
318 if let Err(e) = self.connection.ensure_elevated().await {
319 debug!(error = ?e, "Elevation failed, will run as normal user");
320 }
321 }
322
323 if !requires_elevation && let Err(e) = self.connection.ensure_connected().await {
325 error!(error = ?e, "Failed to ensure SSH connection");
326 return Ok(CallToolResult::error(vec![Content::text(e.to_string())]));
327 }
328
329 self.execute_detachable_foreground_impl(&sanitized, &sanitized, timeout)
330 .await
331 }
332
333 async fn execute_command(
334 &self,
335 command: &str,
336 ) -> std::result::Result<CallToolResult, McpError> {
337 self.execute_command_with_timeout(command, self.timeout)
338 .await
339 }
340
341 async fn execute_background_command(
342 &self,
343 command: &str,
344 log_path: Option<&str>,
345 ) -> std::result::Result<CallToolResult, McpError> {
346 self.execute_background_impl(command, log_path, exec::BackgroundPrivilege::Normal)
347 .await
348 }
349
350 async fn execute_sudo_command_with_timeout(
352 &self,
353 command: &str,
354 timeout: Duration,
355 ) -> std::result::Result<CallToolResult, McpError> {
356 debug!(
357 "sudo_shell tool called: cmd_len={}, background=false, sudo=true, timeout_ms={}",
358 command.len(),
359 timeout.as_millis()
360 );
361
362 let sanitized = match self.sanitize_or_tool_error(command) {
364 Ok(cmd) => cmd,
365 Err(result) => return Ok(result),
366 };
367
368 let sudo_password = self.connection.get_sudo_password();
370 let wrapped_command = wrap_sudo_command(&sanitized, sudo_password);
371 debug!(
372 "Wrapped sudo command (password hidden): sudo -n sh -c '...' or printf '...' | sudo ..."
373 );
374
375 if let Err(e) = self.connection.ensure_connected().await {
376 error!(error = ?e, "Failed to ensure SSH connection");
377 return Ok(CallToolResult::error(vec![Content::text(e.to_string())]));
378 }
379
380 self.execute_detachable_foreground_impl(
381 &wrapped_command,
382 &format!("sudo {sanitized}"),
383 timeout,
384 )
385 .await
386 }
387
388 async fn execute_sudo_command(
389 &self,
390 command: &str,
391 ) -> std::result::Result<CallToolResult, McpError> {
392 self.execute_sudo_command_with_timeout(command, self.timeout)
393 .await
394 }
395
396 async fn execute_background_sudo_command(
397 &self,
398 command: &str,
399 log_path: Option<&str>,
400 ) -> std::result::Result<CallToolResult, McpError> {
401 let sudo_password = self.connection.get_sudo_password();
402 self.execute_background_impl(
403 command,
404 log_path,
405 exec::BackgroundPrivilege::Sudo {
406 password: sudo_password,
407 },
408 )
409 .await
410 }
411
412 fn sanitize_or_tool_error(&self, command: &str) -> std::result::Result<String, CallToolResult> {
413 sanitize_command(command, self.max_chars).map_err(|e| {
414 error!(error = ?e, "Command sanitization failed");
415 CallToolResult::error(vec![Content::text(format!("Error: {}", e))])
416 })
417 }
418
419 fn calltool_from_command_output(output: CommandOutput) -> CallToolResult {
420 let mut result_text = output.stdout;
422 if !output.stderr.is_empty() {
423 if !result_text.is_empty() {
424 result_text.push_str("\n--- stderr ---\n");
425 }
426 result_text.push_str(&output.stderr);
427 }
428
429 if output.exit_code.map(|code| code != 0).unwrap_or(true) {
433 CallToolResult::error(vec![Content::text(result_text)])
434 } else {
435 CallToolResult::success(vec![Content::text(result_text)])
436 }
437 }
438
439 fn shell_tool() -> Tool {
441 tools::shell_tool()
442 }
443
444 fn sudo_shell_tool() -> Tool {
446 tools::sudo_shell_tool()
447 }
448
449 fn transfer_tool() -> Tool {
451 tools::transfer_tool()
452 }
453
454 fn check_process_tool() -> Tool {
456 tools::check_process_tool()
457 }
458
459 fn read_file_tool() -> Tool {
461 tools::read_file_tool()
462 }
463
464 fn apply_patch_tool() -> Tool {
466 tools::apply_patch_tool()
467 }
468
469 fn sudo_apply_patch_tool() -> Tool {
471 tools::sudo_apply_patch_tool()
472 }
473
474 pub fn get_tool_documentation(tool_name: &str) -> Option<&'static str> {
479 tools::get_tool_documentation(tool_name)
480 }
481
482 fn resolve_timeout(&self, timeout_ms: Option<u64>) -> Duration {
484 timeout_ms
485 .map(Duration::from_millis)
486 .unwrap_or(self.timeout)
487 }
488
489 fn parse_tool_params<T: serde::de::DeserializeOwned>(
491 &self,
492 args: serde_json::Map<String, serde_json::Value>,
493 tool_name: &str,
494 ) -> std::result::Result<T, McpError> {
495 serde_json::from_value(serde_json::Value::Object(args))
496 .map_err(|e| McpError::invalid_params(format!("invalid {tool_name} params: {e}"), None))
497 }
498
499 async fn execute_transfer(
501 &self,
502 params: TransferParams,
503 verbose: bool,
504 ) -> std::result::Result<CallToolResult, McpError> {
505 let timeout = self.resolve_timeout(params.timeout_ms);
506 let key_path = self.config.key.clone();
507
508 if let Err(e) = self.connection.ensure_connected().await {
510 let resp = crate::transfer::TransferResponse::error(
511 params,
512 self.transfer.local_root(),
513 &e.to_string(),
514 );
515 let body = resp
516 .to_json(verbose)
517 .unwrap_or_else(|_| "{\"ok\":false,\"error\":\"serialization_error\"}".to_string());
518 return Ok(CallToolResult::success(vec![Content::text(body)]));
519 }
520
521 let resp = self
522 .transfer
523 .run(
524 &self.connection,
525 params,
526 TransferRunContext {
527 timeout,
528 ssh: TransferSshOptions {
529 host: self.config.host.clone(),
530 port: self.config.port,
531 user: self.config.user.clone(),
532 key_path,
533 host_key_checking: self.config.strict_host_key_checking,
534 known_hosts: self.config.known_hosts.clone(),
535 },
536 },
537 )
538 .await;
539 let body = resp
540 .to_json(verbose)
541 .unwrap_or_else(|_| "{\"ok\":false,\"error\":\"serialization_error\"}".to_string());
542 Ok(CallToolResult::success(vec![Content::text(body)]))
543 }
544}
545
546impl ServerHandler for SshMcpServer {
547 fn get_info(&self) -> ServerInfo {
549 ServerInfo::new(ServerCapabilities::builder().enable_tools().build())
550 .with_protocol_version(ProtocolVersion::LATEST)
551 .with_server_info(Implementation::from_build_env())
552 .with_instructions(format!(
553 "SSH MCP Server v{} - Execute commands on {}@{}:{}",
554 env!("CARGO_PKG_VERSION"),
555 self.config.user,
556 self.config.host,
557 self.config.port,
558 ))
559 }
560
561 async fn list_tools(
563 &self,
564 _request: Option<PaginatedRequestParams>,
565 _context: RequestContext<RoleServer>,
566 ) -> std::result::Result<ListToolsResult, McpError> {
567 debug!("list_tools called");
568
569 let mut tools = vec![Self::shell_tool()];
570
571 if !self.config.disable_sudo {
573 tools.push(Self::sudo_shell_tool());
574 tools.push(Self::sudo_apply_patch_tool());
575 }
576 tools.push(Self::check_process_tool());
577 tools.push(Self::transfer_tool());
578 tools.push(Self::read_file_tool());
579 tools.push(Self::apply_patch_tool());
580
581 Ok(ListToolsResult {
582 tools,
583 next_cursor: None,
584 meta: Default::default(),
585 })
586 }
587
588 async fn call_tool(
590 &self,
591 request: CallToolRequestParams,
592 _context: RequestContext<RoleServer>,
593 ) -> std::result::Result<CallToolResult, McpError> {
594 let tool_name: &str = request.name.as_ref();
595 debug!("call_tool called: {:?}", tool_name);
596
597 let args = request.arguments.unwrap_or_default();
598
599 match tool_name {
601 "shell" => {
602 let parsed = self.parse_common_tool_args(&args)?;
603 let timeout = self.resolve_timeout(parsed.timeout_ms);
604
605 if parsed.background {
606 self.execute_background_command(&parsed.command, parsed.log_path.as_deref())
607 .await
608 } else {
609 self.execute_command_with_timeout(&parsed.command, timeout)
610 .await
611 }
612 }
613 "sudo_shell" => {
614 if self.config.disable_sudo {
615 return Err(McpError::invalid_params(
616 "sudo_shell tool is disabled",
617 None,
618 ));
619 }
620
621 let parsed = self.parse_common_tool_args(&args)?;
622 let timeout = self.resolve_timeout(parsed.timeout_ms);
623
624 if parsed.background {
625 self.execute_background_sudo_command(
626 &parsed.command,
627 parsed.log_path.as_deref(),
628 )
629 .await
630 } else {
631 self.execute_sudo_command_with_timeout(&parsed.command, timeout)
632 .await
633 }
634 }
635 "transfer" => {
636 let params: TransferParams = self.parse_tool_params(args, "transfer")?;
637 let verbose = params.verbose;
638 self.execute_transfer(params, verbose).await
639 }
640 "check_process" => {
641 let params: CheckProcessParams = self.parse_tool_params(args, "check_process")?;
642 self.execute_check_process(params).await
643 }
644 "read" => {
645 let params: ReadFileParams = self.parse_tool_params(args, "read")?;
646 self.execute_read_file(params).await
647 }
648 "apply_patch" => {
649 let params: ApplyPatchParams = self.parse_tool_params(args, "apply_patch")?;
650 self.execute_apply_patch(
651 params,
652 FileEditFaultInjection::None,
653 FileEditPrivilege::User,
654 )
655 .await
656 }
657 "sudo_apply_patch" => {
658 if self.config.disable_sudo {
659 return Err(McpError::invalid_params(
660 "sudo_apply_patch tool is disabled",
661 None,
662 ));
663 }
664
665 let params: ApplyPatchParams = self.parse_tool_params(args, "sudo_apply_patch")?;
666 self.execute_apply_patch(
667 params,
668 FileEditFaultInjection::None,
669 FileEditPrivilege::Sudo,
670 )
671 .await
672 }
673 _ => Err(McpError::invalid_params(
674 format!("Unknown tool: {}", tool_name),
675 None,
676 )),
677 }
678 }
679}
680
681#[cfg(test)]
682mod tests {
683 use super::*;
684 use crate::background::response::{
685 BACKGROUND_JSON_SNIPPET_LIMIT_CHARS, background_json_err, background_json_timeout,
686 };
687 use crate::background::wrapper::{build_background_wrapper_script, remote_job_log_path};
688 use crate::server::validation::common::validate_read_file_path;
689 use crate::server::validation::read_file::sanitize_read_file_stderr_snippet;
690
691 fn extract_text_from_result(result: &CallToolResult) -> String {
692 result
693 .content
694 .iter()
695 .filter_map(|c| c.raw.as_text().map(|text| text.text.clone()))
696 .collect::<Vec<_>>()
697 .join("\n")
698 }
699
700 #[test]
704 fn test_server_info() {
705 assert!(!env!("CARGO_PKG_VERSION").is_empty());
707 }
708
709 #[test]
710 fn test_shell_tool_definition() {
711 let tool = SshMcpServer::shell_tool();
712 assert_eq!(tool.name.as_ref(), "shell");
713 assert!(tool.description.is_some());
714 }
715
716 #[test]
717 fn test_sudo_shell_tool_definition() {
718 let tool = SshMcpServer::sudo_shell_tool();
719 assert_eq!(tool.name.as_ref(), "sudo_shell");
720 assert!(tool.description.is_some());
721 }
722
723 #[test]
724 fn test_read_file_tool_definition() {
725 let tool = SshMcpServer::read_file_tool();
726 assert_eq!(tool.name.as_ref(), "read");
727 assert!(tool.description.is_some());
728 }
729
730 #[test]
731 fn test_apply_patch_tool_definition() {
732 let tool = SshMcpServer::apply_patch_tool();
733 assert_eq!(tool.name.as_ref(), "apply_patch");
734 assert!(tool.description.is_some());
735 }
736
737 #[test]
738 fn test_sudo_apply_patch_tool_definition() {
739 let tool = SshMcpServer::sudo_apply_patch_tool();
740 assert_eq!(tool.name.as_ref(), "sudo_apply_patch");
741 assert!(tool.description.is_some());
742 }
743
744 #[test]
745 fn test_build_background_wrapper_escapes_single_quotes_in_user_command() {
746 let remote_log = remote_job_log_path("job-1");
747 let script = build_background_wrapper_script("job-1", "echo 'hello world'", &remote_log);
748 assert!(script.contains("exec sh -c 'set +m; echo '\"'\"'hello world'\"'\"''"));
749 }
750
751 #[test]
752 fn test_build_background_wrapper_is_busybox_friendly() {
753 let remote_log = remote_job_log_path("job-1");
754 let script = build_background_wrapper_script("job-1", "echo test", &remote_log);
755 assert!(!script.contains("dirname --"));
756 assert!(!script.contains("mkdir -p --"));
757 assert!(!script.contains("sh -lc"));
758 assert!(script.contains("exec sh -c"));
759 assert!(!script.contains("nohup"));
760 }
761
762 #[test]
763 fn test_background_wrapper_emits_markers_and_exec() {
764 let remote_log = remote_job_log_path("job-1");
765 let script = build_background_wrapper_script("job-1", "echo test", &remote_log);
766 assert!(script.contains("__SSH_MCP_JOB_ID=job-1"));
767 assert!(script.contains("__SSH_MCP_PID=$$"));
768 assert!(script.contains("__SSH_MCP_LOG=$LOG"));
769 assert!(script.contains("exec sh -c"));
770 }
771
772 #[test]
773 fn test_background_wrapper_does_not_redirect_remote_output() {
774 let remote_log = remote_job_log_path("job-1");
775 let script = build_background_wrapper_script("job-1", "echo test", &remote_log);
776 assert!(!script.contains(">$LOG"));
777 assert!(!script.contains("2>&1"));
778 assert!(!script.contains("$EXIT"));
779 assert!(!script.contains("nohup"));
780 }
781
782 #[test]
783 fn test_validate_background_log_path_rejects_leading_dash() {
784 let err =
785 validate_background_log_path(Path::new("/tmp/ssh-mcp"), "-not-a-path").unwrap_err();
786 assert!(err.contains("start with '-'") || err.contains("start with"));
787 }
788
789 #[test]
790 fn test_validate_background_log_path_rejects_newlines() {
791 assert!(
792 validate_background_log_path(Path::new("/tmp/ssh-mcp"), "/tmp/x\nrm -rf /").is_err()
793 );
794 assert!(
795 validate_background_log_path(Path::new("/tmp/ssh-mcp"), "/tmp/x\rrm -rf /").is_err()
796 );
797 }
798
799 #[test]
800 fn test_validate_read_file_path_requires_absolute() {
801 let err = validate_read_file_path("relative/path").unwrap_err();
802 assert!(err.contains("absolute"));
803 }
804
805 #[test]
806 fn test_validate_read_file_path_rejects_trailing_slash() {
807 let err = validate_read_file_path("/etc/").unwrap_err();
808 assert!(err.contains("must not end with '/'"));
809 }
810
811 #[test]
812 fn test_resolve_read_file_max_bytes_uses_token_limit() {
813 assert_eq!(
814 resolve_read_file_max_bytes(Some(12_000)),
815 12_000 * READ_FILE_BYTES_PER_TOKEN
816 );
817 }
818
819 #[test]
820 fn test_resolve_read_file_max_bytes_none_uses_hard_cap() {
821 assert_eq!(resolve_read_file_max_bytes(None), READ_FILE_HARD_MAX_BYTES);
822 }
823
824 #[test]
825 fn test_resolve_read_file_max_bytes_applies_hard_cap() {
826 let very_large_tokens = READ_FILE_HARD_MAX_BYTES;
827 assert_eq!(
828 resolve_read_file_max_bytes(Some(very_large_tokens)),
829 READ_FILE_HARD_MAX_BYTES
830 );
831 }
832
833 #[test]
834 fn test_estimate_tokens_from_bytes_rounds_up() {
835 assert_eq!(estimate_tokens_from_bytes(0), 0);
836 assert_eq!(estimate_tokens_from_bytes(1), 1);
837 assert_eq!(estimate_tokens_from_bytes(4), 1);
838 assert_eq!(estimate_tokens_from_bytes(5), 2);
839 }
840
841 #[test]
842 fn test_resolve_read_file_line_limit_defaults_to_preview_window() {
843 let preview = resolve_read_file_line_limit(ReadFileMode::Preview, None)
844 .expect("preview lines should resolve");
845 assert_eq!(preview, Some(READ_FILE_DEFAULT_PREVIEW_LINES));
846
847 let head = resolve_read_file_line_limit(ReadFileMode::Head, None)
848 .expect("head lines should resolve");
849 assert_eq!(head, Some(READ_FILE_DEFAULT_PREVIEW_LINES));
850
851 let tail = resolve_read_file_line_limit(ReadFileMode::Tail, None)
852 .expect("tail lines should resolve");
853 assert_eq!(tail, Some(READ_FILE_DEFAULT_PREVIEW_LINES));
854 }
855
856 #[test]
857 fn test_resolve_read_file_line_limit_for_full_ignores_lines() {
858 let full = resolve_read_file_line_limit(ReadFileMode::Full, Some(123))
859 .expect("full mode should ignore lines");
860 assert_eq!(full, None);
861 }
862
863 #[test]
864 fn test_resolve_read_file_line_limit_rejects_zero() {
865 let err = resolve_read_file_line_limit(ReadFileMode::Head, Some(0)).unwrap_err();
866 assert!(err.contains("positive"));
867 }
868
869 #[test]
870 fn test_resolve_read_file_line_limit_rejects_too_large() {
871 let err =
872 resolve_read_file_line_limit(ReadFileMode::Tail, Some(READ_FILE_MAX_LINE_WINDOW + 1))
873 .unwrap_err();
874 assert!(err.contains("<="));
875 }
876
877 #[test]
878 fn test_sanitize_read_file_stderr_snippet_normalizes_whitespace_and_controls() {
879 let stderr = "line1\nline2\t\u{0007}bad\rline3";
880 let snippet = sanitize_read_file_stderr_snippet(stderr)
881 .expect("snippet should be present for non-empty stderr");
882 assert_eq!(snippet, "line1 line2 bad line3");
883 }
884
885 #[test]
886 fn test_background_json_err_omits_unregistered_job_fields() {
887 let long_error = "e".repeat(BACKGROUND_JSON_SNIPPET_LIMIT_CHARS + 10);
888 let long_stderr = "s".repeat(BACKGROUND_JSON_SNIPPET_LIMIT_CHARS + 10);
889
890 let result = background_json_err(&long_error, &long_stderr);
891 let text = extract_text_from_result(&result);
892
893 let value: serde_json::Value =
894 serde_json::from_str(text.trim()).expect("background_json_err should return JSON");
895
896 assert_eq!(value.get("ok").and_then(|v| v.as_bool()), Some(false));
897 assert_eq!(
898 value.get("background").and_then(|v| v.as_bool()),
899 Some(true)
900 );
901 assert_eq!(value.get("truncated").and_then(|v| v.as_bool()), Some(true));
902 assert!(value.get("job_id").is_none());
903 assert!(value.get("log_path").is_none());
904 assert!(value.get("hint").is_none());
905
906 let fields = value
907 .get("truncated_fields")
908 .expect("expected truncated_fields");
909 assert_eq!(fields.get("error").and_then(|v| v.as_bool()), Some(true));
910 assert_eq!(fields.get("stderr").and_then(|v| v.as_bool()), Some(true));
911
912 let error_snippet = value
913 .get("error")
914 .and_then(|v| v.as_str())
915 .expect("expected error field");
916 assert_eq!(
917 error_snippet.chars().count(),
918 BACKGROUND_JSON_SNIPPET_LIMIT_CHARS
919 );
920 let stderr_snippet = value
921 .get("stderr")
922 .and_then(|v| v.as_str())
923 .expect("expected stderr field");
924 assert_eq!(
925 stderr_snippet.chars().count(),
926 BACKGROUND_JSON_SNIPPET_LIMIT_CHARS
927 );
928 }
929
930 #[test]
931 fn test_background_json_timeout_hint_contains_pid_and_check_process_tool() {
932 let result = background_json_timeout(
933 "job-42",
934 4242,
935 "/tmp/ssh-mcp/local.log",
936 &crate::background::response::BackgroundTimeoutSnapshot {
937 state: "running",
938 still_running: true,
939 exit_code: None,
940 state_reason: None,
941 elapsed_time: "00:01",
942 log_exists: true,
943 log_tail: "tail line",
944 tail_lines_used: 50,
945 },
946 );
947 let text = extract_text_from_result(&result);
948
949 let value: serde_json::Value =
950 serde_json::from_str(text.trim()).expect("background_json_timeout should return JSON");
951
952 assert_eq!(value.get("ok").and_then(|v| v.as_bool()), Some(false));
953 assert_eq!(value.get("timeout").and_then(|v| v.as_bool()), Some(true));
954 assert_eq!(
955 value.get("background").and_then(|v| v.as_bool()),
956 Some(true)
957 );
958 assert_eq!(
959 value.get("still_running").and_then(|v| v.as_bool()),
960 Some(true)
961 );
962 assert_eq!(value.get("state").and_then(|v| v.as_str()), Some("running"));
963 assert_eq!(
964 value.get("tail_lines_used").and_then(|v| v.as_u64()),
965 Some(50)
966 );
967 assert_eq!(
968 value.get("elapsed_time").and_then(|v| v.as_str()),
969 Some("00:01")
970 );
971 assert_eq!(
972 value.get("log_tail").and_then(|v| v.as_str()),
973 Some("tail line")
974 );
975
976 let hint = value
977 .get("hint")
978 .and_then(|v| v.as_str())
979 .expect("expected hint field");
980
981 assert!(
983 hint.contains("job_id=job-42"),
984 "hint should contain the actual job_id value; got: '{hint}'"
985 );
986 assert!(
988 hint.contains("check_process"),
989 "hint should mention check_process tool; got: '{hint}'"
990 );
991 assert!(
993 hint.contains("DO NOT restart"),
994 "hint should warn against restarting; got: '{hint}'"
995 );
996 assert!(
998 hint.contains("TIMEOUT_RECOVERY"),
999 "hint should start with TIMEOUT_RECOVERY; got: '{hint}'"
1000 );
1001 assert!(
1003 !hint.contains("<pid>"),
1004 "hint should not contain <pid> placeholder; got: '{hint}'"
1005 );
1006 assert!(
1007 !hint.contains("<log_path>"),
1008 "hint should not contain <log_path> placeholder; got: '{hint}'"
1009 );
1010 }
1011
1012 #[test]
1013 fn test_tool_documentation_available() {
1014 assert!(SshMcpServer::get_tool_documentation("shell").is_some());
1016 assert!(SshMcpServer::get_tool_documentation("sudo_shell").is_some());
1017 assert!(SshMcpServer::get_tool_documentation("transfer").is_some());
1018 assert!(SshMcpServer::get_tool_documentation("read").is_some());
1019 assert!(SshMcpServer::get_tool_documentation("apply_patch").is_some());
1020 assert!(SshMcpServer::get_tool_documentation("sudo_apply_patch").is_some());
1021 assert!(SshMcpServer::get_tool_documentation("write-file").is_none());
1022 assert!(SshMcpServer::get_tool_documentation("replace-in-file").is_none());
1023 assert!(SshMcpServer::get_tool_documentation("unknown").is_none());
1024 }
1025
1026 #[test]
1027 fn test_shell_documentation_content() {
1028 let docs = SshMcpServer::get_tool_documentation("shell").unwrap();
1029 assert!(docs.contains("SHELL TOOL"));
1030 assert!(docs.contains("PARAMETERS:"));
1031 assert!(docs.contains("BACKGROUND MODE:"));
1032 assert!(docs.contains("command"));
1033 assert!(docs.contains("background"));
1034 assert!(docs.contains("still_running"));
1035 }
1036
1037 #[test]
1038 fn test_sudo_shell_documentation_content() {
1039 let docs = SshMcpServer::get_tool_documentation("sudo_shell").unwrap();
1040 assert!(docs.contains("SUDO_SHELL TOOL"));
1041 assert!(docs.contains("sudo"));
1042 }
1043
1044 #[test]
1045 fn test_transfer_documentation_content() {
1046 let docs = SshMcpServer::get_tool_documentation("transfer").unwrap();
1047 assert!(docs.contains("TRANSFER TOOL"));
1048 assert!(docs.contains("put"));
1049 assert!(docs.contains("get"));
1050 assert!(docs.contains("TRANSPORTS:"));
1051 }
1052
1053 #[test]
1054 fn test_read_file_documentation_content() {
1055 let docs = SshMcpServer::get_tool_documentation("read").unwrap();
1056 assert!(docs.contains("READ TOOL"));
1057 assert!(docs.contains("remote_path"));
1058 assert!(docs.contains("mode"));
1059 assert!(docs.contains("UTF-8"));
1060 }
1061
1062 #[test]
1063 fn test_apply_patch_documentation_content() {
1064 let docs = SshMcpServer::get_tool_documentation("apply_patch").unwrap();
1065 assert!(docs.contains("APPLY_PATCH TOOL"));
1066 assert!(docs.contains("Add File"));
1067 assert!(docs.contains("Delete File"));
1068 }
1069
1070 #[test]
1071 fn test_sudo_apply_patch_documentation_content() {
1072 let docs = SshMcpServer::get_tool_documentation("sudo_apply_patch").unwrap();
1073 assert!(docs.contains("SUDO_APPLY_PATCH TOOL"));
1074 assert!(docs.contains("sudo"));
1075 }
1076
1077 #[test]
1078 fn test_compact_tool_descriptions() {
1079 let shell = SshMcpServer::shell_tool();
1081 let sudo_shell = SshMcpServer::sudo_shell_tool();
1082 let transfer = SshMcpServer::transfer_tool();
1083 let read_file = SshMcpServer::read_file_tool();
1084 let apply_patch = SshMcpServer::apply_patch_tool();
1085 let sudo_apply_patch = SshMcpServer::sudo_apply_patch_tool();
1086
1087 if let Some(desc) = shell.description {
1089 assert!(
1090 desc.len() < 100,
1091 "shell description too long: {} chars",
1092 desc.len()
1093 );
1094 }
1095 if let Some(desc) = sudo_shell.description {
1096 assert!(
1097 desc.len() < 100,
1098 "sudo_shell description too long: {} chars",
1099 desc.len()
1100 );
1101 }
1102 if let Some(desc) = transfer.description {
1103 assert!(
1104 desc.len() < 100,
1105 "transfer description too long: {} chars",
1106 desc.len()
1107 );
1108 }
1109 if let Some(desc) = read_file.description {
1110 assert!(
1111 desc.len() < 100,
1112 "read description too long: {} chars",
1113 desc.len()
1114 );
1115 }
1116 if let Some(desc) = apply_patch.description {
1117 assert!(
1118 desc.len() < 100,
1119 "apply_patch description too long: {} chars",
1120 desc.len()
1121 );
1122 }
1123 if let Some(desc) = sudo_apply_patch.description {
1124 assert!(
1125 desc.len() < 100,
1126 "sudo_apply_patch description too long: {} chars",
1127 desc.len()
1128 );
1129 }
1130 }
1131}