Add KillCaller feature to test server for CLI crash recovery testing

acceptance/internal/config.go

@@ -156,6 +156,10 @@ type ServerStub struct {
 	// Configure as "1ms", "2s", "3m", etc.
 	// See [time.ParseDuration] for details.
 	Delay time.Duration
+
+	// If true, kill the caller process instead of returning a response.
+	// Requires DATABRICKS_CLI_TEST_PID=1 to be set in the test environment.
+	KillCaller bool
 }
 
 // FindConfigs finds all the config relevant for this test,

acceptance/internal/prepare_server.go

@@ -209,6 +209,44 @@ func startLocalServer(t *testing.T,
 				}
 			}
 
+			if stub.KillCaller {
+				pid := testserver.ExtractPidFromHeaders(req.Headers)
+				if pid == 0 {
+					t.Errorf("KillCaller configured but test-pid not found in User-Agent")
+					return testserver.Response{
+						StatusCode: http.StatusBadRequest,
+						Body:       "test-pid not found in User-Agent (set DATABRICKS_CLI_TEST_PID=1)",
+					}
+				}
+
+				t.Logf("KillCaller: killing PID %d (pattern: %s)", pid, stub.Pattern)
+
+				process, err := os.FindProcess(pid)
+				if err != nil {
+					t.Errorf("Failed to find process %d: %s", pid, err)
+					return testserver.Response{
+						StatusCode: http.StatusInternalServerError,
+						Body:       fmt.Sprintf("failed to find process: %s", err),
+					}
+				}
+
+				// Use process.Kill() for cross-platform compatibility.
+				// On Unix, this sends SIGKILL. On Windows, this calls TerminateProcess.
+				if err := process.Kill(); err != nil {
+					t.Errorf("Failed to kill process %d: %s", pid, err)
+					return testserver.Response{
+						StatusCode: http.StatusInternalServerError,
+						Body:       fmt.Sprintf("failed to kill process: %s", err),
+					}
+				}
+
+				// Return a response (the CLI will likely be killed before it receives this)
+				return testserver.Response{
+					StatusCode: http.StatusOK,
+					Body:       fmt.Sprintf("killed PID %d", pid),
+				}
+			}
+
 			return stub.Response
 		})
 	}

acceptance/selftest/kill_caller/mid_request/output.txt

@@ -0,0 +1,5 @@
+
+>>> errcode [CLI] workspace list /
+[PROCESS_KILLED]
+
+Exit code: [KILLED]

acceptance/selftest/kill_caller/mid_request/script

@@ -0,0 +1 @@
+trace errcode $CLI workspace list /

acceptance/selftest/kill_caller/mid_request/test.toml

@@ -0,0 +1,30 @@
+Local = true
+Env.DATABRICKS_CLI_TEST_PID = "1"
+
+[[Server]]
+Pattern = "GET /api/2.0/workspace/list"
+KillCaller = true
+
+[[Repls]]
+# macOS bash shows "Killed: 9" (with signal number), Linux shows "Killed"
+# Normalize the whole killed line to a placeholder
+Old = 'script: line \d+:\s+\d+ Killed(: 9)?\s+"\$@"'
+New = '[PROCESS_KILLED]'
+
+[[Repls]]
+# On Windows, there's no "Killed" message - just empty line before Exit code
+# Insert [PROCESS_KILLED] placeholder for consistency
+Old = '(\n>>> errcode [^\n]+\n)\nExit code:'
+New = """${1}[PROCESS_KILLED]
+
+Exit code:"""
+
+[[Repls]]
+# Normalize exit code: 137 on Unix (128 + SIGKILL), 1 on Windows
+Old = 'Exit code: (137|1)'
+New = 'Exit code: [KILLED]'
+
+[[Repls]]
+# Normalize Windows line endings (CRLF -> LF) - must be LAST
+Old = "\r"
+New = ''

acceptance/selftest/kill_caller/output.txt

@@ -0,0 +1,5 @@
+
+>>> errcode [CLI] current-user me
+[PROCESS_KILLED]
+
+Exit code: [KILLED]

acceptance/selftest/kill_caller/script

@@ -0,0 +1 @@
+trace errcode $CLI current-user me

acceptance/selftest/kill_caller/test.toml

@@ -0,0 +1,31 @@
+Local = true
+Env.DATABRICKS_CLI_TEST_PID = "1"
+
+# Kill the CLI when it calls /Me endpoint
+[[Server]]
+Pattern = "GET /api/2.0/preview/scim/v2/Me"
+KillCaller = true
+
+[[Repls]]
+# macOS bash shows "Killed: 9" (with signal number), Linux shows "Killed"
+# Normalize the whole killed line to a placeholder
+Old = 'script: line \d+:\s+\d+ Killed(: 9)?\s+"\$@"'
+New = '[PROCESS_KILLED]'
+
+[[Repls]]
+# On Windows, there's no "Killed" message - just empty line before Exit code
+# Insert [PROCESS_KILLED] placeholder for consistency
+Old = '(\n>>> errcode [^\n]+\n)\nExit code:'
+New = """${1}[PROCESS_KILLED]
+
+Exit code:"""
+
+[[Repls]]
+# Normalize exit code: 137 on Unix (128 + SIGKILL), 1 on Windows
+Old = 'Exit code: (137|1)'
+New = 'Exit code: [KILLED]'
+
+[[Repls]]
+# Normalize Windows line endings (CRLF -> LF) - must be LAST
+Old = "\r"
+New = ''

