fix merge issues

Author: rclarey

experimental/ssh/internal/client/client.go

@@ -585,7 +585,7 @@ func fetchServerErrorLogs(ctx context.Context, client *databricks.WorkspaceClien
 
 // Assemble the SubmitRun request that bootstraps the SSH server.
 // Extracted from submitSSHTunnelJob so this logic can be unit tested.
-func buildSSHServerSubmitRun(version, secretScopeName, jobNotebookPath string, opts ClientOptions) jobs.SubmitRun {
+func buildSSHServerSubmitRun(version, secretScopeName, jobNotebookPath, baseEnvironment string, opts ClientOptions) jobs.SubmitRun {
 	sessionID := opts.SessionIdentifier()
 
 	baseParams := map[string]string{
@@ -632,11 +632,7 @@ func buildSSHServerSubmitRun(version, secretScopeName, jobNotebookPath string, o
 		// base_environment and environment_version are mutually exclusive: a custom
 		// base environment carries its own version, so we don't also set one.
 		var spec compute.Environment
-		if opts.BaseEnvironment != "" {
-			baseEnvironment, err := resolveBaseEnvironment(ctx, client, opts.BaseEnvironment)
-			if err != nil {
-				return 0, err
-			}
+		if baseEnvironment != "" {
 			spec.BaseEnvironment = baseEnvironment
 		} else {
 			spec.EnvironmentVersion = strconv.Itoa(max(opts.EnvironmentVersion, minEnvironmentVersion))
@@ -687,7 +683,15 @@ func submitSSHTunnelJob(ctx context.Context, client *databricks.WorkspaceClient,
 		log.Infof(ctx, "Using accelerator: %s", opts.Accelerator)
 	}
 
-	submitRequest := buildSSHServerSubmitRun(version, secretScopeName, jobNotebookPath, opts)
+	var baseEnvironment string
+	if opts.IsServerlessMode() && opts.BaseEnvironment != "" {
+		baseEnvironment, err = resolveBaseEnvironment(ctx, client, opts.BaseEnvironment)
+		if err != nil {
+			return 0, err
+		}
+	}
+
+	submitRequest := buildSSHServerSubmitRun(version, secretScopeName, jobNotebookPath, baseEnvironment, opts)
 
 	waiter, err := client.Jobs.Submit(ctx, submitRequest)
 	if err != nil {

experimental/ssh/internal/client/submit_internal_test.go

@@ -6,6 +6,7 @@ import (
 
 	"github.com/databricks/databricks-sdk-go/service/compute"
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 )
 
 func TestBuildSSHServerSubmitRun(t *testing.T) {
@@ -18,7 +19,7 @@ func TestBuildSSHServerSubmitRun(t *testing.T) {
 			ServerTimeout:      time.Hour,
 			EnvironmentVersion: 4,
 		}
-		got := buildSSHServerSubmitRun("v1", "scope", notebookPath, opts)
+		got := buildSSHServerSubmitRun("v1", "scope", notebookPath, "", opts)
 
 		// Usage policy flows onto the run and into the base params the server reads.
 		assert.Equal(t, "pol-1", got.BudgetPolicyId)
@@ -38,17 +39,32 @@ func TestBuildSSHServerSubmitRun(t *testing.T) {
 			Accelerator:    "GPU_1xA10",
 			ServerTimeout:  time.Hour,
 		}
-		got := buildSSHServerSubmitRun("v1", "scope", notebookPath, opts)
+		got := buildSSHServerSubmitRun("v1", "scope", notebookPath, "", opts)
 
 		assert.Equal(t, compute.HardwareAcceleratorType("GPU_1xA10"), got.Tasks[0].Compute.HardwareAccelerator)
 	})
 
+	t.Run("serverless with base environment", func(t *testing.T) {
+		opts := ClientOptions{
+			ConnectionName:     "conn",
+			ServerTimeout:      time.Hour,
+			EnvironmentVersion: 4,
+			BaseEnvironment:    "my-env",
+		}
+		got := buildSSHServerSubmitRun("v1", "scope", notebookPath, "workspace-base-environments/dbe_123", opts)
+
+		// A resolved base environment carries its own version, so environment_version is not set.
+		require.Len(t, got.Environments, 1)
+		assert.Equal(t, "workspace-base-environments/dbe_123", got.Environments[0].Spec.BaseEnvironment)
+		assert.Empty(t, got.Environments[0].Spec.EnvironmentVersion)
+	})
+
 	t.Run("dedicated cluster", func(t *testing.T) {
 		opts := ClientOptions{
 			ClusterID:     "abc-123",
 			ServerTimeout: time.Hour,
 		}
-		got := buildSSHServerSubmitRun("v1", "scope", notebookPath, opts)
+		got := buildSSHServerSubmitRun("v1", "scope", notebookPath, "", opts)
 
 		// Usage policy is serverless-only; a dedicated run carries none and targets the cluster.
 		assert.Empty(t, got.BudgetPolicyId)