[SSH] Check port_fwd_serverless_enabled before starting serverless compute

acceptance/ssh/connect-serverless-feature-disabled/script

@@ -1,6 +1,7 @@
 # An administrator disabled port forwarding to serverless compute, so the driver proxy refuses
-# every request to the tunnel server. The CLI must report the proxy's reason as soon as it gets
-# the 403 instead of polling until the server start timeout, and cancel the run it submitted
+# every request to the tunnel server. The setting itself cannot be read here (see test.toml), so
+# the CLI goes ahead and starts the server. It must then report the proxy's reason as soon as it
+# gets the 403 instead of polling until the server start timeout, and cancel the run it submitted
 # since nothing can ever connect to that server.
 # The body repeats the name: with only {"value": false} in it, the API sees an empty message.
 trace $CLI workspace-settings-v2 patch-public-workspace-setting port_fwd_serverless_enabled --json '{"name": "port_fwd_serverless_enabled", "boolean_val": {"value": false}}' >/dev/null

acceptance/ssh/connect-serverless-feature-disabled/test.toml

@@ -14,3 +14,16 @@ GOOS.darwin = false
 GOOS.windows = false
 
 EnvMatrix.DATABRICKS_BUNDLE_ENGINE = ["direct"]
+
+# Production does not serve port_fwd_serverless_enabled through the public API, so the CLI cannot
+# check it up front there and only learns about it from the driver proxy's 403. This reproduces
+# that environment, which the fake workspace's settings handler does not model.
+[[Server]]
+Pattern = "GET /api/2.1/settings/port_fwd_serverless_enabled"
+Response.StatusCode = 404
+Response.Body = '''
+{
+  "error_code": "NOT_FOUND",
+  "message": "Setting port_fwd_serverless_enabled is not enabled for public API"
+}
+'''

acceptance/ssh/connect-serverless-port-forwarding-disabled/out.test.toml

@@ -0,0 +1,5 @@
+Cloud = false
+GOOS.darwin = false
+GOOS.linux = true
+GOOS.windows = false
+EnvMatrix.DATABRICKS_BUNDLE_ENGINE = ["direct"]

acceptance/ssh/connect-serverless-port-forwarding-disabled/output.txt

@@ -0,0 +1,10 @@
+
+>>> [CLI] workspace-settings-v2 patch-public-workspace-setting port_fwd_serverless_enabled --json {"name": "port_fwd_serverless_enabled", "boolean_val": {"value": false}}
+
+=== CLI ssh connect
+
+>>> musterr [CLI] ssh connect --name port-forwarding-disabled -- echo 'Connection successful'
+Connecting to port-forwarding-disabled...
+Error: port forwarding to serverless compute is disabled for this workspace (setting port_fwd_serverless_enabled), and the SSH tunnel needs it. Ask a workspace or account admin to enable it, or connect to a dedicated cluster with --cluster
+
+>>> print_requests.py //api

acceptance/ssh/connect-serverless-port-forwarding-disabled/script

@@ -0,0 +1,13 @@
+# An administrator disabled port forwarding to serverless compute. The CLI reads the setting
+# before anything else, so it refuses the connection without creating a secret scope, uploading
+# binaries or starting compute that nobody could reach.
+# The body repeats the name: with only {"value": false} in it, the API sees an empty message.
+trace $CLI workspace-settings-v2 patch-public-workspace-setting port_fwd_serverless_enabled --json '{"name": "port_fwd_serverless_enabled", "boolean_val": {"value": false}}' >/dev/null
+# Drop the setup request so the check below only sees what ssh connect sent.
+print_requests.py //api/2.1/settings >/dev/null
+
+title "CLI ssh connect\n"
+trace musterr $CLI ssh connect --name port-forwarding-disabled -- "echo 'Connection successful'"
+
+# Expect no requests other than reads: the connection stops at the settings check.
+trace print_requests.py //api

acceptance/ssh/connect-serverless-port-forwarding-disabled/test.toml

