From 7eb5f2f3369797bb07aaa6ea9c651782c01b41dd Mon Sep 17 00:00:00 2001 From: briiiqtt Date: Fri, 5 Jun 2026 16:53:42 +0900 Subject: [PATCH] Shut down the LSP child reliably when the server is orphaned When the host that launched this MCP server goes away, the spawned LSP child process must not be left running. This change makes teardown reliable, concurrency-safe, and bounded across the shutdown paths. main.go - Trigger shutdown on stdin EOF (the primary, cross-platform signal) and on a getppid() reparent-to-PID-1 backstop for the pty/held-pipe case where stdio stays open. The poll is conservative (only the PID-1 reparent case) to avoid killing the child when an intermediate process in a wrapper/subreaper chain exits mid-session; its known gap under a Linux PR_SET_CHILD_SUBREAPER ancestor is documented. - Run cleanup once (cleanupOnce) and bound it with a 5s backstop run in a goroutine so a wedged Close()/Cmd.Wait() can never block process exit, including the orphan case where ServeStdio never returns. - Normalize mcp-go's context.Canceled (its own signal handler) to nil so an intentional shutdown exits 0, while preserving genuine start() failures as exit 1. - Resolve the init-time orphan race with an lspClientMu + shuttingDown handshake so exactly one path owns closing the freshly spawned child. internal/lsp/client.go, transport.go - Make teardown concurrency-safe: closeOnce serializes the full teardown (CloseAllFiles -> shutdown/exit -> stdin close -> Cmd.Wait), with a bounded force-kill backstop and a single-close killDone. - Serialize all stdin writes under stdinMu; close stdin outside the lock so a writer parked on a full pipe is unblocked. - Signal a closed channel so an in-flight Call no longer parks forever. It is signalled both by Close() and by handleMessages on read-loop exit (LSP server death: stdout EOF or crash); a dedicated closedOnce guards the close so both paths are safe against a double-close panic. --- internal/lsp/client.go | 87 ++++++++++++++++++--- internal/lsp/transport.go | 57 ++++++++++++-- main.go | 160 ++++++++++++++++++++++++++++++++------ 3 files changed, 266 insertions(+), 38 deletions(-) diff --git a/internal/lsp/client.go b/internal/lsp/client.go index fc07059d..fb927907 100644 --- a/internal/lsp/client.go +++ b/internal/lsp/client.go @@ -22,6 +22,12 @@ type Client struct { stdout *bufio.Reader stderr io.ReadCloser + // stdinMu serializes writes to stdin. A single LSP message is framed as a + // Content-Length header followed by the body in two separate Write calls, so + // concurrent senders (request/response, notifications, the workspace watcher) + // must not interleave or the wire framing is corrupted. + stdinMu sync.Mutex + // Request ID counter nextID atomic.Int32 @@ -44,6 +50,32 @@ type Client struct { // Files are currently opened by the LSP openFiles map[string]*OpenFileInfo openFilesMu sync.RWMutex + + // closed is closed exactly once, the moment the transport starts going away — + // either because Close() began tearing it down, or because handleMessages + // stopped reading (LSP server died: stdout EOF or crash). An in-flight Call + // selects on it so a request whose response can no longer arrive does not + // block forever, whether the shutdown was initiated by us or forced on us by a + // dead server. It is a best-effort fast-fail for new Call/Notify, not a write + // barrier: stdinMu serializes the actual writes, and a write that loses the + // race with stdin.Close() simply returns an error. + // + // closedOnce guards the channel close so the Close() path and the + // handleMessages-exit path can both signal it without a double-close panic. + // It is distinct from closeOnce, which serializes the full teardown (stdin + // close + Cmd.Wait) and must run only from Close(), never from the read + // goroutine. + closeOnce sync.Once + closedOnce sync.Once + closed chan struct{} +} + +// signalClosed marks the transport as going away, unblocking any in-flight Call. +// Safe to call from multiple goroutines and more than once. +func (c *Client) signalClosed() { + c.closedOnce.Do(func() { + close(c.closed) + }) } func NewClient(command string, args ...string) (*Client, error) { @@ -76,6 +108,7 @@ func NewClient(command string, args ...string) (*Client, error) { serverRequestHandlers: make(map[string]ServerRequestHandler), diagnostics: make(map[protocol.DocumentUri][]protocol.Diagnostic), openFiles: make(map[string]*OpenFileInfo), + closed: make(chan struct{}), } // Start the LSP server process @@ -228,15 +261,38 @@ func (c *Client) InitializeLSPClient(ctx context.Context, workspaceDir string) ( } func (c *Client) Close() error { - // Try to close all open files first + var err error + // Guard against a double Close: closing c.closed (and stdin) twice would + // panic, and Close can race in from more than one shutdown trigger. + c.closeOnce.Do(func() { + err = c.closeLocked() + }) + return err +} + +func (c *Client) closeLocked() error { + // Try to close all open files first (graceful didClose notifications still + // need the transport, so do this before signalling c.closed below). ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) defer cancel() // Attempt to close files but continue shutdown regardless c.CloseAllFiles(ctx) - // Force kill the LSP process if it doesn't exit within timeout - forcedKill := make(chan struct{}) + // From here the transport is going away. Unblock any in-flight Call waiting + // on a response and make new Call/Notify fail fast instead of writing to a + // pipe we are about to close. handleMessages may have already signalled this + // if the server died first, so go through the once-guarded helper. + c.signalClosed() + + // Force kill the LSP process if it doesn't exit within timeout. + // killDone is closed exactly once, by this function, after Cmd.Wait + // returns. The kill goroutine only ever *reads* it (never closes it), so + // the two paths can never double-close the same channel: when the process + // exits on its own we close killDone and the goroutine takes its <-killDone + // arm; when it hangs the goroutine kills it and then observes killDone + // closing once Wait returns. + killDone := make(chan struct{}) go func() { select { case <-time.After(2 * time.Second): @@ -248,25 +304,38 @@ func (c *Client) Close() error { lspLogger.Info("Process killed successfully") } } - close(forcedKill) - case <-forcedKill: - // Channel closed from completion path - return + case <-killDone: + // Process exited on its own before the timeout; nothing to do. } }() - // Close stdin to signal the server + // Close stdin to signal the server. Deliberately NOT under stdinMu: a writer + // blocked on a full pipe holds that lock, and closing the fd here is what + // unblocks it (the in-flight write then returns an error). os.File guards + // concurrent Write/Close internally, so this is a fast-fail, not a race. if err := c.stdin.Close(); err != nil { lspLogger.Error("Failed to close stdin: %v", err) } // Wait for process to exit err := c.Cmd.Wait() - close(forcedKill) // Stop the force kill goroutine + close(killDone) // Stop the force kill goroutine return err } +// isClosed reports whether the transport is going away — either Close() began +// tearing it down or the read loop exited because the server died. It mirrors +// the broadened semantics of the closed channel (see field comment above). +func (c *Client) isClosed() bool { + select { + case <-c.closed: + return true + default: + return false + } +} + type ServerState int const ( diff --git a/internal/lsp/transport.go b/internal/lsp/transport.go index 3cc9107a..0ff3cb89 100644 --- a/internal/lsp/transport.go +++ b/internal/lsp/transport.go @@ -94,8 +94,22 @@ func ReadMessage(r *bufio.Reader) (*Message, error) { return &msg, nil } +// writeMessage serializes a single LSP message onto stdin. The lock prevents +// concurrent senders from interleaving one message's header with another's body. +func (c *Client) writeMessage(msg *Message) error { + c.stdinMu.Lock() + defer c.stdinMu.Unlock() + return WriteMessage(c.stdin, msg) +} + // handleMessages reads and dispatches messages in a loop func (c *Client) handleMessages() { + // Once this loop returns, no further responses can be delivered to waiting + // Calls. Signal c.closed so any in-flight Call unblocks instead of parking + // forever — covers the case where the LSP server dies on its own (stdout EOF + // or crash) and nobody calls Close(). Idempotent and safe to race with + // Close(), which signals the same channel through the same once. + defer c.signalClosed() for { msg, err := ReadMessage(c.stdout) if err != nil { @@ -150,7 +164,7 @@ func (c *Client) handleMessages() { } // Send response back to server - if err := WriteMessage(c.stdin, response); err != nil { + if err := c.writeMessage(response); err != nil { lspLogger.Error("Error sending response to server: %v", err) } @@ -193,6 +207,10 @@ func (c *Client) handleMessages() { // Call makes a request and waits for the response func (c *Client) Call(ctx context.Context, method string, params any, result any) error { + if c.isClosed() { + return fmt.Errorf("LSP client is closed") + } + id := c.nextID.Add(1) lspLogger.Debug("Making call: method=%s id=%v", method, id) @@ -217,14 +235,39 @@ func (c *Client) Call(ctx context.Context, method string, params any, result any }() // Send request - if err := WriteMessage(c.stdin, msg); err != nil { + if err := c.writeMessage(msg); err != nil { return fmt.Errorf("failed to send request: %w", err) } lspLogger.Debug("Waiting for response to request ID: %v", msg.ID) - // Wait for response - resp := <-ch + // Wait for the response, but don't block forever if no response can arrive. + // c.closed is signalled both when Close() tears the transport down and when + // handleMessages stops reading because the LSP server died on its own (stdout + // EOF or crash). Either way, once it is closed no response will ever arrive on + // ch, so we bail out instead of parking forever. + // + // NOTE: deliberately not selecting on ctx.Done() here. Several callers + // (e.g. the diagnostics/definition tools) sleep and then issue calls on a + // context whose deadline may already have elapsed, relying on the response + // still being delivered as long as the LSP is alive. Aborting on ctx here + // would change that long-standing contract. Shutdown bounding is handled by + // the caller-side timeouts in cleanup() plus c.closed, not by ctx. + // + // Prefer an already-delivered response: when both ch and c.closed are ready, + // a bare select picks pseudo-randomly and would discard a valid result. So we + // check ch first, and only fall back to the closed-aware wait if nothing has + // arrived yet. + var resp *Message + select { + case resp = <-ch: + default: + select { + case resp = <-ch: + case <-c.closed: + return fmt.Errorf("LSP client closed while awaiting response to %s", method) + } + } lspLogger.Debug("Received response for request ID: %v", msg.ID) @@ -251,6 +294,10 @@ func (c *Client) Call(ctx context.Context, method string, params any, result any // Notify sends a notification (a request without an ID that doesn't expect a response) func (c *Client) Notify(ctx context.Context, method string, params any) error { + if c.isClosed() { + return fmt.Errorf("LSP client is closed") + } + lspLogger.Debug("Sending notification: method=%s", method) msg, err := NewNotification(method, params) @@ -258,7 +305,7 @@ func (c *Client) Notify(ctx context.Context, method string, params any) error { return fmt.Errorf("failed to create notification: %w", err) } - if err := WriteMessage(c.stdin, msg); err != nil { + if err := c.writeMessage(msg); err != nil { return fmt.Errorf("failed to send notification: %w", err) } diff --git a/main.go b/main.go index f6f3ed5c..a9a53817 100644 --- a/main.go +++ b/main.go @@ -2,12 +2,14 @@ package main import ( "context" + "errors" "flag" "fmt" "os" "os/exec" "os/signal" "path/filepath" + "sync" "syscall" "time" @@ -27,8 +29,16 @@ type config struct { } type mcpServer struct { - config config - lspClient *lsp.Client + config config + // lspClientMu guards lspClient and shuttingDown: start() assigns lspClient + // from the server goroutine while a shutdown trigger may have cleanup() read + // it from another goroutine, so the two must not race on the bare pointer. + lspClientMu sync.RWMutex + lspClient *lsp.Client + // shuttingDown is set by cleanup() under lspClientMu so that initializeLSP, + // if it is mid-flight, can see the shutdown and tear down the LSP child it + // just spawned instead of leaking it as an orphan. + shuttingDown bool mcpServer *server.MCPServer ctx context.Context cancelFunc context.CancelFunc @@ -89,7 +99,22 @@ func (s *mcpServer) initializeLSP() error { if err != nil { return fmt.Errorf("failed to create LSP client: %v", err) } + // NewClient has already spawned the LSP child process. If a shutdown trigger + // fired during that spawn, cleanup() ran before s.lspClient was published and + // could not tear the child down — so close it here to avoid leaking an + // orphan. The lock makes this race-free: cleanup() either reads the assigned + // client (and closes it) or sets shuttingDown before we publish (and we close + // it). Exactly one path owns the teardown. + s.lspClientMu.Lock() + if s.shuttingDown { + s.lspClientMu.Unlock() + if closeErr := client.Close(); closeErr != nil { + coreLogger.Error("Failed to close LSP client spawned during shutdown: %v", closeErr) + } + return fmt.Errorf("shutdown initiated during LSP initialization") + } s.lspClient = client + s.lspClientMu.Unlock() s.workspaceWatcher = watcher.NewWorkspaceWatcher(client) initResult, err := client.InitializeLSPClient(s.ctx, s.config.workspaceDir) @@ -140,11 +165,55 @@ func main() { coreLogger.Fatal("%v", err) } + // Ensure teardown runs at most once, regardless of which trigger fires + // first (OS signal, parent death, or stdin EOF). Without this the racing + // triggers could run cleanup() — and thus the LSP shutdown/kill — twice. + var cleanupOnce sync.Once + shutdown := func() { + cleanupOnce.Do(func() { + cleanup(server, done) + }) + } + + // drainServerErr does a non-blocking read of any result start() has already + // produced, used on the signal/parent-death paths to preserve a genuine + // independent start() failure in the exit code. It normalizes context.Canceled + // to nil: mcp-go's ServeStdio installs its own SIGTERM/SIGINT handler that + // cancels its context, so on a signal start() returns context.Canceled — that + // is our own intentional shutdown, not a failure, and must not become exit 1. + drainServerErr := func(ch <-chan error) error { + select { + case e := <-ch: + if errors.Is(e, context.Canceled) { + return nil + } + return e + default: + return nil + } + } + // Parent process monitoring channel parentDeath := make(chan struct{}) - // Monitor parent process termination - // Claude desktop does not properly kill child processes for MCP servers + // Monitor parent process termination. + // The reliable death signal is the stdio EOF handled below: when the host + // exits, the kernel closes its end of our stdio pipe and server.start() + // returns. That is the primary trigger and works on every platform. This + // getppid() poll is only a best-effort backstop for the rare case where stdio + // stays open (e.g. redirected to a pty or a pipe another process holds) after + // the host is gone — macOS has no PR_SET_PDEATHSIG equivalent, so we cannot + // get a kernel-delivered death signal there. + // + // We only act when the reparent involves PID 1 (reparented to, or originally + // owned by, init): a bare "ppid changed" check would also fire when an + // intermediate process in a wrapper/subreaper spawn chain exits while the real + // host is still alive, killing the LSP child mid-session. The cost of that + // conservatism is a known gap: under a PR_SET_CHILD_SUBREAPER ancestor on + // Linux (systemd user scope, some multiplexers/containers) an orphaned child + // reparents to the subreaper's PID, not 1, so this poll never fires there and + // we fall back entirely on the stdio-EOF path. This poll is a backstop, not a + // guarantee. go func() { ppid := os.Getppid() coreLogger.Debug("Monitoring parent process: %d", ppid) @@ -167,39 +236,82 @@ func main() { } }() - // Handle shutdown triggers + // Run the blocking stdio server in the background so teardown can be driven + // from a single place, no matter which trigger fires first. ServeStdio only + // returns on a clean stdio EOF (host closed the pipe) or an error; it does + // NOT return on parent death while stdio stays open, so the main goroutine + // must not park on it — otherwise an orphaned process would never exit. + serverErrCh := make(chan error, 1) go func() { - select { - case sig := <-sigChan: - coreLogger.Info("Received signal %v in PID: %d", sig, os.Getpid()) - cleanup(server, done) - case <-parentDeath: - coreLogger.Info("Parent death detected, initiating shutdown") - cleanup(server, done) - } + serverErrCh <- server.start() }() - if err := server.start(); err != nil { - coreLogger.Error("Server error: %v", err) - cleanup(server, done) - os.Exit(1) + // Wait for the first teardown trigger: OS signal, parent death, or the + // stdio server returning. + var serverErr error + select { + case sig := <-sigChan: + coreLogger.Info("Received signal %v in PID: %d", sig, os.Getpid()) + // If start() had already failed on its own before the signal arrived, + // preserve that error for the exit code. We have not begun teardown yet, + // so anything already queued here is a genuine independent failure, not a + // side effect of us closing stdin during cleanup. + serverErr = drainServerErr(serverErrCh) + case <-parentDeath: + coreLogger.Info("Parent death detected, initiating shutdown") + serverErr = drainServerErr(serverErrCh) + case serverErr = <-serverErrCh: + if serverErr != nil { + coreLogger.Error("Server error: %v", serverErr) + } else { + coreLogger.Info("MCP stdio closed (EOF), initiating shutdown for PID: %d", os.Getpid()) + } } - <-done - coreLogger.Info("Server shutdown complete for PID: %d", os.Getpid()) + // Run cleanup concurrently and bound how long we wait for it. Running it in + // a goroutine is what lets the timeout below actually fire: cleanup()'s + // s.lspClient.Close() is unbounded, so a wedged LSP must not be able to + // block us inside cleanup(). It also lets us exit even while server.start() + // is still parked in ServeStdio (the orphan/pty case where stdin never EOFs). + go shutdown() + + select { + case <-done: + coreLogger.Info("Server shutdown complete for PID: %d", os.Getpid()) + case <-time.After(5 * time.Second): + coreLogger.Warn("Shutdown wait timed out, forcing exit for PID: %d", os.Getpid()) + } + if serverErr != nil { + os.Exit(1) + } os.Exit(0) } func cleanup(s *mcpServer, done chan struct{}) { coreLogger.Info("Cleanup initiated for PID: %d", os.Getpid()) + // Stop the workspace watcher first. It runs on s.ctx and otherwise keeps + // issuing notifications (writes to the LSP stdin) throughout teardown, racing + // the shutdown/exit/close below. Cancelling s.ctx makes its loop exit. + if s.cancelFunc != nil { + s.cancelFunc() + } + // Create a context with timeout for shutdown operations ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) defer cancel() - if s.lspClient != nil { + // Mark shutdown and snapshot the client under the lock. If start() is still + // mid-init on another goroutine, either we read the client it already + // published (and close it here) or it sees shuttingDown after we set it and + // closes the freshly spawned child itself — so the child is never leaked. + s.lspClientMu.Lock() + s.shuttingDown = true + lspClient := s.lspClient + s.lspClientMu.Unlock() + if lspClient != nil { coreLogger.Info("Closing open files") - s.lspClient.CloseAllFiles(ctx) + lspClient.CloseAllFiles(ctx) // Create a shorter timeout context for the shutdown request shutdownCtx, shutdownCancel := context.WithTimeout(ctx, 500*time.Millisecond) @@ -209,7 +321,7 @@ func cleanup(s *mcpServer, done chan struct{}) { shutdownDone := make(chan struct{}) go func() { coreLogger.Info("Sending shutdown request") - if err := s.lspClient.Shutdown(shutdownCtx); err != nil { + if err := lspClient.Shutdown(shutdownCtx); err != nil { coreLogger.Error("Shutdown request failed: %v", err) } close(shutdownDone) @@ -224,12 +336,12 @@ func cleanup(s *mcpServer, done chan struct{}) { } coreLogger.Info("Sending exit notification") - if err := s.lspClient.Exit(ctx); err != nil { + if err := lspClient.Exit(ctx); err != nil { coreLogger.Error("Exit notification failed: %v", err) } coreLogger.Info("Closing LSP client") - if err := s.lspClient.Close(); err != nil { + if err := lspClient.Close(); err != nil { coreLogger.Error("Failed to close LSP client: %v", err) } }