From 995c155d7b54d78c466268ba3b0f41bb2e1a5806 Mon Sep 17 00:00:00 2001
From: Samuel K <skevetter@pm.me>
Date: Sat, 18 Jul 2026 12:00:45 -0500
Subject: [PATCH] refactor(up): modernize WithSignals with signal.NotifyContext

Replace the hand-rolled context.WithCancel + signal goroutine with
signal.NotifyContext for first-signal cancellation, matching the pattern
already used in cmd/internal. The second-signal force-exit behavior is
preserved via a small watchForForceShutdown helper. No behavior change.
---
 cmd/workspace/up/up.go | 70 ++++++++++++++++++++++--------------------
 1 file changed, 36 insertions(+), 34 deletions(-)

diff --git a/cmd/workspace/up/up.go b/cmd/workspace/up/up.go
index 36c6c30bb..aace0f72c 100644
--- a/cmd/workspace/up/up.go
+++ b/cmd/workspace/up/up.go
@@ -424,46 +424,48 @@ func (cmd *UpCmd) executeDevsyUp(
 	return &workspaceContext{result: result, user: user, workdir: workdir}, nil
 }
 
+var shutdownSignals = []os.Signal{
+	os.Interrupt, syscall.SIGHUP, syscall.SIGTERM, syscall.SIGQUIT,
+}
+
+// WithSignals returns a context cancelled on the first shutdown signal so
+// deferred cleanup can run, and forces exit on a second signal. The returned
+// func must be called to release the signal handler.
 func WithSignals(ctx context.Context) (context.Context, func()) {
-	ctx, cancel := context.WithCancel(ctx)
-	signals := make(chan os.Signal, 1)
-	signal.Notify(signals, os.Interrupt, syscall.SIGHUP, syscall.SIGTERM, syscall.SIGQUIT)
+	ctx, stop := signal.NotifyContext(ctx, shutdownSignals...)
 
-	// done lets cleanup unblock goroutines parked on <-signals; signal.Stop
-	// alone wouldn't wake them.
+	// done lets the force-shutdown watcher exit once cleanup runs; stop()
+	// alone wouldn't wake it.
 	done := make(chan struct{})
+	go watchForForceShutdown(ctx, done)
 
-	go func() {
-		select {
-		case <-signals:
-			cancel()
-		case <-ctx.Done():
-		case <-done:
-		}
-	}()
+	return ctx, func() {
+		stop()
+		close(done)
+	}
+}
 
-	go func() {
-		select {
-		case <-ctx.Done():
-		case <-done:
-			return
-		}
-		// Skip the second-signal wait if cleanup already closed done.
-		select {
-		case <-done:
-			return
-		default:
-		}
-		select {
-		case <-signals:
-			os.Exit(1) // second signal — force shutdown
-		case <-done:
-		}
-	}()
+// watchForForceShutdown waits for the first-signal cancellation, then forces
+// exit if a second signal arrives before cleanup closes done.
+func watchForForceShutdown(ctx context.Context, done <-chan struct{}) {
+	select {
+	case <-ctx.Done():
+	case <-done:
+		return
+	}
+	// Skip the second-signal wait if cleanup already closed done.
+	select {
+	case <-done:
+		return
+	default:
+	}
 
-	return ctx, func() {
-		cancel()
+	signals := make(chan os.Signal, 1)
+	signal.Notify(signals, shutdownSignals...)
+	select {
+	case <-signals:
+		os.Exit(1) // second signal — force shutdown
+	case <-done:
 		signal.Stop(signals)
-		close(done)
 	}
 }