diff --git a/experimental/ssh/internal/client/agentshim.go b/experimental/ssh/internal/client/agentshim.go index d596a4e49d9..8278b846080 100644 --- a/experimental/ssh/internal/client/agentshim.go +++ b/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 diff --git a/experimental/ssh/internal/client/agentshim_signal_test.go b/experimental/ssh/internal/client/agentshim_signal_test.go new file mode 100644 index 00000000000..0ddc6fb98b5 --- /dev/null +++ b/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) +} diff --git a/experimental/ssh/internal/client/agentshim_test.go b/experimental/ssh/internal/client/agentshim_test.go index 283d4288fbe..60160a8368d 100644 --- a/experimental/ssh/internal/client/agentshim_test.go +++ b/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) { diff --git a/libs/process/signal.go b/libs/process/signal.go new file mode 100644 index 00000000000..06380207acb --- /dev/null +++ b/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" +} diff --git a/libs/process/signal_test.go b/libs/process/signal_test.go new file mode 100644 index 00000000000..218af438290 --- /dev/null +++ b/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) +} diff --git a/libs/process/signal_unix_test.go b/libs/process/signal_unix_test.go new file mode 100644 index 00000000000..fe63900f1a1 --- /dev/null +++ b/libs/process/signal_unix_test.go @@ -0,0 +1,130 @@ +//go:build unix + +// The second-signal escape hatch is a Unix signal-delivery behavior: the test +// re-execs itself and sends SIGINT/SIGTERM/SIGHUP, which os/signal does not +// support on Windows (the CLI's own signal handling there is likewise a no-op path). +package process + +import ( + "bufio" + "context" + "os" + "os/exec" + "strings" + "syscall" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// The second-signal behavior can only be observed in a real process: signal +// dispositions are process-wide, and the test binary's own handlers would mask +// them. So the test re-executes itself as a child (Go resets dispositions to +// SIG_DFL across exec, unlike a shell background job, which inherits SIGINT as +// SIG_IGN and would silently invalidate the result). +const signalChildEnv = "TEST_PROCESS_SIGNAL_CHILD" + +// TestMain runs the signal-handler-under-test when re-executed as the child. +func TestMain(m *testing.M) { + if os.Getenv(signalChildEnv) == "" { + os.Exit(m.Run()) + } + + // TestMain has no *testing.T, so t.Context() is unavailable here. + //nolint:gocritic + ctx, stop := NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM, syscall.SIGHUP) + defer stop() + + // Stand in for a caller still tearing down after the first signal (e.g. + // draining a process group during its SIGKILL grace window) — the interval + // in which the user needs the hatch. + os.Stdout.WriteString("READY\n") + <-ctx.Done() + os.Stdout.WriteString("CANCELLED\n") + time.Sleep(30 * time.Second) + os.Stdout.WriteString("SURVIVED\n") + os.Exit(0) +} + +// TestNotifyContextSecondSignalStillKills is the regression test for the +// escape hatch: signal.Notify disables the default disposition process-wide, so +// a handler that only relays the first signal leaves the user unable to abort. +// The first signal must cancel the context, and the second must terminate the +// process outright. +func TestNotifyContextSecondSignalStillKills(t *testing.T) { + for _, sig := range []syscall.Signal{syscall.SIGINT, syscall.SIGTERM, syscall.SIGHUP} { + t.Run(sig.String(), func(t *testing.T) { + cmd := exec.Command(os.Args[0], "-test.run=TestNotifyContextSecondSignalStillKills") + cmd.Env = append(os.Environ(), signalChildEnv+"=1") + stdout, err := cmd.StdoutPipe() + require.NoError(t, err) + require.NoError(t, cmd.Start()) + t.Cleanup(func() { _ = cmd.Process.Kill() }) + + lines := make(chan string, 8) + go func() { + sc := bufio.NewScanner(stdout) + for sc.Scan() { + lines <- sc.Text() + } + close(lines) + }() + await := func(want string) bool { + deadline := time.After(30 * time.Second) + for { + select { + case l, ok := <-lines: + if !ok { + return false + } + if strings.Contains(l, want) { + return true + } + case <-deadline: + return false + } + } + } + + require.True(t, await("READY"), "child never started") + require.NoError(t, cmd.Process.Signal(sig)) + require.True(t, await("CANCELLED"), "first %s did not cancel the context", sig) + + // The second signal must reach the default disposition and kill the + // child rather than being buffered and dropped by the handler. + require.NoError(t, cmd.Process.Signal(sig)) + waitErr := make(chan error, 1) + go func() { waitErr <- cmd.Wait() }() + select { + case err := <-waitErr: + // Killed by the signal, so a non-nil (non-exit-zero) error. + assert.Error(t, err, "child should die by signal, not exit cleanly") + case <-time.After(15 * time.Second): + t.Fatalf("second %s was swallowed: the user has no escape hatch "+ + "while the caller tears down", sig) + } + }) + } +} + +// TestNotifyContextSignalCause checks that a signal cancels the context with a +// *SignalError cause and that stop leaves it in place, so a caller that stops +// the handler before checking can still tell a signal apart from the stop. +func TestNotifyContextSignalCause(t *testing.T) { + ctx, stop := NotifyContext(t.Context(), syscall.SIGHUP) + defer stop() + + require.NoError(t, syscall.Kill(os.Getpid(), syscall.SIGHUP)) + select { + case <-ctx.Done(): + case <-time.After(10 * time.Second): + t.Fatal("SIGHUP did not cancel the context") + } + stop() + + var sigErr *SignalError + require.ErrorAs(t, context.Cause(ctx), &sigErr) + assert.Equal(t, syscall.SIGHUP, sigErr.Signal) +}