From 390869e170ee9d3beed9f9a1727e2c36e3541358 Mon Sep 17 00:00:00 2001 From: Ilia Babanov Date: Mon, 19 Jan 2026 15:23:10 +0100 Subject: [PATCH 1/3] Add IDE flag and support for IDE integration for vscode and cursor --- experimental/ssh/cmd/connect.go | 11 +- experimental/ssh/cmd/setup.go | 21 +++- experimental/ssh/internal/client/client.go | 138 ++++++++++++++++++++- experimental/ssh/internal/setup/setup.go | 16 +-- 4 files changed, 163 insertions(+), 23 deletions(-) diff --git a/experimental/ssh/cmd/connect.go b/experimental/ssh/cmd/connect.go index 453a9db13b..dc0660d88f 100644 --- a/experimental/ssh/cmd/connect.go +++ b/experimental/ssh/cmd/connect.go @@ -32,6 +32,7 @@ For serverless compute: var connectionName string var accelerator string var proxyMode bool + var ide string var serverMetadata string var shutdownDelay time.Duration var maxClients int @@ -42,8 +43,9 @@ For serverless compute: var liteswap string cmd.Flags().StringVar(&clusterID, "cluster", "", "Databricks cluster ID (for dedicated clusters)") - cmd.Flags().StringVar(&connectionName, "name", "", "Connection name (for serverless compute)") - cmd.Flags().StringVar(&accelerator, "accelerator", "", "GPU accelerator type for serverless compute (GPU_1xA10 or GPU_8xH100)") + cmd.Flags().StringVar(&connectionName, "name", "", "Connection name") + cmd.Flags().StringVar(&accelerator, "accelerator", "", "GPU accelerator type (GPU_1xA10 or GPU_8xH100)") + cmd.Flags().StringVar(&ide, "ide", "", "Open remote IDE window (vscode or cursor)") cmd.Flags().DurationVar(&shutdownDelay, "shutdown-delay", defaultShutdownDelay, "Delay before shutting down the server after the last client disconnects") cmd.Flags().IntVar(&maxClients, "max-clients", defaultMaxClients, "Maximum number of SSH clients") cmd.Flags().BoolVar(&autoStartCluster, "auto-start-cluster", true, "Automatically start the cluster if it is not running") @@ -80,7 +82,7 @@ For serverless compute: wsClient := cmdctx.WorkspaceClient(ctx) if !proxyMode && clusterID == "" && connectionName == "" { - return errors.New("please provide --cluster flag with the cluster ID, or --name flag with the serverless connection name") + return errors.New("please provide --cluster flag with the cluster ID, or --name flag with the connection name (for serverless compute)") } if accelerator != "" && connectionName == "" { @@ -89,7 +91,7 @@ For serverless compute: // Remove when we add support for serverless CPU if connectionName != "" && accelerator == "" { - return errors.New("--name flag requires --accelerator to be set (e.g. for now we only support serverless GPU compute)") + return errors.New("--name flag requires --accelerator to be set (for now we only support serverless GPU compute)") } // TODO: validate connectionName if provided @@ -100,6 +102,7 @@ For serverless compute: ConnectionName: connectionName, Accelerator: accelerator, ProxyMode: proxyMode, + IDE: ide, ServerMetadata: serverMetadata, ShutdownDelay: shutdownDelay, MaxClients: maxClients, diff --git a/experimental/ssh/cmd/setup.go b/experimental/ssh/cmd/setup.go index 3e4523904c..81b7863666 100644 --- a/experimental/ssh/cmd/setup.go +++ b/experimental/ssh/cmd/setup.go @@ -1,9 +1,11 @@ package ssh import ( + "fmt" "time" "github.com/databricks/cli/cmd/root" + "github.com/databricks/cli/experimental/ssh/internal/client" "github.com/databricks/cli/experimental/ssh/internal/setup" "github.com/databricks/cli/libs/cmdctx" "github.com/spf13/cobra" @@ -43,16 +45,27 @@ an SSH host configuration to your SSH config file. cmd.RunE = func(cmd *cobra.Command, args []string) error { ctx := cmd.Context() - client := cmdctx.WorkspaceClient(ctx) - opts := setup.SetupOptions{ + wsClient := cmdctx.WorkspaceClient(ctx) + setupOpts := setup.SetupOptions{ HostName: hostName, ClusterID: clusterID, AutoStartCluster: autoStartCluster, SSHConfigPath: sshConfigPath, ShutdownDelay: shutdownDelay, - Profile: client.Config.Profile, + Profile: wsClient.Config.Profile, } - return setup.Setup(ctx, client, opts) + clientOpts := client.ClientOptions{ + ClusterID: setupOpts.ClusterID, + AutoStartCluster: setupOpts.AutoStartCluster, + ShutdownDelay: setupOpts.ShutdownDelay, + Profile: setupOpts.Profile, + } + proxyCommand, err := clientOpts.ToProxyCommand() + if err != nil { + return fmt.Errorf("failed to generate ProxyCommand: %w", err) + } + setupOpts.ProxyCommand = proxyCommand + return setup.Setup(ctx, wsClient, setupOpts) } return cmd diff --git a/experimental/ssh/internal/client/client.go b/experimental/ssh/internal/client/client.go index 839705c4ec..23572bf323 100644 --- a/experimental/ssh/internal/client/client.go +++ b/experimental/ssh/internal/client/client.go @@ -14,6 +14,7 @@ import ( "os/exec" "os/signal" "path/filepath" + "regexp" "strconv" "strings" "syscall" @@ -58,6 +59,8 @@ type ClientOptions struct { // to the cluster and proxy all traffic through stdin/stdout. // In the non proxy mode the CLI spawns an ssh client with the ProxyCommand config. ProxyMode bool + // Open remote IDE window with a specific ssh config (empty, 'vscode', or 'cursor') + IDE string // Expected format: ",,". // If present, the CLI won't attempt to start the server. ServerMetadata string @@ -171,8 +174,7 @@ func Run(ctx context.Context, client *databricks.WorkspaceClient, opts ClientOpt } // Only check cluster state for dedicated clusters - // TODO: we can remove liteswap check when we can start serverless GPU clusters via API. - if !opts.IsServerlessMode() && opts.Liteswap == "" { + if !opts.IsServerlessMode() { err := checkClusterState(ctx, client, opts.ClusterID, opts.AutoStartCluster) if err != nil { return err @@ -250,12 +252,144 @@ func Run(ctx context.Context, client *databricks.WorkspaceClient, opts ClientOpt if opts.ProxyMode { return runSSHProxy(ctx, client, serverPort, clusterID, opts) + } else if opts.IDE != "" { + return runIDE(ctx, client, userName, keyPath, serverPort, clusterID, opts) } else { cmdio.LogString(ctx, fmt.Sprintf("Additional SSH arguments: %v", opts.AdditionalArgs)) return spawnSSHClient(ctx, userName, keyPath, serverPort, clusterID, opts) } } +func runIDE(ctx context.Context, client *databricks.WorkspaceClient, userName, keyPath string, serverPort int, clusterID string, opts ClientOptions) error { + // Validate IDE value + if opts.IDE != "vscode" && opts.IDE != "cursor" { + return fmt.Errorf("invalid IDE value: %s, expected 'vscode' or 'cursor'", opts.IDE) + } + + // Get connection name + connectionName := opts.SessionIdentifier() + if connectionName == "" { + return errors.New("connection name is required for IDE integration") + } + + // Get Databricks user name for the workspace path + currentUser, err := client.CurrentUser.Me(ctx) + if err != nil { + return fmt.Errorf("failed to get current user: %w", err) + } + databricksUserName := currentUser.UserName + + // Ensure SSH config entry exists + configPath, err := getSSHConfigPath() + if err != nil { + return fmt.Errorf("failed to get SSH config path: %w", err) + } + + err = ensureSSHConfigEntry(ctx, configPath, connectionName, userName, keyPath, serverPort, clusterID, opts) + if err != nil { + return fmt.Errorf("failed to ensure SSH config entry: %w", err) + } + + // Determine the IDE command + ideCommand := "code" + if opts.IDE == "cursor" { + ideCommand = "cursor" + } + + // Construct the remote SSH URI + // Format: ssh-remote+@ /Workspace/Users// + remoteURI := fmt.Sprintf("ssh-remote+%s@%s", userName, connectionName) + remotePath := fmt.Sprintf("/Workspace/Users/%s/", databricksUserName) + + cmdio.LogString(ctx, fmt.Sprintf("Launching %s with remote URI: %s and path: %s", opts.IDE, remoteURI, remotePath)) + + // Launch the IDE + ideCmd := exec.CommandContext(ctx, ideCommand, "--remote", remoteURI, remotePath) + ideCmd.Stdout = os.Stdout + ideCmd.Stderr = os.Stderr + + return ideCmd.Run() +} + +func getSSHConfigPath() (string, error) { + homeDir, err := os.UserHomeDir() + if err != nil { + return "", fmt.Errorf("failed to get home directory: %w", err) + } + return filepath.Join(homeDir, ".ssh", "config"), nil +} + +func ensureSSHConfigEntry(ctx context.Context, configPath, hostName, userName, keyPath string, serverPort int, clusterID string, opts ClientOptions) error { + // Ensure SSH directory and config file exist + sshDir := filepath.Dir(configPath) + err := os.MkdirAll(sshDir, 0o700) + if err != nil { + return fmt.Errorf("failed to create SSH directory: %w", err) + } + + _, err = os.Stat(configPath) + if os.IsNotExist(err) { + err = os.WriteFile(configPath, []byte(""), 0o600) + if err != nil { + return fmt.Errorf("failed to create SSH config file: %w", err) + } + } else if err != nil { + return fmt.Errorf("failed to check SSH config file: %w", err) + } + + // Check if the host entry already exists + existingContent, err := os.ReadFile(configPath) + if err != nil { + return fmt.Errorf("failed to read SSH config file: %w", err) + } + + hostPattern := fmt.Sprintf(`(?m)^\s*Host\s+%s\s*$`, regexp.QuoteMeta(hostName)) + matched, err := regexp.Match(hostPattern, existingContent) + if err != nil { + return fmt.Errorf("failed to check for existing host: %w", err) + } + + if matched { + cmdio.LogString(ctx, fmt.Sprintf("SSH config entry for '%s' already exists", hostName)) + return nil + } + + // Generate ProxyCommand with server metadata + optsWithMetadata := opts + optsWithMetadata.ServerMetadata = FormatMetadata(userName, serverPort, clusterID) + + proxyCommand, err := optsWithMetadata.ToProxyCommand() + if err != nil { + return fmt.Errorf("failed to generate ProxyCommand: %w", err) + } + + // Generate host config + hostConfig := fmt.Sprintf(` +Host %s + User %s + ConnectTimeout 360 + StrictHostKeyChecking accept-new + IdentitiesOnly yes + IdentityFile %q + ProxyCommand %s +`, hostName, userName, keyPath, proxyCommand) + + // Append to config file + content := string(existingContent) + if !strings.HasSuffix(content, "\n") && content != "" { + content += "\n" + } + content += hostConfig + + err = os.WriteFile(configPath, []byte(content), 0o600) + if err != nil { + return fmt.Errorf("failed to update SSH config file: %w", err) + } + + cmdio.LogString(ctx, fmt.Sprintf("Added SSH config entry for '%s'", hostName)) + return nil +} + // getServerMetadata retrieves the server metadata from the workspace and validates it via Driver Proxy. // sessionID is the unique identifier for the session (cluster ID for dedicated clusters, connection name for serverless). // For dedicated clusters, clusterID should be the same as sessionID. diff --git a/experimental/ssh/internal/setup/setup.go b/experimental/ssh/internal/setup/setup.go index adfe204427..0d76071a65 100644 --- a/experimental/ssh/internal/setup/setup.go +++ b/experimental/ssh/internal/setup/setup.go @@ -10,7 +10,6 @@ import ( "strings" "time" - "github.com/databricks/cli/experimental/ssh/internal/client" "github.com/databricks/cli/experimental/ssh/internal/keys" "github.com/databricks/cli/libs/cmdio" "github.com/databricks/databricks-sdk-go" @@ -32,6 +31,8 @@ type SetupOptions struct { SSHKeysDir string // Optional auth profile name. If present, will be added as --profile flag to the ProxyCommand Profile string + // Proxy command to use for the SSH connection + ProxyCommand string } func validateClusterAccess(ctx context.Context, client *databricks.WorkspaceClient, clusterID string) error { @@ -62,17 +63,6 @@ func generateHostConfig(opts SetupOptions) (string, error) { return "", fmt.Errorf("failed to get local keys folder: %w", err) } - clientOpts := client.ClientOptions{ - ClusterID: opts.ClusterID, - AutoStartCluster: opts.AutoStartCluster, - ShutdownDelay: opts.ShutdownDelay, - Profile: opts.Profile, - } - proxyCommand, err := clientOpts.ToProxyCommand() - if err != nil { - return "", fmt.Errorf("failed to generate ProxyCommand: %w", err) - } - hostConfig := fmt.Sprintf(` Host %s User root @@ -81,7 +71,7 @@ Host %s IdentitiesOnly yes IdentityFile %q ProxyCommand %s -`, opts.HostName, identityFilePath, proxyCommand) +`, opts.HostName, identityFilePath, opts.ProxyCommand) return hostConfig, nil } From 59d0605144d0180ee3fae651c682a884bc5c40c1 Mon Sep 17 00:00:00 2001 From: Ilia Babanov Date: Thu, 5 Feb 2026 17:16:33 +0100 Subject: [PATCH 2/3] Mark hidden flags for serverless compute, remove manual Http call --- experimental/ssh/cmd/connect.go | 10 +- experimental/ssh/internal/client/client.go | 104 ++------------------- go.mod | 2 +- go.sum | 4 +- 4 files changed, 16 insertions(+), 104 deletions(-) diff --git a/experimental/ssh/cmd/connect.go b/experimental/ssh/cmd/connect.go index dc0660d88f..ce6fc8e2e8 100644 --- a/experimental/ssh/cmd/connect.go +++ b/experimental/ssh/cmd/connect.go @@ -43,13 +43,17 @@ For serverless compute: var liteswap string cmd.Flags().StringVar(&clusterID, "cluster", "", "Databricks cluster ID (for dedicated clusters)") - cmd.Flags().StringVar(&connectionName, "name", "", "Connection name") - cmd.Flags().StringVar(&accelerator, "accelerator", "", "GPU accelerator type (GPU_1xA10 or GPU_8xH100)") - cmd.Flags().StringVar(&ide, "ide", "", "Open remote IDE window (vscode or cursor)") cmd.Flags().DurationVar(&shutdownDelay, "shutdown-delay", defaultShutdownDelay, "Delay before shutting down the server after the last client disconnects") cmd.Flags().IntVar(&maxClients, "max-clients", defaultMaxClients, "Maximum number of SSH clients") cmd.Flags().BoolVar(&autoStartCluster, "auto-start-cluster", true, "Automatically start the cluster if it is not running") + cmd.Flags().StringVar(&connectionName, "name", "", "Connection name (for serverless compute)") + cmd.Flags().MarkHidden("name") + cmd.Flags().StringVar(&accelerator, "accelerator", "", "GPU accelerator type (GPU_1xA10 or GPU_8xH100)") + cmd.Flags().MarkHidden("accelerator") + cmd.Flags().StringVar(&ide, "ide", "", "Open remote IDE window (vscode or cursor)") + cmd.Flags().MarkHidden("ide") + cmd.Flags().BoolVar(&proxyMode, "proxy", false, "ProxyCommand mode") cmd.Flags().MarkHidden("proxy") cmd.Flags().StringVar(&serverMetadata, "metadata", "", "Metadata of the running SSH server (format: ,)") diff --git a/experimental/ssh/internal/client/client.go b/experimental/ssh/internal/client/client.go index 23572bf323..2fcc08f373 100644 --- a/experimental/ssh/internal/client/client.go +++ b/experimental/ssh/internal/client/client.go @@ -1,11 +1,9 @@ package client import ( - "bytes" "context" _ "embed" "encoding/base64" - "encoding/json" "errors" "fmt" "io" @@ -486,11 +484,6 @@ func submitSSHTunnelJob(ctx context.Context, client *databricks.WorkspaceClient, cmdio.LogString(ctx, "Submitting a job to start the ssh server...") - // Use manual HTTP call when hardware_accelerator is needed (SDK doesn't support it yet) - if opts.Accelerator != "" { - return submitSSHTunnelJobManual(ctx, client, jobNotebookPath, baseParams, opts) - } - task := jobs.SubmitTask{ TaskKey: sshServerTaskKey, NotebookTask: &jobs.NotebookTask{ @@ -502,6 +495,12 @@ func submitSSHTunnelJob(ctx context.Context, client *databricks.WorkspaceClient, if opts.IsServerlessMode() { task.EnvironmentKey = serverlessEnvironmentKey + if opts.Accelerator != "" { + cmdio.LogString(ctx, "Using accelerator: "+opts.Accelerator) + task.Compute = &jobs.Compute{ + HardwareAccelerator: compute.HardwareAcceleratorType(opts.Accelerator), + } + } } else { task.ExistingClusterId = opts.ClusterID } @@ -533,97 +532,6 @@ func submitSSHTunnelJob(ctx context.Context, client *databricks.WorkspaceClient, return waitForJobToStart(ctx, client, waiter.RunId, opts.TaskStartupTimeout) } -// submitSSHTunnelJobManual submits a job using manual HTTP call for features not yet supported by the SDK. -// Currently used for hardware_accelerator field which is not yet in the SDK. -func submitSSHTunnelJobManual(ctx context.Context, client *databricks.WorkspaceClient, jobNotebookPath string, baseParams map[string]string, opts ClientOptions) error { - sessionID := opts.SessionIdentifier() - sshTunnelJobName := "ssh-server-bootstrap-" + sessionID - - // Construct the request payload manually to allow custom parameters - task := map[string]any{ - "task_key": sshServerTaskKey, - "notebook_task": map[string]any{ - "notebook_path": jobNotebookPath, - "base_parameters": baseParams, - }, - "timeout_seconds": int(opts.ServerTimeout.Seconds()), - } - - if opts.IsServerlessMode() { - task["environment_key"] = serverlessEnvironmentKey - if opts.Accelerator != "" { - cmdio.LogString(ctx, "Using accelerator: "+opts.Accelerator) - task["compute"] = map[string]any{ - "hardware_accelerator": opts.Accelerator, - } - } - } else { - task["existing_cluster_id"] = opts.ClusterID - } - - submitRequest := map[string]any{ - "run_name": sshTunnelJobName, - "timeout_seconds": int(opts.ServerTimeout.Seconds()), - "tasks": []map[string]any{task}, - } - - if opts.IsServerlessMode() { - submitRequest["environments"] = []map[string]any{ - { - "environment_key": serverlessEnvironmentKey, - "spec": map[string]any{ - "environment_version": "3", - }, - }, - } - } - - requestBody, err := json.Marshal(submitRequest) - if err != nil { - return fmt.Errorf("failed to marshal request body: %w", err) - } - - cmdio.LogString(ctx, "Request body: "+string(requestBody)) - - apiURL := client.Config.Host + "/api/2.1/jobs/runs/submit" - req, err := http.NewRequestWithContext(ctx, "POST", apiURL, bytes.NewReader(requestBody)) - if err != nil { - return fmt.Errorf("failed to create request: %w", err) - } - - req.Header.Set("Content-Type", "application/json") - if err := client.Config.Authenticate(req); err != nil { - return fmt.Errorf("failed to authenticate request: %w", err) - } - - resp, err := http.DefaultClient.Do(req) - if err != nil { - return fmt.Errorf("failed to submit job: %w", err) - } - defer resp.Body.Close() - - responseBody, err := io.ReadAll(resp.Body) - if err != nil { - return fmt.Errorf("failed to read response body: %w", err) - } - - if resp.StatusCode != http.StatusOK { - return fmt.Errorf("failed to submit job, status code %d: %s", resp.StatusCode, string(responseBody)) - } - - var result struct { - RunID int64 `json:"run_id"` - } - if err := json.Unmarshal(responseBody, &result); err != nil { - return fmt.Errorf("failed to parse response: %w", err) - } - - cmdio.LogString(ctx, fmt.Sprintf("Job submitted successfully with run ID: %d", result.RunID)) - - // For manual submissions we still need to poll manually - return waitForJobToStart(ctx, client, result.RunID, opts.TaskStartupTimeout) -} - func spawnSSHClient(ctx context.Context, userName, privateKeyPath string, serverPort int, clusterID string, opts ClientOptions) error { // Create a copy with metadata for the ProxyCommand optsWithMetadata := opts diff --git a/go.mod b/go.mod index 27f170cc04..49042cdbf7 100644 --- a/go.mod +++ b/go.mod @@ -12,7 +12,7 @@ require ( github.com/charmbracelet/bubbletea v1.3.10 // MIT github.com/charmbracelet/huh v0.8.0 github.com/charmbracelet/lipgloss v1.1.0 // MIT - github.com/databricks/databricks-sdk-go v0.100.0 // Apache 2.0 + github.com/databricks/databricks-sdk-go v0.104.0 // Apache 2.0 github.com/fatih/color v1.18.0 // MIT github.com/google/uuid v1.6.0 // BSD-3-Clause github.com/gorilla/mux v1.8.1 // BSD 3-Clause diff --git a/go.sum b/go.sum index 3d6376294c..1d342bb862 100644 --- a/go.sum +++ b/go.sum @@ -67,8 +67,8 @@ github.com/creack/pty v1.1.24 h1:bJrF4RRfyJnbTJqzRLHzcGaZK1NeM5kTC9jGgovnR1s= github.com/creack/pty v1.1.24/go.mod h1:08sCNb52WyoAwi2QDyzUCTgcvVFhUzewun7wtTfvcwE= github.com/cyphar/filepath-securejoin v0.4.1 h1:JyxxyPEaktOD+GAnqIqTf9A8tHyAG22rowi7HkoSU1s= github.com/cyphar/filepath-securejoin v0.4.1/go.mod h1:Sdj7gXlvMcPZsbhwhQ33GguGLDGQL7h7bg04C/+u9jI= -github.com/databricks/databricks-sdk-go v0.100.0 h1:1jODpEVBlKKrtLG/HyTkhLjhWCwdlIeGMMifp9ZweRk= -github.com/databricks/databricks-sdk-go v0.100.0/go.mod h1:hWoHnHbNLjPKiTm5K/7bcIv3J3Pkgo5x9pPzh8K3RVE= +github.com/databricks/databricks-sdk-go v0.104.0 h1:8V8JAebwuQjaMszA1e3F+BtQ816oSUQm2yU8mmbRc28= +github.com/databricks/databricks-sdk-go v0.104.0/go.mod h1:hWoHnHbNLjPKiTm5K/7bcIv3J3Pkgo5x9pPzh8K3RVE= github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c= github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= From b5cc387d1098cd624090d4dc103db330a5ab9485 Mon Sep 17 00:00:00 2001 From: Ilia Babanov Date: Fri, 6 Feb 2026 14:38:58 +0100 Subject: [PATCH 3/3] Separate IDE options into constants --- experimental/ssh/internal/client/client.go | 19 ++++++++++--------- 1 file changed, 10 insertions(+), 9 deletions(-) diff --git a/experimental/ssh/internal/client/client.go b/experimental/ssh/internal/client/client.go index 2fcc08f373..6c4695b6ac 100644 --- a/experimental/ssh/internal/client/client.go +++ b/experimental/ssh/internal/client/client.go @@ -40,6 +40,11 @@ var errServerMetadata = errors.New("server metadata error") const ( sshServerTaskKey = "start_ssh_server" serverlessEnvironmentKey = "ssh_tunnel_serverless" + + VSCodeOption = "vscode" + VSCodeCommand = "code" + CursorOption = "cursor" + CursorCommand = "cursor" ) type ClientOptions struct { @@ -259,12 +264,10 @@ func Run(ctx context.Context, client *databricks.WorkspaceClient, opts ClientOpt } func runIDE(ctx context.Context, client *databricks.WorkspaceClient, userName, keyPath string, serverPort int, clusterID string, opts ClientOptions) error { - // Validate IDE value - if opts.IDE != "vscode" && opts.IDE != "cursor" { - return fmt.Errorf("invalid IDE value: %s, expected 'vscode' or 'cursor'", opts.IDE) + if opts.IDE != VSCodeOption && opts.IDE != CursorOption { + return fmt.Errorf("invalid IDE value: %s, expected '%s' or '%s'", opts.IDE, VSCodeOption, CursorOption) } - // Get connection name connectionName := opts.SessionIdentifier() if connectionName == "" { return errors.New("connection name is required for IDE integration") @@ -288,10 +291,9 @@ func runIDE(ctx context.Context, client *databricks.WorkspaceClient, userName, k return fmt.Errorf("failed to ensure SSH config entry: %w", err) } - // Determine the IDE command - ideCommand := "code" - if opts.IDE == "cursor" { - ideCommand = "cursor" + ideCommand := VSCodeCommand + if opts.IDE == CursorOption { + ideCommand = CursorCommand } // Construct the remote SSH URI @@ -301,7 +303,6 @@ func runIDE(ctx context.Context, client *databricks.WorkspaceClient, userName, k cmdio.LogString(ctx, fmt.Sprintf("Launching %s with remote URI: %s and path: %s", opts.IDE, remoteURI, remotePath)) - // Launch the IDE ideCmd := exec.CommandContext(ctx, ideCommand, "--remote", remoteURI, remotePath) ideCmd.Stdout = os.Stdout ideCmd.Stderr = os.Stderr