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) } }