Add error-category telemetry to SshTunnelEvent

experimental/ssh/internal/client/client.go

@@ -258,7 +258,7 @@ func (o *ClientOptions) ToProxyCommand() (string, error) {
 	return proxyCommand, nil
 }
 
-func Run(ctx context.Context, client *databricks.WorkspaceClient, opts ClientOptions) error {
+func Run(ctx context.Context, client *databricks.WorkspaceClient, opts ClientOptions) (retErr error) {
 	ctx, cancel := context.WithCancel(ctx)
 	defer cancel()
 
@@ -270,6 +270,17 @@ func Run(ctx context.Context, client *databricks.WorkspaceClient, opts ClientOpt
 		cancel()
 	}()
 
+	// Report the outcome of every path below, so a failure is attributable to the step that
+	// caused it. Registered before the first early return -- in particular before the IDE
+	// preconditions, which fail fast on a permanent per-machine condition and are the failures
+	// most worth measuring. Each failing step sets outcome.errorCategory; the returned error is
+	// picked up here via the named return.
+	outcome := connectOutcome{isReconnect: opts.ServerMetadata != ""}
+	defer func() {
+		outcome.err = retErr
+		logSshTunnelEvent(ctx, opts, outcome)
+	}()
+
 	sessionID := opts.SessionIdentifier()
 	if sessionID == "" {
 		return errors.New("either --cluster or --name must be provided")
@@ -281,9 +292,11 @@ func Run(ctx context.Context, client *databricks.WorkspaceClient, opts ClientOpt
 
 	if opts.IDE != "" && !opts.ProxyMode {
 		if err := vscode.CheckIDECommand(opts.IDE); err != nil {
+			outcome.errorCategory = protos.SshTunnelErrorCategoryIDECommandNotOnPath
 			return err
 		}
 		if err := vscode.CheckIDESSHExtension(ctx, opts.IDE, opts.AutoApprove); err != nil {
+			outcome.errorCategory = protos.SshTunnelErrorCategoryIDESSHExtensionMissing
 			return err
 		}
 	}
@@ -299,31 +312,28 @@ func Run(ctx context.Context, client *databricks.WorkspaceClient, opts ClientOpt
 			cmdio.LogString(ctx, vscode.GetManualInstructions(opts.IDE, opts.ConnectionName))
 			cmdio.LogString(ctx, "Use --skip-settings-check to bypass IDE settings verification.")
 			if opts.AutoApprove {
+				outcome.errorCategory = protos.SshTunnelErrorCategoryIDESettingsUpdateDeclined
 				return fmt.Errorf("aborted: IDE settings need to be updated manually: %w", err)
 			}
 			shouldProceed, promptErr := cmdio.AskYesOrNo(ctx, "Do you want to proceed with the connection?")
 			if promptErr != nil {
+				outcome.errorCategory = protos.SshTunnelErrorCategoryIDESettingsUpdateDeclined
 				return fmt.Errorf("failed to prompt user: %w", promptErr)
 			}
 			if !shouldProceed {
+				outcome.errorCategory = protos.SshTunnelErrorCategoryIDESettingsUpdateDeclined
 				return errors.New("aborted: IDE settings need to be updated manually, user declined to proceed")
 			}
 		}
 	}
 
-	isReconnect := opts.ServerMetadata != ""
-	var serverStartTimeMs int64
-	isSuccess := false
-	defer func() {
-		logSshTunnelEvent(ctx, opts, isSuccess, isReconnect, serverStartTimeMs)
-	}()
-
 	// A direct `connect --cluster` bypasses `ssh setup`, which is where the access mode is
 	// normally validated, so validate it here too. Proxy mode is skipped because its
 	// ProxyCommand was generated by `setup` (already validated), and re-checking would add a
 	// Clusters.Get on every (re)connection. Serverless has no cluster to inspect.
 	if !opts.ProxyMode && !opts.IsServerlessMode() {
 		if err := ValidateClusterAccess(ctx, client, opts.ClusterID); err != nil {
+			outcome.errorCategory = protos.SshTunnelErrorCategoryClusterAccessDenied
 			return err
 		}
 	}
@@ -333,27 +343,32 @@ func Run(ctx context.Context, client *databricks.WorkspaceClient, opts ClientOpt
 		cmdio.LogString(ctx, "Checking cluster state...")
 		err := checkClusterState(ctx, client, opts.ClusterID, opts.AutoStartCluster)
 		if err != nil {
+			outcome.errorCategory = protos.SshTunnelErrorCategoryClusterStartFailed
 			return err
 		}
 	}
 
 	secretScopeName, err := keys.CreateKeysSecretScope(ctx, client, sessionID)
 	if err != nil {
+		outcome.errorCategory = protos.SshTunnelErrorCategorySecretScopeFailed
 		return fmt.Errorf("failed to create secret scope: %w", err)
 	}
 
 	privateKeyBytes, publicKeyBytes, err := keys.CheckAndGenerateSSHKeyPairFromSecrets(ctx, client, secretScopeName, opts.ClientPrivateKeyName, opts.ClientPublicKeyName)
 	if err != nil {
+		outcome.errorCategory = protos.SshTunnelErrorCategoryKeyGenerationFailed
 		return fmt.Errorf("failed to get or generate SSH key pair from secrets: %w", err)
 	}
 
 	keyPath, err := keys.GetLocalSSHKeyPath(ctx, sessionID, opts.SSHKeysDir)
 	if err != nil {
+		outcome.errorCategory = protos.SshTunnelErrorCategoryKeyGenerationFailed
 		return fmt.Errorf("failed to get local keys folder: %w", err)
 	}
 
 	err = keys.SaveSSHKeyPair(keyPath, privateKeyBytes, publicKeyBytes)
 	if err != nil {
+		outcome.errorCategory = protos.SshTunnelErrorCategoryKeyGenerationFailed
 		return fmt.Errorf("failed to save SSH key pair locally: %w", err)
 	}
 	log.Infof(ctx, "Using SSH key: %s", keyPath)
@@ -372,15 +387,21 @@ func Run(ctx context.Context, client *databricks.WorkspaceClient, opts ClientOpt
 		err := UploadTunnelReleases(ctx, client, version, opts.ReleasesDir)
 		sp.Close()
 		if err != nil {
+			outcome.errorCategory = protos.SshTunnelErrorCategoryBinaryUploadFailed
 			return fmt.Errorf("failed to upload ssh-tunnel binaries: %w", err)
 		}
 		serverStartTime := time.Now()
 		userName, serverPort, clusterID, err = ensureSSHServerIsRunning(ctx, client, version, secretScopeName, opts)
 		if err != nil {
+			outcome.errorCategory = protos.SshTunnelErrorCategoryServerStartTimeout
 			return fmt.Errorf("failed to ensure that ssh server is running: %w", err)
 		}
-		serverStartTimeMs = time.Since(serverStartTime).Milliseconds()
+		outcome.serverStartTimeMs = time.Since(serverStartTime).Milliseconds()
 	} else {
+		// The failures below are left to fall through to UNKNOWN on purpose: --metadata is a
+		// hidden flag whose value we generated ourselves in ToProxyCommand, so a parse failure
+		// here is a CLI bug rather than a per-environment blocker. Attributing them to
+		// SERVER_START_TIMEOUT would pollute the bucket that tracks unreachable servers.
 		// Metadata format: "<user_name>,<port>,<cluster_id>"
 		metadata := strings.Split(opts.ServerMetadata, ",")
 		if len(metadata) < 2 {
@@ -416,7 +437,9 @@ func Run(ctx context.Context, client *databricks.WorkspaceClient, opts ClientOpt
 		cmdio.LogString(ctx, "Connected!")
 	}
 
-	isSuccess = true
+	// The tunnel is up from here on. A later non-zero exit belongs to the SSH client or the
+	// user's own remote command, so it is not counted as a connection failure.
+	outcome.isSuccess = true
 
 	if opts.ProxyMode {
 		return runSSHProxy(ctx, client, serverPort, clusterID, opts)
@@ -1209,15 +1232,44 @@ func ensureSSHServerIsRunning(ctx context.Context, client *databricks.WorkspaceC
 	return meta.UserName, meta.Port, meta.ClusterID, nil
 }
 
-func logSshTunnelEvent(ctx context.Context, opts ClientOptions, isSuccess, isReconnect bool, serverStartTimeMs int64) {
+// connectOutcome is the observed result of a connection attempt, collected by Run for telemetry.
+type connectOutcome struct {
+	// isSuccess reports whether the tunnel was established. It stays true when the SSH client
+	// itself later exits non-zero, since by then the tunnel was up.
+	isSuccess         bool
+	isReconnect       bool
+	serverStartTimeMs int64
+	// errorCategory is set at the failure site. Empty means the failure was not attributed.
+	errorCategory protos.SshTunnelErrorCategory
+	err           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.
+func (o connectOutcome) category() protos.SshTunnelErrorCategory {
+	if o.isSuccess || o.err == nil {
+		return protos.SshTunnelErrorCategoryUnspecified
+	}
+	if errors.Is(o.err, context.Canceled) {
+		return protos.SshTunnelErrorCategoryUserAborted
+	}
+	if o.errorCategory == "" {
+		return protos.SshTunnelErrorCategoryUnknown
+	}
+	return o.errorCategory
+}
+
+func logSshTunnelEvent(ctx context.Context, opts ClientOptions, outcome connectOutcome) {
 	telemetry.Log(ctx, protos.DatabricksCliLog{
-		SshTunnelEvent: buildSshTunnelEvent(opts, isSuccess, isReconnect, serverStartTimeMs),
+		SshTunnelEvent: buildSshTunnelEvent(opts, outcome),
 	})
 }
 
 // buildSshTunnelEvent maps the connection options and outcome onto the telemetry
 // event. It is separated from logSshTunnelEvent so the field mapping can be unit tested.
-func buildSshTunnelEvent(opts ClientOptions, isSuccess, isReconnect bool, serverStartTimeMs int64) *protos.SshTunnelEvent {
+func buildSshTunnelEvent(opts ClientOptions, outcome connectOutcome) *protos.SshTunnelEvent {
 	computeType := protos.SshTunnelComputeTypeDedicated
 	if opts.IsServerlessMode() {
 		computeType = protos.SshTunnelComputeTypeServerless
@@ -1238,11 +1290,12 @@ func buildSshTunnelEvent(opts ClientOptions, isSuccess, isReconnect bool, server
 		AcceleratorType:    opts.Accelerator,
 		IdeType:            opts.IDE,
 		ClientMode:         clientMode,
-		IsReconnect:        isReconnect,
+		IsReconnect:        outcome.isReconnect,
 		AutoStartCluster:   opts.AutoStartCluster,
-		ServerStartTimeMs:  serverStartTimeMs,
-		IsSuccess:          isSuccess,
+		ServerStartTimeMs:  outcome.serverStartTimeMs,
+		IsSuccess:          outcome.isSuccess,
 		HasBaseEnvironment: opts.BaseEnvironment != "",
 		HasUsagePolicy:     opts.UsagePolicyID != "",
+		ErrorCategory:      outcome.category(),
 	}
 }

experimental/ssh/internal/client/client_internal_test.go

@@ -1,7 +1,9 @@
 package client
 
 import (
+	"context"
 	"errors"
+	"fmt"
 	"strings"
 	"testing"
 	"time"
@@ -474,11 +476,79 @@ func TestBuildSshTunnelEvent(t *testing.T) {
 
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			got := buildSshTunnelEvent(tt.opts, true, true, 1500)
+			got := buildSshTunnelEvent(tt.opts, connectOutcome{
+				isSuccess:         true,
+				isReconnect:       true,
+				serverStartTimeMs: 1500,
+			})
 			tt.want.IsSuccess = true
 			tt.want.IsReconnect = true
 			tt.want.ServerStartTimeMs = 1500
+			tt.want.ErrorCategory = protos.SshTunnelErrorCategoryUnspecified
 			assert.Equal(t, &tt.want, got)
 		})
 	}
 }
+
+func TestConnectOutcomeCategory(t *testing.T) {
+	errFailed := errors.New("failed")
+
+	tests := []struct {
+		name    string
+		outcome connectOutcome
+		want    protos.SshTunnelErrorCategory
+	}{
+		{
+			name:    "success reports no category",
+			outcome: connectOutcome{isSuccess: true},
+			want:    protos.SshTunnelErrorCategoryUnspecified,
+		},
+		{
+			// A non-zero exit after the tunnel is up belongs to the ssh client, not the connection.
+			name:    "error after a successful connection reports no category",
+			outcome: connectOutcome{isSuccess: true, err: errFailed},
+			want:    protos.SshTunnelErrorCategoryUnspecified,
+		},
+		{
+			name:    "attributed failure keeps its category",
+			outcome: connectOutcome{errorCategory: protos.SshTunnelErrorCategoryIDECommandNotOnPath, err: errFailed},
+			want:    protos.SshTunnelErrorCategoryIDECommandNotOnPath,
+		},
+		{
+			name:    "unattributed failure falls back to UNKNOWN",
+			outcome: connectOutcome{err: errFailed},
+			want:    protos.SshTunnelErrorCategoryUnknown,
+		},
+		{
+			name:    "cancellation reports USER_ABORTED",
+			outcome: connectOutcome{err: fmt.Errorf("wrapped: %w", context.Canceled)},
+			want:    protos.SshTunnelErrorCategoryUserAborted,
+		},
+		{
+			// Ctrl-C surfaces as a cancellation from whichever step observed it first, so the
+			// interruption must win over the category that step recorded.
+			name: "cancellation wins over the category set at the failure site",
+			outcome: connectOutcome{
+				errorCategory: protos.SshTunnelErrorCategoryServerStartTimeout,
+				err:           fmt.Errorf("wrapped: %w", context.Canceled),
+			},
+			want: protos.SshTunnelErrorCategoryUserAborted,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			assert.Equal(t, tt.want, tt.outcome.category())
+		})
+	}
+}
+
+func TestBuildSshTunnelEventReportsErrorCategory(t *testing.T) {
+	got := buildSshTunnelEvent(ClientOptions{ConnectionName: "my-conn", IDE: "vscode"}, connectOutcome{
+		errorCategory: protos.SshTunnelErrorCategoryIDECommandNotOnPath,
+		err:           errors.New("failed"),
+	})
+
+	assert.False(t, got.IsSuccess)
+	assert.Equal(t, protos.SshTunnelErrorCategoryIDECommandNotOnPath, got.ErrorCategory)
+}

libs/telemetry/protos/ssh_tunnel.go

@@ -17,6 +17,54 @@ const (
 	SshTunnelClientModeIDE         SshTunnelClientMode = "IDE"
 )
 
+// SshTunnelErrorCategory is a coarse classification of why a connection attempt failed.
+// The categories name the distinct early-return sites of the connect flow so a failure can
+// be attributed without logging the error text, which carries cluster names, paths and user
+// names.
+type SshTunnelErrorCategory string
+
+const (
+	SshTunnelErrorCategoryUnspecified SshTunnelErrorCategory = "TYPE_UNSPECIFIED"
+
+	// The IDE's shell command ("code"/"cursor") is not on PATH. A permanent per-machine
+	// condition rather than a transient failure, so it is distinguished from the rest.
+	SshTunnelErrorCategoryIDECommandNotOnPath SshTunnelErrorCategory = "IDE_COMMAND_NOT_ON_PATH"
+
+	// The required Remote-SSH extension is missing or too old and was not installed.
+	SshTunnelErrorCategoryIDESSHExtensionMissing SshTunnelErrorCategory = "IDE_SSH_EXTENSION_MISSING"
+
+	// IDE settings had to be updated for serverless but the update failed and the user
+	// declined to continue (or --auto-approve turned the failure into an abort).
+	SshTunnelErrorCategoryIDESettingsUpdateDeclined SshTunnelErrorCategory = "IDE_SETTINGS_UPDATE_DECLINED"
+
+	// The cluster is not a dedicated single-user cluster, or it could not be inspected.
+	SshTunnelErrorCategoryClusterAccessDenied SshTunnelErrorCategory = "CLUSTER_ACCESS_DENIED"
+
+	// The cluster was not running and could not be started.
+	SshTunnelErrorCategoryClusterStartFailed SshTunnelErrorCategory = "CLUSTER_START_FAILED"
+
+	// Creating or reading the secret scope holding the SSH keys failed.
+	SshTunnelErrorCategorySecretScopeFailed SshTunnelErrorCategory = "SECRET_SCOPE_FAILED"
+
+	// Generating or persisting the local SSH key pair failed.
+	SshTunnelErrorCategoryKeyGenerationFailed SshTunnelErrorCategory = "KEY_GENERATION_FAILED"
+
+	// Uploading the SSH tunnel binaries to the workspace failed.
+	SshTunnelErrorCategoryBinaryUploadFailed SshTunnelErrorCategory = "BINARY_UPLOAD_FAILED"
+
+	// The SSH server never became reachable: the bootstrap job failed to start, died, or
+	// its metadata never appeared before the timeout.
+	SshTunnelErrorCategoryServerStartTimeout SshTunnelErrorCategory = "SERVER_START_TIMEOUT"
+
+	// The user interrupted the connection (Ctrl-C or a termination signal).
+	SshTunnelErrorCategoryUserAborted SshTunnelErrorCategory = "USER_ABORTED"
+
+	// A failure that does not correspond to any of the categories above. The connect path
+	// attributes every per-environment blocker, so a rise here points at a CLI bug (or a new
+	// failure mode that needs its own category) rather than a user's setup.
+	SshTunnelErrorCategoryUnknown SshTunnelErrorCategory = "UNKNOWN"
+)
+
 // SshTunnelEvent is emitted when a user establishes an SSH tunnel connection
 // via the Databricks CLI.
 type SshTunnelEvent struct {
@@ -53,4 +101,10 @@ type SshTunnelEvent struct {
 	// Whether a serverless usage policy was set via --usage-policy-id.
 	// Only the presence is recorded, not the policy ID itself.
 	HasUsagePolicy bool `json:"has_usage_policy,omitempty"`
+
+	// Why the connection attempt failed, or TYPE_UNSPECIFIED on success. Deliberately
+	// without omitempty: the field is what identifies a failure's cause, so an empty value
+	// must not be silently dropped into an indistinguishable NULL. Every failure path sets
+	// a category, falling back to UNKNOWN.
+	ErrorCategory SshTunnelErrorCategory `json:"error_category"`
 }