From f9b10ea9ca3ac9b01f65334724ff9d21c0457622 Mon Sep 17 00:00:00 2001 From: blackax Date: Tue, 29 Sep 2026 21:42:41 -0700 Subject: [PATCH 1/2] refactor: behaviour-preserving /simplify pass over src/ --- pyproject.toml | 2 +- src/ssh_mcp/config.py | 99 +++++++++++++++----------------------- src/ssh_mcp/formatting.py | 33 +++---------- src/ssh_mcp/healthcheck.py | 9 +--- src/ssh_mcp/server.py | 22 ++------- src/ssh_mcp/ssh.py | 29 +++++------ 6 files changed, 63 insertions(+), 131 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 7558a9a..4edf57f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -187,7 +187,7 @@ constraint-dependencies = [ # fires on a `# noqa: F401` that is perfectly legitimate under this config, # purely because `--isolated` re-derives which rules are on. Under the real # select it fires on exactly two directives — the `# noqa: S310` pair in -# healthcheck.py:152,162 — which are inert only because `S` is not selected. +# healthcheck.py::_check_http — which are inert only because `S` is not selected. # They document a deliberate `urllib.request` decision, so enabling RUF100 # would force deleting intent-documenting comments to satisfy a linter. No # safety is bought; leave it off. diff --git a/src/ssh_mcp/config.py b/src/ssh_mcp/config.py index 1dff695..cacc324 100644 --- a/src/ssh_mcp/config.py +++ b/src/ssh_mcp/config.py @@ -161,8 +161,7 @@ def servers_in_group(self, group_name: str) -> list[ServerConfig]: Raises: KeyError: If group name not found """ - # Validate group exists first - self.get_group(group_name) + self.get_group(group_name) # raises KeyError for an undefined group return [ server for server in self._servers.values() if group_name in server.groups ] @@ -235,47 +234,37 @@ def _load(self) -> None: valid = ", ".join(_valid_keys(Settings)) raise ConfigError(f"{detail}. Valid keys: {valid}") from e - # Warn if known_hosts verification is disabled if not self._settings.known_hosts: logger.warning( "known_hosts verification disabled - connections vulnerable to MITM attacks" ) - # Load groups - if "groups" in config_data: - for group_name, group_data in config_data["groups"].items(): - try: - self._groups[group_name] = GroupConfig( - name=group_name, - **dict(group_data), - ) - except ValidationError as e: - detail = _format_validation_error("groups", group_name, e) - valid = ", ".join( - k for k in _valid_keys(GroupConfig) if k != "name" - ) - raise ConfigError(f"{detail}. Valid keys: {valid}") from e - - # Load servers - if "servers" in config_data: - for server_name, server_data in config_data["servers"].items(): - # Convert groups list to tuple before Pydantic sees it - data: dict[str, Any] = dict(server_data) - if "groups" in data: - data["groups"] = tuple(data["groups"]) - try: - self._servers[server_name] = ServerConfig( - name=server_name, - **data, - ) - except ValidationError as e: - detail = _format_validation_error("servers", server_name, e) - valid = ", ".join( - k for k in _valid_keys(ServerConfig) if k != "name" - ) - raise ConfigError(f"{detail}. Valid keys: {valid}") from e + for group_name, group_data in config_data.get("groups", {}).items(): + try: + self._groups[group_name] = GroupConfig( + name=group_name, + **dict(group_data), + ) + except ValidationError as e: + detail = _format_validation_error("groups", group_name, e) + valid = ", ".join(k for k in _valid_keys(GroupConfig) if k != "name") + raise ConfigError(f"{detail}. Valid keys: {valid}") from e + + for server_name, server_data in config_data.get("servers", {}).items(): + # Convert groups list to tuple before Pydantic sees it + data: dict[str, Any] = dict(server_data) + if "groups" in data: + data["groups"] = tuple(data["groups"]) + try: + self._servers[server_name] = ServerConfig( + name=server_name, + **data, + ) + except ValidationError as e: + detail = _format_validation_error("servers", server_name, e) + valid = ", ".join(k for k in _valid_keys(ServerConfig) if k != "name") + raise ConfigError(f"{detail}. Valid keys: {valid}") from e - # Validate configuration self._validate() def _validate(self) -> None: @@ -289,9 +278,6 @@ def _validate(self) -> None: ValueError: If a circular jump host chain is detected. This is the only fatal check here, so it aborts ``__init__``. """ - server_names = set(self._servers.keys()) - group_names = set(self._groups.keys()) - for server_name, server in self._servers.items(): # A server should reference at least one group (advisory only) if not server.groups: @@ -299,7 +285,7 @@ def _validate(self) -> None: # A group referenced by a server should be defined (advisory only) for group in server.groups: - if group not in group_names: + if group not in self._groups: logger.warning( "Server '%s' references undefined group '%s'", server_name, @@ -307,11 +293,11 @@ def _validate(self) -> None: ) # Server names should not collide with group names (advisory only) - if server_name in group_names: + if server_name in self._groups: logger.warning("Server name '%s' collides with group name", server_name) # jump_host should reference another defined server (advisory only) - if server.jump_host and server.jump_host not in server_names: + if server.jump_host and server.jump_host not in self._servers: logger.warning( "Server '%s' references undefined jump_host '%s'", server_name, @@ -320,20 +306,13 @@ def _validate(self) -> None: # Detect circular jump host chains for server_name, server in self._servers.items(): - if server.jump_host: - visited = {server_name} - current = server.jump_host - path = [server_name, current] - - while current in self._servers: - if current in visited: - cycle_path = " -> ".join(path) - raise ValueError(f"Circular jump host chain: {cycle_path}") - - visited.add(current) - next_jump = self._servers[current].jump_host - if next_jump: - path.append(next_jump) - current = next_jump - else: - break + visited = {server_name} + path = [server_name] + current = server.jump_host + while current and current in self._servers: + path.append(current) + if current in visited: + cycle_path = " -> ".join(path) + raise ValueError(f"Circular jump host chain: {cycle_path}") + visited.add(current) + current = self._servers[current].jump_host diff --git a/src/ssh_mcp/formatting.py b/src/ssh_mcp/formatting.py index 2e3a1de..fcb2183 100644 --- a/src/ssh_mcp/formatting.py +++ b/src/ssh_mcp/formatting.py @@ -33,32 +33,24 @@ def format_server_table(servers: list[ServerConfig], filter_label: str = "") -> if not servers: return "No servers found." - # Calculate column widths max_name = max(len("SERVER"), max(len(s.name) for s in servers)) - max_groups = max( - len("GROUPS"), max(len(", ".join(s.groups)) if s.groups else 0 for s in servers) - ) + max_groups = max(len("GROUPS"), max(len(", ".join(s.groups)) for s in servers)) - # Build header lines = [ f"{'SERVER':<{max_name}} {'GROUPS':<{max_groups}} DESCRIPTION", ] - # Build rows for server in servers: - groups_str = ", ".join(server.groups) if server.groups else "" + groups_str = ", ".join(server.groups) lines.append( f"{server.name:<{max_name}} {groups_str:<{max_groups}} {server.description}" ) - # Build footer lines.append("") count = len(servers) plural = "server" if count == 1 else "servers" - if filter_label: - lines.append(f"Total: {count} {plural} ({filter_label})") - else: - lines.append(f"Total: {count} {plural}") + suffix = f" ({filter_label})" if filter_label else "" + lines.append(f"Total: {count} {plural}{suffix}") return "\n".join(lines) @@ -84,18 +76,15 @@ def format_group_table(groups: list[GroupConfig], server_counts: dict[str, int]) if not groups: return "No groups found." - # Calculate column widths max_name = max(len("GROUP"), max(len(g.name) for g in groups)) max_count = max( len("SERVERS"), max(len(str(server_counts.get(g.name, 0))) for g in groups) ) - # Build header lines = [ f"{'GROUP':<{max_name}} {'SERVERS':<{max_count}} DESCRIPTION", ] - # Build rows for group in groups: count = server_counts.get(group.name, 0) lines.append( @@ -158,11 +147,9 @@ def format_exec_result(result: ExecResult) -> str: """ lines = [f"[{result.server}] $ {result.command}"] - # Add stdout if present if result.stdout: lines.append(result.stdout) - # Add stderr if present if result.stderr: lines.append("") lines.append("STDERR:") @@ -243,21 +230,13 @@ def format_group_results(results: list[ExecResult], group_name: str) -> str: "", ] - # Track success/failure counts succeeded = 0 - failed = 0 - - # Format each result for result in results: if result.error: - failed += 1 lines.append(f"[{result.server}] ERROR: {result.error}") else: - is_success = result.exit_code == 0 - if is_success: + if result.exit_code == 0: succeeded += 1 - else: - failed += 1 exit_code = result.exit_code if result.exit_code is not None else "unknown" lines.append( @@ -272,7 +251,7 @@ def format_group_results(results: list[ExecResult], group_name: str) -> str: lines.append("") - # Add summary + failed = len(results) - succeeded lines.append(f"Summary: {succeeded} succeeded, {failed} failed") return "\n".join(lines) diff --git a/src/ssh_mcp/healthcheck.py b/src/ssh_mcp/healthcheck.py index 0f6a101..797c2d5 100644 --- a/src/ssh_mcp/healthcheck.py +++ b/src/ssh_mcp/healthcheck.py @@ -82,11 +82,6 @@ def _check_stdio() -> tuple[bool, str]: Returns (ok, diagnostic). """ - try: - import ssh_mcp # noqa: F401 - except ImportError as e: - return False, f"import failed: {e}" - try: from ssh_mcp.server import _get_config_path except ImportError as e: @@ -164,9 +159,7 @@ def _check_http() -> tuple[bool, str]: except urllib.error.HTTPError as e: # Any non-5xx status means the server is alive but the request was rejected # (wrong auth, wrong protocol version, etc.) — still healthy. - if e.code < 500: - return True, f"http {e.code}" - return False, f"http {e.code}" + return e.code < 500, f"http {e.code}" except urllib.error.URLError as e: return False, f"connect failed: {type(e.reason).__name__}" except Exception as e: # noqa: BLE001 - an unknown probe failure is 'unhealthy', not a crash of the healthcheck itself diff --git a/src/ssh_mcp/server.py b/src/ssh_mcp/server.py index ca4902a..672ae9e 100644 --- a/src/ssh_mcp/server.py +++ b/src/ssh_mcp/server.py @@ -283,7 +283,6 @@ def _cleanup_connections() -> None: backup that fires only if the lifespan didn't run (e.g. stdio mode or abnormal exit). """ - global _ssh if _ssh is None: return try: @@ -369,9 +368,7 @@ async def wrapper(*args: Any, **kwargs: Any) -> str: with _span(f"mcp.tool.{tool_name}", **{"mcp.tool.name": tool_name}): await _init() return await func(*args, **kwargs) - except ToolError: - raise - except asyncio.CancelledError: + except (ToolError, asyncio.CancelledError): raise except Exception as e: # Redact BOTH the message and the traceback. `exc_info=True` @@ -414,7 +411,6 @@ async def list_servers(group: str | None = None) -> str: registry = _get_registry() if group is not None: - # Filter by group try: servers = registry.servers_in_group(group) except KeyError as e: @@ -424,7 +420,6 @@ async def list_servers(group: str | None = None) -> str: if not servers: return f"No servers found in group '{group}'" else: - # Show all servers servers = registry.all_servers() filter_label = "" @@ -443,11 +438,9 @@ async def list_groups() -> str: groups = registry.all_groups() - # Count servers per group - server_counts = {} - for group in groups: - count = len(registry.servers_in_group(group.name)) - server_counts[group.name] = count + server_counts = { + group.name: len(registry.servers_in_group(group.name)) for group in groups + } return format_group_table(groups, server_counts) @@ -716,8 +709,6 @@ def _build_http_app( Returns a ``Starlette`` instance ready to hand to ``uvicorn.run``. """ - from contextlib import asynccontextmanager - from starlette.applications import Starlette from starlette.routing import Mount @@ -762,7 +753,7 @@ def _build_http_app( "in server.py._build_http_app or pin a known-good mcp." ) from exc - @asynccontextmanager + @contextlib.asynccontextmanager async def _lifespan(_app: Starlette) -> AsyncGenerator[None]: # Step 1: start the MCPServer session manager's task group via the # public session_manager.run() API (see note above). @@ -772,7 +763,6 @@ async def _lifespan(_app: Starlette) -> AsyncGenerator[None]: finally: # Step 2: drain SSH BEFORE exiting the inner lifespan # so close_all() can still dispatch on a live event loop. - global _ssh if _ssh is not None: logger.info("Draining SSH connections on HTTP shutdown") try: @@ -1538,8 +1528,6 @@ def main() -> None: run_healthcheck() # exits 0 or 1 return # unreachable but keeps mypy happy - from ssh_mcp import __version__ - transport = os.environ.get("SSH_MCP_TRANSPORT", "stdio").strip().lower() if transport in ("http", "streamable-http"): diff --git a/src/ssh_mcp/ssh.py b/src/ssh_mcp/ssh.py index 9fc6275..3aec6a4 100644 --- a/src/ssh_mcp/ssh.py +++ b/src/ssh_mcp/ssh.py @@ -1614,10 +1614,8 @@ async def _execute_impl( # Prepend working directory if specified effective_command = command - if working_dir: - effective_command = f"cd {shlex.quote(working_dir)} && {command}" - elif server.default_dir: - effective_command = f"cd {shlex.quote(server.default_dir)} && {command}" + if cd_dir := working_dir or server.default_dir: + effective_command = f"cd {shlex.quote(cd_dir)} && {command}" # Get or create connection conn = await self._get_connection(server_name) @@ -1924,10 +1922,8 @@ async def execute_on_group( # of the configured limit. Behaviour change, stated plainly: # independent execute_on_group() calls now serialise against # each other for their share of the shared pool. - semaphore = self._group_semaphore - async def execute_with_semaphore(server: ServerConfig) -> ExecResult: - async with semaphore: + async with self._group_semaphore: return await self.execute( server.name, command, @@ -2509,7 +2505,6 @@ async def _download_impl( async with sftp.open(remote_path, "rb") as remote_file: with os.fdopen(local_fd, "wb", closefd=True) as local_file: local_fd = None # ownership transferred - offset = 0 while True: # asyncssh's SFTPClientFile.read() is # annotated `-> AnyStr` without the @@ -2521,11 +2516,10 @@ async def _download_impl( # of a blanket ignore; encoding=None # (set automatically by the 'b' in # "rb") guarantees bytes at runtime. - data: bytes = await remote_file.read(block, offset) + data: bytes = await remote_file.read(block, written) if not data: break await asyncio.to_thread(local_file.write, data) - offset += len(data) written += len(data) duration_ms = int((time.monotonic() - start_time) * 1000) @@ -2792,11 +2786,10 @@ async def _create_connection( # Create connection try: - conn = await asyncio.wait_for( + return await asyncio.wait_for( asyncssh.connect(host, **connect_params), timeout=self.settings.command_timeout, ) - return conn # `except asyncssh.PermissionDenied` MUST precede # `except asyncssh.DisconnectError`: PermissionDenied is a # SUBCLASS of DisconnectError (verified against asyncssh 2.24.0 — @@ -2882,11 +2875,11 @@ async def _eviction_loop(self) -> None: idle_threshold = self.settings.connection_idle_timeout # Find idle connections - to_evict = [] - for server_name, last_used in self._last_used.items(): - idle_time = now - last_used - if idle_time > idle_threshold: - to_evict.append((server_name, idle_time)) + to_evict = [ + name + for name, last_used in self._last_used.items() + if now - last_used > idle_threshold + ] logger.info( "Connection pool: %d active, %d locks", @@ -2895,7 +2888,7 @@ async def _eviction_loop(self) -> None: ) # Evict idle connections - for server_name, idle_time in to_evict: + for server_name in to_evict: # Get lock for this server (if it exists) lock = self._locks.get(server_name) if lock is None: From a62331f0458721a4dbe79b2c47f99e1ef22bf540 Mon Sep 17 00:00:00 2001 From: blackax Date: Tue, 29 Sep 2026 22:42:20 -0700 Subject: [PATCH 2/2] refactor: address PR #67 review readability nits --- src/ssh_mcp/formatting.py | 1 + src/ssh_mcp/ssh.py | 8 ++++++-- 2 files changed, 7 insertions(+), 2 deletions(-) diff --git a/src/ssh_mcp/formatting.py b/src/ssh_mcp/formatting.py index fcb2183..74fba3f 100644 --- a/src/ssh_mcp/formatting.py +++ b/src/ssh_mcp/formatting.py @@ -251,6 +251,7 @@ def format_group_results(results: list[ExecResult], group_name: str) -> str: lines.append("") + # Anything not counted as succeeded failed: an error, a non-zero or unknown exit. failed = len(results) - succeeded lines.append(f"Summary: {succeeded} succeeded, {failed} failed") diff --git a/src/ssh_mcp/ssh.py b/src/ssh_mcp/ssh.py index 3aec6a4..e5896b6 100644 --- a/src/ssh_mcp/ssh.py +++ b/src/ssh_mcp/ssh.py @@ -1612,9 +1612,10 @@ async def _execute_impl( # Use server-specific timeout if configured effective_timeout = server.timeout or timeout - # Prepend working directory if specified + # Prepend working_dir if given, else the server's default_dir effective_command = command - if cd_dir := working_dir or server.default_dir: + cd_dir = working_dir or server.default_dir + if cd_dir: effective_command = f"cd {shlex.quote(cd_dir)} && {command}" # Get or create connection @@ -2516,6 +2517,9 @@ async def _download_impl( # of a blanket ignore; encoding=None # (set automatically by the 'b' in # "rb") guarantees bytes at runtime. + # + # `written` doubles as the remote read + # offset: both advance by len(data). data: bytes = await remote_file.read(block, written) if not data: break