Move air and environments setup-local to process.NotifyContext

cmd/air/logs.go

@@ -5,12 +5,14 @@ import (
 	"errors"
 	"fmt"
 	"io"
+	"os"
 	"slices"
 	"strconv"
 
 	"github.com/databricks/cli/cmd/root"
 	"github.com/databricks/cli/libs/cmdctx"
 	"github.com/databricks/cli/libs/flags"
+	"github.com/databricks/cli/libs/process"
 	"github.com/databricks/databricks-sdk-go"
 	"github.com/databricks/databricks-sdk-go/apierr"
 	"github.com/databricks/databricks-sdk-go/service/iam"
@@ -116,7 +118,7 @@ JOB_RUN_ID is the job run ID returned by "air run".`,
 		streamCtx := ctx
 		if downloadTo == "" && root.OutputType(cmd) != flags.OutputJSON {
 			var stop context.CancelFunc
-			streamCtx, stop = notifyInterrupt(ctx)
+			streamCtx, stop = process.NotifyContext(ctx, os.Interrupt)
 			defer stop()
 		}
 

cmd/air/run.go

@@ -6,7 +6,6 @@ import (
 	"fmt"
 	"io"
 	"os"
-	"os/signal"
 	"strconv"
 	"strings"
 	"time"
@@ -16,6 +15,7 @@ import (
 	"github.com/databricks/cli/libs/cmdctx"
 	"github.com/databricks/cli/libs/cmdio"
 	"github.com/databricks/cli/libs/flags"
+	"github.com/databricks/cli/libs/process"
 	"github.com/databricks/cli/libs/shellquote"
 	"github.com/databricks/databricks-sdk-go"
 	"github.com/spf13/cobra"
@@ -213,7 +213,9 @@ The path must be a separate argument: cobra reserves -h as a boolean, so
 			retryTracker: newRetryTracker(logRunStatus{}),
 		}
 
-		watchCtx, stop := notifyInterrupt(ctx)
+		// Catch the interrupt so the log stream unwinds and handleWatchResult
+		// prints resume guidance; the CLI root installs no signal handler of its own.
+		watchCtx, stop := process.NotifyContext(ctx, os.Interrupt)
 		defer stop()
 
 		if !jsonOut {
@@ -288,34 +290,6 @@ func printPostSubmitGuidance(out io.Writer, profile, runID string) {
 	fmt.Fprintln(out, "  "+airLogsCommand(profile, runID))
 }
 
-// notifyInterrupt returns a context cancelled on the first interrupt (Ctrl-C),
-// which lets the log stream unwind and print resume guidance via
-// handleWatchResult; the CLI root installs no signal handler of its own. The
-// caller must defer the returned stop.
-//
-// signal.Notify disables the default SIGINT disposition process-wide 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 hung teardown. Calling signal.Stop
-// before cancel restores SIG_DFL first, so the cancellation is only observable
-// once a second signal is guaranteed to terminate the process.
-func notifyInterrupt(parent context.Context) (context.Context, context.CancelFunc) {
-	ctx, cancel := context.WithCancel(parent)
-	sigCh := make(chan os.Signal, 1)
-	signal.Notify(sigCh, os.Interrupt)
-	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():
-			signal.Stop(sigCh)
-		}
-	}()
-	return ctx, cancel
-}
-
 func handleWatchResult(out io.Writer, profile, runID string, err error) error {
 	if !errors.Is(err, context.Canceled) {
 		return err

cmd/environments/sync.go

@@ -1,9 +1,7 @@
 package environments
 
 import (
-	"context"
 	"os"
-	"os/signal"
 	"path/filepath"
 	"syscall"
 
@@ -13,6 +11,7 @@ import (
 	libslocalenv "github.com/databricks/cli/libs/localenv"
 	"github.com/databricks/cli/libs/log"
 	"github.com/databricks/cli/libs/logdiag"
+	"github.com/databricks/cli/libs/process"
 	"github.com/spf13/cobra"
 )
 
@@ -81,55 +80,19 @@ func addComputeFlags(cmd *cobra.Command) {
 	// consumer relies on, instead of a bare pre-RunE Cobra error.
 }
 
-// watchInterruptSignals cancels ctx on the first SIGINT (Ctrl-C) or SIGTERM (how
-// a supervisor, CI timeout, or VS Code stops the child), which propagates to the
-// uv subprocesses the pipeline spawns so they are reaped instead of orphaned
-// mid-provision. The CLI root installs no signal handler of its own.
-//
-// The returned stop function uninstalls the handler and joins the goroutine; the
-// caller must defer it.
-//
-// The handler must give the *second* signal back to the OS. signal.Notify (like
-// signal.NotifyContext, which wraps it) disables the default disposition for
-// SIGINT/SIGTERM for as long as the channel stays registered, so without the
-// signal.Stop below a second Ctrl-C is merely buffered and dropped: the user
-// would have no way to abort during the process group's SIGKILL grace window.
-// Stopping the relay as soon as the first signal lands restores SIG_DFL, so a
-// second signal terminates the CLI immediately. That matters more here than in
-// most commands because WithProcessGroup moves uv out of the foreground process
-// group, so the tty no longer delivers Ctrl-C to it directly — this handler is
-// the only delivery path.
-func watchInterruptSignals(ctx context.Context, cancel context.CancelFunc) func() {
-	sigCh := make(chan os.Signal, 1)
-	signal.Notify(sigCh, os.Interrupt, syscall.SIGTERM)
-
-	done := make(chan struct{})
-	go func() {
-		defer close(done)
-		// Selecting on ctx.Done() too lets the goroutine exit on the normal (no
-		// signal) path rather than blocking on sigCh for the rest of the process:
-		// signal.Stop unregisters the channel but never closes it.
-		select {
-		case <-sigCh:
-			signal.Stop(sigCh)
-			cancel()
-		case <-ctx.Done():
-		}
-	}()
-
-	return func() {
-		signal.Stop(sigCh)
-		// Wake the goroutine in case neither sigCh nor ctx.Done has fired.
-		cancel()
-		<-done
-	}
-}
-
 // runPipeline builds and runs the setup-local Pipeline.
 func runPipeline(cmd *cobra.Command) error {
-	ctx, cancel := context.WithCancel(cmd.Context())
-	defer cancel()
-	defer watchInterruptSignals(ctx, cancel)()
+	// Cancel on the first SIGINT (Ctrl-C) or SIGTERM (how a supervisor, CI
+	// timeout, or VS Code stops the child), which propagates to the uv
+	// subprocesses the pipeline spawns so they are reaped instead of orphaned
+	// mid-provision. The CLI root installs no signal handler of its own.
+	//
+	// Letting the second signal through matters more here than in most commands
+	// because WithProcessGroup moves uv out of the foreground process group, so
+	// the tty no longer delivers Ctrl-C to it directly — this handler is the only
+	// delivery path.
+	ctx, stop := process.NotifyContext(cmd.Context(), os.Interrupt, syscall.SIGTERM)
+	defer stop()
 
 	cluster, _ := cmd.Flags().GetString("cluster-id")
 	clusterName, _ := cmd.Flags().GetString("cluster-name")