fix: handle context cancellation gracefully during startup/shutdown (#11)

jhaynie · web-flow · commit 2b6085317c28 · 2026-01-27T11:47:43.000-06:00
* fix: handle context cancellation gracefully during startup/shutdown When the gravity process receives SIGTERM during startup or is shut down gracefully, context cancellation errors should not be logged as fatal errors. Changes: - Propagate context.Canceled cleanly from stack.CreateNetworkProvider() - Handle context cancellation in cmd/start.go by returning gracefully instead of calling logger.Fatal() - Change context done during startup from Fatal to Debug log level Fixes: agentuity/sdk#775 * bump go-common * fix: use errors.Is() for context cancellation detection Replace direct equality checks against context.Canceled with errors.Is() to correctly detect wrapped cancellation errors: - cmd/start.go: Use errors.Is(err, context.Canceled) for network provider error - internal/stack/stack.go: Use errors.Is() for proxy error handler and client start This ensures wrapped cancellation errors are properly detected and treated as graceful shutdowns rather than unexpected errors.
diff --git a/cmd/start.go b/cmd/start.go
@@ -6,6 +6,7 @@ import (
 	"crypto/x509"
 	"encoding/base64"
 	"encoding/pem"
+	"errors"
 	"fmt"
 	"net/http"
 	"os"
@@ -156,13 +157,20 @@ var rootCmd = &cobra.Command{
 			return nil
 		})
 		if err != nil {
+			// Don't log fatal error for context cancellation (graceful shutdown)
+			if errors.Is(err, context.Canceled) {
+				logger.Debug("gravity client shutdown due to context cancellation")
+				return
+			}
 			logger.Fatal("failed to create network provider: %v", err)
 		}
 
 		// Wait for provider connection
 		select {
 		case <-ctx.Done():
-			logger.Fatal("context done: %v", ctx.Err())
+			// Context cancelled during startup - graceful shutdown, not an error
+			logger.Debug("context cancelled during startup: %v", ctx.Err())
+			return
 		case <-time.After(time.Second * 10):
 			logger.Error("timed out waiting for provider connection")
 			os.Exit(1)
diff --git a/go.mod b/go.mod
@@ -3,7 +3,7 @@ module github.com/agentuity/gravity
 go 1.25.5
 
 require (
-	github.com/agentuity/go-common v1.0.142
+	github.com/agentuity/go-common v1.0.144
 	github.com/spf13/cobra v1.10.1
 	gvisor.dev/gvisor v0.0.0-20240423190808-9d7a357edefe
 )
diff --git a/go.sum b/go.sum
@@ -1,5 +1,5 @@
-github.com/agentuity/go-common v1.0.142 h1:hcrf9y0pYSne47B1ReM+VZdhuAJxDmJn1mM1VXiBemw=
-github.com/agentuity/go-common v1.0.142/go.mod h1:rvM6VhfltrwrfbjKjNToBQ7kNx3lR6955+n/G1C8soM=
+github.com/agentuity/go-common v1.0.144 h1:tbIqPNIa0xJA+4XVV/Z9Dg6JF5+Saqgu+i/qZAZbBhc=
+github.com/agentuity/go-common v1.0.144/go.mod h1:rvM6VhfltrwrfbjKjNToBQ7kNx3lR6955+n/G1C8soM=
 github.com/cespare/xxhash/v2 v2.3.0 h1:UL815xU9SqsFlibzuggzjXhog7bL6oX9BbNZnL2UFvs=
 github.com/cespare/xxhash/v2 v2.3.0/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs=
 github.com/cockroachdb/errors v1.12.0 h1:d7oCs6vuIMUQRVbi6jWWWEJZahLCfJpnJSVobd1/sUo=
diff --git a/internal/stack/stack.go b/internal/stack/stack.go
@@ -6,6 +6,7 @@ import (
 	"crypto/tls"
 	"crypto/x509"
 	"encoding/pem"
+	"errors"
 	"fmt"
 	"io"
 	"net"
@@ -175,7 +176,7 @@ func StartServer(ctx context.Context, logger _logger.Logger, tlsConfig *tls.Conf
 	proxy.FlushInterval = -1
 	proxy.ErrorHandler = func(w http.ResponseWriter, r *http.Request, err error) {
 		// Suppress expected context cancellation errors (client disconnect, WebSocket close)
-		if ctx.Err() == context.Canceled || r.Context().Err() == context.Canceled {
+		if errors.Is(ctx.Err(), context.Canceled) || errors.Is(r.Context().Err(), context.Canceled) {
 			return
 		}
 		logger.Error("proxy error: %v", err)
@@ -379,6 +380,10 @@ func CreateNetworkProvider(
 	network.client = client
 
 	if err := client.Start(); err != nil {
+		// Don't wrap context.Canceled - let it propagate cleanly for graceful shutdown
+		if errors.Is(err, context.Canceled) {
+			return nil, nil, err
+		}
 		return nil, nil, fmt.Errorf("failed to start gravity client: %w", err)
 	}
 

Original file line number	Diff line number	Diff line change
`@@ -3,7 +3,7 @@ module github.com/agentuity/gravity`
`3`	`3`	`go 1.25.5`
`4`	`4`
`5`	`5`	`require (`
`6`		`- github.com/agentuity/go-common v1.0.142`
	`6`	`+ github.com/agentuity/go-common v1.0.144`
`7`	`7`	`github.com/spf13/cobra v1.10.1`
`8`	`8`	`gvisor.dev/gvisor v0.0.0-20240423190808-9d7a357edefe`
`9`	`9`	`)`