[SSH] Read the serverless port forwarding setting once, without retries

The setting check used the default SDK client, which retries 429 and 503
until its retry budget runs out, so a slow settings service could hold up
a working connection for minutes. It now reads through the no-retry probe
client the AI Gateway check already uses.

The check is also skipped when reconnecting with --metadata. The
ProxyCommand child runs Run again with the server already running, so the
read prevented nothing and was repeated on every IDE reconnect.

Co-authored-by: Isaac <no-reply@databricks.com>

Author: anton-107

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

@@ -355,7 +355,9 @@ func Run(ctx context.Context, client *databricks.WorkspaceClient, opts ClientOpt
 		cmdio.LogString(ctx, fmt.Sprintf("Connecting to %s...", sessionID))
 	}
 
-	if opts.IsServerlessMode() {
+	// 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
 		}
@@ -710,9 +712,16 @@ func metadataResponseError(statusCode int, body []byte) error {
 // (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.
+// (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 {
-	setting, err := client.WorkspaceSettingsV2.GetPublicWorkspaceSetting(ctx, settingsv2.GetPublicWorkspaceSettingRequest{
+	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 {

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"
 
@@ -20,7 +23,6 @@ import (
 	"github.com/databricks/databricks-sdk-go/service/compute"
 	"github.com/databricks/databricks-sdk-go/service/environments"
 	"github.com/databricks/databricks-sdk-go/service/jobs"
-	"github.com/databricks/databricks-sdk-go/service/settingsv2"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/mock"
 	"github.com/stretchr/testify/require"
@@ -863,38 +865,56 @@ func TestMetadataResponseError(t *testing.T) {
 func TestCheckPortForwardingToServerless(t *testing.T) {
 	tests := []struct {
 		name    string
-		setting *settingsv2.Setting
-		readErr error
+		status  int
+		body    string
 		wantErr error
 	}{
 		{
 			name:    "disabled",
-			setting: &settingsv2.Setting{EffectiveBooleanVal: &settingsv2.BooleanMessage{Value: false}},
+			status:  http.StatusOK,
+			body:    `{"effective_boolean_val": {"value": false}}`,
 			wantErr: errPortForwardingDisabled,
 		},
 		{
-			name:    "enabled",
-			setting: &settingsv2.Setting{EffectiveBooleanVal: &settingsv2.BooleanMessage{Value: true}},
+			name:   "enabled",
+			status: http.StatusOK,
+			body:   `{"effective_boolean_val": {"value": true}}`,
 		},
 		{
-			name:    "no effective value",
-			setting: &settingsv2.Setting{},
+			name:   "no effective value",
+			status: http.StatusOK,
+			body:   `{}`,
 		},
 		{
 			// Production does not serve the setting through the public API yet.
-			name:    "read fails",
-			readErr: errors.New("Setting port_fwd_serverless_enabled is not enabled for public API"),
+			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) {
-			ctx := cmdio.MockDiscard(t.Context())
-			m := mocks.NewMockWorkspaceClient(t)
-			m.GetMockWorkspaceSettingsV2API().EXPECT().
-				GetPublicWorkspaceSetting(mock.Anything, settingsv2.GetPublicWorkspaceSettingRequest{Name: portFwdServerlessSetting}).
-				Return(tc.setting, tc.readErr)
+			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()
 
-			assert.Equal(t, tc.wantErr, checkPortForwardingToServerless(ctx, m.WorkspaceClient))
+			ctx := cmdio.MockDiscard(t.Context())
+			assert.Equal(t, tc.wantErr, checkPortForwardingToServerless(ctx, newProbeClient(t, srv.URL)))
+			assert.Equal(t, int32(1), reads.Load())
 		})
 	}
 }