acc: converge ssh/connect-serverless-gpu tunnel test onto a local web…

.github/workflows/push.yml

@@ -262,9 +262,9 @@ jobs:
         with:
           cache-key: test-exp-ssh
 
-      # The proxy tests exercise a real OpenSSH server, which the runner image
-      # ships only the client half of. Install it on Linux (the OS the sshd
-      # test is gated to); other platforms skip that test.
+      # The ssh acceptance test drives a real sshd behind the tunnel; the runner
+      # ships only the client. Install it on Linux (the OS the test is gated to);
+      # elsewhere and in the general test job it skips via SKIP_TEST.
       - name: Install OpenSSH server
         if: runner.os == 'Linux'
         run: sudo apt-get update && sudo apt-get install -y openssh-server

acceptance/ssh/connect-serverless-gpu/known_hosts

@@ -0,0 +1 @@
+# not actually checked by tests; accept-new appends the ephemeral host key here

acceptance/ssh/connect-serverless-gpu/out.stdout.txt

@@ -0,0 +1 @@
+Connection successful

acceptance/ssh/connect-serverless-gpu/script

@@ -7,16 +7,23 @@ if [ -z "${CLOUD_ENV:-}" ]; then
   CLI_RELEASES_DIR="$PWD/releases"
 fi
 
-errcode $CLI ssh connect --name serverless-gpu-test --accelerator=GPU_1xA10 --releases-dir=$CLI_RELEASES_DIR -- "echo 'Connection successful'" >out.stdout.txt 2>LOG.stderr
-
 if [ -n "${CLOUD_ENV:-}" ]; then
-  # On cloud the connection runs the remote command; dump the run on failure.
+  # On cloud serverless GPU compute runs the remote command over a full SSH
+  # handshake; dump the run on failure.
+  errcode $CLI ssh connect --name serverless-gpu-test --accelerator=GPU_1xA10 --releases-dir=$CLI_RELEASES_DIR -- "echo 'Connection successful'" >out.stdout.txt 2>LOG.stderr
   if ! grep -q "Connection successful" out.stdout.txt; then
     run_id=$(cat LOG.stderr | grep -o "Job submitted successfully with run ID: [0-9]*" | grep -o "[0-9]*$")
     trace $CLI jobs get-run "$run_id" > LOG.job
   fi
 else
-  # Locally the handshake can't complete (no compute); just assert the CLI
-  # submitted the right serverless GPU bootstrap job.
+  # No compute locally: the test server backs the driver-proxy /ssh websocket with a
+  # real sshd, so this one run asserts both the bootstrap job and a full handshake +
+  # remote exec over the ws:// tunnel. Needs sshd (only task test-exp-ssh provisions
+  # it); the test server probes the same way, so skip where it's absent.
+  if [ -z "$(command -v sshd || ls /usr/sbin/sshd /usr/local/sbin/sshd /sbin/sshd 2>/dev/null)" ]; then
+    echo "SKIP_TEST sshd (openssh-server) not installed"
+    exit 0
+  fi
+  errcode $CLI ssh connect --name serverless-gpu-test --accelerator=GPU_1xA10 --releases-dir=$CLI_RELEASES_DIR --user-known-hosts-file=known_hosts -- "echo 'Connection successful'" >out.stdout.txt 2>LOG.stderr
   trace print_requests.py //api/2.2/jobs/runs/submit
 fi

acceptance/ssh/connect-serverless-gpu/test.toml

@@ -12,6 +12,14 @@ Ignore = [
   "releases",
 ]
 
+# Linux-only: the local run drives a real sshd, reliable only there (and where
+# task test-exp-ssh provisions it). Absent GOOS keys default to enabled, so the
+# other OSes must be disabled explicitly.
+[GOOS]
+  linux = true
+  darwin = false
+  windows = false
+
 # Serverless GPU is not available in GCP yet
 [CloudEnvs]
   gcp = false

experimental/ssh/internal/client/websockets.go

@@ -4,19 +4,20 @@ import (
 	"context"
 	"fmt"
 	"net/http"
+	"net/url"
 
 	"github.com/databricks/cli/libs/auth"
 	"github.com/databricks/databricks-sdk-go"
 	"github.com/gorilla/websocket"
 )
 
 func createWebsocketConnection(ctx context.Context, client *databricks.WorkspaceClient, connID, clusterID string, serverPort int, liteswap string) (*websocket.Conn, error) {
-	url, err := getProxyURL(ctx, client, connID, clusterID, serverPort)
+	proxyURL, err := getProxyURL(ctx, client, connID, clusterID, serverPort)
 	if err != nil {
 		return nil, fmt.Errorf("failed to get proxy URL: %w", err)
 	}
 
-	req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil)
+	req, err := http.NewRequestWithContext(ctx, http.MethodGet, proxyURL, nil)
 	if err != nil {
 		return nil, fmt.Errorf("failed to create request: %w", err)
 	}
