Skip to main content

ssh_mcp/
server.rs

1//! MCP Server implementation
2//!
3//! This module provides the main MCP server that integrates SSH connection
4//! management with the `shell` and `sudo_shell` tools.
5
6use 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, 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/// SSH MCP Server
69///
70/// The main server implementation that provides MCP tools for remote SSH
71/// command execution.
72#[derive(Clone)]
73pub struct SshMcpServer {
74    /// Server configuration
75    config: Config,
76
77    /// SSH connection manager
78    connection: Arc<SshConnectionManager>,
79
80    /// Command execution timeout
81    timeout: Duration,
82
83    /// Maximum command length
84    max_chars: Option<usize>,
85
86    spooler: Arc<LocalLogSpooler>,
87    job_registry: Arc<JobRegistry>,
88
89    transfer: TransferEngine,
90}
91
92impl SshMcpServer {
93    /// Create a new SSH MCP Server
94    ///
95    /// This sets up the SSH connection manager based on the provided configuration.
96    /// Connection is not established until a tool is actually used.
97    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        // Build SSH configuration
110        let mut ssh_config = SshConfig::new(&config.host, &config.user).with_port(config.port);
111
112        // Add authentication
113        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            // Read the key file
119            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        // Add elevation passwords if provided
126        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        // Add keepalive settings for human-like connection persistence
135        ssh_config = ssh_config
136            .with_keepalive_interval(config.keepalive_interval)
137            .with_keepalive_max(config.keepalive_max);
138
139        // Add reconnect and health probe settings
140        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        // Add host key verification settings
146        ssh_config = ssh_config
147            .with_host_key_checking(config.strict_host_key_checking)
148            .with_known_hosts(config.known_hosts.clone());
149
150        // Add output token limit for OOM protection
151        ssh_config = ssh_config.with_max_output_tokens(config.max_output_tokens);
152
153        // Create connection manager
154        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    /// Get a reference to the SSH connection manager
278    pub fn connection(&self) -> &Arc<SshConnectionManager> {
279        &self.connection
280    }
281
282    /// Close the server and cleanup resources
283    pub async fn shutdown(&self) {
284        info!("Shutting down SSH MCP Server...");
285        self.connection.close().await;
286    }
287
288    /// Execute a command (used by shell tool)
289    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        // Sanitize the command
301        let sanitized = match self.sanitize_or_tool_error(command) {
302            Ok(cmd) => cmd,
303            Err(result) => return Ok(result),
304        };
305
306        // Foreground execution is detachable-by-design:
307        // - Start the command on a dedicated SSH channel
308        // - Stream remote stdout/stderr into a local spool file
309        // - If timeout elapses, return JSON with job_id/pid/log_path while the stream continues
310
311        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![ContentBlock::text(
316                    e.to_string(),
317                )]));
318            }
319
320            if let Err(e) = self.connection.ensure_elevated().await {
321                debug!(error = ?e, "Elevation failed, will run as normal user");
322            }
323        }
324
325        // Ensure connection is established for detached foreground execution path.
326        if !requires_elevation && let Err(e) = self.connection.ensure_connected().await {
327            error!(error = ?e, "Failed to ensure SSH connection");
328            return Ok(CallToolResult::error(vec![ContentBlock::text(
329                e.to_string(),
330            )]));
331        }
332
333        self.execute_detachable_foreground_impl(&sanitized, &sanitized, timeout)
334            .await
335    }
336
337    async fn execute_command(
338        &self,
339        command: &str,
340    ) -> std::result::Result<CallToolResult, McpError> {
341        self.execute_command_with_timeout(command, self.timeout)
342            .await
343    }
344
345    async fn execute_background_command(
346        &self,
347        command: &str,
348        log_path: Option<&str>,
349    ) -> std::result::Result<CallToolResult, McpError> {
350        self.execute_background_impl(command, log_path, exec::BackgroundPrivilege::Normal)
351            .await
352    }
353
354    /// Execute a command with sudo (used by sudo_shell tool)
355    async fn execute_sudo_command_with_timeout(
356        &self,
357        command: &str,
358        timeout: Duration,
359    ) -> std::result::Result<CallToolResult, McpError> {
360        debug!(
361            "sudo_shell tool called: cmd_len={}, background=false, sudo=true, timeout_ms={}",
362            command.len(),
363            timeout.as_millis()
364        );
365
366        // Sanitize the command
367        let sanitized = match self.sanitize_or_tool_error(command) {
368            Ok(cmd) => cmd,
369            Err(result) => return Ok(result),
370        };
371
372        // Wrap the command with sudo
373        let sudo_password = self.connection.get_sudo_password();
374        let wrapped_command = wrap_sudo_command(&sanitized, sudo_password);
375        debug!(
376            "Wrapped sudo command (password hidden): sudo -n sh -c '...' or printf '...' | sudo ..."
377        );
378
379        if let Err(e) = self.connection.ensure_connected().await {
380            error!(error = ?e, "Failed to ensure SSH connection");
381            return Ok(CallToolResult::error(vec![ContentBlock::text(
382                e.to_string(),
383            )]));
384        }
385
386        self.execute_detachable_foreground_impl(
387            &wrapped_command,
388            &format!("sudo {sanitized}"),
389            timeout,
390        )
391        .await
392    }
393
394    async fn execute_sudo_command(
395        &self,
396        command: &str,
397    ) -> std::result::Result<CallToolResult, McpError> {
398        self.execute_sudo_command_with_timeout(command, self.timeout)
399            .await
400    }
401
402    async fn execute_background_sudo_command(
403        &self,
404        command: &str,
405        log_path: Option<&str>,
406    ) -> std::result::Result<CallToolResult, McpError> {
407        let sudo_password = self.connection.get_sudo_password();
408        self.execute_background_impl(
409            command,
410            log_path,
411            exec::BackgroundPrivilege::Sudo {
412                password: sudo_password,
413            },
414        )
415        .await
416    }
417
418    fn sanitize_or_tool_error(&self, command: &str) -> std::result::Result<String, CallToolResult> {
419        sanitize_command(command, self.max_chars).map_err(|e| {
420            error!(error = ?e, "Command sanitization failed");
421            CallToolResult::error(vec![ContentBlock::text(format!("Error: {}", e))])
422        })
423    }
424
425    fn calltool_from_command_output(output: CommandOutput) -> CallToolResult {
426        // Combine stdout and stderr for the response
427        let mut result_text = output.stdout;
428        if !output.stderr.is_empty() {
429            if !result_text.is_empty() {
430                result_text.push_str("\n--- stderr ---\n");
431            }
432            result_text.push_str(&output.stderr);
433        }
434
435        // Check for error exit code.
436        // exit_code=None means the SSH channel was torn down without delivering
437        // an exit status or exit signal — treat as error, not success.
438        if output.exit_code.map(|code| code != 0).unwrap_or(true) {
439            CallToolResult::error(vec![ContentBlock::text(result_text)])
440        } else {
441            CallToolResult::success(vec![ContentBlock::text(result_text)])
442        }
443    }
444
445    /// Build shell tool definition (compact)
446    fn shell_tool() -> Tool {
447        tools::shell_tool()
448    }
449
450    /// Build sudo_shell tool definition (compact)
451    fn sudo_shell_tool() -> Tool {
452        tools::sudo_shell_tool()
453    }
454
455    /// Build transfer tool definition (compact)
456    fn transfer_tool() -> Tool {
457        tools::transfer_tool()
458    }
459
460    /// Build check_process tool definition
461    fn check_process_tool() -> Tool {
462        tools::check_process_tool()
463    }
464
465    /// Build read tool definition
466    fn read_file_tool() -> Tool {
467        tools::read_file_tool()
468    }
469
470    /// Build apply_patch tool definition
471    fn apply_patch_tool() -> Tool {
472        tools::apply_patch_tool()
473    }
474
475    /// Build sudo_apply_patch tool definition
476    fn sudo_apply_patch_tool() -> Tool {
477        tools::sudo_apply_patch_tool()
478    }
479
480    /// Get extended documentation for a tool by name
481    ///
482    /// Returns the full documentation text that was removed from compact tool definitions
483    /// to save tokens in the MCP protocol.
484    pub fn get_tool_documentation(tool_name: &str) -> Option<&'static str> {
485        tools::get_tool_documentation(tool_name)
486    }
487
488    /// Resolve timeout duration from optional milliseconds, falling back to server default.
489    fn resolve_timeout(&self, timeout_ms: Option<u64>) -> Duration {
490        timeout_ms
491            .map(Duration::from_millis)
492            .unwrap_or(self.timeout)
493    }
494
495    /// Parse tool parameters from JSON with standardized error handling.
496    fn parse_tool_params<T: serde::de::DeserializeOwned>(
497        &self,
498        args: serde_json::Map<String, serde_json::Value>,
499        tool_name: &str,
500    ) -> std::result::Result<T, McpError> {
501        serde_json::from_value(serde_json::Value::Object(args))
502            .map_err(|e| McpError::invalid_params(format!("invalid {tool_name} params: {e}"), None))
503    }
504
505    /// Execute transfer tool with connection management and JSON serialization.
506    async fn execute_transfer(
507        &self,
508        params: TransferParams,
509        verbose: bool,
510    ) -> std::result::Result<CallToolResult, McpError> {
511        let timeout = self.resolve_timeout(params.timeout_ms);
512        let key_path = self.config.key.clone();
513
514        // Ensure connection is established (so errors are deterministic).
515        if let Err(e) = self.connection.ensure_connected().await {
516            let resp = crate::transfer::TransferResponse::error(
517                params,
518                self.transfer.local_root(),
519                &e.to_string(),
520            );
521            let body = resp
522                .to_json(verbose)
523                .unwrap_or_else(|_| "{\"ok\":false,\"error\":\"serialization_error\"}".to_string());
524            return Ok(CallToolResult::success(vec![ContentBlock::text(body)]));
525        }
526
527        let resp = self
528            .transfer
529            .run(
530                &self.connection,
531                params,
532                TransferRunContext {
533                    timeout,
534                    ssh: TransferSshOptions {
535                        host: self.config.host.clone(),
536                        port: self.config.port,
537                        user: self.config.user.clone(),
538                        key_path,
539                        host_key_checking: self.config.strict_host_key_checking,
540                        known_hosts: self.config.known_hosts.clone(),
541                    },
542                },
543            )
544            .await;
545        let body = resp
546            .to_json(verbose)
547            .unwrap_or_else(|_| "{\"ok\":false,\"error\":\"serialization_error\"}".to_string());
548        Ok(CallToolResult::success(vec![ContentBlock::text(body)]))
549    }
550}
551
552impl ServerHandler for SshMcpServer {
553    /// Return server information
554    fn get_info(&self) -> ServerInfo {
555        ServerInfo::new(ServerCapabilities::builder().enable_tools().build())
556            .with_protocol_version(ProtocolVersion::LATEST)
557            .with_server_info(Implementation::from_build_env())
558            .with_instructions(format!(
559                "SSH MCP Server v{} - Execute commands on {}@{}:{}",
560                env!("CARGO_PKG_VERSION"),
561                self.config.user,
562                self.config.host,
563                self.config.port,
564            ))
565    }
566
567    /// List available tools
568    async fn list_tools(
569        &self,
570        _request: Option<PaginatedRequestParams>,
571        _context: RequestContext<RoleServer>,
572    ) -> std::result::Result<ListToolsResult, McpError> {
573        debug!("list_tools called");
574
575        let mut tools = vec![Self::shell_tool()];
576
577        // Docs/expected order: shell, optional sudo tools, check_process, transfer, read, apply_patch.
578        if !self.config.disable_sudo {
579            tools.push(Self::sudo_shell_tool());
580            tools.push(Self::sudo_apply_patch_tool());
581        }
582        tools.push(Self::check_process_tool());
583        tools.push(Self::transfer_tool());
584        tools.push(Self::read_file_tool());
585        tools.push(Self::apply_patch_tool());
586
587        Ok(ListToolsResult {
588            tools,
589            next_cursor: None,
590            meta: Default::default(),
591        })
592    }
593
594    /// Call a tool
595    async fn call_tool(
596        &self,
597        request: CallToolRequestParams,
598        context: RequestContext<RoleServer>,
599    ) -> std::result::Result<CallToolResult, McpError> {
600        let tool_name: &str = request.name.as_ref();
601        debug!("call_tool called: {:?}", tool_name);
602
603        let args = request.arguments.unwrap_or_default();
604
605        // Route to the appropriate tool
606        match tool_name {
607            "shell" => {
608                let parsed = self.parse_common_tool_args(&args)?;
609                let timeout = self.resolve_timeout(parsed.timeout_ms);
610
611                if parsed.background {
612                    self.execute_background_command(&parsed.command, parsed.log_path.as_deref())
613                        .await
614                } else {
615                    self.execute_command_with_timeout(&parsed.command, timeout)
616                        .await
617                }
618            }
619            "sudo_shell" => {
620                if self.config.disable_sudo {
621                    return Err(McpError::invalid_params(
622                        "sudo_shell tool is disabled",
623                        None,
624                    ));
625                }
626
627                let parsed = self.parse_common_tool_args(&args)?;
628                let timeout = self.resolve_timeout(parsed.timeout_ms);
629
630                if parsed.background {
631                    self.execute_background_sudo_command(
632                        &parsed.command,
633                        parsed.log_path.as_deref(),
634                    )
635                    .await
636                } else {
637                    self.execute_sudo_command_with_timeout(&parsed.command, timeout)
638                        .await
639                }
640            }
641            "transfer" => {
642                let params: TransferParams = self.parse_tool_params(args, "transfer")?;
643                let verbose = params.verbose;
644                self.execute_transfer(params, verbose).await
645            }
646            "check_process" => {
647                let params: args::CheckProcessToolArgs =
648                    self.parse_tool_params(args, "check_process")?;
649                self.execute_check_process(params.check, params.wait_for, context.ct.cancelled())
650                    .await
651            }
652            "read" => {
653                let params: ReadFileParams = self.parse_tool_params(args, "read")?;
654                self.execute_read_file(params).await
655            }
656            "apply_patch" => {
657                let params: ApplyPatchParams = self.parse_tool_params(args, "apply_patch")?;
658                self.execute_apply_patch(
659                    params,
660                    FileEditFaultInjection::None,
661                    FileEditPrivilege::User,
662                )
663                .await
664            }
665            "sudo_apply_patch" => {
666                if self.config.disable_sudo {
667                    return Err(McpError::invalid_params(
668                        "sudo_apply_patch tool is disabled",
669                        None,
670                    ));
671                }
672
673                let params: ApplyPatchParams = self.parse_tool_params(args, "sudo_apply_patch")?;
674                self.execute_apply_patch(
675                    params,
676                    FileEditFaultInjection::None,
677                    FileEditPrivilege::Sudo,
678                )
679                .await
680            }
681            _ => Err(McpError::invalid_params(
682                format!("Unknown tool: {}", tool_name),
683                None,
684            )),
685        }
686    }
687}
688
689#[cfg(test)]
690mod tests {
691    use super::*;
692    use crate::background::response::{
693        BACKGROUND_JSON_SNIPPET_LIMIT_CHARS, background_json_err, background_json_timeout,
694    };
695    use crate::background::wrapper::{build_background_wrapper_script, remote_job_log_path};
696    use crate::server::validation::common::validate_read_file_path;
697    use crate::server::validation::read_file::sanitize_read_file_stderr_snippet;
698
699    fn extract_text_from_result(result: &CallToolResult) -> String {
700        result
701            .content
702            .iter()
703            .filter_map(|c| c.as_text().map(|text| text.text.clone()))
704            .collect::<Vec<_>>()
705            .join("\n")
706    }
707
708    #[test]
709    fn test_server_info() {
710        // Verify the package version is defined
711        assert!(!env!("CARGO_PKG_VERSION").is_empty());
712    }
713
714    #[test]
715    fn test_shell_tool_definition() {
716        let tool = SshMcpServer::shell_tool();
717        assert_eq!(tool.name.as_ref(), "shell");
718        assert!(tool.description.is_some());
719    }
720
721    #[test]
722    fn test_sudo_shell_tool_definition() {
723        let tool = SshMcpServer::sudo_shell_tool();
724        assert_eq!(tool.name.as_ref(), "sudo_shell");
725        assert!(tool.description.is_some());
726    }
727
728    #[test]
729    fn test_read_file_tool_definition() {
730        let tool = SshMcpServer::read_file_tool();
731        assert_eq!(tool.name.as_ref(), "read");
732        assert!(tool.description.is_some());
733    }
734
735    #[test]
736    fn test_apply_patch_tool_definition() {
737        let tool = SshMcpServer::apply_patch_tool();
738        assert_eq!(tool.name.as_ref(), "apply_patch");
739        assert!(tool.description.is_some());
740    }
741
742    #[test]
743    fn test_sudo_apply_patch_tool_definition() {
744        let tool = SshMcpServer::sudo_apply_patch_tool();
745        assert_eq!(tool.name.as_ref(), "sudo_apply_patch");
746        assert!(tool.description.is_some());
747    }
748
749    #[test]
750    fn test_build_background_wrapper_escapes_single_quotes_in_user_command() {
751        let remote_log = remote_job_log_path("job-1");
752        let script = build_background_wrapper_script("job-1", "echo 'hello world'", &remote_log);
753        assert!(script.contains("exec sh -c 'set +m; echo '\"'\"'hello world'\"'\"''"));
754    }
755
756    #[test]
757    fn test_build_background_wrapper_is_busybox_friendly() {
758        let remote_log = remote_job_log_path("job-1");
759        let script = build_background_wrapper_script("job-1", "echo test", &remote_log);
760        assert!(!script.contains("dirname --"));
761        assert!(!script.contains("mkdir -p --"));
762        assert!(!script.contains("sh -lc"));
763        assert!(script.contains("exec sh -c"));
764        assert!(!script.contains("nohup"));
765    }
766
767    #[test]
768    fn test_background_wrapper_emits_markers_and_exec() {
769        let remote_log = remote_job_log_path("job-1");
770        let script = build_background_wrapper_script("job-1", "echo test", &remote_log);
771        assert!(script.contains("__SSH_MCP_JOB_ID=job-1"));
772        assert!(script.contains("__SSH_MCP_PID=$$"));
773        assert!(script.contains("__SSH_MCP_LOG=$LOG"));
774        assert!(script.contains("exec sh -c"));
775    }
776
777    #[test]
778    fn test_background_wrapper_does_not_redirect_remote_output() {
779        let remote_log = remote_job_log_path("job-1");
780        let script = build_background_wrapper_script("job-1", "echo test", &remote_log);
781        assert!(!script.contains(">$LOG"));
782        assert!(!script.contains("2>&1"));
783        assert!(!script.contains("$EXIT"));
784        assert!(!script.contains("nohup"));
785    }
786
787    #[test]
788    fn test_validate_background_log_path_rejects_leading_dash() {
789        let err =
790            validate_background_log_path(Path::new("/tmp/ssh-mcp"), "-not-a-path").unwrap_err();
791        assert!(err.contains("start with '-'") || err.contains("start with"));
792    }
793
794    #[test]
795    fn test_validate_background_log_path_rejects_newlines() {
796        assert!(
797            validate_background_log_path(Path::new("/tmp/ssh-mcp"), "/tmp/x\nrm -rf /").is_err()
798        );
799        assert!(
800            validate_background_log_path(Path::new("/tmp/ssh-mcp"), "/tmp/x\rrm -rf /").is_err()
801        );
802    }
803
804    #[test]
805    fn test_validate_read_file_path_requires_absolute() {
806        let err = validate_read_file_path("relative/path").unwrap_err();
807        assert!(err.contains("absolute"));
808    }
809
810    #[test]
811    fn test_validate_read_file_path_rejects_trailing_slash() {
812        let err = validate_read_file_path("/etc/").unwrap_err();
813        assert!(err.contains("must not end with '/'"));
814    }
815
816    #[test]
817    fn test_resolve_read_file_max_bytes_uses_token_limit() {
818        assert_eq!(
819            resolve_read_file_max_bytes(Some(12_000)),
820            12_000 * READ_FILE_BYTES_PER_TOKEN
821        );
822    }
823
824    #[test]
825    fn test_resolve_read_file_max_bytes_none_uses_hard_cap() {
826        assert_eq!(resolve_read_file_max_bytes(None), READ_FILE_HARD_MAX_BYTES);
827    }
828
829    #[test]
830    fn test_resolve_read_file_max_bytes_applies_hard_cap() {
831        let very_large_tokens = READ_FILE_HARD_MAX_BYTES;
832        assert_eq!(
833            resolve_read_file_max_bytes(Some(very_large_tokens)),
834            READ_FILE_HARD_MAX_BYTES
835        );
836    }
837
838    #[test]
839    fn test_estimate_tokens_from_bytes_rounds_up() {
840        assert_eq!(estimate_tokens_from_bytes(0), 0);
841        assert_eq!(estimate_tokens_from_bytes(1), 1);
842        assert_eq!(estimate_tokens_from_bytes(4), 1);
843        assert_eq!(estimate_tokens_from_bytes(5), 2);
844    }
845
846    #[test]
847    fn test_resolve_read_file_line_limit_defaults_to_preview_window() {
848        let preview = resolve_read_file_line_limit(ReadFileMode::Preview, None)
849            .expect("preview lines should resolve");
850        assert_eq!(preview, Some(READ_FILE_DEFAULT_PREVIEW_LINES));
851
852        let head = resolve_read_file_line_limit(ReadFileMode::Head, None)
853            .expect("head lines should resolve");
854        assert_eq!(head, Some(READ_FILE_DEFAULT_PREVIEW_LINES));
855
856        let tail = resolve_read_file_line_limit(ReadFileMode::Tail, None)
857            .expect("tail lines should resolve");
858        assert_eq!(tail, Some(READ_FILE_DEFAULT_PREVIEW_LINES));
859    }
860
861    #[test]
862    fn test_resolve_read_file_line_limit_for_full_ignores_lines() {
863        let full = resolve_read_file_line_limit(ReadFileMode::Full, Some(123))
864            .expect("full mode should ignore lines");
865        assert_eq!(full, None);
866    }
867
868    #[test]
869    fn test_resolve_read_file_line_limit_rejects_zero() {
870        let err = resolve_read_file_line_limit(ReadFileMode::Head, Some(0)).unwrap_err();
871        assert!(err.contains("positive"));
872    }
873
874    #[test]
875    fn test_resolve_read_file_line_limit_rejects_too_large() {
876        let err =
877            resolve_read_file_line_limit(ReadFileMode::Tail, Some(READ_FILE_MAX_LINE_WINDOW + 1))
878                .unwrap_err();
879        assert!(err.contains("<="));
880    }
881
882    #[test]
883    fn test_sanitize_read_file_stderr_snippet_normalizes_whitespace_and_controls() {
884        let stderr = "line1\nline2\t\u{0007}bad\rline3";
885        let snippet = sanitize_read_file_stderr_snippet(stderr)
886            .expect("snippet should be present for non-empty stderr");
887        assert_eq!(snippet, "line1 line2 bad line3");
888    }
889
890    #[test]
891    fn test_background_json_err_omits_unregistered_job_fields() {
892        let long_error = "e".repeat(BACKGROUND_JSON_SNIPPET_LIMIT_CHARS + 10);
893        let long_stderr = "s".repeat(BACKGROUND_JSON_SNIPPET_LIMIT_CHARS + 10);
894
895        let result = background_json_err(&long_error, &long_stderr);
896        let text = extract_text_from_result(&result);
897
898        let value: serde_json::Value =
899            serde_json::from_str(text.trim()).expect("background_json_err should return JSON");
900
901        assert_eq!(value.get("ok").and_then(|v| v.as_bool()), Some(false));
902        assert_eq!(
903            value.get("background").and_then(|v| v.as_bool()),
904            Some(true)
905        );
906        assert_eq!(value.get("truncated").and_then(|v| v.as_bool()), Some(true));
907        assert!(value.get("job_id").is_none());
908        assert!(value.get("log_path").is_none());
909        assert!(value.get("hint").is_none());
910
911        let fields = value
912            .get("truncated_fields")
913            .expect("expected truncated_fields");
914        assert_eq!(fields.get("error").and_then(|v| v.as_bool()), Some(true));
915        assert_eq!(fields.get("stderr").and_then(|v| v.as_bool()), Some(true));
916
917        let error_snippet = value
918            .get("error")
919            .and_then(|v| v.as_str())
920            .expect("expected error field");
921        assert_eq!(
922            error_snippet.chars().count(),
923            BACKGROUND_JSON_SNIPPET_LIMIT_CHARS
924        );
925        let stderr_snippet = value
926            .get("stderr")
927            .and_then(|v| v.as_str())
928            .expect("expected stderr field");
929        assert_eq!(
930            stderr_snippet.chars().count(),
931            BACKGROUND_JSON_SNIPPET_LIMIT_CHARS
932        );
933    }
934
935    #[test]
936    fn test_background_json_timeout_hint_contains_pid_and_check_process_tool() {
937        let result = background_json_timeout(
938            "job-42",
939            4242,
940            "/tmp/ssh-mcp/local.log",
941            &crate::background::response::BackgroundTimeoutSnapshot {
942                state: "running",
943                still_running: true,
944                exit_code: None,
945                state_reason: None,
946                elapsed_time: "00:01",
947                log_exists: true,
948                log_tail: "tail line",
949                tail_lines_used: 50,
950            },
951        );
952        let text = extract_text_from_result(&result);
953
954        let value: serde_json::Value =
955            serde_json::from_str(text.trim()).expect("background_json_timeout should return JSON");
956
957        assert_eq!(value.get("ok").and_then(|v| v.as_bool()), Some(false));
958        assert_eq!(value.get("timeout").and_then(|v| v.as_bool()), Some(true));
959        assert_eq!(
960            value.get("background").and_then(|v| v.as_bool()),
961            Some(true)
962        );
963        assert_eq!(
964            value.get("still_running").and_then(|v| v.as_bool()),
965            Some(true)
966        );
967        assert_eq!(value.get("state").and_then(|v| v.as_str()), Some("running"));
968        assert_eq!(
969            value.get("tail_lines_used").and_then(|v| v.as_u64()),
970            Some(50)
971        );
972        assert_eq!(
973            value.get("elapsed_time").and_then(|v| v.as_str()),
974            Some("00:01")
975        );
976        assert_eq!(
977            value.get("log_tail").and_then(|v| v.as_str()),
978            Some("tail line")
979        );
980
981        let hint = value
982            .get("hint")
983            .and_then(|v| v.as_str())
984            .expect("expected hint field");
985
986        // Hint should contain the actual job_id value
987        assert!(
988            hint.contains("job_id=job-42"),
989            "hint should contain the actual job_id value; got: '{hint}'"
990        );
991        // Hint should mention the check_process tool
992        assert!(
993            hint.contains("check_process"),
994            "hint should mention check_process tool; got: '{hint}'"
995        );
996        // Hint should warn against restarting
997        assert!(
998            hint.contains("DO NOT restart"),
999            "hint should warn against restarting; got: '{hint}'"
1000        );
1001        // Hint should use TIMEOUT_RECOVERY prefix
1002        assert!(
1003            hint.contains("TIMEOUT_RECOVERY"),
1004            "hint should start with TIMEOUT_RECOVERY; got: '{hint}'"
1005        );
1006        assert!(
1007            hint.contains("MCP client deadlines may be shorter than timeout_ms"),
1008            "hint should distinguish the client deadline from timeout_ms; got: '{hint}'"
1009        );
1010        assert!(
1011            hint.contains("background=true"),
1012            "hint should recommend explicit background mode; got: '{hint}'"
1013        );
1014        // Hint should NOT contain old placeholders
1015        assert!(
1016            !hint.contains("<pid>"),
1017            "hint should not contain <pid> placeholder; got: '{hint}'"
1018        );
1019        assert!(
1020            !hint.contains("<log_path>"),
1021            "hint should not contain <log_path> placeholder; got: '{hint}'"
1022        );
1023    }
1024
1025    #[test]
1026    fn test_tool_documentation_available() {
1027        // Verify that extended documentation is available for all tools
1028        assert!(SshMcpServer::get_tool_documentation("shell").is_some());
1029        assert!(SshMcpServer::get_tool_documentation("sudo_shell").is_some());
1030        assert!(SshMcpServer::get_tool_documentation("transfer").is_some());
1031        assert!(SshMcpServer::get_tool_documentation("read").is_some());
1032        assert!(SshMcpServer::get_tool_documentation("apply_patch").is_some());
1033        assert!(SshMcpServer::get_tool_documentation("sudo_apply_patch").is_some());
1034        assert!(SshMcpServer::get_tool_documentation("write-file").is_none());
1035        assert!(SshMcpServer::get_tool_documentation("replace-in-file").is_none());
1036        assert!(SshMcpServer::get_tool_documentation("unknown").is_none());
1037    }
1038
1039    #[test]
1040    fn test_shell_documentation_content() {
1041        let docs = SshMcpServer::get_tool_documentation("shell").unwrap();
1042        assert!(docs.contains("SHELL TOOL"));
1043        assert!(docs.contains("PARAMETERS:"));
1044        assert!(docs.contains("BACKGROUND MODE:"));
1045        assert!(docs.contains("command"));
1046        assert!(docs.contains("background"));
1047        assert!(docs.contains("still_running"));
1048        assert!(docs.contains("not the full tool-call deadline"));
1049        assert!(docs.contains("client may stop waiting earlier"));
1050    }
1051
1052    #[test]
1053    fn test_sudo_shell_documentation_content() {
1054        let docs = SshMcpServer::get_tool_documentation("sudo_shell").unwrap();
1055        assert!(docs.contains("SUDO_SHELL TOOL"));
1056        assert!(docs.contains("sudo"));
1057        assert!(docs.contains("not the full tool-call deadline"));
1058    }
1059
1060    #[test]
1061    fn test_transfer_documentation_content() {
1062        let docs = SshMcpServer::get_tool_documentation("transfer").unwrap();
1063        assert!(docs.contains("TRANSFER TOOL"));
1064        assert!(docs.contains("put"));
1065        assert!(docs.contains("get"));
1066        assert!(docs.contains("TRANSPORTS:"));
1067    }
1068
1069    #[test]
1070    fn test_read_file_documentation_content() {
1071        let docs = SshMcpServer::get_tool_documentation("read").unwrap();
1072        assert!(docs.contains("READ TOOL"));
1073        assert!(docs.contains("remote_path"));
1074        assert!(docs.contains("mode"));
1075        assert!(docs.contains("UTF-8"));
1076    }
1077
1078    #[test]
1079    fn test_apply_patch_documentation_content() {
1080        let docs = SshMcpServer::get_tool_documentation("apply_patch").unwrap();
1081        assert!(docs.contains("APPLY_PATCH TOOL"));
1082        assert!(docs.contains("Add File"));
1083        assert!(docs.contains("Delete File"));
1084    }
1085
1086    #[test]
1087    fn test_sudo_apply_patch_documentation_content() {
1088        let docs = SshMcpServer::get_tool_documentation("sudo_apply_patch").unwrap();
1089        assert!(docs.contains("SUDO_APPLY_PATCH TOOL"));
1090        assert!(docs.contains("sudo"));
1091    }
1092
1093    #[test]
1094    fn test_compact_tool_descriptions() {
1095        // Verify that tool descriptions are compact (not verbose)
1096        let shell = SshMcpServer::shell_tool();
1097        let sudo_shell = SshMcpServer::sudo_shell_tool();
1098        let transfer = SshMcpServer::transfer_tool();
1099        let read_file = SshMcpServer::read_file_tool();
1100        let apply_patch = SshMcpServer::apply_patch_tool();
1101        let sudo_apply_patch = SshMcpServer::sudo_apply_patch_tool();
1102
1103        // Descriptions should be present but concise (under 100 chars)
1104        if let Some(desc) = shell.description {
1105            assert!(
1106                desc.len() < 100,
1107                "shell description too long: {} chars",
1108                desc.len()
1109            );
1110        }
1111        if let Some(desc) = sudo_shell.description {
1112            assert!(
1113                desc.len() < 100,
1114                "sudo_shell description too long: {} chars",
1115                desc.len()
1116            );
1117        }
1118        if let Some(desc) = transfer.description {
1119            assert!(
1120                desc.len() < 100,
1121                "transfer description too long: {} chars",
1122                desc.len()
1123            );
1124        }
1125        if let Some(desc) = read_file.description {
1126            assert!(
1127                desc.len() < 100,
1128                "read description too long: {} chars",
1129                desc.len()
1130            );
1131        }
1132        if let Some(desc) = apply_patch.description {
1133            assert!(
1134                desc.len() < 100,
1135                "apply_patch description too long: {} chars",
1136                desc.len()
1137            );
1138        }
1139        if let Some(desc) = sudo_apply_patch.description {
1140            assert!(
1141                desc.len() < 100,
1142                "sudo_apply_patch description too long: {} chars",
1143                desc.len()
1144            );
1145        }
1146    }
1147}