diff --git a/src/remoteshell_mcp/ssh_client.py b/src/remoteshell_mcp/ssh_client.py index 6261567..892197a 100644 --- a/src/remoteshell_mcp/ssh_client.py +++ b/src/remoteshell_mcp/ssh_client.py @@ -35,7 +35,8 @@ def __init__( port: int = 22, password: Optional[str] = None, private_key: Optional[str] = None, - timeout: int = 30 + timeout: int = 30, + keepalive_interval: int = 30 ): """ Initialize SSH client. @@ -47,6 +48,7 @@ def __init__( password: Password for authentication (optional) key_path: Path to SSH private key file (optional) timeout: Connection timeout in seconds + keepalive_interval: Seconds between SSH keepalive packets (0 disables) """ self.host = host self.user = user @@ -54,6 +56,7 @@ def __init__( self.password = password self.private_key = private_key self.timeout = timeout + self.keepalive_interval = keepalive_interval self._client: Optional[SSHClient] = None self._sftp: Optional[paramiko.SFTPClient] = None @@ -115,7 +118,13 @@ def connect(self) -> None: # Connect self._client.connect(**connect_kwargs) - + + # Enable transport-level keepalive so idle connections are not + # dropped by firewalls or sshd idle timeouts. + transport = self._client.get_transport() + if transport is not None: + transport.set_keepalive(self.keepalive_interval) + except AuthenticationException as e: raise SSHConnectionError(f"Authentication failed: {e}") except SSHException as e: @@ -175,21 +184,30 @@ def execute_command( """ self.ensure_connected() + if working_dir: + command = f"cd {working_dir} && {command}" + try: - # Add working directory change if specified - if working_dir: - command = f"cd {working_dir} && {command}" - - # Execute command - stdin, stdout, stderr = self._client.exec_command( + stdin, stdout, stderr = self._exec_command_with_retry( command, timeout=timeout or self.timeout ) - # Read output - stdout_data = stdout.read().decode('utf-8', errors='replace') - stderr_data = stderr.read().decode('utf-8', errors='replace') - exit_code = stdout.channel.recv_exit_status() + channel = stdout.channel + try: + # Read output + stdout_data = stdout.read().decode('utf-8', errors='replace') + stderr_data = stderr.read().decode('utf-8', errors='replace') + exit_code = channel.recv_exit_status() + finally: + # Always close the channel so the server-side session slot is + # released. Without this, open sessions accumulate until the + # server's MaxSessions limit is hit and it rejects new + # channels (e.g. "Administratively prohibited"). + try: + channel.close() + except Exception: + pass return { "stdout": stdout_data, @@ -201,6 +219,25 @@ def execute_command( except Exception as e: raise SSHCommandError(f"Command execution failed: {e}") + def _exec_command_with_retry( + self, + command: str, + timeout: int + ): + """Open an exec channel, reconnecting once if the session open fails. + + Safe to retry: if opening the channel fails, the command has not + started executing on the remote host. + """ + try: + return self._client.exec_command(command, timeout=timeout) + except SSHException: + # The transport or session table is broken (stale connection, + # MaxSessions exhaustion, ...). Reconnect once and retry. + self.disconnect() + self.connect() + return self._client.exec_command(command, timeout=timeout) + def _get_sftp(self) -> paramiko.SFTPClient: """Get or create SFTP client.""" self.ensure_connected() diff --git a/tests/test_ssh_client.py b/tests/test_ssh_client.py new file mode 100644 index 0000000..d70c77a --- /dev/null +++ b/tests/test_ssh_client.py @@ -0,0 +1,144 @@ +"""Tests for RemoteSSHClient session/channel handling without real SSH.""" + +from __future__ import annotations + +from unittest.mock import MagicMock, patch + +import pytest + +from remoteshell_mcp import ssh_client +from remoteshell_mcp.ssh_client import RemoteSSHClient, SSHCommandError + + +def _make_fake_exec(open_channels: list): + """Build a fake exec_command that tracks open channels like a transport.""" + from paramiko.ssh_exception import SSHException + + max_sessions = 2 # mimic a small sshd MaxSessions limit + + def fake_exec_command(self, command, timeout=None, environment=None): + if len(open_channels) >= max_sessions: + raise SSHException("ChannelException: Administratively prohibited") + chan = MagicMock() + chan.closed = False + + def close(): + chan.closed = True + if chan in open_channels: + open_channels.remove(chan) + + chan.close.side_effect = close + stdin = MagicMock() + stdout = MagicMock() + stdout.read.return_value = b"ok" + stdout.channel = chan + stderr = MagicMock() + stderr.read.return_value = b"" + chan.recv_exit_status.return_value = 0 + open_channels.append(chan) + return stdin, stdout, stderr + + return fake_exec_command + + +@pytest.fixture +def client(): + c = RemoteSSHClient(host="h", user="u", password="p") + c._client = MagicMock() + c._client.exec_command = MagicMock( + wraps=_make_fake_exec(open_channels=[]).__get__(c._client) + ) + # Pretend the connection is healthy so ensure_connected() is a no-op. + with patch.object(RemoteSSHClient, "is_connected", return_value=True): + yield c + + +def test_execute_command_closes_channel(client: RemoteSSHClient): + open_channels = [] + + def fake_exec(command, timeout=None, environment=None): + chan = MagicMock() + stdout = MagicMock() + stdout.read.return_value = b"ok" + stdout.channel = chan + stderr = MagicMock() + stderr.read.return_value = b"" + chan.recv_exit_status.return_value = 0 + open_channels.append(chan) + return MagicMock(), stdout, stderr + + client._client.exec_command.side_effect = fake_exec + + for _ in range(10): + result = client.execute_command("echo hi") + assert result["success"] is True + # Channel must be closed after each execution. + assert open_channels[-1].close.called + + assert len(open_channels) == 10 + + +def test_many_commands_do_not_exhaust_sessions(client: RemoteSSHClient): + """Simulate sshd MaxSessions=2: unclosed channels would fail run 3+.""" + open_channels: list = [] + client._client.exec_command.side_effect = _make_fake_exec(open_channels).__get__( + client._client + ) + + for _ in range(20): + result = client.execute_command("echo hi") + assert result["exit_code"] == 0 + + assert open_channels == [] + + +def test_connect_enables_keepalive(): + with patch.object(ssh_client.SSHClient, "connect"): + with patch.object( + ssh_client.SSHClient, "get_transport", return_value=MagicMock() + ) as gt: + c = RemoteSSHClient(host="h", user="u", password="p", keepalive_interval=15) + c.connect() + gt.assert_called_once() + transport = gt.return_value + transport.set_keepalive.assert_called_once_with(15) + + +def test_exec_retry_reconnects_on_ssh_exception(client: RemoteSSHClient): + """A broken session open triggers one reconnect and a safe retry.""" + from paramiko.ssh_exception import SSHException + + open_channels: list = [] + good = _make_fake_exec(open_channels).__get__(client._client) + calls = {"n": 0} + + def flaky(command, timeout=None, environment=None): + calls["n"] += 1 + if calls["n"] == 1: + raise SSHException("Socket is closed") + return good(command, timeout=timeout) + + client._client.exec_command.side_effect = flaky + + with patch.object(RemoteSSHClient, "disconnect") as disconnect: + with patch.object(RemoteSSHClient, "connect") as connect: + result = client.execute_command("echo hi") + + assert result["success"] is True + assert calls["n"] == 2 + disconnect.assert_called_once() + connect.assert_called_once() + + +def test_exec_retry_gives_up_after_second_failure(client: RemoteSSHClient): + from paramiko.ssh_exception import SSHException + + def always_fail(command, timeout=None, environment=None): + raise SSHException("boom") + + client._client.exec_command.side_effect = always_fail + + with patch.object(RemoteSSHClient, "disconnect"): + with patch.object(RemoteSSHClient, "connect"): + with pytest.raises(SSHCommandError): + client.execute_command("echo hi") diff --git a/uv.lock b/uv.lock index 84baf96..6864ed3 100644 --- a/uv.lock +++ b/uv.lock @@ -1187,7 +1187,7 @@ wheels = [ [[package]] name = "remoteshell-mcp" -version = "0.2.0" +version = "1.1.1" source = { editable = "." } dependencies = [ { name = "fastmcp" },