@@ -28,7 +29,6 @@ func createWebsocketConnection(ctx context.Context, client *databricks.Workspace
 		return nil, fmt.Errorf("failed to authenticate: %w", err)
 	}
 
-	req.URL.Scheme = "wss"
 	// websocket connection manages lifecycle of the response object, no need to close the body
 	conn, _, err := websocket.DefaultDialer.Dial(req.URL.String(), req.Header) // nolint:bodyclose
 	if err != nil {
@@ -43,10 +43,28 @@ func getProxyURL(ctx context.Context, client *databricks.WorkspaceClient, connID
 	if err != nil {
 		return "", fmt.Errorf("failed to get current workspace ID: %w", err)
 	}
-	host := client.Config.Host
+	return buildProxyWebsocketURL(client.Config.Host, workspaceID, clusterID, serverPort, connID)
+}
+
+// buildProxyWebsocketURL builds the driver-proxy websocket URL for an SSH tunnel.
+//
+// The scheme follows the host (http -> ws, else wss) instead of being hardcoded to
+// wss, so the tunnel is also diallable against the plaintext local test server.
+func buildProxyWebsocketURL(host, workspaceID, clusterID string, serverPort int, connID string) (string, error) {
+	u, err := url.Parse(host)
+	if err != nil {
+		return "", fmt.Errorf("failed to parse host %q: %w", host, err)
+	}
+	switch u.Scheme {
+	case "http":
+		u.Scheme = "ws"
+	default:
+		u.Scheme = "wss"
+	}
 	// The /driver-proxy-api/o/<workspace-id>/... path is a legacy URL form on
 	// the driver-proxy endpoint and uses an "o" path segment regardless of
 	// whether the workspace ID itself is the legacy or new shape.
-	url := fmt.Sprintf("%s/driver-proxy-api/o/%s/%s/%d/ssh?id=%s", host, workspaceID, clusterID, serverPort, connID)
-	return url, nil
+	u.Path = fmt.Sprintf("/driver-proxy-api/o/%s/%s/%d/ssh", workspaceID, clusterID, serverPort)
+	u.RawQuery = url.Values{"id": {connID}}.Encode()
+	return u.String(), nil
 }

experimental/ssh/internal/client/websockets_test.go

@@ -0,0 +1,35 @@
+package client
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func TestBuildProxyWebsocketURL(t *testing.T) {
+	tests := []struct {
+		name string
+		host string
+		want string
+	}{
+		{
+			name: "https host is dialed over wss",
+			host: "https://my-workspace.cloud.databricks.test",
+			want: "wss://my-workspace.cloud.databricks.test/driver-proxy-api/o/900800700600/1234-567890-abc/7772/ssh?id=conn-1",
+		},
+		{
+			name: "plaintext http host is dialed over ws",
+			host: "http://127.0.0.1:8080",
+			want: "ws://127.0.0.1:8080/driver-proxy-api/o/900800700600/1234-567890-abc/7772/ssh?id=conn-1",
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			got, err := buildProxyWebsocketURL(tt.host, "900800700600", "1234-567890-abc", 7772, "conn-1")
+			require.NoError(t, err)
+			assert.Equal(t, tt.want, got)
+		})
+	}
+}

libs/testserver/handlers.go

@@ -743,16 +743,19 @@ func AddDefaultHandlers(server *Server) {
 		return req.Workspace.SecretsList(req)
 	})
 
-	// SSH tunnel server behind the driver proxy: /metadata returns the remote
-	// login user, /logs is a best-effort error tail fetched on failure.
+	// SSH tunnel server behind the driver proxy: /metadata returns the remote login
+	// user, /logs is a best-effort error tail, /ssh hijacks the connection for a
+	// websocket (so it's raw, not a JSON response).
 	server.Handle("GET", "/driver-proxy-api/o/{workspace_id}/{cluster_id}/{port}/metadata", func(req Request) any {
-		return Response{Body: sshTunnelRemoteUser}
+		return Response{Body: sshTunnelRemoteUser()}
 	})
 
 	server.Handle("GET", "/driver-proxy-api/o/{workspace_id}/{cluster_id}/{port}/logs", func(req Request) any {
 		return Response{Body: ""}
 	})
 
+	server.HandleRaw("GET", "/driver-proxy-api/o/{workspace_id}/{cluster_id}/{port}/ssh", server.sshTunnelHandler)
+
 	// Secrets ACLs:
 	server.Handle("GET", "/api/2.0/secrets/acls/get", func(req Request) any {
 		return req.Workspace.SecretsAclsGet(req)

libs/testserver/jobs.go

@@ -534,7 +534,6 @@ const (
 	sshTunnelBootstrapNotebook  = "ssh-server-bootstrap"
 	sshTunnelServerPort         = 7772
 	sshTunnelClusterID          = "1234-567890-serverless"
-	sshTunnelRemoteUser         = "spark"
 )
 
 // writeSSHTunnelMetadata publishes the metadata.json a real tunnel server would