Split the SSH extension telemetry category by failure outcome

experimental/ssh/internal/client/client.go

@@ -280,6 +280,12 @@ func Run(ctx context.Context, client *databricks.WorkspaceClient, opts ClientOpt
 	outcome := connectOutcome{isReconnect: opts.ServerMetadata != ""}
 	defer func() {
 		outcome.err = retErr
+		// A cancelled context is the only trace a Ctrl-C leaves: exec.CommandContext kills the
+		// child and reports *exec.ExitError, which does not wrap context.Canceled, so a step
+		// that shells out cannot recognise the interruption itself. This defer is registered
+		// after `defer cancel()` and so runs before it (LIFO), which means ctx is cancelled
+		// here only by the signal handler or the caller, never by Run's own cleanup.
+		outcome.ctxErr = ctx.Err()
 		logSshTunnelEvent(ctx, opts, outcome)
 	}()
 
@@ -298,7 +304,7 @@ func Run(ctx context.Context, client *databricks.WorkspaceClient, opts ClientOpt
 			return err
 		}
 		if err := vscode.CheckIDESSHExtension(ctx, opts.IDE, opts.AutoApprove); err != nil {
-			outcome.errorCategory = protos.SshTunnelErrorCategoryIDESSHExtensionMissing
+			outcome.errorCategory = sshExtensionErrorCategory(err)
 			return err
 		}
 	}
@@ -1245,17 +1251,41 @@ type connectOutcome struct {
 	// errorCategory is set at the failure site. Empty means the failure was not attributed.
 	errorCategory protos.SshTunnelErrorCategory
 	err           error
+	// ctxErr is the connect context's error when the outcome is logged. Tracked apart from err
+	// because a step that shells out reports a killed child as *exec.ExitError, which carries no
+	// trace of the cancellation. Only a cancellation counts as the user giving up: a deadline
+	// would be a timeout, so category() matches the cause rather than testing for non-nil.
+	ctxErr error
 }
 
