From 463472d4598e7970bc4a86c39cf429906d47c342 Mon Sep 17 00:00:00 2001 From: Russell Clarey Date: Wed, 14 Oct 2026 11:56:52 +0000 Subject: [PATCH 1/4] Report USER_ABORTED when agent-shim setup is interrupted Nothing on the agent-shim path handled signals, so Ctrl-C, SIGTERM, or SIGHUP during setup terminated the process before the failure event was logged or uploaded, and USER_ABORTED was unreachable. Cancel the context on the first signal instead so the in-flight step aborts and the event is reported. Co-authored-by: Isaac --- experimental/ssh/internal/client/agentshim.go | 40 +++++++++++ .../internal/client/agentshim_signal_test.go | 68 +++++++++++++++++++ 2 files changed, 108 insertions(+) create mode 100644 experimental/ssh/internal/client/agentshim_signal_test.go diff --git a/experimental/ssh/internal/client/agentshim.go b/experimental/ssh/internal/client/agentshim.go index d596a4e49d9..ba714edd77a 100644 --- a/experimental/ssh/internal/client/agentshim.go +++ b/experimental/ssh/internal/client/agentshim.go @@ -10,10 +10,12 @@ import ( "net/http" "os" "os/exec" + "os/signal" "path/filepath" "regexp" "runtime" "strings" + "syscall" "time" "github.com/databricks/cli/internal/build" @@ -118,6 +120,14 @@ func agentByName(name string) (agentSpec, bool) { } func RunAgentShim(ctx context.Context, client *databricks.WorkspaceClient, agentName string, agentArgs []string) (retErr error) { + // Without a handler, Ctrl-C, SIGTERM, or SIGHUP (the SSH connection dropping) would + // terminate the process outright, skipping the failure event below and cmd/root's upload. + // Cancelling ctx instead aborts the in-flight step, so the launch fails as USER_ABORTED. + // Set up before the deferred event so that defer runs first (LIFO) and sees ctx cancelled + // only by a signal or the caller, never by stopSignals itself. + ctx, stopSignals := notifyAbort(ctx) + 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. @@ -169,9 +179,39 @@ func RunAgentShim(ctx context.Context, client *databricks.WorkspaceClient, agent launched = true logAgentShimEvent(ctx, agentName, outcome, time.Since(start)) uploadAgentShimTelemetry(ctx, agentName, time.Since(start)) + // Nothing below observes ctx, so restore the default signal handling: a signal from here + // on terminates the process instead of cancelling a context no step checks any more. + stopSignals() return execProcess(argv0, argv, environ) } +// notifyAbort returns a context cancelled on the first SIGINT, SIGTERM, or SIGHUP. The returned +// stop uninstalls the handler before returning; the caller must defer it. +// +// signal.Notify disables the default disposition for these signals for as long as the channel +// stays registered, so a second Ctrl-C would merely be buffered and dropped, leaving no way to +// abort a step that ignores cancellation. Calling signal.Stop before cancel restores the +// default first, so once the cancellation is observable a second signal terminates the process. +func notifyAbort(parent context.Context) (context.Context, func()) { + ctx, cancel := context.WithCancel(parent) + sigCh := make(chan os.Signal, 1) + signal.Notify(sigCh, os.Interrupt, syscall.SIGTERM, syscall.SIGHUP) + go func() { + // Selecting on ctx.Done() too lets the goroutine exit on the normal (no signal) path + // rather than parking on sigCh for the life of the process. + select { + case <-sigCh: + signal.Stop(sigCh) + cancel() + case <-ctx.Done(): + } + }() + return ctx, func() { + signal.Stop(sigCh) + cancel() + } +} + // agentShimOutcome is the observed result of an agent-shim launch, collected by // RunAgentShim for telemetry. type agentShimOutcome struct { 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..2e9cb6a872f --- /dev/null +++ b/experimental/ssh/internal/client/agentshim_signal_test.go @@ -0,0 +1,68 @@ +//go:build unix + +// The test interrupts the shim by sending a signal to its own process, which os/signal cannot +// deliver to self on Windows. The shim itself only runs on the Linux driver. +package client + +import ( + "encoding/json" + "io" + "net/http" + "net/http/httptest" + "os" + "strings" + "syscall" + "testing" + + "github.com/databricks/cli/libs/cmdctx" + "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" +) + +// 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 := 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"): + // 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() + 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("{}")) + } + })) + defer srv.Close() + + ctx := telemetry.WithNewLogger(t.Context()) + ctx = cmdctx.SetConfigUsed(ctx, &config.Config{Host: srv.URL, Token: "test-token"}) + + err := RunAgentShim(ctx, newProbeClient(t, srv.URL), "claude", nil) + require.Error(t, err) + + // cmd/root uploads on the command's own context, which the signal leaves intact. + 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) + }) + } +} From c8d1b364d0a02ad98c65fcd92062b18b45914826 Mon Sep 17 00:00:00 2001 From: Russell Clarey Date: Wed, 14 Oct 2026 11:59:58 +0000 Subject: [PATCH 2/4] Add process.NotifyContext and use it in agent-shim agent-shim needs to cancel on the first signal and then restore the previous disposition, so a second one gets its usual handling. environments setup-local and air each carry a near-copy of that logic. Add one helper in libs/process, modelled on signal.NotifyContext, and use it in place of agent-shim's own copy. The second-signal regression test from environments is copied alongside it, now covering SIGINT, SIGTERM, and SIGHUP. Moving environments and air over to the helper is left to a follow-up. Co-authored-by: Isaac --- experimental/ssh/internal/client/agentshim.go | 31 +---- libs/process/signal.go | 48 ++++++++ libs/process/signal_test.go | 42 +++++++ libs/process/signal_unix_test.go | 110 ++++++++++++++++++ 4 files changed, 202 insertions(+), 29 deletions(-) create mode 100644 libs/process/signal.go create mode 100644 libs/process/signal_test.go create mode 100644 libs/process/signal_unix_test.go diff --git a/experimental/ssh/internal/client/agentshim.go b/experimental/ssh/internal/client/agentshim.go index ba714edd77a..3ae520a98a1 100644 --- a/experimental/ssh/internal/client/agentshim.go +++ b/experimental/ssh/internal/client/agentshim.go @@ -10,7 +10,6 @@ import ( "net/http" "os" "os/exec" - "os/signal" "path/filepath" "regexp" "runtime" @@ -24,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" @@ -125,7 +125,7 @@ func RunAgentShim(ctx context.Context, client *databricks.WorkspaceClient, agent // Cancelling ctx instead aborts the in-flight step, so the launch fails as USER_ABORTED. // Set up before the deferred event so that defer runs first (LIFO) and sees ctx cancelled // only by a signal or the caller, never by stopSignals itself. - ctx, stopSignals := notifyAbort(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 @@ -185,33 +185,6 @@ func RunAgentShim(ctx context.Context, client *databricks.WorkspaceClient, agent return execProcess(argv0, argv, environ) } -// notifyAbort returns a context cancelled on the first SIGINT, SIGTERM, or SIGHUP. The returned -// stop uninstalls the handler before returning; the caller must defer it. -// -// signal.Notify disables the default disposition for these signals for as long as the channel -// stays registered, so a second Ctrl-C would merely be buffered and dropped, leaving no way to -// abort a step that ignores cancellation. Calling signal.Stop before cancel restores the -// default first, so once the cancellation is observable a second signal terminates the process. -func notifyAbort(parent context.Context) (context.Context, func()) { - ctx, cancel := context.WithCancel(parent) - sigCh := make(chan os.Signal, 1) - signal.Notify(sigCh, os.Interrupt, syscall.SIGTERM, syscall.SIGHUP) - go func() { - // Selecting on ctx.Done() too lets the goroutine exit on the normal (no signal) path - // rather than parking on sigCh for the life of the process. - select { - case <-sigCh: - signal.Stop(sigCh) - cancel() - case <-ctx.Done(): - } - }() - return ctx, func() { - signal.Stop(sigCh) - cancel() - } -} - // agentShimOutcome is the observed result of an agent-shim launch, collected by // RunAgentShim for telemetry. type agentShimOutcome struct { diff --git a/libs/process/signal.go b/libs/process/signal.go new file mode 100644 index 00000000000..869216dfa07 --- /dev/null +++ b/libs/process/signal.go @@ -0,0 +1,48 @@ +package process + +import ( + "context" + "os" + "os/signal" +) + +// 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. +// +// The returned stop uninstalls the handler, cancels the context, and waits for +// the handler goroutine to exit. 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.WithCancel(parent) + sigCh := make(chan os.Signal, 1) + signal.Notify(sigCh, sigs...) + + done := make(chan struct{}) + go func() { + defer close(done) + // Selecting on ctx.Done() too lets the goroutine exit on the normal (no + // signal) path: signal.Stop unregisters the channel but never closes it. + select { + case <-sigCh: + case <-ctx.Done(): + } + signal.Stop(sigCh) + cancel() + }() + + return ctx, func() { + signal.Stop(sigCh) + cancel() + <-done + } +} diff --git a/libs/process/signal_test.go b/libs/process/signal_test.go new file mode 100644 index 00000000000..5c9e6da07f9 --- /dev/null +++ b/libs/process/signal_test.go @@ -0,0 +1,42 @@ +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) +} + +// 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..4410dda28c4 --- /dev/null +++ b/libs/process/signal_unix_test.go @@ -0,0 +1,110 @@ +//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) + } + }) + } +} From 10419068f24c5f9eebebd2393121bc8fc202bb47 Mon Sep 17 00:00:00 2001 From: Russell Clarey Date: Wed, 14 Oct 2026 12:00:51 +0000 Subject: [PATCH 3/4] Handle late signals and cancelled probes in agent-shim A signal that arrived while a step ignored ctx (warm bootstrap, file I/O) was swallowed and the agent launched anyway. Restore the default signal handling before committing to the launch, so a later signal kills the shim as it would without a handler, then check whether a signal already came in. Checking first would leave a window in which a signal is consumed by the handler and the agent still launches. stopSignals cancels ctx, so process.NotifyContext now cancels with a *SignalError cause on a signal, and its stop still reports a signal that arrived before it. The success event uses the caller's context. probeAIGateway now returns the cancellation before classifying its results, so an interrupted probe is no longer reported as a disabled gateway or a transient error to retry. Co-authored-by: Isaac --- experimental/ssh/internal/client/agentshim.go | 42 +++++++++++---- libs/process/signal.go | 52 ++++++++++++++----- 2 files changed, 72 insertions(+), 22 deletions(-) diff --git a/experimental/ssh/internal/client/agentshim.go b/experimental/ssh/internal/client/agentshim.go index 3ae520a98a1..8278b846080 100644 --- a/experimental/ssh/internal/client/agentshim.go +++ b/experimental/ssh/internal/client/agentshim.go @@ -120,11 +120,13 @@ func agentByName(name string) (agentSpec, bool) { } func RunAgentShim(ctx context.Context, client *databricks.WorkspaceClient, agentName string, agentArgs []string) (retErr error) { - // Without a handler, Ctrl-C, SIGTERM, or SIGHUP (the SSH connection dropping) would - // terminate the process outright, skipping the failure event below and cmd/root's upload. - // Cancelling ctx instead aborts the in-flight step, so the launch fails as USER_ABORTED. - // Set up before the deferred event so that defer runs first (LIFO) and sees ctx cancelled - // only by a signal or the caller, never by stopSignals itself. + // 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() @@ -174,14 +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)) - // Nothing below observes ctx, so restore the default signal handling: a signal from here - // on terminates the process instead of cancelling a context no step checks any more. - stopSignals() + logAgentShimEvent(parentCtx, agentName, outcome, time.Since(start)) + uploadAgentShimTelemetry(parentCtx, agentName, time.Since(start)) return execProcess(argv0, argv, environ) } @@ -309,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, "/") @@ -323,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) { @@ -333,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/libs/process/signal.go b/libs/process/signal.go index 869216dfa07..06380207acb 100644 --- a/libs/process/signal.go +++ b/libs/process/signal.go @@ -4,6 +4,7 @@ import ( "context" "os" "os/signal" + "sync" ) // NotifyContext is like signal.NotifyContext, but stops relaying as soon as @@ -19,30 +20,57 @@ import ( // SIGHUP or SIGINT if the process started with it ignored (e.g. under nohup); // signal.Stop ignores it again. // -// The returned stop uninstalls the handler, cancels the context, and waits for -// the handler goroutine to exit. It is safe to call more than once; the caller -// must defer it. +// 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.WithCancel(parent) + 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() too lets the goroutine exit on the normal (no - // signal) path: signal.Stop unregisters the channel but never closes it. + // 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 <-sigCh: + 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: + } } - signal.Stop(sigCh) - cancel() }() - return ctx, func() { + return ctx, sync.OnceFunc(func() { signal.Stop(sigCh) - cancel() + 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" } From 2006cfc8df9fad85e3f14329d504391b058e44fb Mon Sep 17 00:00:00 2001 From: Russell Clarey Date: Wed, 14 Oct 2026 12:06:23 +0000 Subject: [PATCH 4/4] Test agent-shim interruption paths Cover a probe run on a cancelled context, and a signal during a warm-path bootstrap step that swallows the cancellation: the shim must not exec the agent and must report USER_ABORTED. Isolate HOME and provide the exec ID and runtime detection cmd/root sets up, so a regression fails by launching the stub agent rather than by panicking earlier. Co-authored-by: Isaac --- .../internal/client/agentshim_signal_test.go | 124 +++++++++++++----- .../ssh/internal/client/agentshim_test.go | 12 ++ libs/process/signal_test.go | 2 + libs/process/signal_unix_test.go | 20 +++ 4 files changed, 125 insertions(+), 33 deletions(-) diff --git a/experimental/ssh/internal/client/agentshim_signal_test.go b/experimental/ssh/internal/client/agentshim_signal_test.go index 2e9cb6a872f..0ddc6fb98b5 100644 --- a/experimental/ssh/internal/client/agentshim_signal_test.go +++ b/experimental/ssh/internal/client/agentshim_signal_test.go @@ -1,20 +1,25 @@ //go:build unix -// The test interrupts the shim by sending a signal to its own process, which os/signal cannot +// 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" @@ -22,6 +27,52 @@ import ( "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. @@ -29,40 +80,47 @@ 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 := 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"): - // 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() - 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("{}")) - } - })) - defer srv.Close() - - ctx := telemetry.WithNewLogger(t.Context()) - ctx = cmdctx.SetConfigUsed(ctx, &config.Config{Host: srv.URL, Token: "test-token"}) + 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.Error(t, err) - - // cmd/root uploads on the command's own context, which the signal leaves intact. - 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) + 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_test.go b/libs/process/signal_test.go index 5c9e6da07f9..218af438290 100644 --- a/libs/process/signal_test.go +++ b/libs/process/signal_test.go @@ -28,6 +28,8 @@ func TestNotifyContextStopsWithoutSignal(t *testing.T) { 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 diff --git a/libs/process/signal_unix_test.go b/libs/process/signal_unix_test.go index 4410dda28c4..fe63900f1a1 100644 --- a/libs/process/signal_unix_test.go +++ b/libs/process/signal_unix_test.go @@ -108,3 +108,23 @@ func TestNotifyContextSecondSignalStillKills(t *testing.T) { }) } } + +// 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) +}