Handle user cancellation in ssh agent-shim

experimental/ssh/internal/client/agentshim.go

@@ -14,6 +14,7 @@ import (
 	"regexp"
 	"runtime"
 	"strings"
+	"syscall"
 	"time"
 
 	"github.com/databricks/cli/internal/build"
@@ -22,6 +23,7 @@ import (
 	"github.com/databricks/cli/libs/dbr"
 	"github.com/databricks/cli/libs/env"
 	"github.com/databricks/cli/libs/log"
+	"github.com/databricks/cli/libs/process"
 	"github.com/databricks/cli/libs/telemetry"
 	"github.com/databricks/cli/libs/telemetry/protos"
 	"github.com/databricks/databricks-sdk-go"
@@ -118,6 +120,16 @@ func agentByName(name string) (agentSpec, bool) {
 }
 
 func RunAgentShim(ctx context.Context, client *databricks.WorkspaceClient, agentName string, agentArgs []string) (retErr error) {
+	// Cancel ctx on Ctrl-C, SIGTERM, or SIGHUP (the SSH connection dropping) so the in-flight
+	// step aborts and the launch is reported as USER_ABORTED, instead of the signal killing the
+	// process before any event is logged. The event defer below runs before the deferred
+	// stopSignals (LIFO), so a failed step sees ctx cancelled only by a signal or the caller. Once
+	// the launch is committed,
+	// late signals get the default handling again.
+	parentCtx := ctx
+	ctx, stopSignals := process.NotifyContext(ctx, os.Interrupt, syscall.SIGTERM, syscall.SIGHUP)
+	defer stopSignals()
+
 	// Report the outcome of the launch. Setup runs from here to the exec below, so time it
 	// from here too. Each failing step sets outcome.errorCategory; the returned error is
 	// picked up via the named return.
@@ -164,11 +176,22 @@ func RunAgentShim(ctx context.Context, client *databricks.WorkspaceClient, agent
 		return ensureCategory(protos.SshAgentShimErrorCategoryAgentLaunchFailed, err)
 	}
 
+	// From here a signal should kill the shim outright, as it would without a handler. A
+	// signal during a step that ignores ctx (warm bootstrap, file I/O) only cancelled it, so
+	// check after stopping the handler, which leaves no window for one to be swallowed.
+	// stopSignals cancels ctx, so tell a signal apart by its cause, and the success path
+	// uses the caller's context.
+	stopSignals()
+	_, signalled := errors.AsType[*process.SignalError](context.Cause(ctx))
+	if signalled || parentCtx.Err() != nil {
+		return fmt.Errorf("agent launch interrupted: %w", ctx.Err())
+	}
+
 	// The agent is ready to launch. execProcess replaces this process, so cmd/root's
 	// end-of-command telemetry upload never runs; emit and flush the success event first.
 	launched = true
-	logAgentShimEvent(ctx, agentName, outcome, time.Since(start))
-	uploadAgentShimTelemetry(ctx, agentName, time.Since(start))
+	logAgentShimEvent(parentCtx, agentName, outcome, time.Since(start))
+	uploadAgentShimTelemetry(parentCtx, agentName, time.Since(start))
 	return execProcess(argv0, argv, environ)
 }
 
@@ -296,6 +319,12 @@ type gatewayProbe struct {
 	resourceAvailable bool
 }
 
+// gatewayCheckInterrupted reports a cancelled probe as such. Classifying its results would
+// misreport a Ctrl-C as a disabled gateway or a transient error to retry.
+func gatewayCheckInterrupted(err error) error {
+	return fmt.Errorf("the Unity AI Gateway check was interrupted: %w", err)
+}
+
 func probeAIGateway(ctx context.Context, wsclient *databricks.WorkspaceClient) error {
 	host := strings.TrimRight(wsclient.Config.Host, "/")
 
@@ -310,6 +339,9 @@ func probeAIGateway(ctx context.Context, wsclient *databricks.WorkspaceClient) e
 	}
 
 	modelSvc := probeModelServices(ctx, apiClient)
+	if err := ctx.Err(); err != nil {
+		return gatewayCheckInterrupted(err)
+	}
 	// A 401 (or a 400 "invalid token") can't be rescued by trying another API,
 	// so surface it before the fallback probe.
 	if !modelSvc.reachable && looksLikeDefinitiveAuthFailure(modelSvc.err) {
@@ -320,6 +352,9 @@ func probeAIGateway(ctx context.Context, wsclient *databricks.WorkspaceClient) e
 	}
 
 	legacy := probeLegacyEndpoints(ctx, apiClient)
+	if err := ctx.Err(); err != nil {
+		return gatewayCheckInterrupted(err)
+	}
 	switch {
 	case legacy.reachable:
 		// The legacy endpoints API answered, so the gateway is enabled even if no

experimental/ssh/internal/client/agentshim_signal_test.go

@@ -0,0 +1,126 @@
+//go:build unix
+
+// The tests interrupt the shim by sending a signal to their own process, which os/signal cannot
+// deliver to self on Windows. The shim itself only runs on the Linux driver.
+package client
+
+import (
+	"context"
+	"encoding/json"
+	"io"
+	"net/http"
+	"net/http/httptest"
+	"os"
+	"os/exec"
+	"path/filepath"
+	"strings"
+	"syscall"
+	"testing"
+
+	"github.com/databricks/cli/libs/cmdctx"
+	"github.com/databricks/cli/libs/dbr"
+	"github.com/databricks/cli/libs/env"
+	"github.com/databricks/cli/libs/telemetry"
+	"github.com/databricks/cli/libs/telemetry/protos"
+	"github.com/databricks/databricks-sdk-go/config"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+// newAgentShimServer serves the gateway probe with modelServices and records the telemetry
+// upload into body.
+func newAgentShimServer(t *testing.T, modelServices http.HandlerFunc, body *telemetry.RequestBody) *httptest.Server {
+	t.Helper()
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		switch {
+		case strings.HasPrefix(r.URL.Path, "/api/2.1/unity-catalog/model-services"):
+			modelServices(w, r)
+		case r.URL.Path == "/telemetry-ext":
+			b, err := io.ReadAll(r.Body)
+			assert.NoError(t, err)
+			assert.NoError(t, json.Unmarshal(b, body))
+			_ = json.NewEncoder(w).Encode(telemetry.ResponseBody{NumProtoSuccess: 1})
+		default:
+			w.WriteHeader(http.StatusNotFound)
+			_, _ = w.Write([]byte("{}"))
+		}
+	}))
+	t.Cleanup(srv.Close)
+	return srv
+}
+
+// newAgentShimContext returns a context with the telemetry and auth setup cmd/root provides,
+// and an isolated home so the shim never touches the real one.
+func newAgentShimContext(t *testing.T, host string) context.Context {
+	ctx := cmdctx.GenerateExecId(t.Context())
+	ctx = dbr.MockRuntime(ctx, dbr.Environment{})
+	ctx = telemetry.WithNewLogger(ctx)
+	ctx = cmdctx.SetConfigUsed(ctx, &config.Config{Host: host, Token: "test-token"})
+	return env.WithUserHomeDir(ctx, t.TempDir())
+}
+
+// assertUserAbortedEvent uploads the buffered telemetry as cmd/root does, on the command's own
+// context (which the signal leaves intact), and checks it holds one USER_ABORTED failure.
+func assertUserAbortedEvent(t *testing.T, ctx context.Context, body *telemetry.RequestBody) {
+	t.Helper()
+	require.NoError(t, telemetry.Upload(ctx, protos.ExecutionContext{}))
+	require.Len(t, body.ProtoLogs, 1)
+	var logged protos.FrontendLog
+	require.NoError(t, json.Unmarshal([]byte(body.ProtoLogs[0]), &logged))
+	event := logged.Entry.DatabricksCliLog.SshAgentShimEvent
+	require.NotNil(t, event)
+	assert.False(t, event.IsSuccess)
+	assert.Equal(t, protos.SshAgentShimErrorCategoryUserAborted, event.ErrorCategory)
+}
+
+// TestRunAgentShimSignalReportsUserAborted interrupts a launch mid-setup and checks the
+// failure event still reaches the telemetry upload as USER_ABORTED. Without a signal handler
+// in RunAgentShim the signal would terminate the test process instead.
+func TestRunAgentShimSignalReportsUserAborted(t *testing.T) {
+	for _, sig := range []syscall.Signal{syscall.SIGINT, syscall.SIGTERM, syscall.SIGHUP} {
+		t.Run(sig.String(), func(t *testing.T) {
+			var body telemetry.RequestBody
+			srv := newAgentShimServer(t, func(w http.ResponseWriter, r *http.Request) {
+				// The gateway probe is in flight, so RunAgentShim's handler is installed:
+				// interrupt it and hold the request until the cancellation aborts it.
+				assert.NoError(t, syscall.Kill(os.Getpid(), sig))
+				<-r.Context().Done()
+			}, &body)
+			ctx := newAgentShimContext(t, srv.URL)
+
+			err := RunAgentShim(ctx, newProbeClient(t, srv.URL), "claude", nil)
+			require.ErrorIs(t, err, context.Canceled)
+			assertUserAbortedEvent(t, ctx, &body)
+		})
+	}
+}
+
+// TestRunAgentShimSignalDuringIgnoredStepSkipsLaunch interrupts the launch during a step that
+// swallows the cancellation: the warm-path `npm config set` run by bootstrapToolchain. The
+// probe has already succeeded, so only the check before the exec can catch the signal, and
+// the agent must not launch.
+func TestRunAgentShimSignalDuringIgnoredStepSkipsLaunch(t *testing.T) {
+	const modelServicesBody = `{"model_services":[{"name":"model-services/system.ai.gpt-5"}]}`
+
+	var body telemetry.RequestBody
+	srv := newAgentShimServer(t, jsonHandler(http.StatusOK, modelServicesBody), &body)
+	ctx := newAgentShimContext(t, srv.URL)
+
+	// Stub the toolchain so bootstrap takes the warm path. npm interrupts the shim and waits to
+	// be killed by the cancellation, so ctx is cancelled by the time bootstrap returns; the
+	// sleep is bounded so a regression there can't hang the test. ucode stands in for the
+	// agent: if the shim execs it, it leaves the marker and fails the run. PATH holds only the
+	// stubs, so they use shell builtins and an absolute path to sleep.
+	sleep, err := exec.LookPath("sleep")
+	require.NoError(t, err)
+	binDir := t.TempDir()
+	marker := filepath.Join(t.TempDir(), "launched")
+	writeFakeExec(t, binDir, "npm", "kill -INT $PPID\nexec \""+sleep+"\" 30\n")
+	writeFakeExec(t, binDir, "ucode", ": > \""+marker+"\"\nexit 1\n")
+	t.Setenv("PATH", binDir)
+
+	err = RunAgentShim(ctx, newProbeClient(t, srv.URL), "claude", nil)
+	require.ErrorIs(t, err, context.Canceled)
+	assert.NoFileExists(t, marker, "the agent must not launch after a signal")
+	assertUserAbortedEvent(t, ctx, &body)
+}

experimental/ssh/internal/client/agentshim_test.go

@@ -539,6 +539,18 @@ func TestProbeAIGateway(t *testing.T) {
 		assert.Contains(t, err.Error(), "transient error")
 		assert.NotContains(t, err.Error(), "not enabled")
 	})
