feat: add bidi communcation

Raphael Amorim created

Change summary

internal/lsp/client.go    | 119 ++++++++++++++++++++++++++++++++++------
internal/lsp/methods.go   |  86 +++++++++++++++++++++--------
internal/lsp/transport.go |  25 +++++++-
3 files changed, 184 insertions(+), 46 deletions(-)

Detailed changes

internal/lsp/client.go 🔗

@@ -276,35 +276,50 @@ func (c *Client) InitializeLSPClient(ctx context.Context, workspaceDir string) (
 }
 
 func (c *Client) Close() error {
-	// Try to close all open files first
+	// Try graceful shutdown first
 	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
 	defer cancel()
 
-	// Attempt to close files but continue shutdown regardless
+	// Attempt to close all files
 	c.CloseAllFiles(ctx)
 
+	// Do full shutdown following LSP spec
+	if err := c.Shutdown(ctx); err != nil {
+		slog.Warn("LSP shutdown failed during close", "name", c.name, "error", err)
+	}
+
 	// Close stdin to signal the server
-	if err := c.stdin.Close(); err != nil {
-		return fmt.Errorf("failed to close stdin: %w", err)
+	if c.stdin != nil {
+		if err := c.stdin.Close(); err != nil {
+			slog.Warn("Failed to close stdin", "name", c.name, "error", err)
+		}
 	}
 
-	// Use a channel to handle the Wait with timeout
-	done := make(chan error, 1)
-	go func() {
-		done <- c.Cmd.Wait()
-	}()
+	// Terminate the process if still running
+	if c.Cmd != nil && c.Cmd.Process != nil {
+		// Use a channel to handle the Wait with timeout
+		done := make(chan error, 1)
+		go func() {
+			done <- c.Cmd.Wait()
+		}()
 
-	// Wait for process to exit with timeout
-	select {
-	case err := <-done:
-		return err
-	case <-time.After(2 * time.Second):
-		// If we timeout, try to kill the process
-		if err := c.Cmd.Process.Kill(); err != nil {
-			return fmt.Errorf("failed to kill process: %w", err)
+		// Wait for process to exit with timeout
+		select {
+		case err := <-done:
+			if err != nil {
+				slog.Debug("LSP process exited with error", "name", c.name, "error", err)
+			}
+			return nil
+		case <-time.After(2 * time.Second):
+			// If we timeout, try to kill the process
+			if err := c.Cmd.Process.Kill(); err != nil {
+				return fmt.Errorf("failed to kill process: %w", err)
+			}
+			return fmt.Errorf("process killed after timeout")
 		}
-		return fmt.Errorf("process killed after timeout")
 	}
+
+	return nil
 }
 
 type ServerState int
@@ -876,3 +891,71 @@ func (c *Client) ClearDiagnosticsForURI(uri protocol.DocumentURI) {
 	defer c.diagnosticsMu.Unlock()
 	delete(c.diagnostics, uri)
 }
+
+// startHealthCheckAndCleanup runs periodic health checks and cleans up stale handlers
+func (c *Client) startHealthCheckAndCleanup(ctx context.Context) {
+	healthTicker := time.NewTicker(30 * time.Second)
+	cleanupTicker := time.NewTicker(60 * time.Second)
+	defer healthTicker.Stop()
+	defer cleanupTicker.Stop()
+
+	for {
+		select {
+		case <-ctx.Done():
+			return
+		case <-c.shutdownChan:
+			return
+		case <-healthTicker.C:
+			// Perform health check
+			if c.GetServerState() == StateReady {
+				// Try a simple ping to check if server is still responsive
+				pingCtx, cancel := context.WithTimeout(ctx, 5*time.Second)
+				err := c.pingServerByType(pingCtx, c.detectServerType())
+				cancel()
+				if err != nil {
+					slog.Warn("LSP server health check failed", "name", c.name, "error", err)
+					c.SetServerState(StateError)
+				}
+			}
+		case <-cleanupTicker.C:
+			// Clean up stale pending requests
+			c.cleanupStaleHandlers()
+		}
+	}
+}
+
+// cleanupStaleHandlers removes handlers for requests that have been pending too long
+func (c *Client) cleanupStaleHandlers() {
+	threshold := time.Now().Add(-5 * time.Minute)
+	var staleIDs []int32
+
+	// Find stale requests
+	c.pendingRequestsMu.RLock()
+	for id, timestamp := range c.pendingRequests {
+		if timestamp.Before(threshold) {
+			staleIDs = append(staleIDs, id)
+		}
+	}
+	c.pendingRequestsMu.RUnlock()
+
+	if len(staleIDs) == 0 {
+		return
+	}
+
+	// Clean up stale handlers
+	c.handlersMu.Lock()
+	c.pendingRequestsMu.Lock()
+	for _, id := range staleIDs {
+		if ch, exists := c.handlers[id]; exists {
+			close(ch)
+			delete(c.handlers, id)
+		}
+		delete(c.pendingRequests, id)
+	}
+	c.pendingRequestsMu.Unlock()
+	c.handlersMu.Unlock()
+
+	if len(staleIDs) > 0 {
+		slog.Debug("Cleaned up stale LSP handlers", "count", len(staleIDs), "name", c.name)
+	}
+}

internal/lsp/methods.go 🔗

@@ -6,6 +6,7 @@ import (
 	"log/slog"
 	"time"
 
+	"github.com/charmbracelet/crush/internal/config"
 	"github.com/charmbracelet/crush/internal/lsp/protocol"
 )
 
@@ -246,53 +247,90 @@ func (c *Client) Initialize(ctx context.Context, params protocol.ParamInitialize
 func (c *Client) Shutdown(ctx context.Context) error {
 	var shutdownErr error
 	c.shutdownOnce.Do(func() {
-		// Signal shutdown to goroutines
-		close(c.shutdownChan)
+		cfg := config.Get()
+
+		// Step 1: Send shutdown request (expects a response)
+		if cfg.Options.DebugLSP {
+			slog.Debug("Sending LSP shutdown request", "name", c.name)
+		}
 
-		// Send shutdown request to server
-		shutdownErr = c.Call(ctx, "shutdown", nil, nil)
+		shutdownCtx, cancel := context.WithTimeout(ctx, 2*time.Second)
+		defer cancel()
 
-		// Clean up handlers map to prevent memory leaks
-		c.handlersMu.Lock()
-		for id, ch := range c.handlers {
-			close(ch)
-			delete(c.handlers, id)
+		if err := c.Call(shutdownCtx, "shutdown", nil, nil); err != nil {
+			// Log but don't fail - server might already be shutting down
+			slog.Warn("LSP shutdown request failed", "name", c.name, "error", err)
+			shutdownErr = err
 		}
-		c.handlersMu.Unlock()
 
-		// Clean up open files map
-		c.openFilesMu.Lock()
-		for uri := range c.openFiles {
-			delete(c.openFiles, uri)
+		// Step 2: Send exit notification (no response expected)
+		if cfg.Options.DebugLSP {
+			slog.Debug("Sending LSP exit notification", "name", c.name)
 		}
-		c.openFilesMu.Unlock()
 
-		// Clean up diagnostics map
-		c.diagnosticsMu.Lock()
-		for uri := range c.diagnostics {
-			delete(c.diagnostics, uri)
+		if err := c.Notify(ctx, "exit", nil); err != nil {
+			slog.Warn("LSP exit notification failed", "name", c.name, "error", err)
 		}
-		c.diagnosticsMu.Unlock()
 
-		// Wait for goroutines to finish with timeout
+		// Step 3: Signal our goroutines to stop
+		close(c.shutdownChan)
+
+		// Step 4: Clean up resources
+		c.cleanupResources()
+
+		// Step 5: Wait for goroutines with timeout
 		done := make(chan struct{})
 		go func() {
 			<-c.stderrDone
 			<-c.messageHandlerDone
+			<-c.healthCheckDone
 			close(done)
 		}()
 
 		select {
 		case <-done:
-			// Goroutines finished cleanly
+			if cfg.Options.DebugLSP {
+				slog.Debug("LSP goroutines stopped cleanly", "name", c.name)
+			}
 		case <-time.After(2 * time.Second):
-			// Timeout waiting for goroutines
-			slog.Warn("Timeout waiting for LSP goroutines to finish", "name", c.name)
+			slog.Warn("Timeout waiting for LSP goroutines", "name", c.name)
 		}
 	})
 	return shutdownErr
 }
 
+// cleanupResources closes all handlers and cleans up internal state
+func (c *Client) cleanupResources() {
+	// Close all pending request handlers
+	c.handlersMu.Lock()
+	for id, ch := range c.handlers {
+		close(ch)
+		delete(c.handlers, id)
+	}
+	c.handlersMu.Unlock()
+
+	// Clear pending requests tracking
+	c.pendingRequestsMu.Lock()
+	for id := range c.pendingRequests {
+		delete(c.pendingRequests, id)
+	}
+	c.pendingRequestsMu.Unlock()
+
+	// Clear open files
+	c.openFilesMu.Lock()
+	for uri := range c.openFiles {
+		delete(c.openFiles, uri)
+	}
+	c.openFilesMu.Unlock()
+
+	// Clear diagnostics
+	c.diagnosticsMu.Lock()
+	for uri := range c.diagnostics {
+		delete(c.diagnostics, uri)
+	}
+	c.diagnosticsMu.Unlock()
+}
+
 // WillSaveWaitUntil sends a textDocument/willSaveWaitUntil request to the LSP server.
 // A document will save request is sent from the client to the server before the document is actually saved. The request can return an array of TextEdits which will be applied to the text document before it is saved. Please note that clients might drop results if computing the text edits took too long or if a server constantly fails on this request. This is done to keep the save fast and reliable.
 func (c *Client) WillSaveWaitUntil(ctx context.Context, params protocol.WillSaveTextDocumentParams) ([]protocol.TextEdit, error) {

internal/lsp/transport.go 🔗

@@ -8,6 +8,7 @@ import (
 	"io"
 	"log/slog"
 	"strings"
+	"time"
 
 	"github.com/charmbracelet/crush/internal/config"
 )
@@ -186,9 +187,17 @@ func (c *Client) handleMessages() {
 	}
 }
 
-// Call makes a request and waits for the response
+// Call makes a request and waits for the response with timeout
 func (c *Client) Call(ctx context.Context, method string, params any, result any) error {
 	cfg := config.Get()
+
+	// Add default timeout if context doesn't have one
+	if _, hasDeadline := ctx.Deadline(); !hasDeadline {
+		var cancel context.CancelFunc
+		ctx, cancel = context.WithTimeout(ctx, 30*time.Second)
+		defer cancel()
+	}
+
 	id := c.nextID.Add(1)
 
 	if cfg.Options.DebugLSP {
@@ -200,16 +209,24 @@ func (c *Client) Call(ctx context.Context, method string, params any, result any
 		return fmt.Errorf("failed to create request: %w", err)
 	}
 
-	// Create response channel
+	// Create response channel and track request
 	ch := make(chan *Message, 1)
 	c.handlersMu.Lock()
 	c.handlers[id] = ch
 	c.handlersMu.Unlock()
 
+	c.pendingRequestsMu.Lock()
+	c.pendingRequests[id] = time.Now()
+	c.pendingRequestsMu.Unlock()
+
 	defer func() {
 		c.handlersMu.Lock()
 		delete(c.handlers, id)
 		c.handlersMu.Unlock()
+
+		c.pendingRequestsMu.Lock()
+		delete(c.pendingRequests, id)
+		c.pendingRequestsMu.Unlock()
 	}()
 
 	// Send request
@@ -221,10 +238,10 @@ func (c *Client) Call(ctx context.Context, method string, params any, result any
 		slog.Debug("Request sent", "method", method, "id", id)
 	}
 
-	// Wait for response
+	// Wait for response with context timeout
 	select {
 	case <-ctx.Done():
-		return ctx.Err()
+		return fmt.Errorf("request timeout or cancelled for %s: %w", method, ctx.Err())
 	case resp := <-ch:
 		if cfg.Options.DebugLSP {
 			slog.Debug("Received response", "id", id)