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 +- internal/appconfig/config_alias_test.go | 3 +- internal/appconfig/config_env_model_test.go | 9 ++-- internal/appconfig/config_features_test.go | 15 +++--- internal/appconfig/config_load.go | 18 +++++-- internal/appconfig/config_test.go | 47 ++++++++++++------ internal/appconfig/custom_validation_more_test.go | 5 +- internal/askcli/command_edit.go | 11 +++-- internal/askcli/command_edit_test.go | 4 +- internal/editor/editor.go | 23 +++++---- internal/editor/editor_test.go | 51 ++++++++++++-------- internal/hexaiaction/custom_action_test.go | 2 +- internal/hexaiaction/custom_exec_more_test.go | 4 +- internal/hexaiaction/custom_exec_test.go | 4 +- internal/hexaiaction/prompts_simplify_test.go | 2 +- internal/hexaiaction/run.go | 4 +- internal/hexaiaction/run_more_test.go | 8 ++-- internal/hexaiaction/run_seam_test.go | 2 +- internal/hexaiaction/run_test.go | 2 +- internal/hexaicli/editor_integration_test.go | 6 ++- internal/hexaicli/run_editor_behavior_test.go | 2 +- internal/hexaicli/runner.go | 8 ++-- internal/hexaicli/runner_test.go | 6 +-- internal/hexailsp/dependencies.go | 3 +- internal/hexailsp/run.go | 31 +++++++----- internal/hexailsp/run_more_test.go | 16 +++---- internal/hexailsp/run_test.go | 19 ++++---- internal/hexaimcp/run.go | 33 ++++++++----- internal/hexaimcp/run_test.go | 29 ++++++------ internal/lsp/chat_commands.go | 4 +- internal/lsp/chat_commands_test.go | 5 +- internal/lsp/server.go | 30 +++++++++++- internal/lsp/server_test.go | 58 +++++++++++++++++++++++ internal/mcp/handlers_prompt_test.go | 5 +- internal/mcp/server.go | 16 ++++++- internal/mcp/server_test.go | 21 +++++++- internal/runtimeconfig/store.go | 9 ++-- internal/runtimeconfig/store_test.go | 11 +++-- internal/tmuxedit/run.go | 5 +- 42 files changed, 404 insertions(+), 184 deletions(-) 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) { diff --git a/internal/appconfig/config_alias_test.go b/internal/appconfig/config_alias_test.go index da7909e..764611b 100644 --- a/internal/appconfig/config_alias_test.go +++ b/internal/appconfig/config_alias_test.go @@ -1,6 +1,7 @@ package appconfig import ( + "context" "log" "os" "path/filepath" @@ -28,7 +29,7 @@ codex = "gpt-5-codex" if err := os.WriteFile(path, []byte(toml), 0o644); err != nil { t.Fatalf("write: %v", err) } - cfg := Load(log.New(os.Stderr, "test ", 0)) + cfg := Load(context.Background(), log.New(os.Stderr, "test ", 0)) if cfg.OpenAIModel != "gpt-5-codex" { t.Fatalf("expected alias to resolve to gpt-5-codex, got %q", cfg.OpenAIModel) } diff --git a/internal/appconfig/config_env_model_test.go b/internal/appconfig/config_env_model_test.go index e10fa5d..9856a4e 100644 --- a/internal/appconfig/config_env_model_test.go +++ b/internal/appconfig/config_env_model_test.go @@ -1,6 +1,7 @@ package appconfig import ( + "context" "log" "os" "testing" @@ -12,14 +13,14 @@ func TestEnv_GenericModelOverrideAndPrecedence(t *testing.T) { t.Setenv("HEXAI_MODEL", "gpt-5-codex") t.Setenv("HEXAI_PROVIDER", "openai") // No provider-specific env set yet: HEXAI_MODEL should flow into OpenAIModel - cfg := Load(log.New(os.Stderr, "test ", 0)) + cfg := Load(context.Background(), log.New(os.Stderr, "test ", 0)) if cfg.OpenAIModel != "gpt-5-codex" { t.Fatalf("expected OpenAIModel=gpt-5-codex via HEXAI_MODEL, got %q", cfg.OpenAIModel) } // Now set a provider-specific model; it should win over HEXAI_MODEL t.Setenv("HEXAI_OPENAI_MODEL", "gpt-5-thinking") - cfg2 := Load(log.New(os.Stderr, "test ", 0)) + cfg2 := Load(context.Background(), log.New(os.Stderr, "test ", 0)) if cfg2.OpenAIModel != "gpt-5-thinking" { t.Fatalf("expected OpenAIModel from HEXAI_OPENAI_MODEL to win, got %q", cfg2.OpenAIModel) } @@ -30,7 +31,7 @@ func TestEnv_ModelForce_OverridesProviderSpecific(t *testing.T) { t.Setenv("HEXAI_OPENAI_MODEL", "gpt-5-main") t.Setenv("HEXAI_MODEL_FORCE", "gpt-5-codex") t.Setenv("HEXAI_PROVIDER", "openai") - cfg := Load(log.New(os.Stderr, "test ", 0)) + cfg := Load(context.Background(), log.New(os.Stderr, "test ", 0)) if cfg.OpenAIModel != "gpt-5-codex" { t.Fatalf("expected OpenAIModel forced to gpt-5-codex, got %q", cfg.OpenAIModel) } @@ -43,7 +44,7 @@ func TestEnv_SurfaceModelOverrides(t *testing.T) { t.Setenv("HEXAI_MODEL_CLI", "gpt-cli") t.Setenv("HEXAI_TEMPERATURE_CLI", "0.22") t.Setenv("HEXAI_PROVIDER_CLI", "ollama") - cfg := Load(log.New(os.Stderr, "test ", 0)) + cfg := Load(context.Background(), log.New(os.Stderr, "test ", 0)) if len(cfg.CompletionConfigs) != 1 { t.Fatalf("expected single completion entry, got %+v", cfg.CompletionConfigs) } diff --git a/internal/appconfig/config_features_test.go b/internal/appconfig/config_features_test.go index 123283a..77beb07 100644 --- a/internal/appconfig/config_features_test.go +++ b/internal/appconfig/config_features_test.go @@ -3,6 +3,7 @@ package appconfig import ( + "context" "os" "path/filepath" "reflect" @@ -11,7 +12,7 @@ import ( func TestIgnoreConfig_Defaults(t *testing.T) { clearHexaiEnv(t) - cfg := Load(nil) + cfg := Load(context.Background(), nil) if cfg.IgnoreGitignore == nil || !*cfg.IgnoreGitignore { t.Error("expected IgnoreGitignore default true") } @@ -33,7 +34,7 @@ gitignore = false extra_patterns = ["*.min.js", "dist/**"] lsp_notify_ignored = false `) - cfg := LoadWithOptions(newLogger(), LoadOptions{ConfigPath: cfgPath, ProjectRoot: dir}) + cfg := LoadWithOptions(context.Background(), newLogger(), LoadOptions{ConfigPath: cfgPath, ProjectRoot: dir}) if cfg.IgnoreGitignore == nil || *cfg.IgnoreGitignore { t.Error("expected IgnoreGitignore false from file") } @@ -58,7 +59,7 @@ lsp_notify_ignored = true withEnv(t, "HEXAI_IGNORE_GITIGNORE", "false") withEnv(t, "HEXAI_IGNORE_LSP_NOTIFY", "0") withEnv(t, "HEXAI_IGNORE_EXTRA_PATTERNS", "*.bak,*.tmp") - cfg := LoadWithOptions(newLogger(), LoadOptions{ConfigPath: cfgPath, ProjectRoot: dir}) + cfg := LoadWithOptions(context.Background(), newLogger(), LoadOptions{ConfigPath: cfgPath, ProjectRoot: dir}) if cfg.IgnoreGitignore == nil || *cfg.IgnoreGitignore { t.Error("expected IgnoreGitignore false from env override") } @@ -90,7 +91,7 @@ gitignore = true gitignore = false extra_patterns = ["build/**"] `) - cfg := LoadWithOptions(newLogger(), LoadOptions{ConfigPath: cfgPath, ProjectRoot: projectDir}) + cfg := LoadWithOptions(context.Background(), newLogger(), LoadOptions{ConfigPath: cfgPath, ProjectRoot: projectDir}) if cfg.IgnoreGitignore == nil || *cfg.IgnoreGitignore { t.Error("expected project override to set IgnoreGitignore false") } @@ -108,7 +109,7 @@ func TestIgnoreConfig_DisableGitignore(t *testing.T) { [ignore] gitignore = false `) - cfg := LoadWithOptions(newLogger(), LoadOptions{ConfigPath: cfgPath, ProjectRoot: dir}) + cfg := LoadWithOptions(context.Background(), newLogger(), LoadOptions{ConfigPath: cfgPath, ProjectRoot: dir}) if cfg.IgnoreGitignore == nil || *cfg.IgnoreGitignore { t.Error("expected IgnoreGitignore false") } @@ -149,7 +150,7 @@ clear_keys = "C-u" newline_keys = "S-Enter" submit_keys = "Enter" `) - cfg := LoadWithOptions(newLogger(), LoadOptions{ConfigPath: cfgPath, ProjectRoot: dir}) + cfg := LoadWithOptions(context.Background(), newLogger(), LoadOptions{ConfigPath: cfgPath, ProjectRoot: dir}) if cfg.TmuxEditPopupWidth != "90%" { t.Errorf("PopupWidth = %q, want 90%%", cfg.TmuxEditPopupWidth) } @@ -212,7 +213,7 @@ func TestTmuxEditConfig_SkipsEmptyName(t *testing.T) { name = "" display_name = "Empty" `) - cfg := LoadWithOptions(newLogger(), LoadOptions{ConfigPath: cfgPath, ProjectRoot: dir}) + cfg := LoadWithOptions(context.Background(), newLogger(), LoadOptions{ConfigPath: cfgPath, ProjectRoot: dir}) if len(cfg.TmuxEditAgents) != 0 { t.Errorf("got %d agents, want 0 (empty name should be skipped)", len(cfg.TmuxEditAgents)) } diff --git a/internal/appconfig/config_load.go b/internal/appconfig/config_load.go index 3ed1654..1e4c6b7 100644 --- a/internal/appconfig/config_load.go +++ b/internal/appconfig/config_load.go @@ -1,6 +1,7 @@ package appconfig import ( + "context" "fmt" "log" "os" @@ -15,15 +16,26 @@ import ( const ProjectConfigFilename = ".hexaiconfig.toml" // Load reads configuration from a file and merges with defaults. -// It respects the XDG Base Directory Specification. -func Load(logger *log.Logger) App { return LoadWithOptions(logger, LoadOptions{}) } +// It respects the XDG Base Directory Specification. The context lets callers +// abort the (blocking) file reads on shutdown/cancellation. +func Load(ctx context.Context, logger *log.Logger) App { + return LoadWithOptions(ctx, logger, LoadOptions{}) +} // LoadWithOptions reads configuration and applies the requested loading options. -func LoadWithOptions(logger *log.Logger, opts LoadOptions) App { +// ctx is honored before performing the blocking file I/O so a cancelled caller +// gets a clean default config back instead of stalling on disk access. +func LoadWithOptions(ctx context.Context, logger *log.Logger, opts LoadOptions) App { cfg := newDefaultConfig() if logger == nil { return cfg // Return defaults if no logger is provided (e.g. in tests) } + // Respect cancellation up front: config loading touches several files and + // there is no value in starting that work once the caller has given up. + if ctx != nil && ctx.Err() != nil { + logger.Printf("config load cancelled: %v", ctx.Err()) + return cfg + } // Step 1: Load global config file configPath := strings.TrimSpace(opts.ConfigPath) diff --git a/internal/appconfig/config_test.go b/internal/appconfig/config_test.go index d4e7739..06e73a1 100644 --- a/internal/appconfig/config_test.go +++ b/internal/appconfig/config_test.go @@ -2,6 +2,7 @@ package appconfig import ( "bytes" + "context" "errors" "io" "log" @@ -14,6 +15,24 @@ import ( func newLogger() *log.Logger { return log.New(io.Discard, "", 0) } +// TestLoadWithOptions_CancelledContextReturnsDefaults verifies that a +// pre-cancelled context short-circuits config loading (skipping the blocking +// file reads) and returns the default config instead. +func TestLoadWithOptions_CancelledContextReturnsDefaults(t *testing.T) { + dir := t.TempDir() + cfgPath := filepath.Join(dir, "config.toml") + writeFile(t, cfgPath, "[core]\nmax_tokens = 9999\n") + + ctx, cancel := context.WithCancel(context.Background()) + cancel() + + cfg := LoadWithOptions(ctx, newLogger(), LoadOptions{ConfigPath: cfgPath}) + def := newDefaultConfig() + if cfg.MaxTokens != def.MaxTokens { + t.Fatalf("expected default config on cancelled ctx, got MaxTokens=%d", cfg.MaxTokens) + } +} + func writeFile(t *testing.T, path, content string) { t.Helper() if err := os.MkdirAll(filepath.Dir(path), 0o755); err != nil { @@ -46,7 +65,7 @@ func withEnv(t *testing.T, k, v string) { } func TestLoad_Defaults_NoLogger(t *testing.T) { - cfg := Load(nil) + cfg := Load(context.Background(), nil) if cfg.MaxTokens == 0 || cfg.ContextMode == "" || cfg.ContextWindowLines == 0 || cfg.MaxContextTokens == 0 { t.Fatalf("expected defaults populated, got %+v", cfg) } @@ -60,7 +79,7 @@ func TestLoad_Defaults_WithLogger_NoFile_NoEnv(t *testing.T) { dir := t.TempDir() t.Setenv("XDG_CONFIG_HOME", dir) logger := newLogger() - cfg := LoadWithOptions(logger, LoadOptions{ProjectRoot: dir}) + cfg := LoadWithOptions(context.Background(), logger, LoadOptions{ProjectRoot: dir}) def := newDefaultConfig() if cfg.MaxTokens != def.MaxTokens || cfg.ContextMode != def.ContextMode || cfg.ContextWindowLines != def.ContextWindowLines { t.Fatalf("expected defaults; got %+v want %+v", cfg, def) @@ -186,7 +205,7 @@ temperature = 0.0 withEnv(t, "HEXAI_PROVIDER_CLI", "ollama") logger := newLogger() - cfg := LoadWithOptions(logger, LoadOptions{ProjectRoot: dir}) + cfg := LoadWithOptions(context.Background(), logger, LoadOptions{ProjectRoot: dir}) // Check overrides if cfg.MaxTokens != 321 || cfg.ContextMode != "always-full" || cfg.ContextWindowLines != 77 || cfg.MaxContextTokens != 888 { @@ -255,7 +274,7 @@ temperature = 0.0 } { t.Setenv(k, "") } - cfg2 := LoadWithOptions(logger, LoadOptions{ProjectRoot: dir}) + cfg2 := LoadWithOptions(context.Background(), logger, LoadOptions{ProjectRoot: dir}) if cfg2.MaxTokens != 123 || cfg2.ContextMode != "file-on-new-func" || cfg2.ContextWindowLines != 50 || cfg2.MaxContextTokens != 999 || cfg2.LogPreviewLimit != 0 { t.Fatalf("file merge not applied: %+v", cfg2) } @@ -434,7 +453,7 @@ temperature = 0.0 // Ensure no env override interferes with manual_invoke_min_prefix in this test t.Setenv("HEXAI_MANUAL_INVOKE_MIN_PREFIX", "") logger := newLogger() - cfg := LoadWithOptions(logger, LoadOptions{ProjectRoot: dir}) + cfg := LoadWithOptions(context.Background(), logger, LoadOptions{ProjectRoot: dir}) if cfg.MaxTokens != 111 || cfg.ContextMode != "window" || cfg.ContextWindowLines != 42 || cfg.MaxContextTokens != 777 { t.Fatalf("sectioned basics wrong: %+v", cfg) @@ -495,7 +514,7 @@ explain_system = "CLI-EXPLAIN" ` writeFile(t, cfgPath, content) - cfg := Load(newLogger()) + cfg := Load(context.Background(), newLogger()) // completion if cfg.PromptCompletionSystemGeneral != "SYS-GENERAL" || cfg.PromptCompletionSystemParams != "SYS-PARAMS" || cfg.PromptCompletionSystemInline != "SYS-INLINE" { @@ -557,7 +576,7 @@ user = "Diagnostics to resolve (selection only):\n{{diagnostics}}\n\nSelected co custom_menu_hotkey = "a" ` writeFile(t, cfgPath, content) - cfg := Load(newLogger()) + cfg := Load(context.Background(), newLogger()) if err := cfg.Validate(); err != nil { t.Fatalf("validate: %v", err) } @@ -592,7 +611,7 @@ id = "DUP" title = "B" instruction = "y" `) - cfg := Load(newLogger()) + cfg := Load(context.Background(), newLogger()) if err := cfg.Validate(); err == nil || !strings.Contains(err.Error(), "duplicate custom action id") { t.Fatalf("expected duplicate id error, got %v", err) } @@ -616,7 +635,7 @@ title = "B" instruction = "y" hotkey = "E" `) - cfg := Load(newLogger()) + cfg := Load(context.Background(), newLogger()) if err := cfg.Validate(); err == nil || !strings.Contains(err.Error(), "duplicate custom action hotkey") { t.Fatalf("expected duplicate hotkey error, got %v", err) } @@ -635,7 +654,7 @@ title = "A" instruction = "x" scope = "bad" `) - cfg := Load(newLogger()) + cfg := Load(context.Background(), newLogger()) if err := cfg.Validate(); err == nil || !strings.Contains(err.Error(), "invalid scope") { t.Fatalf("expected invalid scope error, got %v", err) } @@ -650,7 +669,7 @@ func TestTmuxMenuHotkey_Clash_Error(t *testing.T) { [tmux] custom_menu_hotkey = "r" `) - cfg := Load(newLogger()) + cfg := Load(context.Background(), newLogger()) if err := cfg.Validate(); err == nil || !strings.Contains(err.Error(), "clashes with built-in") { t.Fatalf("expected clash error, got %v", err) } @@ -727,7 +746,7 @@ name = "anthropic" // Load using explicit ProjectRoot (avoids needing to chdir) logger := newLogger() - cfg := LoadWithOptions(logger, LoadOptions{ProjectRoot: projectDir}) + cfg := LoadWithOptions(context.Background(), logger, LoadOptions{ProjectRoot: projectDir}) // Project config should override global values if cfg.MaxTokens != 8000 { @@ -768,7 +787,7 @@ max_tokens = 8000 withEnv(t, "HEXAI_MAX_TOKENS", "9999") logger := newLogger() - cfg := LoadWithOptions(logger, LoadOptions{ProjectRoot: projectDir}) + cfg := LoadWithOptions(context.Background(), logger, LoadOptions{ProjectRoot: projectDir}) if cfg.MaxTokens != 9999 { t.Fatalf("expected env max_tokens=9999 to override project, got %d", cfg.MaxTokens) @@ -799,7 +818,7 @@ max_tokens = 2000 } logger := newLogger() - cfg := LoadWithOptions(logger, LoadOptions{}) + cfg := LoadWithOptions(context.Background(), logger, LoadOptions{}) // Should get global config values, not defaults if cfg.MaxTokens != 2000 { diff --git a/internal/appconfig/custom_validation_more_test.go b/internal/appconfig/custom_validation_more_test.go index 36212aa..255963b 100644 --- a/internal/appconfig/custom_validation_more_test.go +++ b/internal/appconfig/custom_validation_more_test.go @@ -1,6 +1,7 @@ package appconfig import ( + "context" "path/filepath" "strings" "testing" @@ -19,7 +20,7 @@ instruction = "x" id = "no-title" instruction = "x" `) - cfg := Load(newLogger()) + cfg := Load(context.Background(), newLogger()) if err := cfg.Validate(); err == nil || (!strings.Contains(err.Error(), "missing required field id") && !strings.Contains(err.Error(), "missing required field title")) { t.Fatalf("expected missing field error, got %v", err) } @@ -40,7 +41,7 @@ hotkey = "too" [tmux] custom_menu_hotkey = "ab" `) - cfg := Load(newLogger()) + cfg := Load(context.Background(), newLogger()) if err := cfg.Validate(); err == nil || (!strings.Contains(err.Error(), "hotkey must be a single character") && !strings.Contains(err.Error(), "invalid tmux.custom_menu_hotkey")) { t.Fatalf("expected invalid hotkey error, got %v", err) } diff --git a/internal/askcli/command_edit.go b/internal/askcli/command_edit.go index 27c9106..a82d575 100644 --- a/internal/askcli/command_edit.go +++ b/internal/askcli/command_edit.go @@ -10,9 +10,10 @@ import ( // captureFromEditor opens the user's editor on a temporary file pre-filled with // the given initial content and returns its trimmed contents after the editor -// exits. It is a variable so tests can stub it. -var captureFromEditor = func(initial []byte) (string, error) { - return editor.OpenTempAndEdit(initial) +// exits. ctx is forwarded to the editor subprocess so it can be cancelled with +// the surrounding command. It is a variable so tests can stub it. +var captureFromEditor = func(ctx context.Context, initial []byte) (string, error) { + return editor.OpenTempAndEdit(ctx, initial) } // handleEdit opens the configured editor on a temporary file. With no selector @@ -22,7 +23,7 @@ func (d *Dispatcher) handleEdit(ctx context.Context, args []string, stdout, stde if len(args) >= 2 { return d.editTaskDescription(ctx, args[1], stdout, stderr) } - description, err := captureFromEditor(nil) + description, err := captureFromEditor(ctx, nil) if err != nil { writeInfoError(stderr, err) return 1, nil @@ -43,7 +44,7 @@ func (d *Dispatcher) editTaskDescription(ctx context.Context, selector string, s return code, nil } - description, err := captureFromEditor([]byte(tasks[0].Description)) + description, err := captureFromEditor(ctx, []byte(tasks[0].Description)) if err != nil { writeInfoError(stderr, err) return 1, nil diff --git a/internal/askcli/command_edit_test.go b/internal/askcli/command_edit_test.go index 80cbd56..4db19c1 100644 --- a/internal/askcli/command_edit_test.go +++ b/internal/askcli/command_edit_test.go @@ -12,7 +12,7 @@ import ( func stubEditorCapture(t *testing.T, content string, err error) { t.Helper() old := captureFromEditor - captureFromEditor = func(initial []byte) (string, error) { + captureFromEditor = func(_ context.Context, initial []byte) (string, error) { return content, err } t.Cleanup(func() { captureFromEditor = old }) @@ -60,7 +60,7 @@ func TestHandleEdit_ExistingTaskModifiesDescription(t *testing.T) { var initialContent []byte old := captureFromEditor - captureFromEditor = func(initial []byte) (string, error) { + captureFromEditor = func(_ context.Context, initial []byte) (string, error) { initialContent = initial return "updated description", nil } diff --git a/internal/editor/editor.go b/internal/editor/editor.go index 722e336..9a9e737 100644 --- a/internal/editor/editor.go +++ b/internal/editor/editor.go @@ -1,6 +1,7 @@ package editor import ( + "context" "errors" "os" "os/exec" @@ -21,9 +22,11 @@ func Resolve() (string, error) { } // RunEditor is the seam that invokes the editor on the given file path. -// Override in tests to avoid launching a real editor. -var RunEditor = func(editor, path string) error { - cmd := exec.Command(editor, path) +// Override in tests to avoid launching a real editor. It uses +// exec.CommandContext so a cancelled ctx (e.g. process shutdown) kills the +// editor subprocess instead of leaving it blocking on terminal input. +var RunEditor = func(ctx context.Context, editor, path string) error { + cmd := exec.CommandContext(ctx, editor, path) cmd.Stdin = os.Stdin cmd.Stdout = os.Stdout cmd.Stderr = os.Stderr @@ -32,8 +35,9 @@ var RunEditor = func(editor, path string) error { // OpenTempAndEdit creates a temporary .md file, writes initial content if provided, // opens it in the resolved editor, then reads the final content and removes the file. -// Returns the trimmed content. -func OpenTempAndEdit(initial []byte) (string, error) { +// Returns the trimmed content. ctx is forwarded to the editor subprocess so it +// can be cancelled along with the surrounding command. +func OpenTempAndEdit(ctx context.Context, initial []byte) (string, error) { ed, err := Resolve() if err != nil { return "", err @@ -59,7 +63,7 @@ func OpenTempAndEdit(initial []byte) (string, error) { if err := f.Close(); err != nil { return "", err } - if err := RunEditor(ed, path); err != nil { + if err := RunEditor(ctx, ed, path); err != nil { return "", err } b, err := os.ReadFile(filepath.Clean(path)) @@ -70,8 +74,9 @@ func OpenTempAndEdit(initial []byte) (string, error) { } // OpenFile ensures the parent directory exists, then opens path in the editor -// from Resolve() (HEXAI_EDITOR or EDITOR). -func OpenFile(path string) error { +// from Resolve() (HEXAI_EDITOR or EDITOR). ctx is forwarded to the editor +// subprocess so it can be cancelled with the surrounding command. +func OpenFile(ctx context.Context, path string) error { ed, err := Resolve() if err != nil { return err @@ -86,5 +91,5 @@ func OpenFile(path string) error { return mkErr } } - return RunEditor(ed, path) + return RunEditor(ctx, ed, path) } diff --git a/internal/editor/editor_test.go b/internal/editor/editor_test.go index 403d165..f2c1e22 100644 --- a/internal/editor/editor_test.go +++ b/internal/editor/editor_test.go @@ -1,6 +1,7 @@ package editor import ( + "context" "errors" "os" "path/filepath" @@ -14,19 +15,31 @@ func TestRunEditor_Default(t *testing.T) { if err := os.WriteFile(tmp, []byte("hello"), 0o600); err != nil { t.Fatal(err) } - if err := RunEditor("true", tmp); err != nil { + if err := RunEditor(context.Background(), "true", tmp); err != nil { t.Fatalf("RunEditor with 'true': %v", err) } } // TestRunEditor_Default_BadCommand verifies RunEditor returns an error for a nonexistent command. func TestRunEditor_Default_BadCommand(t *testing.T) { - err := RunEditor("nonexistent-editor-cmd-12345", "/dev/null") + err := RunEditor(context.Background(), "nonexistent-editor-cmd-12345", "/dev/null") if err == nil { t.Fatal("expected error for nonexistent editor command") } } +// TestRunEditor_CancelledContextKillsProcess verifies that a cancelled context +// terminates the editor subprocess (via exec.CommandContext) so RunEditor +// returns an error rather than blocking forever. +func TestRunEditor_CancelledContextKillsProcess(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + cancel() // already cancelled: CommandContext refuses to start the process + // "sleep 60" would otherwise block; the cancelled ctx kills it immediately. + if err := RunEditor(ctx, "sleep", "60"); err == nil { + t.Fatal("expected error when context is cancelled") + } +} + func TestResolve_EnvPriority(t *testing.T) { t.Setenv("HEXAI_EDITOR", "ed1") t.Setenv("EDITOR", "ed2") @@ -67,12 +80,12 @@ func TestOpenTempAndEdit_UsesRunEditor(t *testing.T) { // Ensure Resolve() succeeds t.Setenv("HEXAI_EDITOR", "dummy") var capturedPath string - RunEditor = func(editor, path string) error { + RunEditor = func(_ context.Context, editor, path string) error { capturedPath = path // simulate user writing content return os.WriteFile(path, []byte("Hello\nWorld\n"), 0o600) } - out, err := OpenTempAndEdit([]byte("# Start\n\n")) + out, err := OpenTempAndEdit(context.Background(), []byte("# Start\n\n")) if err != nil { t.Fatalf("OpenTempAndEdit: %v", err) } @@ -88,7 +101,7 @@ func TestOpenTempAndEdit_UsesRunEditor(t *testing.T) { func TestOpenTempAndEdit_NoEditor(t *testing.T) { t.Setenv("HEXAI_EDITOR", "") t.Setenv("EDITOR", "") - _, err := OpenTempAndEdit(nil) + _, err := OpenTempAndEdit(context.Background(), nil) if err == nil { t.Fatal("expected error when no editor is set") } @@ -99,11 +112,11 @@ func TestOpenTempAndEdit_NilInitial(t *testing.T) { old := RunEditor t.Cleanup(func() { RunEditor = old }) t.Setenv("HEXAI_EDITOR", "dummy") - RunEditor = func(editor, path string) error { + RunEditor = func(_ context.Context, editor, path string) error { // simulate user writing content into a file that started empty return os.WriteFile(path, []byte("result"), 0o600) } - out, err := OpenTempAndEdit(nil) + out, err := OpenTempAndEdit(context.Background(), nil) if err != nil { t.Fatalf("OpenTempAndEdit with nil initial: %v", err) } @@ -118,10 +131,10 @@ func TestOpenTempAndEdit_EmptyInitial(t *testing.T) { old := RunEditor t.Cleanup(func() { RunEditor = old }) t.Setenv("HEXAI_EDITOR", "dummy") - RunEditor = func(editor, path string) error { + RunEditor = func(_ context.Context, editor, path string) error { return os.WriteFile(path, []byte(" trimmed "), 0o600) } - out, err := OpenTempAndEdit([]byte{}) + out, err := OpenTempAndEdit(context.Background(), []byte{}) if err != nil { t.Fatalf("OpenTempAndEdit with empty initial: %v", err) } @@ -136,10 +149,10 @@ func TestOpenTempAndEdit_EditorError(t *testing.T) { t.Cleanup(func() { RunEditor = old }) t.Setenv("HEXAI_EDITOR", "dummy") editorErr := errors.New("editor crashed") - RunEditor = func(editor, path string) error { + RunEditor = func(_ context.Context, editor, path string) error { return editorErr } - _, err := OpenTempAndEdit([]byte("some content")) + _, err := OpenTempAndEdit(context.Background(), []byte("some content")) if err == nil { t.Fatal("expected error when editor fails") } @@ -153,11 +166,11 @@ func TestOpenTempAndEdit_EditorDeletesFile(t *testing.T) { old := RunEditor t.Cleanup(func() { RunEditor = old }) t.Setenv("HEXAI_EDITOR", "dummy") - RunEditor = func(editor, path string) error { + RunEditor = func(_ context.Context, editor, path string) error { // simulate the editor deleting the file return os.Remove(path) } - _, err := OpenTempAndEdit([]byte("content")) + _, err := OpenTempAndEdit(context.Background(), []byte("content")) if err == nil { t.Fatal("expected error when temp file is deleted by editor") } @@ -169,11 +182,11 @@ func TestOpenTempAndEdit_TempFileCleanup(t *testing.T) { t.Cleanup(func() { RunEditor = old }) t.Setenv("HEXAI_EDITOR", "dummy") var capturedPath string - RunEditor = func(editor, path string) error { + RunEditor = func(_ context.Context, editor, path string) error { capturedPath = path return os.WriteFile(path, []byte("done"), 0o600) } - _, err := OpenTempAndEdit(nil) + _, err := OpenTempAndEdit(context.Background(), nil) if err != nil { t.Fatalf("OpenTempAndEdit: %v", err) } @@ -189,12 +202,12 @@ func TestOpenFile_CreatesParentAndInvokesEditor(t *testing.T) { t.Setenv("HEXAI_EDITOR", "dummy") target := filepath.Join(t.TempDir(), "nested", "config.toml") var gotEditor, gotPath string - RunEditor = func(editorCmd, path string) error { + RunEditor = func(_ context.Context, editorCmd, path string) error { gotEditor = editorCmd gotPath = path return nil } - if err := OpenFile(target); err != nil { + if err := OpenFile(context.Background(), target); err != nil { t.Fatalf("OpenFile: %v", err) } if gotEditor != "dummy" { @@ -211,7 +224,7 @@ func TestOpenFile_CreatesParentAndInvokesEditor(t *testing.T) { func TestOpenFile_NoEditor(t *testing.T) { t.Setenv("HEXAI_EDITOR", "") t.Setenv("EDITOR", "") - err := OpenFile(filepath.Join(t.TempDir(), "x.toml")) + err := OpenFile(context.Background(), filepath.Join(t.TempDir(), "x.toml")) if err == nil { t.Fatal("expected error when no editor is set") } @@ -219,7 +232,7 @@ func TestOpenFile_NoEditor(t *testing.T) { func TestOpenFile_EmptyPath(t *testing.T) { t.Setenv("HEXAI_EDITOR", "true") - err := OpenFile(" ") + err := OpenFile(context.Background(), " ") if err == nil { t.Fatal("expected error for empty path") } diff --git a/internal/hexaiaction/custom_action_test.go b/internal/hexaiaction/custom_action_test.go index 74ac350..e2f7902 100644 --- a/internal/hexaiaction/custom_action_test.go +++ b/internal/hexaiaction/custom_action_test.go @@ -30,7 +30,7 @@ func TestActionCustom_UsesEditorPrompt(t *testing.T) { runner.newClient = func(_ appconfig.App) (actionClient, error) { return llmFake2{}, nil } oldRunEd := editor.RunEditor - editor.RunEditor = func(_ string, path string) error { + editor.RunEditor = func(_ context.Context, _ string, path string) error { return os.WriteFile(path, []byte("make it done"), 0o600) } t.Cleanup(func() { editor.RunEditor = oldRunEd }) diff --git a/internal/hexaiaction/custom_exec_more_test.go b/internal/hexaiaction/custom_exec_more_test.go index d826b39..73d9839 100644 --- a/internal/hexaiaction/custom_exec_more_test.go +++ b/internal/hexaiaction/custom_exec_more_test.go @@ -19,7 +19,7 @@ func (c *capDoer) Chat(_ context.Context, msgs []llm.Message, _ ...llm.RequestOp func (*capDoer) DefaultModel() string { return "m" } func TestExecuteAction_Custom_DoesNotMutateProvidedSelection(t *testing.T) { - cfg := appconfig.Load(nil) + cfg := appconfig.Load(context.Background(), nil) parts := InputParts{Selection: "code"} selectedCustom := &appconfig.CustomAction{ID: "x", Title: "X", Instruction: "Do it"} _, _ = executeAction(context.Background(), ActionCustom, parts, &cfg, fakeDoer{"OK"}, nil, selectedCustom) @@ -29,7 +29,7 @@ func TestExecuteAction_Custom_DoesNotMutateProvidedSelection(t *testing.T) { } func TestRunCustom_UserTemplate_InjectsDiagnostics(t *testing.T) { - cfg := appconfig.Load(nil) + cfg := appconfig.Load(context.Background(), nil) parts := InputParts{Selection: "code", Diagnostics: []string{"L1", "L2"}} ca := appconfig.CustomAction{ID: "y", Title: "Y", User: "{{diagnostics}}\n{{selection}}"} cap := &capDoer{} diff --git a/internal/hexaiaction/custom_exec_test.go b/internal/hexaiaction/custom_exec_test.go index 1a5b99e..11670fd 100644 --- a/internal/hexaiaction/custom_exec_test.go +++ b/internal/hexaiaction/custom_exec_test.go @@ -9,7 +9,7 @@ import ( ) func TestExecuteAction_CustomConfigured_Instruction(t *testing.T) { - cfg := appconfig.Load(nil) + cfg := appconfig.Load(context.Background(), nil) parts := InputParts{Selection: "code"} selectedCustom := &appconfig.CustomAction{ID: "x", Title: "X", Instruction: "Do it"} out, err := executeAction(context.Background(), ActionCustom, parts, &cfg, fakeDoer{"OK"}, nil, selectedCustom) @@ -19,7 +19,7 @@ func TestExecuteAction_CustomConfigured_Instruction(t *testing.T) { } func TestExecuteAction_CustomConfigured_User(t *testing.T) { - cfg := appconfig.Load(nil) + cfg := appconfig.Load(context.Background(), nil) parts := InputParts{Selection: "sel"} selectedCustom := &appconfig.CustomAction{ID: "y", Title: "Y", User: "Apply to: {{selection}}"} out, err := executeAction(context.Background(), ActionCustom, parts, &cfg, fakeDoer{"OK2"}, nil, selectedCustom) diff --git a/internal/hexaiaction/prompts_simplify_test.go b/internal/hexaiaction/prompts_simplify_test.go index 4cd831d..b13bf13 100644 --- a/internal/hexaiaction/prompts_simplify_test.go +++ b/internal/hexaiaction/prompts_simplify_test.go @@ -16,7 +16,7 @@ func (simplifyClient) Chat(_ context.Context, _ []llm.Message, _ ...llm.RequestO func (simplifyClient) DefaultModel() string { return "m" } func TestRunSimplify_Smoke(t *testing.T) { - cfg := appconfig.Load(nil) + cfg := appconfig.Load(context.Background(), nil) out, err := runSimplify(context.Background(), &cfg, simplifyClient{}, "code") if err != nil { t.Fatalf("runSimplify: %v", err) diff --git a/internal/hexaiaction/run.go b/internal/hexaiaction/run.go index 92aa72d..f0ce7ee 100644 --- a/internal/hexaiaction/run.go +++ b/internal/hexaiaction/run.go @@ -114,7 +114,7 @@ func (tmuxActionStatusSink) SetLLMStart(provider, model string) error { } func loadActionConfig(ctx context.Context, logger *log.Logger) appconfig.App { - return appconfig.LoadWithOptions(logger, appconfig.LoadOptions{ConfigPath: configPathFromContext(ctx)}) + return appconfig.LoadWithOptions(ctx, logger, appconfig.LoadOptions{ConfigPath: configPathFromContext(ctx)}) } type actionPlan struct { @@ -381,7 +381,7 @@ func handleCustomAction(ctx context.Context, parts InputParts, cfg actionConfig, } func handleCustomPromptAction(ctx context.Context, parts InputParts, cfg actionConfig, client chatDoer, stderr io.Writer) (string, error) { - prompt, err := editor.OpenTempAndEdit(nil) + prompt, err := editor.OpenTempAndEdit(ctx, nil) if err != nil || strings.TrimSpace(prompt) == "" { _, _ = fmt.Fprintln(stderr, logging.AnsiBase+"hexai-tmux-action: custom prompt canceled or empty; echoing input"+logging.AnsiReset) return parts.Selection, nil diff --git a/internal/hexaiaction/run_more_test.go b/internal/hexaiaction/run_more_test.go index 67ffa96..97bcbe9 100644 --- a/internal/hexaiaction/run_more_test.go +++ b/internal/hexaiaction/run_more_test.go @@ -31,7 +31,7 @@ func TestRun_MissingAPIKey(t *testing.T) { func TestRun_NoInput_IsActionable(t *testing.T) { runner := NewRunner() - runner.loadConfig = func(context.Context, *log.Logger) appconfig.App { return appconfig.Load(nil) } + runner.loadConfig = func(context.Context, *log.Logger) appconfig.App { return appconfig.Load(context.Background(), nil) } runner.newClient = func(appconfig.App) (actionClient, error) { return llmFake{}, nil } runner.chooseAction = func(appconfig.App) (actionChoice, error) { return actionChoice{kind: ActionSkip}, nil @@ -62,7 +62,7 @@ func TestHandleDiagnosticsActionInvokesLLM(t *testing.T) { t.Setenv("HEXAI_TMUX_STATUS", "0") parts := InputParts{Diagnostics: []string{"warn1"}, Selection: "code"} client := &stubChatDoer{} - cfg := appconfig.Load(nil) + cfg := appconfig.Load(context.Background(), nil) if _, err := handleDiagnosticsAction(context.Background(), parts, &cfg, client); err != nil { t.Fatalf("handleDiagnosticsAction: %v", err) } @@ -84,7 +84,7 @@ func TestHandleSimplifyActionPassesSelection(t *testing.T) { t.Setenv("HEXAI_TMUX_STATUS", "0") parts := InputParts{Selection: "value := 1"} client := &stubChatDoer{} - cfg := appconfig.Load(nil) + cfg := appconfig.Load(context.Background(), nil) if _, err := handleSimplifyAction(context.Background(), parts, &cfg, client); err != nil { t.Fatalf("handleSimplifyAction: %v", err) } @@ -107,7 +107,7 @@ func TestHandleCustomActionUsesProvidedCustom(t *testing.T) { sel := appconfig.CustomAction{ID: "custom", Title: "Do", Instruction: "do it"} parts := InputParts{Selection: "text"} client := &stubChatDoer{} - cfg := appconfig.Load(nil) + cfg := appconfig.Load(context.Background(), nil) if _, err := handleCustomAction(context.Background(), parts, &cfg, client, &sel); err != nil { t.Fatalf("handleCustomAction: %v", err) } diff --git a/internal/hexaiaction/run_seam_test.go b/internal/hexaiaction/run_seam_test.go index 8fb8533..d37f179 100644 --- a/internal/hexaiaction/run_seam_test.go +++ b/internal/hexaiaction/run_seam_test.go @@ -65,7 +65,7 @@ func TestRun_WithInjectedConfigAndStatusSink(t *testing.T) { t.Setenv("XDG_CONFIG_HOME", t.TempDir()) sink := &recordingActionStatusSink{} runner := NewRunner() - runner.loadConfig = func(context.Context, *log.Logger) appconfig.App { return appconfig.Load(nil) } + runner.loadConfig = func(context.Context, *log.Logger) appconfig.App { return appconfig.Load(context.Background(), nil) } runner.newClient = func(_ appconfig.App) (actionClient, error) { return llmFake{}, nil } runner.chooseAction = func(_ appconfig.App) (actionChoice, error) { return actionChoice{kind: ActionSkip}, nil diff --git a/internal/hexaiaction/run_test.go b/internal/hexaiaction/run_test.go index b927d57..574874f 100644 --- a/internal/hexaiaction/run_test.go +++ b/internal/hexaiaction/run_test.go @@ -26,7 +26,7 @@ func TestExecuteAction_Skip(t *testing.T) { } func TestExecuteAction_Rewrite_Document_GoTest(t *testing.T) { - cfg := appconfig.Load(nil) // defaults + cfg := appconfig.Load(context.Background(), nil) // defaults // Use fenced output to exercise StripFences client := fakeDoer{"```\nDONE\n```"} diff --git a/internal/hexaicli/editor_integration_test.go b/internal/hexaicli/editor_integration_test.go index 7d53d1f..e5580be 100644 --- a/internal/hexaicli/editor_integration_test.go +++ b/internal/hexaicli/editor_integration_test.go @@ -28,7 +28,9 @@ func TestRun_NoArgs_OpensEditor(t *testing.T) { newClientFromApp = func(_ appconfig.App) (llm.Client, error) { return cliFake{}, nil } t.Cleanup(func() { newClientFromApp = oldNew }) oldRun := editor.RunEditor - editor.RunEditor = func(_ string, path string) error { return os.WriteFile(path, []byte("PROMPT"), 0o600) } + editor.RunEditor = func(_ context.Context, _ string, path string) error { + return os.WriteFile(path, []byte("PROMPT"), 0o600) + } t.Cleanup(func() { editor.RunEditor = oldRun }) t.Setenv("HEXAI_EDITOR", "dummy") @@ -50,7 +52,7 @@ func TestRun_WithArgs_DoesNotOpenEditor(t *testing.T) { // Stub editor and detect if called (should not be) called := false oldRun := editor.RunEditor - editor.RunEditor = func(_ string, _ string) error { called = true; return nil } + editor.RunEditor = func(_ context.Context, _ string, _ string) error { called = true; return nil } t.Cleanup(func() { editor.RunEditor = oldRun }) var stdout, stderr bytes.Buffer if err := Run(context.Background(), []string{"ARG"}, bytes.NewBufferString("SEL"), &stdout, &stderr); err != nil { diff --git a/internal/hexaicli/run_editor_behavior_test.go b/internal/hexaicli/run_editor_behavior_test.go index a934473..99a2f2d 100644 --- a/internal/hexaicli/run_editor_behavior_test.go +++ b/internal/hexaicli/run_editor_behavior_test.go @@ -25,7 +25,7 @@ func TestRun_DoesNotOpenEditorWhenStdinPresent(t *testing.T) { // Guard: make editor invocation fatal if called oldRunEd := editor.RunEditor defer func() { editor.RunEditor = oldRunEd }() - editor.RunEditor = func(_ string, _ string) error { + editor.RunEditor = func(_ context.Context, _ string, _ string) error { t.Fatalf("editor should not be invoked when stdin has content") return nil } diff --git a/internal/hexaicli/runner.go b/internal/hexaicli/runner.go index eaae1cd..3929001 100644 --- a/internal/hexaicli/runner.go +++ b/internal/hexaicli/runner.go @@ -18,7 +18,7 @@ import ( type cliConfigLoader func(context.Context, *log.Logger) appconfig.App -type cliEditorOpener func([]byte) (string, error) +type cliEditorOpener func(context.Context, []byte) (string, error) type cliClientFactory func(appconfig.App) (llm.Client, error) @@ -85,7 +85,7 @@ func (r *Runner) Run(ctx context.Context, args []string, stdin io.Reader, stdout } cfgPath = p } - if err := editor.OpenFile(cfgPath); err != nil { + if err := editor.OpenFile(ctx, cfgPath); err != nil { _, _ = fmt.Fprintf(stderr, logging.AnsiBase+"hexai %s: %v"+logging.AnsiReset+"\n", sub, err) return err } @@ -127,7 +127,7 @@ func (r *Runner) Run(ctx context.Context, args []string, stdin io.Reader, stdout input, rerr := readInput(stdin, args) if rerr != nil && len(args) == 0 { - if prompt, eerr := runner.openEditor(nil); eerr == nil && strings.TrimSpace(prompt) != "" { + if prompt, eerr := runner.openEditor(ctx, nil); eerr == nil && strings.TrimSpace(prompt) != "" { args = []string{prompt} input, rerr = readInput(stdin, args) } @@ -182,5 +182,5 @@ func normalizeRunner(r *Runner) Runner { } func loadConfigFromContext(ctx context.Context, logger *log.Logger) appconfig.App { - return appconfig.LoadWithOptions(logger, appconfig.LoadOptions{ConfigPath: configPathFromContext(ctx)}) + return appconfig.LoadWithOptions(ctx, logger, appconfig.LoadOptions{ConfigPath: configPathFromContext(ctx)}) } diff --git a/internal/hexaicli/runner_test.go b/internal/hexaicli/runner_test.go index 009af54..8b52b89 100644 --- a/internal/hexaicli/runner_test.go +++ b/internal/hexaicli/runner_test.go @@ -40,7 +40,7 @@ func TestRunner_UsesInjectedDependencies(t *testing.T) { PromptConfig: appconfig.PromptConfig{PromptCLIDefaultSystem: "SYS"}, } }, - openEditor: func([]byte) (string, error) { return "PROMPT", nil }, + openEditor: func(context.Context, []byte) (string, error) { return "PROMPT", nil }, newClient: func(appconfig.App) (client llm.Client, err error) { return &fakeClient{name: "fake", model: "m", resp: "OUT"}, nil }, @@ -67,7 +67,7 @@ func TestRunner_ConfigSubcommand_OpensConfigFromContext(t *testing.T) { t.Cleanup(func() { editor.RunEditor = old }) t.Setenv("EDITOR", "true") var gotPath string - editor.RunEditor = func(_, path string) error { + editor.RunEditor = func(_ context.Context, _, path string) error { gotPath = path return nil } @@ -89,7 +89,7 @@ func TestRunner_ConfigSubcommand_UsesXDGWhenNoOverride(t *testing.T) { xdg := t.TempDir() t.Setenv("XDG_CONFIG_HOME", xdg) var gotPath string - editor.RunEditor = func(_, path string) error { + editor.RunEditor = func(_ context.Context, _, path string) error { gotPath = path return nil } diff --git a/internal/hexailsp/dependencies.go b/internal/hexailsp/dependencies.go index 7e025d4..d664b4e 100644 --- a/internal/hexailsp/dependencies.go +++ b/internal/hexailsp/dependencies.go @@ -1,6 +1,7 @@ package hexailsp import ( + "context" "log" "codeberg.org/snonux/hexai/internal/appconfig" @@ -12,7 +13,7 @@ import ( "codeberg.org/snonux/hexai/internal/runtimeconfig" ) -type configLoader func(*log.Logger, appconfig.LoadOptions) appconfig.App +type configLoader func(context.Context, *log.Logger, appconfig.LoadOptions) appconfig.App type clientBuilder func(appconfig.App, llm.Client) llm.Client diff --git a/internal/hexailsp/run.go b/internal/hexailsp/run.go index 25d1929..8b3e840 100644 --- a/internal/hexailsp/run.go +++ b/internal/hexailsp/run.go @@ -3,6 +3,7 @@ package hexailsp import ( + "context" "fmt" "io" "log" @@ -20,7 +21,11 @@ import ( ) // ServerRunner is the minimal interface satisfied by lsp.Server. -type ServerRunner interface{ Run() error } +// Run takes a context so the serve loop is cancelled when the process is +// shutting down (e.g. on SIGINT/SIGTERM from the top-level caller). +type ServerRunner interface { + Run(ctx context.Context) error +} // ConfigurableServerRunner supports runtime option updates. type ConfigurableServerRunner interface { @@ -48,19 +53,21 @@ type ServerFactory func(r io.Reader, w io.Writer, logger *log.Logger, opts lsp.S // Run configures logging, loads config, builds the LLM client and runs the LSP server. // It is thin and delegates to RunWithFactory for testability. -func Run(logPath string, stdin io.Reader, stdout io.Writer, stderr io.Writer) error { - return RunWithConfig(logPath, "", stdin, stdout, stderr) +func Run(ctx context.Context, logPath string, stdin io.Reader, stdout io.Writer, stderr io.Writer) error { + return RunWithConfig(ctx, logPath, "", stdin, stdout, stderr) } // RunWithConfig is like Run but accepts an explicit config file path. -func RunWithConfig(logPath string, configPath string, stdin io.Reader, stdout io.Writer, stderr io.Writer) error { +// ctx is threaded through config loading and the LSP serve loop so the whole +// run is cancellable from the process entry point. +func RunWithConfig(ctx context.Context, logPath string, configPath string, stdin io.Reader, stdout io.Writer, stderr io.Writer) error { if err := llm.RegisterAllProviders(); err != nil { return fmt.Errorf("failed to register LLM providers: %w", err) } - return runWithConfigDependencies(logPath, configPath, stdin, stdout, stderr, defaultRunDependencies()) + return runWithConfigDependencies(ctx, logPath, configPath, stdin, stdout, stderr, defaultRunDependencies()) } -func runWithConfigDependencies(logPath string, configPath string, stdin io.Reader, stdout io.Writer, stderr io.Writer, deps runDependencies) error { +func runWithConfigDependencies(ctx context.Context, logPath string, configPath string, stdin io.Reader, stdout io.Writer, stderr io.Writer, deps runDependencies) error { deps = normalizeRunDependencies(deps) logger := log.New(stderr, "hexai-lsp-server ", log.LstdFlags|log.Lmsgprefix) if strings.TrimSpace(logPath) != "" { @@ -77,23 +84,23 @@ func runWithConfigDependencies(logPath string, configPath string, stdin io.Reade } logging.Bind(logger) loadOpts := appconfig.LoadOptions{ConfigPath: configPath} - cfg := deps.loadConfig(logger, loadOpts) + cfg := deps.loadConfig(ctx, logger, loadOpts) if err := cfg.Validate(); err != nil { return fmt.Errorf("invalid config: %w", err) } if cfg.StatsWindowMinutes > 0 { stats.SetWindow(time.Duration(cfg.StatsWindowMinutes) * time.Minute) } - return runWithDependencies(logPath, configPath, stdin, stdout, logger, cfg, nil, nil, deps) + return runWithDependencies(ctx, logPath, configPath, stdin, stdout, logger, cfg, nil, nil, deps) } // RunWithFactory is the testable entrypoint. When client is nil, it is built from cfg+env. // When factory is nil, lsp.NewServer is used. -func RunWithFactory(logPath string, configPath string, stdin io.Reader, stdout io.Writer, logger *log.Logger, cfg appconfig.App, client llm.Client, factory ServerFactory) error { - return runWithDependencies(logPath, configPath, stdin, stdout, logger, cfg, client, factory, defaultRunDependencies()) +func RunWithFactory(ctx context.Context, logPath string, configPath string, stdin io.Reader, stdout io.Writer, logger *log.Logger, cfg appconfig.App, client llm.Client, factory ServerFactory) error { + return runWithDependencies(ctx, logPath, configPath, stdin, stdout, logger, cfg, client, factory, defaultRunDependencies()) } -func runWithDependencies(logPath string, configPath string, stdin io.Reader, stdout io.Writer, logger *log.Logger, cfg appconfig.App, client llm.Client, factory ServerFactory, deps runDependencies) error { +func runWithDependencies(ctx context.Context, logPath string, configPath string, stdin io.Reader, stdout io.Writer, logger *log.Logger, cfg appconfig.App, client llm.Client, factory ServerFactory, deps runDependencies) error { deps = normalizeRunDependencies(deps) normalizeLoggingConfig(&cfg) if err := cfg.Validate(); err != nil { @@ -128,7 +135,7 @@ func runWithDependencies(logPath string, configPath string, stdin io.Reader, std configurable.ApplyOptions(opts) }) } - if err := server.Run(); err != nil { + if err := server.Run(ctx); err != nil { return fmt.Errorf("server error: %w", err) } return nil diff --git a/internal/hexailsp/run_more_test.go b/internal/hexailsp/run_more_test.go index d0f17b5..b0b99fd 100644 --- a/internal/hexailsp/run_more_test.go +++ b/internal/hexailsp/run_more_test.go @@ -15,11 +15,11 @@ import ( type recRunner struct{ ran bool } -func (r *recRunner) Run() error { r.ran = true; return nil } +func (r *recRunner) Run(context.Context) error { r.ran = true; return nil } type applyRunner struct{ opts []lsp.ServerOptions } -func (r *applyRunner) Run() error { return nil } +func (r *applyRunner) Run(context.Context) error { return nil } func (r *applyRunner) ApplyOptions(opts lsp.ServerOptions) { r.opts = append(r.opts, opts) } type stubClient struct{} @@ -43,13 +43,13 @@ func TestRunWithFactory_BuildsOptionsAndClient(t *testing.T) { } var in, out bytes.Buffer logger := log.New(&out, "", 0) - cfg := appconfig.Load(logger) + cfg := appconfig.Load(context.Background(), logger) // Use ollama to avoid API keys cfg.Provider = "ollama" cfg.MaxTokens = 123 cfg.PromptCodeActionRewriteSystem = "RSYS" cfg.PromptCodeActionRewriteUser = "RUSER" - if err := RunWithFactory("", "", &in, &out, logger, cfg, nil, factory); err != nil { + if err := RunWithFactory(context.Background(), "", "", &in, &out, logger, cfg, nil, factory); err != nil { t.Fatalf("RunWithFactory error: %v", err) } if captured.Config == nil { @@ -76,10 +76,10 @@ func TestRunWithFactory_SubscriptionAppliesUpdates(t *testing.T) { runner.opts = append(runner.opts, opts) return runner } - cfg := appconfig.Load(nil) + cfg := appconfig.Load(context.Background(), nil) cfg.StatsWindowMinutes = 0 cfg.ContextMode = " WINDOW " - if err := RunWithFactory("", "", &in, &out, logger, cfg, stubClient{}, factory); err != nil { + if err := RunWithFactory(context.Background(), "", "", &in, &out, logger, cfg, stubClient{}, factory); err != nil { t.Fatalf("RunWithFactory error: %v", err) } if capturedStore == nil { @@ -115,8 +115,8 @@ func TestRunWithDependencies_UsesInjectedClientBuilderAndStatusSink(t *testing.T captured = opts return &recRunner{} } - cfg := appconfig.Load(nil) - if err := runWithDependencies("", "", bytes.NewBuffer(nil), bytes.NewBuffer(nil), log.New(io.Discard, "", 0), cfg, nil, factory, runDependencies{ + cfg := appconfig.Load(context.Background(), nil) + if err := runWithDependencies(context.Background(), "", "", bytes.NewBuffer(nil), bytes.NewBuffer(nil), log.New(io.Discard, "", 0), cfg, nil, factory, runDependencies{ buildClient: func(appconfig.App, llm.Client) llm.Client { buildCalls++ return stubClient{} diff --git a/internal/hexailsp/run_test.go b/internal/hexailsp/run_test.go index b061f17..fa78436 100644 --- a/internal/hexailsp/run_test.go +++ b/internal/hexailsp/run_test.go @@ -3,6 +3,7 @@ package hexailsp import ( "bytes" + "context" "io" "log" "os" @@ -30,7 +31,7 @@ type fakeServer struct { opts lsp.ServerOptions } -func (f *fakeServer) Run() error { f.ran = true; return nil } +func (f *fakeServer) Run(context.Context) error { f.ran = true; return nil } func TestRunWithFactory_UsesDefaultsAndCallsServer(t *testing.T) { old := os.Getenv("OPENAI_API_KEY") @@ -39,7 +40,7 @@ func TestRunWithFactory_UsesDefaultsAndCallsServer(t *testing.T) { var stderr bytes.Buffer logger := log.New(&stderr, "hexai-lsp-server ", 0) - cfg := appconfig.Load(nil) // defaults + cfg := appconfig.Load(context.Background(), nil) // defaults // Pin provider to openai: the in-code default is now ollama, which would // happily build a client without a key and short-circuit the missing-key // assertion below. Load(nil) returns raw defaults and ignores env vars, @@ -50,7 +51,7 @@ func TestRunWithFactory_UsesDefaultsAndCallsServer(t *testing.T) { gotOpts = opts return &fakeServer{opts: opts} } - if err := RunWithFactory("", "", bytes.NewBuffer(nil), bytes.NewBuffer(nil), logger, cfg, nil, factory); err != nil { + if err := RunWithFactory(context.Background(), "", "", bytes.NewBuffer(nil), bytes.NewBuffer(nil), logger, cfg, nil, factory); err != nil { t.Fatalf("RunWithFactory error: %v", err) } if gotOpts.Config == nil { @@ -82,7 +83,7 @@ func TestRunWithFactory_BuildsClientWhenKeysPresent(t *testing.T) { var stderr bytes.Buffer logger := log.New(&stderr, "hexai-lsp-server ", 0) - cfg := appconfig.Load(nil) // defaults + cfg := appconfig.Load(context.Background(), nil) // defaults // Pin provider to openai (the in-code default is now ollama). Load(nil) // returns raw defaults and ignores env vars, so set this on the struct. cfg.Provider = "openai" @@ -91,7 +92,7 @@ func TestRunWithFactory_BuildsClientWhenKeysPresent(t *testing.T) { got = opts.Client return &fakeServer{opts: opts} } - if err := RunWithFactory("", "", bytes.NewBuffer(nil), bytes.NewBuffer(nil), logger, cfg, nil, factory); err != nil { + if err := RunWithFactory(context.Background(), "", "", bytes.NewBuffer(nil), bytes.NewBuffer(nil), logger, cfg, nil, factory); err != nil { t.Fatalf("RunWithFactory error: %v", err) } if got == nil { @@ -103,7 +104,7 @@ func TestRun_RespectsLogPathFlag(t *testing.T) { tmp := t.TempDir() logFile := filepath.Join(tmp, "hexai-lsp-server.log") // Run with real Run but nil env key so client disabled; ensure no panic and file created - if err := Run(logFile, bytes.NewBuffer(nil), bytes.NewBuffer(nil), bytes.NewBuffer(nil)); err != nil { + if err := Run(context.Background(), logFile, bytes.NewBuffer(nil), bytes.NewBuffer(nil), bytes.NewBuffer(nil)); err != nil { t.Fatalf("Run error: %v", err) } if _, err := os.Stat(logFile); err != nil { @@ -126,7 +127,7 @@ func TestRunWithFactory_NormalizesContextMode_AndSetsPreviewLimit(t *testing.T) gotOpts = opts return &fakeServer{opts: opts} } - if err := RunWithFactory("", "", bytes.NewBuffer(nil), bytes.NewBuffer(nil), logger, cfg, nil, factory); err != nil { + if err := RunWithFactory(context.Background(), "", "", bytes.NewBuffer(nil), bytes.NewBuffer(nil), logger, cfg, nil, factory); err != nil { t.Fatalf("RunWithFactory error: %v", err) } if gotOpts.Config == nil { @@ -155,13 +156,13 @@ func TestRunWithFactory_LogContextFlag(t *testing.T) { } return &fakeServer{opts: opts} } - if err := RunWithFactory("/tmp/some.log", "", bytes.NewBuffer(nil), bytes.NewBuffer(nil), logger, cfg, nil, factory); err != nil { + if err := RunWithFactory(context.Background(), "/tmp/some.log", "", bytes.NewBuffer(nil), bytes.NewBuffer(nil), logger, cfg, nil, factory); err != nil { t.Fatalf("RunWithFactory error: %v", err) } if !got1.LogContext { t.Fatalf("expected LogContext true when logPath is non-empty") } - if err := RunWithFactory("", "", bytes.NewBuffer(nil), bytes.NewBuffer(nil), logger, cfg, nil, factory); err != nil { + if err := RunWithFactory(context.Background(), "", "", bytes.NewBuffer(nil), bytes.NewBuffer(nil), logger, cfg, nil, factory); err != nil { t.Fatalf("RunWithFactory error: %v", err) } if got2.LogContext { diff --git a/internal/hexaimcp/run.go b/internal/hexaimcp/run.go index 74eb476..7c487c3 100644 --- a/internal/hexaimcp