-// category returns the error category to report. A cancelled context means the user
-// interrupted the attempt, whichever call happened to observe it first, so it wins over the
-// category recorded at the failure site. An unattributed failure is reported as UNKNOWN so
-// that it stays countable.
+// sshExtensionErrorCategory attributes a Remote SSH extension check failure to the outcome that
+// caused it. The four are kept apart because they imply different fixes, and because only the
+// first two can occur under --auto-approve, which the IDE button always passes -- so a shift
+// between them and the consent outcomes distinguishes button traffic from direct CLI use.
+func sshExtensionErrorCategory(err error) protos.SshTunnelErrorCategory {
+	switch {
+	case errors.Is(err, vscode.ErrSSHExtensionListFailed):
+		return protos.SshTunnelErrorCategoryIDESSHExtensionListFailed
+	case errors.Is(err, vscode.ErrSSHExtensionInstallFailed):
+		return protos.SshTunnelErrorCategoryIDESSHExtensionInstallFailed
+	case errors.Is(err, vscode.ErrSSHExtensionInstallDeclined):
+		return protos.SshTunnelErrorCategoryIDESSHExtensionInstallDeclined
+	case errors.Is(err, vscode.ErrSSHExtensionInstallUnavailable):
+		return protos.SshTunnelErrorCategoryIDESSHExtensionInstallUnavailable
+	}
+	// CheckIDESSHExtension wraps a sentinel on every failure path, so this is only reachable if
+	// a new one is added without a category. UNKNOWN keeps it countable; see category() below.
+	return protos.SshTunnelErrorCategoryUnknown
+}
+
+// category returns the error category to report. An interrupted attempt means the user gave
+// up, whichever call happened to observe it first, so it wins over the category recorded at
+// the failure site. An unattributed failure is reported as UNKNOWN so that it stays countable.
 func (o connectOutcome) category() protos.SshTunnelErrorCategory {
 	if o.isSuccess || o.err == nil {
 		return protos.SshTunnelErrorCategoryUnspecified
 	}
-	if errors.Is(o.err, context.Canceled) {
+	if errors.Is(o.ctxErr, context.Canceled) || errors.Is(o.err, context.Canceled) {
 		return protos.SshTunnelErrorCategoryUserAborted
 	}
 	if o.errorCategory == "" {

experimental/ssh/internal/client/client_internal_test.go

@@ -12,6 +12,7 @@ import (
 	"time"
 
 	"github.com/databricks/cli/experimental/ssh/internal/sshconfig"
+	"github.com/databricks/cli/experimental/ssh/internal/vscode"
 	"github.com/databricks/cli/libs/cmdio"
 	"github.com/databricks/cli/libs/telemetry/protos"
 	"github.com/databricks/databricks-sdk-go/experimental/mocks"
@@ -549,6 +550,38 @@ func TestConnectOutcomeCategory(t *testing.T) {
 			},
 			want: protos.SshTunnelErrorCategoryUserAborted,
 		},
+		{
+			// A step that shells out reports a killed child as *exec.ExitError, which does not
+			// wrap context.Canceled, so the cancelled context is the only evidence left. Without
+			// this branch a Ctrl-C during the extension install counts as a rejected install and
+			// pollutes the bucket that is supposed to mean "the marketplace or a policy blocked
+			// it" -- one of only two reachable under --auto-approve.
+			name: "interruption wins over the category set at the failure site",
+			outcome: connectOutcome{
+				ctxErr:        context.Canceled,
+				errorCategory: protos.SshTunnelErrorCategoryIDESSHExtensionInstallFailed,
+				err:           errors.New("signal: killed"),
+			},
+			want: protos.SshTunnelErrorCategoryUserAborted,
+		},
+		{
+			// Only a cancellation is the user giving up. No ancestor of the connect context
+			// carries a deadline today, so this is unreachable; matching the cause rather than
+			// testing ctxErr for non-nil keeps it that way if one is ever added.
+			name: "an expired deadline is not a user abort",
+			outcome: connectOutcome{
+				ctxErr:        context.DeadlineExceeded,
+				errorCategory: protos.SshTunnelErrorCategoryServerStartTimeout,
+				err:           errFailed,
+			},
+			want: protos.SshTunnelErrorCategoryServerStartTimeout,
+		},
+		{
+			// Interrupting an established tunnel is not a connection failure.
+			name:    "interruption after a successful connection reports no category",
+			outcome: connectOutcome{isSuccess: true, ctxErr: context.Canceled, err: errFailed},
+			want:    protos.SshTunnelErrorCategoryUnspecified,
+		},
 	}
 
 	for _, tt := range tests {
@@ -558,6 +591,50 @@ func TestConnectOutcomeCategory(t *testing.T) {
 	}
 }
 
+// The four Remote SSH extension outcomes were reported as one category until they were split,
+// which left the largest IDE-mode failure bucket unattributable. Pin the mapping, including the
+// wrapping, since CheckIDESSHExtension returns its sentinels wrapped in a message.
+func TestSshExtensionErrorCategory(t *testing.T) {
+	tests := []struct {
+		name string
+		err  error
+		want protos.SshTunnelErrorCategory
+	}{
+		{
+			name: "list failure",
+			err:  fmt.Errorf("%w in VS Code: %w", vscode.ErrSSHExtensionListFailed, errors.New("exit 4")),
+			want: protos.SshTunnelErrorCategoryIDESSHExtensionListFailed,
+		},
+		{
+			name: "install failure",
+			err:  fmt.Errorf("%w: %w", vscode.ErrSSHExtensionInstallFailed, errors.New("exit 3")),
+			want: protos.SshTunnelErrorCategoryIDESSHExtensionInstallFailed,
+		},
+		{
+			name: "user declined the install",
+			err:  fmt.Errorf("%w: install it with ...", vscode.ErrSSHExtensionInstallDeclined),
+			want: protos.SshTunnelErrorCategoryIDESSHExtensionInstallDeclined,
+		},
+		{
+			name: "no way to ask for consent",
+			err:  fmt.Errorf("%w: install it with ...", vscode.ErrSSHExtensionInstallUnavailable),
+			want: protos.SshTunnelErrorCategoryIDESSHExtensionInstallUnavailable,
+		},
+		{
+			// Only reachable if a new failure path forgets its sentinel.
+			name: "unsentinelled failure falls back to UNKNOWN",
+			err:  errors.New("something else"),
+			want: protos.SshTunnelErrorCategoryUnknown,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			assert.Equal(t, tt.want, sshExtensionErrorCategory(tt.err))
+		})
+	}
+}
+
 func TestBuildSshTunnelEventReportsErrorCategory(t *testing.T) {
 	got := buildSshTunnelEvent(ClientOptions{ConnectionName: "my-conn", IDE: "vscode"}, connectOutcome{
 		errorCategory: protos.SshTunnelErrorCategoryIDECommandNotOnPath,

experimental/ssh/internal/vscode/run.go

@@ -2,6 +2,7 @@ package vscode
 
 import (
 	"context"
+	"errors"
 	"fmt"
 	"os"
 	"os/exec"
@@ -118,15 +119,47 @@ func isExtensionVersionAtLeast(version, minVersion string) bool {
 	return semver.IsValid(v) && semver.Compare(v, "v"+minVersion) >= 0
 }
 
+// The ways CheckIDESSHExtension can fail. Callers match these with errors.Is to attribute a
+// failure without matching on message text. They are separate because they call for different
+// fixes: a list failure means the check never ran, an install failure points at the marketplace
+// or a policy blocking it, and the two consent outcomes cannot happen under --auto-approve.
+var (
+	ErrSSHExtensionListFailed         = errors.New("could not list installed extensions")
+	ErrSSHExtensionInstallFailed      = errors.New("could not install the Remote SSH extension")
+	ErrSSHExtensionInstallDeclined    = errors.New("install of the Remote SSH extension declined")
+	ErrSSHExtensionInstallUnavailable = errors.New("cannot prompt to install the Remote SSH extension")
+)
+
+// consentError renders only the message written for the user while still matching its sentinel
+// via errors.Is. The two exec failures wrap theirs with %w instead: their message ends in the
+// underlying error, so the sentinel reads as a natural prefix. A consent message is a full
+// sentence that already states the problem, and prefixing it would lead with an internal tag
+// and then restate what follows.
+type consentError struct {
+	msg      string
+	sentinel error
+}
+
+func (e *consentError) Error() string { return e.msg }
+
+func (e *consentError) Unwrap() error { return e.sentinel }
+
+// consentErrorf formats a user-facing message and tags it with sentinel.
+func consentErrorf(sentinel error, format string, args ...any) error {
+	return &consentError{msg: fmt.Sprintf(format, args...), sentinel: sentinel}
+}
+
 // CheckIDESSHExtension verifies that the required Remote SSH extension is installed
 // with a compatible version, and offers to install/update it if not.
 // When autoApprove is true, the extension is installed without asking.
+//
+// Every returned error wraps one of the Err* sentinels above.
 func CheckIDESSHExtension(ctx context.Context, option string, autoApprove bool) error {
 	ide := getIDE(option)
 
 	out, err := exec.CommandContext(ctx, ide.Command, "--list-extensions", "--show-versions").Output()
 	if err != nil {
-		return fmt.Errorf("failed to list %s extensions: %w", ide.Name, err)
+		return fmt.Errorf("%w in %s: %w", ErrSSHExtensionListFailed, ide.Name, err)
 	}
 
 	version, found := parseExtensionVersion(string(out), ide.SSHExtensionID)
@@ -144,16 +177,18 @@ func CheckIDESSHExtension(ctx context.Context, option string, autoApprove bool)
 
 	if !autoApprove {
 		if !cmdio.IsPromptSupported(ctx) {
-			return fmt.Errorf("%s Install it with: %s --install-extension %s, or pass --auto-approve",
+			return consentErrorf(ErrSSHExtensionInstallUnavailable,
+				"%s Install it with: %s --install-extension %s, or pass --auto-approve",
 				msg, ide.Command, ide.SSHExtensionID)
 		}
 
 		shouldInstall, err := cmdio.AskYesOrNo(ctx, msg+" Would you like to install it?")
 		if err != nil {
-			return fmt.Errorf("failed to prompt user: %w", err)
+			return fmt.Errorf("%w: %w", ErrSSHExtensionInstallUnavailable, err)
 		}
 		if !shouldInstall {
-			return fmt.Errorf("%s Install it with: %s --install-extension %s",
+			return consentErrorf(ErrSSHExtensionInstallDeclined,
+				"%s Install it with: %s --install-extension %s",
 				msg, ide.Command, ide.SSHExtensionID)
 		}
 	} else {
@@ -165,7 +200,7 @@ func CheckIDESSHExtension(ctx context.Context, option string, autoApprove bool)
 	installCmd.Stdout = os.Stdout
 	installCmd.Stderr = os.Stderr
 	if err := installCmd.Run(); err != nil {
-		return fmt.Errorf("failed to install extension %q: %w", ide.SSHExtensionName, err)
+		return fmt.Errorf("%w in %s: %w", ErrSSHExtensionInstallFailed, ide.Name, err)
 	}
 	return nil
 }

experimental/ssh/internal/vscode/run_test.go

@@ -2,6 +2,7 @@ package vscode
 
 import (
 	"fmt"
+	"io"
 	"os"
 	"path/filepath"
 	"runtime"
@@ -232,6 +233,37 @@ func createFakeIDEExecutable(t *testing.T, dir, command, output string) {
 	}
 }
 
+// createFailingIDEExecutable writes a fake IDE command that exits non-zero for every
+// invocation, so "--list-extensions" fails and the check never learns what is installed.
+func createFailingIDEExecutable(t *testing.T, dir, command string) {
+	t.Helper()
+	if runtime.GOOS == "windows" {
+		err := os.WriteFile(filepath.Join(dir, command+".cmd"), []byte("@echo off\nexit /b 4\n"), 0o755)
+		require.NoError(t, err)
+	} else {
+		err := os.WriteFile(filepath.Join(dir, command), []byte("#!/bin/sh\nexit 4\n"), 0o755)
+		require.NoError(t, err)
+	}
+}
+
+// createIDEExecutableFailingInstall writes a fake IDE command that lists extensions
+// successfully but rejects "--install-extension", as a marketplace or policy block would.
+func createIDEExecutableFailingInstall(t *testing.T, dir, command, output string) {
+	t.Helper()
+	if runtime.GOOS == "windows" {
+		payloadPath := filepath.Join(dir, command+"-payload.txt")
+		err := os.WriteFile(payloadPath, []byte(output), 0o644)
+		require.NoError(t, err)
+		script := fmt.Sprintf("@echo off\nif \"%%1\"==\"--install-extension\" exit /b 3\ntype \"%s\"\n", payloadPath)
+		err = os.WriteFile(filepath.Join(dir, command+".cmd"), []byte(script), 0o755)
+		require.NoError(t, err)
+	} else {
+		script := fmt.Sprintf("#!/bin/sh\nfor a in \"$@\"; do\n  [ \"$a\" = \"--install-extension\" ] && exit 3\ndone\nprintf '%%s' '%s'\n", output)
+		err := os.WriteFile(filepath.Join(dir, command), []byte(script), 0o755)
+		require.NoError(t, err)
+	}
+}
+
 func TestCheckIDESSHExtension_UpToDate(t *testing.T) {
 	tmpDir := t.TempDir()
 	t.Setenv("PATH", tmpDir)
@@ -268,6 +300,8 @@ func TestCheckIDESSHExtension_Missing(t *testing.T) {
 	require.Error(t, err)
 	assert.Contains(t, err.Error(), `"Remote - SSH"`)
 	assert.Contains(t, err.Error(), "not installed")
+	// The test context is not a TTY, so consent cannot be asked for.
+	assert.ErrorIs(t, err, ErrSSHExtensionInstallUnavailable)
 }
 
 func TestCheckIDESSHExtension_Outdated(t *testing.T) {
@@ -282,6 +316,7 @@ func TestCheckIDESSHExtension_Outdated(t *testing.T) {
 	require.Error(t, err)
 	assert.Contains(t, err.Error(), "0.100.0")
 	assert.Contains(t, err.Error(), ">= 0.120.0")
+	assert.ErrorIs(t, err, ErrSSHExtensionInstallUnavailable)
 }
 
 func TestCheckIDESSHExtension_Cursor(t *testing.T) {
@@ -318,5 +353,67 @@ func TestCheckIDESSHExtension_NoPrompt_WithoutAutoApprove_Errors(t *testing.T) {
 
 	err := CheckIDESSHExtension(ctx, VSCodeOption, false)
 	require.Error(t, err)
-	assert.Contains(t, err.Error(), "--install-extension")
+	assert.ErrorIs(t, err, ErrSSHExtensionInstallUnavailable)
+	// The sentinel is a telemetry tag, so it must stay out of the message: what the user needs
+	// to read first is the problem, not that the CLI had no way to ask about it.
+	assert.Equal(t, `Required extension "Remote - SSH" is not installed in VS Code. `+
+		"Install it with: code --install-extension ms-vscode-remote.remote-ssh, or pass --auto-approve",
+		err.Error())
+}
+
+// Declining at the prompt is the outcome that separates "the user said no" from "there was no
+// way to ask", which is the distinction the split exists to make. Drive the real prompt rather
+// than asserting against a hand-built error, so the wiring is covered and not just the mapping.
+func TestCheckIDESSHExtension_Declined(t *testing.T) {
+	tmpDir := t.TempDir()
+	t.Setenv("PATH", tmpDir)
+	ctx, tst := cmdio.SetupTest(t.Context(), cmdio.TestOptions{PromptSupported: true})
+	defer tst.Done()
+
+	createFakeIDEExecutable(t, tmpDir, "code", "ms-python.python@2024.1.1\n")
+
+	// Drain stderr, where the prompt is written, or AskYesOrNo blocks on the pipe.
+	go func() { _, _ = io.Copy(io.Discard, tst.Stderr) }()
+	go func() {
+		_, _ = tst.Stdin.WriteString("n\n")
+		_ = tst.Stdin.Flush()
+	}()
+
+	err := CheckIDESSHExtension(ctx, VSCodeOption, false)
+	require.Error(t, err)
+	assert.ErrorIs(t, err, ErrSSHExtensionInstallDeclined)
+	assert.NotErrorIs(t, err, ErrSSHExtensionInstallUnavailable)
+	assert.Equal(t, `Required extension "Remote - SSH" is not installed in VS Code. `+
+		"Install it with: code --install-extension ms-vscode-remote.remote-ssh", err.Error())
+}
+
+// A command that is on PATH but whose --list-extensions fails is reported separately from a
+// missing extension: nothing was learned about what is installed, so it is not an install
+// problem. CheckIDECommand passes here, since the command does resolve.
+func TestCheckIDESSHExtension_ListFails(t *testing.T) {
+	tmpDir := t.TempDir()
+	t.Setenv("PATH", tmpDir)
+	ctx, _ := cmdio.NewTestContextWithStdout(t.Context())
+
+	createFailingIDEExecutable(t, tmpDir, "code")
+
+	err := CheckIDESSHExtension(ctx, VSCodeOption, true)
+	require.Error(t, err)
+	assert.ErrorIs(t, err, ErrSSHExtensionListFailed)
+	assert.NotErrorIs(t, err, ErrSSHExtensionInstallFailed)
+}
+
+// With --auto-approve there is no prompt, so a missing extension goes straight to an install.
+// A rejected install is the one outcome the IDE button can produce on this path.
+func TestCheckIDESSHExtension_AutoApprove_InstallFails(t *testing.T) {
+	tmpDir := t.TempDir()
+	t.Setenv("PATH", tmpDir)
+	ctx, _ := cmdio.NewTestContextWithStdout(t.Context())
+
+	createIDEExecutableFailingInstall(t, tmpDir, "code", "ms-python.python@2024.1.1\n")
+
+	err := CheckIDESSHExtension(ctx, VSCodeOption, true)
+	require.Error(t, err)
+	assert.ErrorIs(t, err, ErrSSHExtensionInstallFailed)
+	assert.NotErrorIs(t, err, ErrSSHExtensionInstallUnavailable)
 }