@@ -0,0 +1,11 @@
+Cloud = false
+
+# Record requests so we can assert that nothing is created before the CLI gives up.
+RecordRequests = true
+
+# The rest of acceptance/ssh only runs on Linux; nothing here needs another OS.
+GOOS.linux = true
+GOOS.darwin = false
+GOOS.windows = false
+
+EnvMatrix.DATABRICKS_BUNDLE_ENGINE = ["direct"]

experimental/ssh/internal/client/agentshim.go

@@ -360,7 +360,7 @@ func probeAIGateway(ctx context.Context, wsclient *databricks.WorkspaceClient) e
 }
 
 // newProbeAPIClient builds a low-level client from cfg with the SDK's retries
-// turned off, so the gateway preflight is one-shot. Clearing ErrorRetriable and
+// turned off, so a preflight check is one-shot. Clearing ErrorRetriable and
 // TransientErrors makes a retriable status (429, 503, ...) surface on the first
 // attempt rather than looping until the retry budget expires; it mirrors
 // client.New, which builds the retrying client the rest of the CLI uses.

experimental/ssh/internal/client/client.go

@@ -38,6 +38,7 @@ import (
 	"github.com/databricks/databricks-sdk-go/service/environments"
 	"github.com/databricks/databricks-sdk-go/service/iam"
 	"github.com/databricks/databricks-sdk-go/service/jobs"
+	"github.com/databricks/databricks-sdk-go/service/settingsv2"
 	"github.com/databricks/databricks-sdk-go/service/workspace"
 	"github.com/gorilla/websocket"
 )
@@ -47,6 +48,14 @@ var sshServerBootstrapScript string
 
 var errServerMetadata = errors.New("server metadata error")
 
+// portFwdServerlessSetting is the workspace setting an administrator turns off to disable port
+// forwarding to serverless compute, which the tunnel needs.
+const portFwdServerlessSetting = "port_fwd_serverless_enabled"
+
+var errPortForwardingDisabled = errors.New("port forwarding to serverless compute is disabled for this workspace " +
+	"(setting " + portFwdServerlessSetting + "), and the SSH tunnel needs it. " +
+	"Ask a workspace or account admin to enable it, or connect to a dedicated cluster with --cluster")
+
 // errDriverProxyForbidden is a 403 from the driver proxy. Unlike the other non-200 responses, which a
 // server that is still starting produces too, retrying cannot change it: the workspace forbids the
 // connection, e.g. because an administrator disabled port forwarding to serverless compute.
@@ -346,6 +355,14 @@ func Run(ctx context.Context, client *databricks.WorkspaceClient, opts ClientOpt
 		cmdio.LogString(ctx, fmt.Sprintf("Connecting to %s...", sessionID))
 	}
 
