fix(lsp): fallback to Kill() on timeout (#2349)

Tai Groot created 1 week ago

Change summary

internal/lsp/client.go  | 31 +++++++++++++++++++++++++------
internal/lsp/manager.go |  4 ++--
2 files changed, 27 insertions(+), 8 deletions(-)

Detailed changes

internal/lsp/client.go

@@ -125,19 +125,38 @@ func (c *Client) Initialize(ctx context.Context, workspaceDir string) (*protocol
 	return result, nil
 }
 
+// closeTimeout is the maximum time to wait for a graceful LSP shutdown.
+const closeTimeout = 5 * time.Second
+
 // Kill kills the client without doing anything else.
 func (c *Client) Kill() { c.client.Kill() }
 
-// Close closes all open files in the client, then the client.
+// Close closes all open files in the client, then shuts down gracefully.
+// If shutdown takes longer than closeTimeout, it falls back to Kill().
 func (c *Client) Close(ctx context.Context) error {
 	c.CloseAllFiles(ctx)
 
-	// Shutdown and exit the client
-	if err := c.client.Shutdown(ctx); err != nil {
-		slog.Warn("Failed to shutdown LSP client", "error", err)
-	}
+	// Use a timeout to prevent hanging on unresponsive LSP servers.
+	// jsonrpc2's send lock doesn't respect context cancellation, so we
+	// need to fall back to Kill() which closes the underlying connection.
+	closeCtx, cancel := context.WithTimeout(ctx, closeTimeout)
+	defer cancel()
 
-	return c.client.Exit()
+	done := make(chan error, 1)
+	go func() {
+		if err := c.client.Shutdown(closeCtx); err != nil {
+			slog.Warn("Failed to shutdown LSP client", "error", err)
+		}
+		done <- c.client.Exit()
+	}()
+
+	select {
+	case err := <-done:
+		return err
+	case <-closeCtx.Done():
+		c.client.Kill()
+		return closeCtx.Err()
+	}
 }
 
 // createPowernapClient creates a new powernap client with the current configuration.

internal/lsp/manager.go

@@ -205,7 +205,7 @@ func (s *Manager) startServer(ctx context.Context, name, filepath string, server
 	if existing, ok := s.clients.Get(name); ok {
 		switch existing.GetServerState() {
 		case StateReady, StateStarting, StateDisabled:
-			client.Close(ctx)
+			_ = client.Close(ctx)
 			s.callback(name, existing)
 			return
 		}
@@ -228,7 +228,7 @@ func (s *Manager) startServer(ctx context.Context, name, filepath string, server
 
 	if _, err := client.Initialize(initCtx, s.cfg.WorkingDir()); err != nil {
 		slog.Error("LSP client initialization failed", "name", name, "error", err)
-		client.Close(ctx)
+		_ = client.Close(ctx)
 		s.clients.Del(name)
 		return
 	}