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/hexailsp/dependencies.go | 3 ++- internal/hexailsp/run.go | 31 +++++++++++++++++++------------ internal/hexailsp/run_more_test.go | 16 ++++++++-------- internal/hexailsp/run_test.go | 19 ++++++++++--------- 4 files changed, 39 insertions(+), 30 deletions(-) (limited to 'internal/hexailsp') 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 { -- cgit v1.2.3