diff --git a/mkdocs.yml b/mkdocs.yml index 89cd591efc..10d57b53a2 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -360,6 +360,7 @@ nav: - dstack ps: docs/reference/cli/dstack/ps.md - dstack stop: docs/reference/cli/dstack/stop.md - dstack attach: docs/reference/cli/dstack/attach.md + - dstack connect: docs/reference/cli/dstack/connect.md - dstack login: docs/reference/cli/dstack/login.md - dstack logs: docs/reference/cli/dstack/logs.md - dstack metrics: docs/reference/cli/dstack/metrics.md diff --git a/mkdocs/docs/guides/cli-api.md b/mkdocs/docs/guides/cli-api.md index b451dc281b..cb708ce188 100644 --- a/mkdocs/docs/guides/cli-api.md +++ b/mkdocs/docs/guides/cli-api.md @@ -215,6 +215,33 @@ $ dstack attach <run name> Use `--job JOB_NUMBER` with `dstack attach` to attach to another job. Ports declared in the run configuration are forwarded while attached. +### Connect without attaching + +When the server has the [SSH proxy](server-deployment.md#ssh-proxy) enabled, +the owner of a running dev environment can connect without keeping an attach +process running: + +
+ +```shell +$ dstack connect <run name> +``` + +
+ +Interactive SSH is the default. Use `--ide` to open the IDE configured by the +dev environment, or `--vscode`, `--cursor`, `--windsurf`, or `--zed` to require +a specific configured IDE. + +`dstack connect` refreshes the run, validates the complete proxy connection +record, and updates a local SSH alias using the current user's built-in key. It +does not fall back to direct host SSH. Use `--expect-sshproxy HOST[:PORT]` to +require a particular proxy endpoint. + +Unlike `dstack attach`, this command does not forward ports or stream logs. See +the [`dstack connect` reference](../reference/cli/dstack/connect.md) for the +connection and validation details. + ??? info "User SSH keys" The server stores a built-in SSH key pair for each user. @@ -269,7 +296,6 @@ available in the CLI or for integrations that need to call the server directly. Out of scope for this guide: - Git repo workflows - File upload workflows -- Direct SSH/sshproxy workflows - Backward compatibility between older clients and servers --> diff --git a/mkdocs/docs/guides/server-deployment.md b/mkdocs/docs/guides/server-deployment.md index caf85f9da7..fd526ce0ca 100644 --- a/mkdocs/docs/guides/server-deployment.md +++ b/mkdocs/docs/guides/server-deployment.md @@ -372,7 +372,19 @@ This lets you: - Reach runs on SSH fleets with a head node without giving users the head node's SSH key. - Let users connect to runs without `dstack attach`. This requires uploading their public SSH key(s) to the `dstack` server. - +After the proxy is deployed and configured, the owner of a running dev +environment can open a direct SSH or IDE session without keeping `dstack +attach` running: + +```shell +dstack connect <run name> +dstack connect --ide <run name> +``` + +The command uses the current user's built-in SSH key and fails if the server +does not return a valid proxy connection. It does not forward ports or stream +logs. See the [`dstack connect` reference](../reference/cli/dstack/connect.md) +for the supported IDEs and security checks. ### Deployment diff --git a/mkdocs/docs/reference/api/python/index.md b/mkdocs/docs/reference/api/python/index.md index 96e39467e3..e6c3076aea 100644 --- a/mkdocs/docs/reference/api/python/index.md +++ b/mkdocs/docs/reference/api/python/index.md @@ -128,6 +128,14 @@ finally: show_root_toc_entry: false heading_level: 4 +### `dstack.api.RunDirectConnection` { #dstack.api.RunDirectConnection data-toc-label="RunDirectConnection" } + +::: dstack.api.RunDirectConnection + options: + show_root_heading: false + show_root_toc_entry: false + heading_level: 4 + ### `dstack.api.Resources` { #dstack.api.Resources data-toc-label="Resources" } #SCHEMA# dstack.api.Resources diff --git a/mkdocs/docs/reference/cli/dstack/connect.md b/mkdocs/docs/reference/cli/dstack/connect.md new file mode 100644 index 0000000000..6929d06f52 --- /dev/null +++ b/mkdocs/docs/reference/cli/dstack/connect.md @@ -0,0 +1,96 @@ +--- +title: connect +description: Connect directly to a running dev environment through the SSH proxy +--- + +# dstack connect + +This command connects directly to a running dev environment through the +server's configured SSH proxy. It can open an interactive SSH session or the +IDE configured by the dev environment without keeping `dstack attach` running. + +## Usage + +
+ +```shell +$ dstack connect --help +#GENERATE# +``` + +
+ +Interactive SSH is the default mode: + +```shell +dstack connect <run name> +``` + +Use `--ide` to open the IDE selected by the dev environment configuration: + +```shell +dstack connect --ide <run name> +``` + +The supported IDEs are VS Code, Cursor, Windsurf, and Zed. You can request one +explicitly with `--vscode`, `--cursor`, `--windsurf`, or `--zed`. An explicit +option fails if it does not match the run configuration. On Linux, the Zed +launcher can be named either `zed` or `zeditor`. + +Use `--replica` and `--job` to select a job. Without `--replica`, the command +uses any running replica with the selected job number. + +## Pin the SSH proxy + +Automation can require an exact proxy endpoint: + +```shell +dstack connect --expect-sshproxy proxy.example.com:2222 <run name> +``` + +The port defaults to `22`. Bracket IPv6 addresses, for example +`[2001:db8::1]:2222`. If the run reports a different endpoint, the command +fails before downloading a key, updating local SSH configuration, or launching +a client. + +## Requirements and behavior + +`dstack connect` is intentionally limited to a running dev environment owned by +the current user. The server must have the SSH proxy enabled and provide a +complete, internally consistent SSH and IDE connection record for the selected +job. The command never falls back to direct host SSH. + +The CLI validates the proxy host, port, upstream identifier, SSH command, IDE +URL, and remote working directory before changing local state. It then stores a +local `dstack-direct-*` alias in `~/.dstack/ssh/config` using the current user's +built-in SSH key. The selected job's dynamic proxy identifier stays in the +`user@alias` target, so resolving the connection again after a retry does not +reuse a stale proxy session. + +Client processes are launched with argument vectors, without a shell. Remote +paths in IDE URLs are percent-encoded before they are passed to the local IDE. + +This command does not forward ports or stream logs. Use +[`dstack attach`](attach.md) when those attached-session features are required. + +## Python API + +Use `Run.get_direct_connection()` to perform the same refresh, validation, key +setup, and SSH alias update without launching a client: + +```python +import subprocess + +from dstack.api import Client + +client = Client.from_config() +run = client.runs.get("my-dev-environment") +if run is None: + raise RuntimeError("Run not found") + +connection = run.get_direct_connection() +subprocess.run(connection.ssh_command, check=True) +``` + +`RunDirectConnection.ssh_command` and `ide_command` are argument tuples. Pass +them directly to a subprocess API; do not join them into a shell command. diff --git a/src/dstack/_internal/cli/commands/connect.py b/src/dstack/_internal/cli/commands/connect.py new file mode 100644 index 0000000000..a36ac6d6dd --- /dev/null +++ b/src/dstack/_internal/cli/commands/connect.py @@ -0,0 +1,217 @@ +import argparse +import shutil +import subprocess +from typing import Literal +from urllib.parse import urlsplit + +from dstack._internal.cli.commands import APIBaseCommand +from dstack._internal.cli.services.completion import RunNameCompleter +from dstack._internal.core.errors import CLIError +from dstack._internal.utils.ssh import find_ssh_client +from dstack.api import Client, RunDirectConnection + +ConnectionMode = Literal["ssh", "ide", "vscode", "cursor", "windsurf", "zed"] + +_IDE_EXECUTABLES = { + "vscode": ("code",), + "cursor": ("cursor",), + "windsurf": ("windsurf",), + "zed": ("zed", "zeditor"), +} + +_IDE_NAMES = { + "vscode": "VS Code", + "cursor": "Cursor", + "windsurf": "Windsurf", + "zed": "Zed", +} + + +def sshproxy_endpoint(value: str) -> tuple[str, int]: + """Parse HOST[:PORT], including bracketed IPv6, without resolving it.""" + if "%" in value or any(c.isspace() or c == "\0" for c in value): + raise argparse.ArgumentTypeError("Invalid SSH proxy endpoint") + parsed = urlsplit(f"ssh://{value}") + if ( + parsed.hostname is None + or parsed.username is not None + or parsed.password is not None + or parsed.path + or parsed.query + or parsed.fragment + ): + raise argparse.ArgumentTypeError("Expected SSH proxy endpoint in HOST[:PORT] format") + try: + parsed_port = parsed.port + except ValueError as e: + raise argparse.ArgumentTypeError("Invalid SSH proxy port") from e + if parsed_port is None: + if value.endswith(":"): + raise argparse.ArgumentTypeError("Invalid SSH proxy port") + port = 22 + else: + port = parsed_port + if not 1 <= port <= 65535: + raise argparse.ArgumentTypeError("Invalid SSH proxy port") + return parsed.hostname, port + + +def resolve_direct_connection( + api: Client, + run_name: str, + *, + replica_num: int | None = None, + job_num: int = 0, + expected_sshproxy: tuple[str, int] | None = None, +) -> RunDirectConnection: + """Resolve and validate a run's direct SSH-proxy connection.""" + run = api.runs.get(run_name) + if run is None: + raise CLIError(f"Run {run_name} not found") + + expected_hostname = None + expected_port = None + if expected_sshproxy is not None: + expected_hostname, expected_port = expected_sshproxy + return run.get_direct_connection( + replica_num=replica_num, + job_num=job_num, + expected_sshproxy_hostname=expected_hostname, + expected_sshproxy_port=expected_port, + ) + + +def launch_direct_connection( + connection: RunDirectConnection, + *, + mode: ConnectionMode = "ssh", +) -> None: + """Launch a validated direct connection without invoking a shell.""" + if mode == "ssh": + command = list(connection.ssh_command) + if not command or command[0] != "ssh": + raise CLIError("The direct SSH command is invalid") + ssh_client = find_ssh_client() + if ssh_client is None: + raise CLIError("SSH client not found") + executable = str(ssh_client) + client_name = "SSH" + else: + requested_ide = None if mode == "ide" else mode + if requested_ide is not None and requested_ide not in _IDE_EXECUTABLES: + raise CLIError(f"Unsupported IDE: {requested_ide}") + if connection.ide is None or connection.ide_command is None: + if requested_ide is None: + raise CLIError("The dev environment is not configured with an IDE") + raise CLIError( + f"The dev environment is not configured for {_IDE_NAMES[requested_ide]}" + ) + if requested_ide is not None and connection.ide != requested_ide: + requested_name = _IDE_NAMES[requested_ide] + configured_name = connection.ide_name or connection.ide + raise CLIError( + f"The dev environment is configured for {configured_name}, not {requested_name}" + ) + + executable_names = _IDE_EXECUTABLES.get(connection.ide) + if executable_names is None: + raise CLIError(f"Unsupported configured IDE: {connection.ide}") + command = list(connection.ide_command) + if not command or command[0] != executable_names[0]: + raise CLIError("The direct IDE command is invalid") + executable = next( + (path for name in executable_names if (path := shutil.which(name)) is not None), + None, + ) + client_name = connection.ide_name or _IDE_NAMES[connection.ide] + if executable is None: + names = " or ".join(f"`{name}`" for name in executable_names) + raise CLIError( + f"{client_name} CLI not found. Install the {names} command and try again" + ) + + command[0] = executable + result = subprocess.run(command, check=False) + if result.returncode != 0: + raise CLIError(f"{client_name} client exited with status {result.returncode}") + + +class ConnectCommand(APIBaseCommand): + NAME = "connect" + DESCRIPTION = "Connect directly to a running dev environment" + + def _register(self): + super()._register() + mode = self._parser.add_mutually_exclusive_group() + mode.add_argument( + "--ssh", + action="store_const", + const="ssh", + dest="mode", + help="Open an interactive SSH shell (default)", + ) + mode.add_argument( + "--ide", + action="store_const", + const="ide", + dest="mode", + help="Open the configured IDE", + ) + mode.add_argument( + "--vscode", + action="store_const", + const="vscode", + dest="mode", + help="Open the dev environment in VS Code", + ) + mode.add_argument( + "--cursor", + action="store_const", + const="cursor", + dest="mode", + help="Open the dev environment in Cursor", + ) + mode.add_argument( + "--windsurf", + action="store_const", + const="windsurf", + dest="mode", + help="Open the dev environment in Windsurf", + ) + mode.add_argument( + "--zed", + action="store_const", + const="zed", + dest="mode", + help="Open the dev environment in Zed", + ) + self._parser.set_defaults(mode="ssh") + self._parser.add_argument( + "--replica", + help="The replica number. Defaults to any running replica.", + type=int, + ) + self._parser.add_argument( + "--job", + help="The job number inside the replica. Defaults to 0.", + type=int, + default=0, + ) + self._parser.add_argument( + "--expect-sshproxy", + metavar="HOST[:PORT]", + type=sshproxy_endpoint, + help="Fail unless the run uses this exact SSH proxy endpoint", + ) + self._parser.add_argument("run_name").completer = RunNameCompleter() # type: ignore[attr-defined] + + def _command(self, args: argparse.Namespace): + super()._command(args) + connection = resolve_direct_connection( + self.api, + args.run_name, + replica_num=args.replica, + job_num=args.job, + expected_sshproxy=args.expect_sshproxy, + ) + launch_direct_connection(connection, mode=args.mode) diff --git a/src/dstack/_internal/cli/main.py b/src/dstack/_internal/cli/main.py index 98e3b70f5c..95351d60ed 100644 --- a/src/dstack/_internal/cli/main.py +++ b/src/dstack/_internal/cli/main.py @@ -7,6 +7,7 @@ from dstack._internal.cli.commands.apply import ApplyCommand from dstack._internal.cli.commands.attach import AttachCommand from dstack._internal.cli.commands.completion import CompletionCommand +from dstack._internal.cli.commands.connect import ConnectCommand from dstack._internal.cli.commands.delete import DeleteCommand from dstack._internal.cli.commands.event import EventCommand from dstack._internal.cli.commands.export import ExportCommand @@ -67,6 +68,7 @@ def main(): subparsers = parser.add_subparsers(metavar="COMMAND") ApplyCommand.register(subparsers) AttachCommand.register(subparsers) + ConnectCommand.register(subparsers) DeleteCommand.register(subparsers) EventCommand.register(subparsers) ExportCommand.register(subparsers) diff --git a/src/dstack/api/__init__.py b/src/dstack/api/__init__.py index 0e6c6ebe3c..428415c1a3 100644 --- a/src/dstack/api/__init__.py +++ b/src/dstack/api/__init__.py @@ -23,7 +23,7 @@ from dstack._internal.core.services.ssh.ports import PortUsedError from dstack.api._public import BackendCollection, Client, RepoCollection, RunCollection from dstack.api._public.backends import Backend -from dstack.api._public.runs import Run, RunStatus +from dstack.api._public.runs import Run, RunDirectConnection, RunStatus Service = _ServiceConfiguration Task = _TaskConfiguration diff --git a/src/dstack/api/_public/runs.py b/src/dstack/api/_public/runs.py index 3fe274c98d..dbb4554085 100644 --- a/src/dstack/api/_public/runs.py +++ b/src/dstack/api/_public/runs.py @@ -1,5 +1,8 @@ import base64 +import hashlib +import ipaddress import queue +import re import tempfile import threading import time @@ -7,10 +10,11 @@ from collections.abc import Iterator from contextlib import contextmanager from copy import copy +from dataclasses import dataclass from datetime import datetime from pathlib import Path from typing import BinaryIO, Dict, Iterable, List, Optional, Union -from urllib.parse import urlencode, urlparse +from urllib.parse import quote, urlencode, urlparse from websocket import WebSocketApp @@ -32,6 +36,7 @@ from dstack._internal.core.models.repos.virtual import VirtualRepo from dstack._internal.core.models.runs import ( Job, + JobConnectionInfo, JobSpec, JobStatus, RunPlan, @@ -49,12 +54,51 @@ from dstack._internal.utils.common import get_or_error, make_proxy_url from dstack._internal.utils.files import create_file_archive from dstack._internal.utils.logging import get_logger -from dstack._internal.utils.path import PathLike +from dstack._internal.utils.path import FilePath, PathLike +from dstack._internal.utils.ssh import ( + build_ssh_command, + build_ssh_url_authority, + include_ssh_config, + update_ssh_config, +) from dstack.api.server import APIClient logger = get_logger(__name__) +@dataclass(frozen=True) +class RunDirectConnection: + """A validated, proxy-only connection to a running dev environment. + + The command values are argument vectors and are safe to pass directly to + :mod:`subprocess` without invoking a shell. The local SSH alias is refreshed + for every resolution, so a retried job uses the current proxy upstream ID. + """ + + run_name: str + replica_num: int + job_num: int + upstream_id: str + sshproxy_hostname: str + sshproxy_port: int + ssh_alias: str + ssh_command: tuple[str, ...] + ide: Optional[str] + ide_name: Optional[str] + ide_command: Optional[tuple[str, ...]] + + +@dataclass(frozen=True) +class ValidatedDirectConnectionInfo: + """Server-provided direct connection data validated without local side effects.""" + + hostname: str + port: int + upstream_id: str + ide_name: Optional[str] + ide_command: Optional[tuple[str, ...]] + + class Run(ABC): """ Attributes: @@ -408,6 +452,99 @@ def attach( return True + def get_direct_connection( + self, + replica_num: Optional[int] = None, + job_num: int = 0, + expected_sshproxy_hostname: Optional[str] = None, + expected_sshproxy_port: Optional[int] = None, + ) -> RunDirectConnection: + """Prepare a direct SSH-proxy connection to a running dev environment. + + The run is refreshed on every call. Direct connections are owner-only, + require a running dev environment, and never fall back to legacy host SSH. + The returned commands are argument vectors; callers must not join them + into a shell command. + + Args: + replica_num: Replica number, or ``None`` for any running replica. + job_num: Job number inside the replica. + expected_sshproxy_hostname: Require this exact SSH proxy hostname when set. + expected_sshproxy_port: Require this SSH proxy port. Defaults to 22 when an expected + hostname is set. + + Returns: + A validated connection with SSH and, when configured, IDE commands. + + Raises: + dstack.api.ClientError: If the run cannot be connected to safely. + """ + self.refresh() + if self.status != RunStatus.RUNNING: + raise ClientError( + f"Direct connection requires a running run; {self.name} is {self.status.value}" + ) + if self._run.run_spec.configuration.type != "dev-environment": + raise ClientError("Direct connection is only supported for dev environments") + + current_user = self._api_client.users.get_my_user() + if self._run.user != current_user.username: + raise ClientError("Direct connection is only available to the run owner") + + job = self._find_job(replica_num=replica_num, job_num=job_num) + if job is None or not job.job_submissions: + replica_repr = replica_num if replica_num is not None else "" + raise ClientError(f"Failed to find running replica={replica_repr} job={job_num}") + if job.job_submissions[-1].status != JobStatus.RUNNING: + replica_repr = replica_num if replica_num is not None else "" + raise ClientError(f"Failed to find running replica={replica_repr} job={job_num}") + + ssh_alias = _get_direct_connection_alias( + server_url=self._api_client.base_url, + project=self._project, + token_hash=self._api_client.get_token_hash(), + run_name=self.name, + ) + ide = getattr(self._run.run_spec.configuration, "ide", None) + validated = validate_direct_connection_info( + job=job, + ssh_alias=ssh_alias, + ide=ide, + expected_sshproxy_hostname=expected_sshproxy_hostname, + expected_sshproxy_port=expected_sshproxy_port, + ) + + # Do not modify local SSH state until all server-provided connection + # fields, including the optional IDE URL, have been validated. + config_manager = ConfigManager() + user_key = UserSSHKeyManager( + self._api_client, config_manager.dstack_ssh_dir + ).get_user_key() + include_ssh_config(config_manager.dstack_ssh_config_path) + update_ssh_config( + config_manager.dstack_ssh_config_path, + ssh_alias, + { + "HostName": validated.hostname, + "Port": validated.port, + "IdentityFile": FilePath(user_key.private_key_path), + "IdentitiesOnly": "yes", + }, + ) + return RunDirectConnection( + run_name=self.name, + replica_num=job.job_spec.replica_num, + job_num=job.job_spec.job_num, + upstream_id=validated.upstream_id, + sshproxy_hostname=validated.hostname, + sshproxy_port=validated.port, + ssh_alias=ssh_alias, + ssh_command=("ssh", f"{validated.upstream_id}@{ssh_alias}"), + ide=ide, + ide_name=validated.ide_name, + ide_command=validated.ide_command, + ) + def detach(self): """ Stop the SSH tunnel to the instance and update SSH config @@ -746,6 +883,162 @@ def _reserve_ports( return PortsLock(ports).acquire() +_HOSTNAME_RE = re.compile( + r"(?=.{1,253}\.?$)(?:[A-Za-z0-9](?:[A-Za-z0-9-]{0,61}[A-Za-z0-9])?\.)*" + r"[A-Za-z0-9](?:[A-Za-z0-9-]{0,61}[A-Za-z0-9])?\.?" +) + + +def _validate_sshproxy_hostname(hostname: Optional[str]) -> str: + if hostname is None or "%" in hostname: + raise ClientError("The server returned an invalid SSH proxy hostname") + try: + ipaddress.ip_address(hostname) + except ValueError: + if _HOSTNAME_RE.fullmatch(hostname) is None: + raise ClientError("The server returned an invalid SSH proxy hostname") + return hostname + + +def _validate_sshproxy_port(port: Optional[int]) -> int: + if port is None: + return 22 + if isinstance(port, bool) or not 1 <= port <= 65535: + raise ClientError("The server returned an invalid SSH proxy port") + return port + + +def _validate_sshproxy_upstream_id(upstream_id: Optional[str]) -> str: + # The server contract deliberately treats this value as an extensible string. + # Keep the accepted subset safe for both OpenSSH's `user@host` syntax and URL userinfo. + if ( + upstream_id is None + or re.fullmatch(r"[A-Za-z0-9][A-Za-z0-9._-]{0,127}", upstream_id) is None + ): + raise ClientError("The server returned an invalid SSH proxy upstream ID") + return upstream_id + + +def _get_direct_connection_alias( + *, server_url: str, project: str, token_hash: str, run_name: str +) -> str: + target = f"{server_url}\0{project}\0{token_hash}\0{run_name}".encode() + digest = hashlib.sha256(target).hexdigest()[:24] + return f"dstack-direct-{digest}" + + +def validate_direct_connection_info( + *, + job: Job, + ssh_alias: str, + ide: Optional[str], + expected_sshproxy_hostname: Optional[str] = None, + expected_sshproxy_port: Optional[int] = None, +) -> ValidatedDirectConnectionInfo: + """Validate direct connection metadata without API calls or local state changes.""" + if not job.job_submissions or job.job_submissions[-1].status != JobStatus.RUNNING: + raise ClientError("Direct connection requires a running job") + connection_info = job.job_connection_info + if connection_info is None: + raise ClientError("The server did not provide SSH proxy connection information") + hostname = _validate_sshproxy_hostname(connection_info.sshproxy_hostname) + port = _validate_sshproxy_port(connection_info.sshproxy_port) + if expected_sshproxy_hostname is None and expected_sshproxy_port is not None: + raise ClientError("An expected SSH proxy port requires an expected hostname") + if expected_sshproxy_hostname is not None: + expected_hostname = _validate_sshproxy_hostname(expected_sshproxy_hostname) + expected_port = _validate_sshproxy_port(expected_sshproxy_port) + if hostname != expected_hostname or port != expected_port: + raise ClientError("The run SSH proxy does not match the expected endpoint") + upstream_id = _validate_sshproxy_upstream_id(connection_info.sshproxy_upstream_id) + expected_proxy_command = build_ssh_command( + username=upstream_id, + hostname=hostname, + port=connection_info.sshproxy_port, + ) + if connection_info.proxied_ssh_command != expected_proxy_command: + raise ClientError("The server returned inconsistent SSH proxy connection information") + ide_name, ide_command = _get_ide_connection( + job=job, + connection_info=connection_info, + hostname=hostname, + port=connection_info.sshproxy_port, + upstream_id=upstream_id, + ssh_alias=ssh_alias, + ide=ide, + ) + return ValidatedDirectConnectionInfo( + hostname=hostname, + port=port, + upstream_id=upstream_id, + ide_name=ide_name, + ide_command=ide_command, + ) + + +_IDE_NAMES = { + "vscode": "VS Code", + "cursor": "Cursor", + "windsurf": "Windsurf", + "zed": "Zed", +} + +_IDE_EXECUTABLES = { + "vscode": "code", + "cursor": "cursor", + "windsurf": "windsurf", + "zed": "zed", +} + + +def _get_ide_connection( + *, + job: Job, + connection_info: JobConnectionInfo, + hostname: str, + port: Optional[int], + upstream_id: str, + ssh_alias: str, + ide: Optional[str], +) -> tuple[Optional[str], Optional[tuple[str, ...]]]: + if ide is None: + if connection_info.ide_name is not None or connection_info.proxied_ide_url is not None: + raise ClientError("The server returned unexpected IDE connection information") + return None, None + ide_name = _IDE_NAMES.get(ide) + executable = _IDE_EXECUTABLES.get(ide) + if ide_name is None or executable is None: + raise ClientError(f"Direct connection does not support the configured IDE: {ide}") + + submission = job.job_submissions[-1] + runtime_data = submission.job_runtime_data + if runtime_data is None or runtime_data.working_dir is None: + raise ClientError("The server did not provide the dev environment working directory") + working_dir = runtime_data.working_dir + if not working_dir.startswith("/") or any(c in working_dir for c in ("\0", "\r", "\n")): + raise ClientError("The server returned an invalid dev environment working directory") + authority = build_ssh_url_authority( + username=upstream_id, + hostname=hostname, + port=port, + ) + if ide == "zed": + expected_url = f"zed://ssh/{authority}{working_dir}" + else: + expected_url = f"{ide}://vscode-remote/ssh-remote+{authority}{working_dir}" + if connection_info.ide_name != ide_name or connection_info.proxied_ide_url != expected_url: + raise ClientError(f"The server returned inconsistent {ide_name} connection information") + + # Keep the dynamic upstream ID in the local URL so a retried job cannot + # reuse a stale remote session keyed only by the stable SSH host alias. + encoded_working_dir = quote(working_dir, safe="/~:@-._") + alias_authority = build_ssh_url_authority(username=upstream_id, hostname=ssh_alias) + if ide == "zed": + return ide_name, (executable, f"ssh://{alias_authority}{encoded_working_dir}") + folder_uri = f"vscode-remote://ssh-remote+{alias_authority}{encoded_working_dir}" + return ide_name, (executable, "--folder-uri", folder_uri) + + @contextmanager def _prepare_code_file(repo: Repo) -> Iterator[tuple[BinaryIO, str]]: with tempfile.TemporaryFile("w+b") as fp: diff --git a/src/tests/_internal/cli/commands/test_connect.py b/src/tests/_internal/cli/commands/test_connect.py new file mode 100644 index 0000000000..b039c6b7a8 --- /dev/null +++ b/src/tests/_internal/cli/commands/test_connect.py @@ -0,0 +1,345 @@ +import argparse +from types import SimpleNamespace +from unittest.mock import Mock, patch + +import pytest + +from dstack._internal.cli.commands.connect import ( + ConnectCommand, + launch_direct_connection, + sshproxy_endpoint, +) +from dstack._internal.core.errors import CLIError +from dstack.api import RunDirectConnection + + +def _parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser() + subparsers = parser.add_subparsers() + ConnectCommand.register(subparsers) + return parser + + +def _invoke(args: list[str], run: Mock): + client = Mock() + client.runs.get.return_value = run + parsed_args = _parser().parse_args(["connect", "--project", "main", *args]) + parsed_args.extra_args = [] + with patch("dstack._internal.cli.commands.Client.from_config", return_value=client): + parsed_args.func(parsed_args) + return client + + +def _connection( + *, + ide: str | None = "zed", + ide_name: str | None = "Zed", + ide_command: tuple[str, ...] | None = ( + "zed", + "ssh://upstream@dstack-direct-safealias/workspace", + ), + ssh_command: tuple[str, ...] = ("ssh", "upstream@dstack-direct-safealias"), +) -> RunDirectConnection: + return RunDirectConnection( + run_name="dev-run", + replica_num=0, + job_num=0, + upstream_id="upstream", + sshproxy_hostname="sshproxy.example.com", + sshproxy_port=22, + ssh_alias="dstack-direct-safealias", + ssh_command=ssh_command, + ide=ide, + ide_name=ide_name, + ide_command=ide_command, + ) + + +class TestConnectCommand: + def test_launches_ssh_with_an_argument_vector(self, monkeypatch: pytest.MonkeyPatch) -> None: + run = Mock() + run.get_direct_connection.return_value = _connection() + subprocess_run = Mock(return_value=SimpleNamespace(returncode=0)) + monkeypatch.setattr( + "dstack._internal.cli.commands.connect.find_ssh_client", + lambda: "/usr/bin/ssh", + ) + monkeypatch.setattr( + "dstack._internal.cli.commands.connect.subprocess.run", + subprocess_run, + ) + + client = _invoke( + [ + "--ssh", + "--replica", + "2", + "--job", + "1", + "--expect-sshproxy", + "sshproxy.example.com:2222", + "dev-run", + ], + run, + ) + + client.runs.get.assert_called_once_with("dev-run") + run.get_direct_connection.assert_called_once_with( + replica_num=2, + job_num=1, + expected_sshproxy_hostname="sshproxy.example.com", + expected_sshproxy_port=2222, + ) + subprocess_run.assert_called_once_with( + ["/usr/bin/ssh", "upstream@dstack-direct-safealias"], + check=False, + ) + + @pytest.mark.parametrize( + ("flag", "ide", "ide_name", "ide_command", "expected_command"), + [ + ( + "--vscode", + "vscode", + "VS Code", + ( + "code", + "--folder-uri", + "vscode-remote://ssh-remote+upstream@safealias/workspace", + ), + [ + "/usr/local/bin/code", + "--folder-uri", + "vscode-remote://ssh-remote+upstream@safealias/workspace", + ], + ), + ( + "--cursor", + "cursor", + "Cursor", + ( + "cursor", + "--folder-uri", + "vscode-remote://ssh-remote+upstream@safealias/workspace", + ), + [ + "/usr/local/bin/cursor", + "--folder-uri", + "vscode-remote://ssh-remote+upstream@safealias/workspace", + ], + ), + ( + "--windsurf", + "windsurf", + "Windsurf", + ( + "windsurf", + "--folder-uri", + "vscode-remote://ssh-remote+upstream@safealias/workspace", + ), + [ + "/usr/local/bin/windsurf", + "--folder-uri", + "vscode-remote://ssh-remote+upstream@safealias/workspace", + ], + ), + ( + "--zed", + "zed", + "Zed", + ("zed", "ssh://upstream@safealias/workspace"), + ["/usr/local/bin/zed", "ssh://upstream@safealias/workspace"], + ), + ], + ) + def test_launches_an_explicit_ide_with_an_argument_vector( + self, + flag: str, + ide: str, + ide_name: str, + ide_command: tuple[str, ...], + expected_command: list[str], + monkeypatch: pytest.MonkeyPatch, + ) -> None: + run = Mock() + run.get_direct_connection.return_value = _connection( + ide=ide, + ide_name=ide_name, + ide_command=ide_command, + ) + subprocess_run = Mock(return_value=SimpleNamespace(returncode=0)) + monkeypatch.setattr( + "dstack._internal.cli.commands.connect.shutil.which", + lambda executable: f"/usr/local/bin/{executable}", + ) + monkeypatch.setattr( + "dstack._internal.cli.commands.connect.subprocess.run", + subprocess_run, + ) + + _invoke([flag, "dev-run"], run) + + subprocess_run.assert_called_once_with(expected_command, check=False) + + def test_launches_the_configured_ide(self, monkeypatch: pytest.MonkeyPatch) -> None: + run = Mock() + run.get_direct_connection.return_value = _connection( + ide="cursor", + ide_name="Cursor", + ide_command=( + "cursor", + "--folder-uri", + "vscode-remote://ssh-remote+upstream@safealias/workspace", + ), + ) + subprocess_run = Mock(return_value=SimpleNamespace(returncode=0)) + monkeypatch.setattr( + "dstack._internal.cli.commands.connect.shutil.which", + lambda executable: f"/usr/local/bin/{executable}", + ) + monkeypatch.setattr( + "dstack._internal.cli.commands.connect.subprocess.run", + subprocess_run, + ) + + _invoke(["--ide", "dev-run"], run) + + subprocess_run.assert_called_once_with( + [ + "/usr/local/bin/cursor", + "--folder-uri", + "vscode-remote://ssh-remote+upstream@safealias/workspace", + ], + check=False, + ) + + def test_rejects_an_explicit_ide_that_is_not_configured(self) -> None: + run = Mock() + run.get_direct_connection.return_value = _connection( + ide=None, + ide_name=None, + ide_command=None, + ) + + with pytest.raises(CLIError, match="not configured for Zed"): + _invoke(["--zed", "dev-run"], run) + + def test_rejects_configured_ide_mode_for_an_ssh_only_run(self) -> None: + run = Mock() + run.get_direct_connection.return_value = _connection( + ide=None, + ide_name=None, + ide_command=None, + ) + + with pytest.raises(CLIError, match="not configured with an IDE"): + _invoke(["--ide", "dev-run"], run) + + def test_uses_zeditor_when_zed_binary_name_is_unavailable( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + run = Mock() + run.get_direct_connection.return_value = _connection( + ide_command=("zed", "ssh://upstream@safealias/workspace"), + ) + subprocess_run = Mock(return_value=SimpleNamespace(returncode=0)) + monkeypatch.setattr( + "dstack._internal.cli.commands.connect.shutil.which", + lambda executable: "/usr/bin/zeditor" if executable == "zeditor" else None, + ) + monkeypatch.setattr( + "dstack._internal.cli.commands.connect.subprocess.run", + subprocess_run, + ) + + _invoke(["--zed", "dev-run"], run) + + subprocess_run.assert_called_once_with( + ["/usr/bin/zeditor", "ssh://upstream@safealias/workspace"], + check=False, + ) + + def test_rejects_an_explicit_ide_that_does_not_match_the_configuration(self) -> None: + run = Mock() + run.get_direct_connection.return_value = _connection() + + with pytest.raises(CLIError, match="configured for Zed, not Cursor"): + _invoke(["--cursor", "dev-run"], run) + + def test_rejects_a_missing_run(self) -> None: + with pytest.raises(CLIError, match="Run missing not found"): + _invoke(["missing"], None) + + def test_rejects_an_invalid_ssh_command(self) -> None: + connection = _connection(ssh_command=("other", "unsafe")) + + with pytest.raises(CLIError, match="direct SSH command is invalid"): + launch_direct_connection(connection) + + def test_rejects_an_invalid_ide_command(self) -> None: + connection = _connection(ide_command=("other", "unsafe")) + + with pytest.raises(CLIError, match="direct IDE command is invalid"): + launch_direct_connection(connection, mode="zed") + + def test_rejects_nonzero_client_exit(self, monkeypatch: pytest.MonkeyPatch) -> None: + run = Mock() + run.get_direct_connection.return_value = _connection( + ide=None, + ide_name=None, + ide_command=None, + ) + monkeypatch.setattr( + "dstack._internal.cli.commands.connect.find_ssh_client", + lambda: "/usr/bin/ssh", + ) + monkeypatch.setattr( + "dstack._internal.cli.commands.connect.subprocess.run", + lambda *_args, **_kwargs: SimpleNamespace(returncode=255), + ) + + with pytest.raises(CLIError, match="status 255"): + _invoke(["dev-run"], run) + + def test_help_uses_standard_project_and_run_shape( + self, capsys: pytest.CaptureFixture[str] + ) -> None: + with pytest.raises(SystemExit) as exc_info: + _parser().parse_args(["connect", "--help"]) + + assert exc_info.value.code == 0 + output = capsys.readouterr().out + assert "--project NAME" in output + for option in ("--ssh", "--ide", "--vscode", "--cursor", "--windsurf", "--zed"): + assert option in output + assert "--expect-sshproxy HOST[:PORT]" in output + assert "run_name" in output + + +@pytest.mark.parametrize( + ("value", "expected"), + [ + ("sshproxy.example.com", ("sshproxy.example.com", 22)), + ("sshproxy.example.com:2222", ("sshproxy.example.com", 2222)), + ("[2001:db8::1]:2222", ("2001:db8::1", 2222)), + ], +) +def test_sshproxy_endpoint(value: str, expected: tuple[str, int]) -> None: + assert sshproxy_endpoint(value) == expected + + +@pytest.mark.parametrize( + "value", + [ + "user@host", + "host/path", + "host:0", + "host:", + "host:70000", + "host\nother", + "[fe80::1%h]:22", + ], +) +def test_sshproxy_endpoint_rejects_invalid_value(value: str) -> None: + with pytest.raises(argparse.ArgumentTypeError): + sshproxy_endpoint(value) diff --git a/src/tests/api/test_direct_connect.py b/src/tests/api/test_direct_connect.py new file mode 100644 index 0000000000..d746f8314a --- /dev/null +++ b/src/tests/api/test_direct_connect.py @@ -0,0 +1,457 @@ +from pathlib import Path +from types import SimpleNamespace +from unittest.mock import Mock +from uuid import UUID, uuid4 + +import pytest + +from dstack._internal.core.errors import ClientError +from dstack._internal.core.models.runs import JobConnectionInfo, JobStatus, RunStatus +from dstack._internal.core.services.ssh.key_manager import UserSSHKey +from dstack._internal.utils.path import FilePath +from dstack._internal.utils.ssh import build_ssh_command, build_ssh_url_authority +from dstack.api._public import runs as public_runs +from dstack.api._public.runs import Run + +_IDE_NAMES = { + "vscode": "VS Code", + "cursor": "Cursor", + "windsurf": "Windsurf", + "zed": "Zed", +} + + +def _run_model( + *, + job_id: UUID, + run_name: str = "dev-run", + upstream_id: str | None = None, + owner: str = "owner", + status: RunStatus = RunStatus.RUNNING, + job_status: JobStatus = JobStatus.RUNNING, + configuration_type: str = "dev-environment", + ide: str | None = "zed", + hostname: str = "sshproxy.example.com", + port: int | None = 2222, + proxied_ssh_command: list[str] | None = None, + proxied_ide_url: str | None = None, + working_dir: str = "/workspace/project", +): + if upstream_id is None: + upstream_id = job_id.hex + if proxied_ssh_command is None: + proxied_ssh_command = build_ssh_command( + username=upstream_id, + hostname=hostname, + port=port, + ) + ide_name = _IDE_NAMES.get(ide) if ide is not None else None + if proxied_ide_url is None and ide is not None: + authority = build_ssh_url_authority( + username=upstream_id, + hostname=hostname, + port=port, + ) + if ide == "zed": + proxied_ide_url = f"zed://ssh/{authority}{working_dir}" + else: + proxied_ide_url = f"{ide}://vscode-remote/ssh-remote+{authority}{working_dir}" + connection_info = JobConnectionInfo( + ide_name=ide_name, + attached_ide_url=None, + proxied_ide_url=proxied_ide_url, + attached_ssh_command=["ssh", "ignored"], + proxied_ssh_command=proxied_ssh_command, + sshproxy_hostname=hostname, + sshproxy_port=port, + sshproxy_upstream_id=upstream_id, + ) + job = SimpleNamespace( + job_spec=SimpleNamespace(replica_num=0, job_num=0), + job_submissions=[ + SimpleNamespace( + status=job_status, + job_runtime_data=SimpleNamespace(working_dir=working_dir), + ) + ], + job_connection_info=connection_info, + ) + return SimpleNamespace( + run_spec=SimpleNamespace( + run_name=run_name, + configuration=SimpleNamespace(type=configuration_type, ide=ide), + ), + jobs=[job], + status=status, + user=owner, + ) + + +def _api_client(run_model, *, username: str = "owner"): + return SimpleNamespace( + base_url="https://dstack.example.com", + get_token_hash=lambda: "token-hash", + runs=SimpleNamespace(get=Mock(return_value=run_model)), + users=SimpleNamespace(get_my_user=Mock(return_value=SimpleNamespace(username=username))), + ) + + +@pytest.fixture +def local_connection_config(monkeypatch: pytest.MonkeyPatch, tmp_path: Path): + config_manager = SimpleNamespace( + dstack_ssh_dir=tmp_path, + dstack_ssh_config_path=tmp_path / "config", + ) + config_manager_factory = Mock(return_value=config_manager) + key_manager = Mock() + key_manager.return_value.get_user_key.return_value = UserSSHKey( + public_key="ssh-ed25519 public", + private_key_path=tmp_path / "private-key", + ) + include = Mock() + update = Mock() + monkeypatch.setattr(public_runs, "ConfigManager", config_manager_factory) + monkeypatch.setattr(public_runs, "UserSSHKeyManager", key_manager) + monkeypatch.setattr(public_runs, "include_ssh_config", include) + monkeypatch.setattr(public_runs, "update_ssh_config", update) + return SimpleNamespace( + config_manager_factory=config_manager_factory, + key_manager=key_manager, + include=include, + update=update, + ) + + +def _assert_no_local_connection_state(state: SimpleNamespace) -> None: + state.config_manager_factory.assert_not_called() + state.key_manager.assert_not_called() + state.include.assert_not_called() + state.update.assert_not_called() + + +class TestRunGetDirectConnection: + def test_returns_validated_argv_and_configures_current_job_alias( + self, + local_connection_config: SimpleNamespace, + ) -> None: + job_id = uuid4() + run_model = _run_model(job_id=job_id) + api_client = _api_client(run_model) + run = Run(api_client=api_client, project="main", run=run_model) + + connection = run.get_direct_connection() + + assert connection.upstream_id == job_id.hex + assert connection.ssh_command == ("ssh", f"{job_id.hex}@{connection.ssh_alias}") + assert connection.ide == "zed" + assert connection.ide_name == "Zed" + assert connection.ide_command == ( + "zed", + f"ssh://{job_id.hex}@{connection.ssh_alias}/workspace/project", + ) + _, alias, raw_options = local_connection_config.update.call_args.args + assert alias == connection.ssh_alias + options = dict(raw_options) + identity_file = options.pop("IdentityFile") + assert isinstance(identity_file, FilePath) + assert identity_file.path == ( + local_connection_config.config_manager_factory.return_value.dstack_ssh_dir + / "private-key" + ) + assert options == { + "HostName": "sshproxy.example.com", + "Port": 2222, + "IdentitiesOnly": "yes", + } + local_connection_config.include.assert_called_once() + api_client.runs.get.assert_called_once_with("main", "dev-run") + + def test_refreshes_job_id_on_every_resolution( + self, + local_connection_config: SimpleNamespace, + ) -> None: + stale_model = _run_model(job_id=uuid4()) + current_job_id = uuid4() + current_model = _run_model(job_id=current_job_id) + api_client = _api_client(current_model) + run = Run(api_client=api_client, project="main", run=stale_model) + + connection = run.get_direct_connection() + + assert connection.upstream_id == current_job_id.hex + assert connection.ssh_command == ( + "ssh", + f"{current_job_id.hex}@{connection.ssh_alias}", + ) + + def test_defaults_a_missing_proxy_port_to_22( + self, + local_connection_config: SimpleNamespace, + ) -> None: + job_id = uuid4() + run_model = _run_model(job_id=job_id, port=None) + run = Run(api_client=_api_client(run_model), project="main", run=run_model) + + connection = run.get_direct_connection() + + assert connection.sshproxy_port == 22 + assert connection.ssh_command == ("ssh", f"{job_id.hex}@{connection.ssh_alias}") + assert local_connection_config.update.call_args.args[2]["Port"] == 22 + + def test_expected_sshproxy_is_checked_before_local_state_changes( + self, + local_connection_config: SimpleNamespace, + ) -> None: + run_model = _run_model(job_id=uuid4()) + run = Run(api_client=_api_client(run_model), project="main", run=run_model) + + with pytest.raises(ClientError, match="expected endpoint"): + run.get_direct_connection( + expected_sshproxy_hostname="other.example.com", + expected_sshproxy_port=2222, + ) + + _assert_no_local_connection_state(local_connection_config) + + def test_rejects_non_owner_before_local_state_changes( + self, + local_connection_config: SimpleNamespace, + ) -> None: + run_model = _run_model(job_id=uuid4(), owner="another-user") + run = Run(api_client=_api_client(run_model), project="main", run=run_model) + + with pytest.raises(ClientError, match="run owner"): + run.get_direct_connection() + + _assert_no_local_connection_state(local_connection_config) + + @pytest.mark.parametrize( + ("status", "configuration_type", "message"), + [ + (RunStatus.PENDING, "dev-environment", "requires a running run"), + (RunStatus.RUNNING, "task", "only supported for dev environments"), + ], + ) + def test_requires_running_dev_environment( + self, + status: RunStatus, + configuration_type: str, + message: str, + local_connection_config: SimpleNamespace, + ) -> None: + run_model = _run_model( + job_id=uuid4(), + status=status, + configuration_type=configuration_type, + ) + run = Run(api_client=_api_client(run_model), project="main", run=run_model) + + with pytest.raises(ClientError, match=message): + run.get_direct_connection() + + _assert_no_local_connection_state(local_connection_config) + + def test_rejects_explicit_replica_without_a_running_job( + self, + local_connection_config: SimpleNamespace, + ) -> None: + run_model = _run_model(job_id=uuid4(), job_status=JobStatus.TERMINATED) + run = Run(api_client=_api_client(run_model), project="main", run=run_model) + + with pytest.raises(ClientError, match="Failed to find running replica"): + run.get_direct_connection(replica_num=0) + + _assert_no_local_connection_state(local_connection_config) + + def test_requires_server_proxy_connection_info( + self, + local_connection_config: SimpleNamespace, + ) -> None: + run_model = _run_model(job_id=uuid4()) + run_model.jobs[0].job_connection_info = None + run = Run(api_client=_api_client(run_model), project="main", run=run_model) + + with pytest.raises(ClientError, match="did not provide SSH proxy"): + run.get_direct_connection() + + _assert_no_local_connection_state(local_connection_config) + + @pytest.mark.parametrize( + "mutate", + [ + lambda jci: setattr(jci, "sshproxy_hostname", "-oProxyCommand=bad"), + lambda jci: setattr(jci, "sshproxy_hostname", "fe80::1%h"), + lambda jci: setattr(jci, "sshproxy_port", True), + lambda jci: setattr(jci, "sshproxy_port", 70000), + lambda jci: setattr(jci, "sshproxy_upstream_id", "bad@upstream"), + lambda jci: setattr(jci, "proxied_ssh_command", ["ssh", "other.example.com"]), + lambda jci: setattr(jci, "ide_name", "Other IDE"), + lambda jci: setattr(jci, "proxied_ide_url", "zed://ssh/untrusted/workspace"), + ], + ) + def test_rejects_inconsistent_or_unsafe_proxy_target( + self, + mutate, + local_connection_config: SimpleNamespace, + ) -> None: + run_model = _run_model(job_id=uuid4()) + mutate(run_model.jobs[0].job_connection_info) + run = Run(api_client=_api_client(run_model), project="main", run=run_model) + + with pytest.raises(ClientError): + run.get_direct_connection() + + _assert_no_local_connection_state(local_connection_config) + + @pytest.mark.parametrize( + "working_dir", + ["workspace/project", "/workspace/project\n-oProxyCommand=bad"], + ) + def test_rejects_an_unsafe_working_directory( + self, + working_dir: str, + local_connection_config: SimpleNamespace, + ) -> None: + run_model = _run_model(job_id=uuid4(), working_dir=working_dir) + run = Run(api_client=_api_client(run_model), project="main", run=run_model) + + with pytest.raises(ClientError, match="invalid dev environment working directory"): + run.get_direct_connection() + + _assert_no_local_connection_state(local_connection_config) + + @pytest.mark.parametrize( + ("ide", "ide_name", "executable"), + [ + ("vscode", "VS Code", "code"), + ("cursor", "Cursor", "cursor"), + ("windsurf", "Windsurf", "windsurf"), + ], + ) + def test_returns_validated_vscode_family_command( + self, + ide: str, + ide_name: str, + executable: str, + local_connection_config: SimpleNamespace, + ) -> None: + job_id = uuid4() + run_model = _run_model(job_id=job_id, ide=ide) + run = Run(api_client=_api_client(run_model), project="main", run=run_model) + + connection = run.get_direct_connection() + + assert connection.ide == ide + assert connection.ide_name == ide_name + assert connection.ide_command == ( + executable, + "--folder-uri", + f"vscode-remote://ssh-remote+{job_id.hex}@{connection.ssh_alias}/workspace/project", + ) + + def test_does_not_offer_an_ide_for_an_ssh_only_configuration( + self, + local_connection_config: SimpleNamespace, + ) -> None: + run_model = _run_model(job_id=uuid4(), ide=None) + run = Run(api_client=_api_client(run_model), project="main", run=run_model) + + connection = run.get_direct_connection() + + assert connection.ide is None + assert connection.ide_name is None + assert connection.ide_command is None + + def test_rejects_unexpected_ide_info_for_an_ssh_only_configuration( + self, + local_connection_config: SimpleNamespace, + ) -> None: + run_model = _run_model(job_id=uuid4(), ide=None) + run_model.jobs[0].job_connection_info.ide_name = "Zed" + run_model.jobs[0].job_connection_info.proxied_ide_url = "zed://ssh/unexpected/workspace" + run = Run(api_client=_api_client(run_model), project="main", run=run_model) + + with pytest.raises(ClientError, match="unexpected IDE connection information"): + run.get_direct_connection() + + _assert_no_local_connection_state(local_connection_config) + + def test_uses_distinct_aliases_for_concurrent_runs_in_one_project( + self, + local_connection_config: SimpleNamespace, + ) -> None: + first_model = _run_model(job_id=uuid4(), run_name="first-run") + second_model = _run_model(job_id=uuid4(), run_name="second-run") + first = Run(api_client=_api_client(first_model), project="main", run=first_model) + second = Run(api_client=_api_client(second_model), project="main", run=second_model) + + first_connection = first.get_direct_connection() + second_connection = second.get_direct_connection() + + assert first_connection.ssh_alias != second_connection.ssh_alias + + def test_accepts_an_extensible_safe_upstream_id( + self, + local_connection_config: SimpleNamespace, + ) -> None: + upstream_id = "runner.v2_job-01" + run_model = _run_model(job_id=uuid4(), upstream_id=upstream_id) + run = Run(api_client=_api_client(run_model), project="main", run=run_model) + + connection = run.get_direct_connection() + + assert connection.upstream_id == upstream_id + assert connection.ssh_command == ("ssh", f"{upstream_id}@{connection.ssh_alias}") + + @pytest.mark.parametrize("upstream_id", ["-leading-dash", "bad@id", "a" * 129]) + def test_rejects_an_unsafe_upstream_id( + self, + upstream_id: str, + local_connection_config: SimpleNamespace, + ) -> None: + run_model = _run_model(job_id=uuid4(), upstream_id=upstream_id) + run = Run(api_client=_api_client(run_model), project="main", run=run_model) + + with pytest.raises(ClientError, match="invalid SSH proxy upstream ID"): + run.get_direct_connection() + + _assert_no_local_connection_state(local_connection_config) + + def test_percent_encodes_the_zed_remote_path( + self, + local_connection_config: SimpleNamespace, + ) -> None: + job_id = uuid4() + run_model = _run_model( + job_id=job_id, + working_dir="/workspace/a project/#literal%value", + ) + run = Run(api_client=_api_client(run_model), project="main", run=run_model) + + connection = run.get_direct_connection() + + assert connection.ide_command == ( + "zed", + f"ssh://{job_id.hex}@{connection.ssh_alias}/workspace/a%20project/%23literal%25value", + ) + + def test_percent_encodes_a_vscode_family_remote_path( + self, + local_connection_config: SimpleNamespace, + ) -> None: + job_id = uuid4() + run_model = _run_model( + job_id=job_id, + ide="cursor", + working_dir="/workspace/a project/#literal%value", + ) + run = Run(api_client=_api_client(run_model), project="main", run=run_model) + + connection = run.get_direct_connection() + + assert connection.ide_command == ( + "cursor", + "--folder-uri", + f"vscode-remote://ssh-remote+{job_id.hex}@{connection.ssh_alias}" + "/workspace/a%20project/%23literal%25value", + )