+	// Skipped when reconnecting with --metadata (the ProxyCommand child): the server is already
+	// running, so there is no compute left to spare, and every IDE reconnect would repeat the read.
+	if opts.IsServerlessMode() && opts.ServerMetadata == "" {
+		if err := checkPortForwardingToServerless(ctx, client); err != nil {
+			return err
+		}
+	}
+
 	if opts.IDE != "" && !opts.ProxyMode {
 		if err := vscode.CheckIDECommand(opts.IDE); err != nil {
 			outcome.errorCategory = protos.SshTunnelErrorCategoryIDECommandNotOnPath
@@ -690,6 +707,33 @@ func metadataResponseError(statusCode int, body []byte) error {
 	return fmt.Errorf("%w: %s", errDriverProxyForbidden, message)
 }
 
+// checkPortForwardingToServerless refuses a serverless connection before any compute is started when
+// the workspace setting that the driver proxy enforces for the tunnel is off. The proxy's gate
+// (PortForwardingToServerlessGate in data-plane-router) also depends on an enforcement flag and a port
+// allow-list that the CLI cannot see, and the setting is not readable through the public API in every
+// environment. A read that fails therefore lets the connection go ahead, and the proxy's own 403
+// (errDriverProxyForbidden) remains the backstop. The read goes through a client with retries turned
+// off, so a throttled or unavailable settings service cannot hold up the connection for the SDK's
+// retry budget.
+func checkPortForwardingToServerless(ctx context.Context, client *databricks.WorkspaceClient) error {
+	apiClient, err := newProbeAPIClient(client.Config)
+	if err != nil {
+		log.Debugf(ctx, "Could not build a client to read the %s setting, continuing: %v", portFwdServerlessSetting, err)
+		return nil
+	}
+	setting, err := settingsv2.NewWorkspaceSettingsV2(apiClient).GetPublicWorkspaceSetting(ctx, settingsv2.GetPublicWorkspaceSettingRequest{
+		Name: portFwdServerlessSetting,
+	})
+	if err != nil {
+		log.Debugf(ctx, "Could not read the %s setting, continuing: %v", portFwdServerlessSetting, err)
+		return nil
+	}
+	if setting.EffectiveBooleanVal != nil && !setting.EffectiveBooleanVal.Value {
+		return errPortForwardingDisabled
+	}
+	return nil
+}
+
 // newDriverProxyRequest builds an authenticated GET request to one of the SSH server's
 // HTTP endpoints behind the workspace driver proxy.
 func newDriverProxyRequest(ctx context.Context, client *databricks.WorkspaceClient, clusterID string, port int, endpoint, liteswap string) (*http.Request, error) {

experimental/ssh/internal/client/client_internal_test.go

@@ -5,9 +5,12 @@ import (
 	"encoding/json"
 	"errors"
 	"fmt"
+	"net/http"
+	"net/http/httptest"
 	"slices"
 	"strconv"
 	"strings"
+	"sync/atomic"
 	"testing"
 	"time"
 
@@ -858,3 +861,60 @@ func TestMetadataResponseError(t *testing.T) {
 		})
 	}
 }
+
+func TestCheckPortForwardingToServerless(t *testing.T) {
+	tests := []struct {
+		name    string
+		status  int
+		body    string
+		wantErr error
+	}{
+		{
+			name:    "disabled",
+			status:  http.StatusOK,
+			body:    `{"effective_boolean_val": {"value": false}}`,
+			wantErr: errPortForwardingDisabled,
+		},
+		{
+			name:   "enabled",
+			status: http.StatusOK,
+			body:   `{"effective_boolean_val": {"value": true}}`,
+		},
+		{
+			name:   "no effective value",
+			status: http.StatusOK,
+			body:   `{}`,
+		},
+		{
+			// Production does not serve the setting through the public API yet.
+			name:   "read fails",
+			status: http.StatusBadRequest,
+			body:   `{"error_code": "INVALID_PARAMETER_VALUE", "message": "Setting port_fwd_serverless_enabled is not enabled for public API"}`,
+		},
+		{
+			// A retriable status is not retried, so a throttled service cannot hold up the connection.
+			name:   "throttled",
+			status: http.StatusTooManyRequests,
+			body:   `{"error_code": "TOO_MANY_REQUESTS", "message": "slow down"}`,
+		},
+	}
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			var reads atomic.Int32
+			srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+				if r.URL.Path != "/api/2.1/settings/"+portFwdServerlessSetting {
+					w.WriteHeader(http.StatusNotFound)
+					return
+				}
+				reads.Add(1)
+				w.WriteHeader(tc.status)
+				_, _ = w.Write([]byte(tc.body))
+			}))
+			defer srv.Close()
+
+			ctx := cmdio.MockDiscard(t.Context())
+			assert.Equal(t, tc.wantErr, checkPortForwardingToServerless(ctx, newProbeClient(t, srv.URL)))
+			assert.Equal(t, int32(1), reads.Load())
+		})
+	}
+}

libs/testserver/handlers.go

