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 --- internal/hexaimcp/run.go | 33 ++++++++++++++++++++------------- internal/hexaimcp/run_test.go | 29 +++++++++++++++-------------- 2 files changed, 35 insertions(+), 27 deletions(-) (limited to 'internal/hexaimcp') diff --git a/internal/hexaimcp/run.go b/internal/hexaimcp/run.go index 74eb476..7c487c3 100644 --- a/internal/hexaimcp/run.go +++ b/internal/hexaimcp/run.go @@ -2,6 +2,7 @@ package hexaimcp import ( + "context" "fmt" "io" "log" @@ -25,8 +26,10 @@ type MCPOverrides struct { } // ServerRunner interface allows dependency injection for testing. +// Run takes a context so the serve loop is cancelled when the process shuts +// down (e.g. on SIGINT/SIGTERM from the top-level caller). type ServerRunner interface { - Run() error + Run(ctx context.Context) error } // ServerFactory creates a server instance (testable). @@ -45,14 +48,17 @@ func defaultServerFactory(r io.Reader, w io.Writer, logger *log.Logger, store pr // Run starts the MCP server with the given configuration and overrides. // This is the main entry point called from cmd/hexai-mcp-server/main.go. -func Run(logPath, configPath string, overrides MCPOverrides, stdin io.Reader, stdout io.Writer, stderr io.Writer) error { - return RunWithFactory(logPath, configPath, overrides, stdin, stdout, stderr, defaultServerFactory) +// ctx is threaded through config loading and the serve loop so the run is +// cancellable from the process entry point. +func Run(ctx context.Context, logPath, configPath string, overrides MCPOverrides, stdin io.Reader, stdout io.Writer, stderr io.Writer) error { + return RunWithFactory(ctx, logPath, configPath, overrides, stdin, stdout, stderr, defaultServerFactory) } // RunWithFactory allows test injection of server factory. // Overrides are applied to the loaded config before use, allowing CLI flags // to take precedence over config file and environment variable settings. func RunWithFactory( + ctx context.Context, logPath string, configPath string, overrides MCPOverrides, @@ -76,14 +82,14 @@ func RunWithFactory( logger.Printf("WARNING: hexai-mcp-server is DEPRECATED and experimental - not actively maintained") // Load configuration and apply CLI overrides - cfg := loadConfig(logger, configPath) + cfg := loadConfig(ctx, logger, configPath) applyOverrides(&cfg, overrides) - return runServer(cfg, logger, stdin, stdout, factory) + return runServer(ctx, cfg, logger, stdin, stdout, factory) } // runServer creates the prompt store, syncer, and runs the MCP server. -func runServer(cfg appconfig.App, logger *log.Logger, stdin io.Reader, stdout io.Writer, factory ServerFactory) error { +func runServer(ctx context.Context, cfg appconfig.App, logger *log.Logger, stdin io.Reader, stdout io.Writer, factory ServerFactory) error { // Determine prompts directory from config (overrides already applied) promptsDir, err := getPromptsDir(cfg) if err != nil { @@ -105,7 +111,7 @@ func runServer(cfg appconfig.App, logger *log.Logger, stdin io.Reader, stdout io // Create and run server server := factory(stdin, stdout, logger, store, syncer) - if err := server.Run(); err != nil { + if err := server.Run(ctx); err != nil { return fmt.Errorf("server error: %w", err) } @@ -136,13 +142,13 @@ func setupLogger(logPath string) (*log.Logger, error) { } // loadConfig loads the hexai configuration. -// Returns default config if loading fails. -func loadConfig(logger *log.Logger, configPath string) appconfig.App { +// Returns default config if loading fails or ctx is already cancelled. +func loadConfig(ctx context.Context, logger *log.Logger, configPath string) appconfig.App { opts := appconfig.LoadOptions{ ConfigPath: configPath, IgnoreEnv: false, } - return appconfig.LoadWithOptions(logger, opts) + return appconfig.LoadWithOptions(ctx, logger, opts) } // applyOverrides applies CLI flag overrides to the loaded config. @@ -212,8 +218,9 @@ func createSyncer(cfg appconfig.App, logger *log.Logger) (*slashcommands.Syncer, } // RunBackfill performs a one-time sync of all prompts and exits. -// Overrides are applied to the loaded config before use. -func RunBackfill(logPath, configPath string, overrides MCPOverrides) error { +// Overrides are applied to the loaded config before use. ctx lets a cancelled +// caller abort config loading before the (blocking) sync work begins. +func RunBackfill(ctx context.Context, logPath, configPath string, overrides MCPOverrides) error { logger, err := setupLogger(logPath) if err != nil { return fmt.Errorf("cannot setup logger: %w", err) @@ -227,7 +234,7 @@ func RunBackfill(logPath, configPath string, overrides MCPOverrides) error { logger.Printf("hexai-mcp-server backfill starting") // Load configuration and apply CLI overrides - cfg := loadConfig(logger, configPath) + cfg := loadConfig(ctx, logger, configPath) applyOverrides(&cfg, overrides) // Force enable sync for backfill diff --git a/internal/hexaimcp/run_test.go b/internal/hexaimcp/run_test.go index 6bedbfc..09f4f87 100644 --- a/internal/hexaimcp/run_test.go +++ b/internal/hexaimcp/run_test.go @@ -3,6 +3,7 @@ package hexaimcp import ( "bytes" + "context" "encoding/json" "fmt" "io" @@ -22,7 +23,7 @@ type mockServerRunner struct { runFunc func() error } -func (m *mockServerRunner) Run() error { +func (m *mockServerRunner) Run(context.Context) error { if m.runFunc != nil { return m.runFunc() } @@ -66,7 +67,7 @@ func TestFullProtocolFlow(t *testing.T) { overrides := MCPOverrides{PromptsDir: tmpDir} // Note: This will hang waiting for more input, which is expected - _ = RunWithFactory("", "", overrides, inBuf, outBuf, errBuf, serverFactory) + _ = RunWithFactory(context.Background(), "", "", overrides, inBuf, outBuf, errBuf, serverFactory) }() // Give server time to process @@ -231,14 +232,14 @@ func TestLoadConfig(t *testing.T) { logger := log.New(io.Discard, "", 0) t.Run("loads default config when path empty", func(t *testing.T) { - cfg := loadConfig(logger, "") + cfg := loadConfig(context.Background(), logger, "") // Should return a valid config (may be defaults) // Just verify it returns without panic _ = cfg }) t.Run("loads config with nonexistent path", func(t *testing.T) { - cfg := loadConfig(logger, "/nonexistent/config.yaml") + cfg := loadConfig(context.Background(), logger, "/nonexistent/config.yaml") // Should return default config without error // Just verify it returns without panic _ = cfg @@ -281,7 +282,7 @@ func TestRun(t *testing.T) { // Pass prompts dir via overrides instead of environment variable overrides := MCPOverrides{PromptsDir: tmpDir} - err := RunWithFactory(logPath, "", overrides, inBuf, outBuf, errBuf, mockFactory) + err := RunWithFactory(context.Background(), logPath, "", overrides, inBuf, outBuf, errBuf, mockFactory) if err != nil { t.Fatalf("RunWithFactory() error = %v", err) } @@ -312,7 +313,7 @@ func TestRunWithFactory_ServerError(t *testing.T) { // Pass prompts dir via overrides instead of environment variable overrides := MCPOverrides{PromptsDir: tmpDir} - err := RunWithFactory(logPath, "", overrides, inBuf, outBuf, errBuf, mockFactory) + err := RunWithFactory(context.Background(), logPath, "", overrides, inBuf, outBuf, errBuf, mockFactory) if err == nil { t.Fatal("RunWithFactory() expected error, got nil") } @@ -331,7 +332,7 @@ func TestRunWithFactory_LoggerError(t *testing.T) { return &mockServerRunner{} } - err := RunWithFactory(badLogPath, "", MCPOverrides{}, &bytes.Buffer{}, &bytes.Buffer{}, &bytes.Buffer{}, mockFactory) + err := RunWithFactory(context.Background(), badLogPath, "", MCPOverrides{}, &bytes.Buffer{}, &bytes.Buffer{}, &bytes.Buffer{}, mockFactory) if err == nil { t.Fatal("expected error for invalid log path, got nil") } @@ -353,7 +354,7 @@ func TestRunWithFactory_StderrLogger(t *testing.T) { overrides := MCPOverrides{PromptsDir: tmpDir} // Empty logPath causes logger to write to stderr (no file to close) - err := RunWithFactory("", "", overrides, &bytes.Buffer{}, &bytes.Buffer{}, &bytes.Buffer{}, mockFactory) + err := RunWithFactory(context.Background(), "", "", overrides, &bytes.Buffer{}, &bytes.Buffer{}, &bytes.Buffer{}, mockFactory) if err != nil { t.Fatalf("RunWithFactory() error = %v", err) } @@ -371,7 +372,7 @@ func TestRun_CallsDefaultFactory(t *testing.T) { // Run with empty stdin — the real server hits EOF and exits cleanly. // This exercises the full Run -> RunWithFactory -> defaultServerFactory path. - err := Run(logPath, "", overrides, &bytes.Buffer{}, &bytes.Buffer{}, &bytes.Buffer{}) + err := Run(context.Background(), logPath, "", overrides, &bytes.Buffer{}, &bytes.Buffer{}, &bytes.Buffer{}) // The server may return nil or an error depending on how it handles EOF; // the important thing is that Run() itself does not panic. _ = err @@ -524,7 +525,7 @@ func TestRunBackfill_FullHappyPath(t *testing.T) { // RunBackfill should succeed: config sets MCPSlashCommandDir, prompts // dir exists, and SyncAll on an empty store is a no-op. - err := RunBackfill(logPath, cfgPath, overrides) + err := RunBackfill(context.Background(), logPath, cfgPath, overrides) if err != nil { t.Fatalf("RunBackfill() error = %v", err) } @@ -550,7 +551,7 @@ func TestRunBackfill_CreateSyncerError(t *testing.T) { t.Fatalf("cannot write config: %v", err) } - err := RunBackfill(logPath, cfgPath, MCPOverrides{}) + err := RunBackfill(context.Background(), logPath, cfgPath, MCPOverrides{}) if err == nil { t.Fatal("expected error for invalid slash command dir, got nil") } @@ -580,7 +581,7 @@ func TestRunBackfill_StderrLogger(t *testing.T) { overrides := MCPOverrides{PromptsDir: promptsDir} // Empty logPath — logger writes to stderr, defer close is a no-op - err := RunBackfill("", cfgPath, overrides) + err := RunBackfill(context.Background(), "", cfgPath, overrides) if err != nil { t.Fatalf("RunBackfill() error = %v", err) } @@ -589,7 +590,7 @@ func TestRunBackfill_StderrLogger(t *testing.T) { // TestRunBackfill_LoggerError verifies RunBackfill returns an error when // the log path is invalid. func TestRunBackfill_LoggerError(t *testing.T) { - err := RunBackfill("/dev/null/impossible/test.log", "", MCPOverrides{}) + err := RunBackfill(context.Background(), "/dev/null/impossible/test.log", "", MCPOverrides{}) if err == nil { t.Fatal("expected error for invalid log path, got nil") } @@ -617,7 +618,7 @@ func TestRunBackfill_NoCmdDir(t *testing.T) { defer os.Setenv("HEXAI_MCP_SLASHCOMMAND_DIR", oldEnv) os.Setenv("HEXAI_MCP_SLASHCOMMAND_DIR", "") - err := RunBackfill(logPath, emptyCfgPath, MCPOverrides{}) + err := RunBackfill(context.Background(), logPath, emptyCfgPath, MCPOverrides{}) if err == nil { t.Fatal("expected error for empty slash command dir, got nil") } -- cgit v1.2.3