From e95f3fdf0a66ba05ba2c8fb7e755e107f9cf7991 Mon Sep 17 00:00:00 2001 From: Paul Buetow Date: Thu, 11 Jun 2026 08:34:41 +0300 Subject: Thread context.Context through blocking I/O entry points Accept ctx as the first parameter on the blocking I/O entry points and propagate it to downstream blocking calls so the work is cancellable from the process entry point: - appconfig.Load / LoadWithOptions: honor ctx before the blocking file reads, returning defaults on a cancelled context. - LSP: lsp.Server.Run(ctx) ties the serve loop to the caller context via a new watchParentContext bridge (cancels the server context, aborting in-flight LLM work). Threaded through hexailsp.Run/RunWithConfig/ RunWithFactory and runtimeconfig.Store.Reload. - MCP: mcp.Server.Run(ctx) stops accepting requests once ctx is cancelled; threaded through hexaimcp.Run/RunWithFactory/RunBackfill. - editor: RunEditor/OpenTempAndEdit/OpenFile take ctx and use exec.CommandContext so a cancelled context kills the editor subprocess; threaded through hexaicli, hexaiaction and askcli call sites. Top-level callers (cmd/hexai-lsp-server, cmd/hexai-mcp-server) now build a signal-cancelled context (SIGINT/SIGTERM) so shutdown tears the run down cleanly. Updated comments to explain the cancellation flow and added cancellation tests for the LSP/MCP loops, editor, and config load. All tests pass with -race; cross-package coverage 86.2%. Co-Authored-By: Claude Opus 4.8 --- cmd/hexai-lsp-server/main.go | 10 +++++++++- cmd/hexai-mcp-server/main.go | 16 ++++++++++++---- cmd/hexai-mcp-server/main_test.go | 27 +++++++++++++++------------ cmd/hexai/app_runner.go | 4 +++- 4 files changed, 39 insertions(+), 18 deletions(-) (limited to 'cmd') diff --git a/cmd/hexai-lsp-server/main.go b/cmd/hexai-lsp-server/main.go index 03e2546..ed3af42 100644 --- a/cmd/hexai-lsp-server/main.go +++ b/cmd/hexai-lsp-server/main.go @@ -2,12 +2,15 @@ package main import ( + "context" "flag" "fmt" "log" "os" + "os/signal" "path/filepath" "strings" + "syscall" "codeberg.org/snonux/hexai/internal" "codeberg.org/snonux/hexai/internal/appconfig" @@ -26,8 +29,13 @@ func main() { return } + // Cancel the run when the process receives an interrupt/terminate signal so + // the LSP serve loop and any in-flight LLM requests are torn down cleanly. + ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM) + defer stop() + path := strings.TrimSpace(*configPath) - if err := hexailsp.RunWithConfig(*logPath, path, os.Stdin, os.Stdout, os.Stderr); err != nil { + if err := hexailsp.RunWithConfig(ctx, *logPath, path, os.Stdin, os.Stdout, os.Stderr); err != nil { log.Fatalf("server error: %v", err) } } diff --git a/cmd/hexai-mcp-server/main.go b/cmd/hexai-mcp-server/main.go index ac88178..03150bb 100644 --- a/cmd/hexai-mcp-server/main.go +++ b/cmd/hexai-mcp-server/main.go @@ -2,11 +2,14 @@ package main import ( + "context" "flag" "fmt" "io" "os" + "os/signal" "path/filepath" + "syscall" "codeberg.org/snonux/hexai/internal" "codeberg.org/snonux/hexai/internal/appconfig" @@ -100,7 +103,11 @@ func runMain(args []string, stdin io.Reader, stdout, stderr io.Writer) int { syncAll: *syncAll, showVersion: *showVersion, } - if err := run(opts, stdin, stdout, stderr); err != nil { + // Cancel the run on interrupt/terminate so the serve loop exits cleanly. + ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM) + defer stop() + + if err := run(ctx, opts, stdin, stdout, stderr); err != nil { fmt.Fprintf(stderr, "error: %v\n", err) return 1 } @@ -109,7 +116,8 @@ func runMain(args []string, stdin io.Reader, stdout, stderr io.Writer) int { // run executes the MCP server logic with the given options and I/O streams. // CLI flag values are passed via MCPOverrides instead of environment variables. -func run(opts mcpOptions, stdin io.Reader, stdout, stderr io.Writer) error { +// ctx is threaded into the server/backfill so they stop on shutdown signals. +func run(ctx context.Context, opts mcpOptions, stdin io.Reader, stdout, stderr io.Writer) error { if opts.showVersion { fmt.Fprintln(stdout, internal.Version) return nil @@ -119,10 +127,10 @@ func run(opts mcpOptions, stdin io.Reader, stdout, stderr io.Writer) error { // Handle backfill operation if opts.syncAll { - return runBackfill(opts.logPath, opts.configPath, overrides) + return runBackfill(ctx, opts.logPath, opts.configPath, overrides) } - return runMCP(opts.logPath, opts.configPath, overrides, stdin, stdout, stderr) + return runMCP(ctx, opts.logPath, opts.configPath, overrides, stdin, stdout, stderr) } // defaultLogPath returns the default MCP log file path in the state directory. diff --git a/cmd/hexai-mcp-server/main_test.go b/cmd/hexai-mcp-server/main_test.go index 33f662d..b2a3895 100644 --- a/cmd/hexai-mcp-server/main_test.go +++ b/cmd/hexai-mcp-server/main_test.go @@ -2,6 +2,7 @@ package main import ( "bytes" + "context" "errors" "io" "strings" @@ -37,7 +38,7 @@ func TestDefaultLogPath(t *testing.T) { func TestRun_ShowVersion(t *testing.T) { var stdout bytes.Buffer opts := mcpOptions{showVersion: true} - if err := run(opts, nil, &stdout, nil); err != nil { + if err := run(context.Background(), opts, nil, &stdout, nil); err != nil { t.Fatalf("run --version: %v", err) } got := strings.TrimSpace(stdout.String()) @@ -70,7 +71,7 @@ func TestRun_SyncAll(t *testing.T) { var gotLog, gotConfig string var gotOverrides hexaimcp.MCPOverrides - runBackfill = func(logPath, configPath string, overrides hexaimcp.MCPOverrides) error { + runBackfill = func(_ context.Context, logPath, configPath string, overrides hexaimcp.MCPOverrides) error { gotLog = logPath gotConfig = configPath gotOverrides = overrides @@ -85,7 +86,7 @@ func TestRun_SyncAll(t *testing.T) { slashCommandSync: true, slashCommandDir: "/tmp/cmds", } - if err := run(opts, nil, nil, nil); err != nil { + if err := run(context.Background(), opts, nil, nil, nil); err != nil { t.Fatalf("run syncAll: %v", err) } if gotLog != "/tmp/test.log" { @@ -110,10 +111,10 @@ func TestRun_SyncAllError(t *testing.T) { t.Cleanup(func() { runBackfill = old }) wantErr := errors.New("backfill failed") - runBackfill = func(_, _ string, _ hexaimcp.MCPOverrides) error { return wantErr } + runBackfill = func(_ context.Context, _, _ string, _ hexaimcp.MCPOverrides) error { return wantErr } opts := mcpOptions{syncAll: true} - if err := run(opts, nil, nil, nil); !errors.Is(err, wantErr) { + if err := run(context.Background(), opts, nil, nil, nil); !errors.Is(err, wantErr) { t.Fatalf("expected backfill error, got: %v", err) } } @@ -123,13 +124,13 @@ func TestRun_MCPServer(t *testing.T) { t.Cleanup(func() { runMCP = old }) called := false - runMCP = func(logPath, configPath string, overrides hexaimcp.MCPOverrides, stdin io.Reader, stdout, stderr io.Writer) error { + runMCP = func(_ context.Context, logPath, configPath string, overrides hexaimcp.MCPOverrides, stdin io.Reader, stdout, stderr io.Writer) error { called = true return nil } opts := mcpOptions{logPath: "/tmp/mcp.log"} - if err := run(opts, nil, nil, nil); err != nil { + if err := run(context.Background(), opts, nil, nil, nil); err != nil { t.Fatalf("run MCP: %v", err) } if !called { @@ -142,9 +143,11 @@ func TestRun_MCPServerError(t *testing.T) { t.Cleanup(func() { runMCP = old }) wantErr := errors.New("server failed") - runMCP = func(_, _ string, _ hexaimcp.MCPOverrides, _ io.Reader, _, _ io.Writer) error { return wantErr } + runMCP = func(_ context.Context, _, _ string, _ hexaimcp.MCPOverrides, _ io.Reader, _, _ io.Writer) error { + return wantErr + } - if err := run(mcpOptions{}, nil, nil, nil); !errors.Is(err, wantErr) { + if err := run(context.Background(), mcpOptions{}, nil, nil, nil); !errors.Is(err, wantErr) { t.Fatalf("expected server error, got: %v", err) } } @@ -173,7 +176,7 @@ func TestRunMain_SyncAllSuccess(t *testing.T) { t.Cleanup(func() { runBackfill = old }) var gotLog string - runBackfill = func(logPath string, _ string, _ hexaimcp.MCPOverrides) error { + runBackfill = func(_ context.Context, logPath string, _ string, _ hexaimcp.MCPOverrides) error { gotLog = logPath return nil } @@ -193,7 +196,7 @@ func TestRunMain_SyncAllSuccess(t *testing.T) { func TestRunMain_ServerErrorReturnsOne(t *testing.T) { old := runMCP t.Cleanup(func() { runMCP = old }) - runMCP = func(string, string, hexaimcp.MCPOverrides, io.Reader, io.Writer, io.Writer) error { + runMCP = func(context.Context, string, string, hexaimcp.MCPOverrides, io.Reader, io.Writer, io.Writer) error { return errors.New("mcp boom") } @@ -212,7 +215,7 @@ func TestRunMain_BadFlagReturnsTwo(t *testing.T) { old := runMCP t.Cleanup(func() { runMCP = old }) called := false - runMCP = func(string, string, hexaimcp.MCPOverrides, io.Reader, io.Writer, io.Writer) error { + runMCP = func(context.Context, string, string, hexaimcp.MCPOverrides, io.Reader, io.Writer, io.Writer) error { called = true return nil } diff --git a/cmd/hexai/app_runner.go b/cmd/hexai/app_runner.go index ab1766b..e9b7e44 100644 --- a/cmd/hexai/app_runner.go +++ b/cmd/hexai/app_runner.go @@ -69,7 +69,9 @@ func normalizeAppRunner(r appRunner) appRunner { func loadAppConfig(configPath string) appconfig.App { logger := log.New(io.Discard, "", 0) - return appconfig.LoadWithOptions(logger, appconfig.LoadOptions{ConfigPath: configPath}) + // Config is loaded before the signal-aware CLI context exists, so use a + // background context here; the load is a quick local file read. + return appconfig.LoadWithOptions(context.Background(), logger, appconfig.LoadOptions{ConfigPath: configPath}) } func parseAppArgs(cfg appconfig.App, configPath string, args []string, stderr io.Writer) (parsedAppArgs, error) { -- cgit v1.2.3