From 4ccf6f20f594178f5c3d1537a1c964c5e37091e8 Mon Sep 17 00:00:00 2001
From: Evgenii Kniazev <evgenii.kniazev@databricks.com>
Date: Thu, 8 Jan 2026 15:51:57 +0000
Subject: [PATCH] Handle shutdown signals gracefully in MCP server

Add signal handling for SIGINT and SIGTERM to prevent crashes
during MCP server shutdown. When Claude Code SDK terminates the
MCP server process, it now exits cleanly instead of crashing.

This fixes the "Command failed with exit code 1" error that occurs
when Claude Agent SDK tries to shut down the MCP server.

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
---
 experimental/apps-mcp/cmd/apps_mcp.go | 38 ++++++++++++++++++++++++---
 1 file changed, 35 insertions(+), 3 deletions(-)

diff --git a/experimental/apps-mcp/cmd/apps_mcp.go b/experimental/apps-mcp/cmd/apps_mcp.go
index 83da91447c..8def4e9996 100644
--- a/experimental/apps-mcp/cmd/apps_mcp.go
+++ b/experimental/apps-mcp/cmd/apps_mcp.go
@@ -1,6 +1,11 @@
 package mcp
 
 import (
+	"context"
+	"os"
+	"os/signal"
+	"syscall"
+
 	mcplib "github.com/databricks/cli/experimental/apps-mcp/lib"
 	"github.com/databricks/cli/experimental/apps-mcp/lib/server"
 	"github.com/databricks/cli/libs/log"
@@ -25,7 +30,13 @@ The server communicates via stdio using the Model Context Protocol.`,
 		Example: `  # Start MCP server with required warehouse
   databricks experimental apps-mcp --warehouse-id abc123`,
 		RunE: func(cmd *cobra.Command, args []string) error {
-			ctx := cmd.Context()
+			// Create cancellable context for graceful shutdown
+			ctx, cancel := context.WithCancel(cmd.Context())
+			defer cancel()
+
+			// Handle shutdown signals (SIGINT, SIGTERM)
+			sigCh := make(chan os.Signal, 1)
+			signal.Notify(sigCh, syscall.SIGINT, syscall.SIGTERM)
 
 			// Build MCP config from flags
 			cfg := &mcplib.Config{}
@@ -41,8 +52,29 @@ The server communicates via stdio using the Model Context Protocol.`,
 				return err
 			}
 
-			// Run server
-			return srv.Run(ctx)
+			// Run server in goroutine so we can handle signals
+			errCh := make(chan error, 1)
+			go func() {
+				errCh <- srv.Run(ctx)
+			}()
+
+			// Wait for either server error or shutdown signal
+			select {
+			case err := <-errCh:
+				// Server stopped (EOF, error, or context cancelled)
+				if shutdownErr := srv.Shutdown(ctx); shutdownErr != nil {
+					log.Warnf(ctx, "Shutdown error: %v", shutdownErr)
+				}
+				return err
+			case sig := <-sigCh:
+				// Received shutdown signal - exit gracefully
+				log.Infof(ctx, "Received signal %v, shutting down gracefully", sig)
+				cancel() // Cancel context to stop server.Run()
+				if shutdownErr := srv.Shutdown(ctx); shutdownErr != nil {
+					log.Warnf(ctx, "Shutdown error: %v", shutdownErr)
+				}
+				return nil
+			}
 		},
 	}