@@ -1070,6 +1070,10 @@ func AddDefaultHandlers(server *Server) {
 
 	server.HandleRaw("GET", "/driver-proxy-api/o/{workspace_id}/{cluster_id}/{port}/ssh", server.sshTunnelHandler)
 
+	server.Handle("GET", "/api/2.1/settings/{name}", func(req Request) any {
+		return req.Workspace.WorkspaceSettingGet(req)
+	})
+
 	server.Handle("PATCH", "/api/2.1/settings/{name}", func(req Request) any {
 		return req.Workspace.WorkspaceSettingPatch(req)
 	})

libs/testserver/settings.go

@@ -11,6 +11,24 @@ import (
 // which the driver proxy enforces for the SSH tunnel.
 const portFwdServerlessSetting = "port_fwd_serverless_enabled"
 
+func (s *FakeWorkspace) WorkspaceSettingGet(req Request) Response {
+	defer s.LockUnlock()()
+
+	name := req.Vars["name"]
+	if setting, ok := s.WorkspaceSettings[name]; ok {
+		return Response{Body: setting}
+	}
+	// The setting's own default is true: port forwarding stays on until an administrator turns it off.
+	if name == portFwdServerlessSetting {
+		enabled := &settingsv2.BooleanMessage{Value: true}
+		return Response{Body: settingsv2.Setting{Name: name, EffectiveBooleanVal: enabled}}
+	}
+	return Response{
+		StatusCode: 404,
+		Body:       map[string]string{"error_code": "RESOURCE_DOES_NOT_EXIST", "message": fmt.Sprintf("Setting %s does not exist", name)},
+	}
+}
+
 func (s *FakeWorkspace) WorkspaceSettingPatch(req Request) Response {
 	var setting settingsv2.Setting
 	if err := json.Unmarshal(req.Body, &setting); err != nil {

libs/testserver/settings_test.go

@@ -3,7 +3,9 @@ package testserver
 import (
 	"testing"
 
+	"github.com/databricks/databricks-sdk-go/service/settingsv2"
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 )
 
 func patchSetting(workspace *FakeWorkspace, body string) Response {
@@ -30,3 +32,35 @@ func TestWorkspaceSettingPatchDisablesPortForwarding(t *testing.T) {
 	assert.Equal(t, 0, response.StatusCode)
 	assert.True(t, workspace.portForwardingToServerlessDisabled())
 }
+
+func TestWorkspaceSettingGetDefaultsPortForwardingToEnabled(t *testing.T) {
+	workspace := NewFakeWorkspace("http://test.test", "token")
+
+	response := workspace.WorkspaceSettingGet(Request{Vars: map[string]string{"name": portFwdServerlessSetting}})
+
+	assert.Equal(t, 0, response.StatusCode)
+	setting, ok := response.Body.(settingsv2.Setting)
+	require.True(t, ok)
+	require.NotNil(t, setting.EffectiveBooleanVal)
+	assert.True(t, setting.EffectiveBooleanVal.Value)
+}
+
+func TestWorkspaceSettingGetReturnsPatchedValue(t *testing.T) {
+	workspace := NewFakeWorkspace("http://test.test", "token")
+	patchSetting(workspace, `{"name": "port_fwd_serverless_enabled", "boolean_val": {"value": false}}`)
+
+	response := workspace.WorkspaceSettingGet(Request{Vars: map[string]string{"name": portFwdServerlessSetting}})
+
+	setting, ok := response.Body.(settingsv2.Setting)
+	require.True(t, ok)
+	require.NotNil(t, setting.EffectiveBooleanVal)
+	assert.False(t, setting.EffectiveBooleanVal.Value)
+}
+
+func TestWorkspaceSettingGetUnknownSetting(t *testing.T) {
+	workspace := NewFakeWorkspace("http://test.test", "token")
+
+	response := workspace.WorkspaceSettingGet(Request{Vars: map[string]string{"name": "no_such_setting"}})
+
+	assert.Equal(t, 404, response.StatusCode)
+}