+
+	t.Run("cancelled context is reported as interrupted, not classified", func(t *testing.T) {
+		unavailable := jsonHandler(http.StatusServiceUnavailable, `{"message":"try later"}`)
+		srv := newGatewayServer(t, unavailable, unavailable)
+		defer srv.Close()
+		ctx, cancel := context.WithCancel(t.Context())
+		cancel()
+		err := probeAIGateway(ctx, newProbeClient(t, srv.URL))
+		require.ErrorIs(t, err, context.Canceled)
+		assert.NotContains(t, err.Error(), "not enabled")
+		assert.NotContains(t, err.Error(), "Retry")
+	})
 }
 
 func TestProbeModelServicesPageError(t *testing.T) {

libs/process/signal.go

@@ -0,0 +1,76 @@
+package process
+
+import (
+	"context"
+	"os"
+	"os/signal"
+	"sync"
+)
+
+// NotifyContext is like signal.NotifyContext, but stops relaying as soon as
+// the first signal arrives, so a second one is handled as it was before the
+// call: for a signal the process doesn't otherwise handle, it terminates.
+//
+// signal.Notify overrides the previous disposition for as long as the channel
+// stays registered, so with signal.NotifyContext a second Ctrl-C is merely
+// buffered and dropped, leaving the user no way to abort a step that ignores
+// cancellation or a slow teardown. Calling signal.Stop before cancel restores
+// the previous disposition first, so by the time the cancellation is
+// observable a second signal gets it. Note that Notify also stops ignoring
+// SIGHUP or SIGINT if the process started with it ignored (e.g. under nohup);
+// signal.Stop ignores it again.
+//
+// A signal cancels the context with a *SignalError cause, so a caller can tell
+// it apart from its own call to stop.
+//
+// The returned stop uninstalls the handler, waits for the handler goroutine to
+// exit, and cancels the context. A signal relayed before stop uninstalls the
+// handler is still reported as the cause rather than dropped, so after stop
+// returns the caller can check whether one came in without racing a late
+// signal. It is safe to call more than once; the caller must defer it.
+func NotifyContext(parent context.Context, sigs ...os.Signal) (context.Context, context.CancelFunc) {
+	ctx, cancel := context.WithCancelCause(parent)
+	sigCh := make(chan os.Signal, 1)
+	signal.Notify(sigCh, sigs...)
+
+	stopping := make(chan struct{})
+	done := make(chan struct{})
+	go func() {
+		defer close(done)
+		// Selecting on ctx.Done() and stopping too lets the goroutine exit on the
+		// normal (no signal) path: signal.Stop unregisters the channel but never
+		// closes it.
+		select {
+		case sig := <-sigCh:
+			signal.Stop(sigCh)
+			cancel(&SignalError{Signal: sig})
+		case <-ctx.Done():
+			signal.Stop(sigCh)
+		case <-stopping:
+			// stop has already called signal.Stop, so a signal that arrived
+			// before then is buffered in sigCh. Report it rather than drop it.
+			select {
+			case sig := <-sigCh:
+				cancel(&SignalError{Signal: sig})
+			default:
+			}
+		}
+	}()
+
+	return ctx, sync.OnceFunc(func() {
+		signal.Stop(sigCh)
+		close(stopping)
+		<-done
+		cancel(nil)
+	})
+}
+
+// SignalError is the cause NotifyContext cancels its context with when a
+// signal arrives.
+type SignalError struct {
+	Signal os.Signal
+}
+
+func (e *SignalError) Error() string {
+	return e.Signal.String() + " signal received"
+}

libs/process/signal_test.go

@@ -0,0 +1,44 @@
+package process
+
+import (
+	"context"
+	"os"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+)
+
+// TestNotifyContextStopsWithoutSignal covers the no-signal path: stop() must
+// return (joining its goroutine) rather than leaving it parked on the channel
+// for the life of the process, and must be safe to call again, as a caller that
+// stops early and also defers stop does.
+func TestNotifyContextStopsWithoutSignal(t *testing.T) {
+	ctx, stop := NotifyContext(t.Context(), os.Interrupt)
+
+	returned := make(chan struct{})
+	go func() {
+		stop()
+		stop()
+		close(returned)
+	}()
+	select {
+	case <-returned:
+	case <-time.After(10 * time.Second):
+		t.Fatal("stop() blocked: the signal goroutine was never joined")
+	}
+	assert.ErrorIs(t, ctx.Err(), context.Canceled)
+	var sigErr *SignalError
+	assert.NotErrorAs(t, context.Cause(ctx), &sigErr, "stop must not look like a signal")
+}
+
+// TestNotifyContextParentCancel checks that cancelling the parent cancels the
+// returned context and that stop still returns afterwards.
+func TestNotifyContextParentCancel(t *testing.T) {
+	parent, cancel := context.WithCancel(t.Context())
+	ctx, stop := NotifyContext(parent, os.Interrupt)
+	cancel()
+	<-ctx.Done()
+	stop()
+	assert.ErrorIs(t, ctx.Err(), context.Canceled)
+}