cmd/pipelines/root/root.go

@@ -10,6 +10,7 @@ import (
 
 	"github.com/databricks/cli/internal/build"
 	"github.com/databricks/cli/libs/log"
+	"github.com/databricks/cli/libs/testserver"
 	"github.com/spf13/cobra"
 )
 
@@ -72,6 +73,7 @@ func New(ctx context.Context) *cobra.Command {
 		ctx = withCommandInUserAgent(ctx, cmd)
 		ctx = withCommandExecIdInUserAgent(ctx)
 		ctx = withUpstreamInUserAgent(ctx)
+		ctx = testserver.InjectPidToUserAgent(ctx)
 		cmd.SetContext(ctx)
 		return nil
 	}

cmd/root/root.go

@@ -17,6 +17,7 @@ import (
 	"github.com/databricks/cli/libs/log"
 	"github.com/databricks/cli/libs/telemetry"
 	"github.com/databricks/cli/libs/telemetry/protos"
+	"github.com/databricks/cli/libs/testserver"
 	"github.com/spf13/cobra"
 )
 
@@ -79,6 +80,7 @@ func New(ctx context.Context) *cobra.Command {
 		ctx = withCommandInUserAgent(ctx, cmd)
 		ctx = withCommandExecIdInUserAgent(ctx)
 		ctx = withUpstreamInUserAgent(ctx)
+		ctx = testserver.InjectPidToUserAgent(ctx)
 		cmd.SetContext(ctx)
 		return nil
 	}

libs/testserver/server.go

@@ -9,15 +9,46 @@ import (
 	"net/http"
 	"net/http/httptest"
 	"net/url"
+	"os"
 	"reflect"
+	"regexp"
+	"strconv"
 	"strings"
 	"sync"
 
+	"github.com/databricks/cli/internal/testutil"
+	"github.com/databricks/cli/libs/env"
+	"github.com/databricks/databricks-sdk-go/useragent"
 	"github.com/gorilla/mux"
+)
 
-	"github.com/databricks/cli/internal/testutil"
+const (
+	TestPidEnvVar = "DATABRICKS_CLI_TEST_PID"
+	testPidKey    = "test-pid"
 )
 
+var testPidRegex = regexp.MustCompile(testPidKey + `/(\d+)`)
+
+func InjectPidToUserAgent(ctx context.Context) context.Context {
+	if env.Get(ctx, TestPidEnvVar) != "1" {
+		return ctx
+	}
+	return useragent.InContext(ctx, testPidKey, strconv.Itoa(os.Getpid()))
+}
+
+func ExtractPidFromHeaders(headers http.Header) int {
+	ua := headers.Get("User-Agent")
+	matches := testPidRegex.FindStringSubmatch(ua)
+	if len(matches) < 2 {
+		return 0
+	}
+	pid, err := strconv.Atoi(matches[1])
+	if err != nil {
+		return 0
+	}
+	return pid
+}
+
 type Server struct {
 	*httptest.Server
 	Router *mux.Router
@@ -274,6 +305,9 @@ func (s *Server) Handle(method, path string, handler HandlerFunc) {
 				StatusCode: 500,
 				Body:       []byte("INJECTED"),
 			}
+		} else if bytes.Contains(request.Body, []byte("KILL_CALLER")) {
+			s.handleKillCaller(&request, w)
+			return
 		} else {
 			respAny := handler(request)
 			if respAny == nil && request.Context.Err() != nil {
@@ -322,3 +356,35 @@ func isNil(i any) bool {
 		return false
 	}
 }
+
+func (s *Server) handleKillCaller(request *Request, w http.ResponseWriter) {
+	pid := ExtractPidFromHeaders(request.Headers)
+	if pid == 0 {
+		s.t.Errorf("KILL_CALLER requested but test-pid not found in User-Agent")
+		w.WriteHeader(http.StatusBadRequest)
+		_, _ = fmt.Fprint(w, "test-pid not found in User-Agent (set DATABRICKS_CLI_TEST_PID=1)")
+		return
+	}
+
+	s.t.Logf("KILL_CALLER: killing PID %d", pid)
+
+	process, err := os.FindProcess(pid)
+	if err != nil {
+		s.t.Errorf("Failed to find process %d: %s", pid, err)
+		w.WriteHeader(http.StatusInternalServerError)
+		_, _ = fmt.Fprintf(w, "failed to find process: %s", err)
+		return
+	}
+
+	// Use process.Kill() for cross-platform compatibility.
+	// On Unix, this sends SIGKILL. On Windows, this calls TerminateProcess.
+	if err := process.Kill(); err != nil {
+		s.t.Errorf("Failed to kill process %d: %s", pid, err)
+		w.WriteHeader(http.StatusInternalServerError)
+		_, _ = fmt.Fprintf(w, "failed to kill process: %s", err)
+		return
+	}
+
+	w.WriteHeader(http.StatusOK)
+	_, _ = fmt.Fprintf(w, "killed PID %d", pid)
+}