client.go

  1package lsp
  2
  3import (
  4	"context"
  5	"encoding/json"
  6	"fmt"
  7	"log/slog"
  8	"maps"
  9	"os"
 10	"path/filepath"
 11	"strings"
 12	"sync"
 13	"sync/atomic"
 14	"time"
 15
 16	"github.com/charmbracelet/crush/internal/config"
 17	"github.com/charmbracelet/crush/internal/csync"
 18	"github.com/charmbracelet/crush/internal/home"
 19	powernap "github.com/charmbracelet/x/powernap/pkg/lsp"
 20	"github.com/charmbracelet/x/powernap/pkg/lsp/protocol"
 21	"github.com/charmbracelet/x/powernap/pkg/transport"
 22)
 23
 24// DiagnosticCounts holds the count of diagnostics by severity.
 25type DiagnosticCounts struct {
 26	Error       int
 27	Warning     int
 28	Information int
 29	Hint        int
 30}
 31
 32type Client struct {
 33	client *powernap.Client
 34	name   string
 35	debug  bool
 36
 37	// Working directory this LSP is scoped to.
 38	workDir string
 39
 40	// File types this LSP server handles (e.g., .go, .rs, .py)
 41	fileTypes []string
 42
 43	// Configuration for this LSP client
 44	config config.LSPConfig
 45
 46	// Original context and resolver for recreating the client
 47	ctx      context.Context
 48	resolver config.VariableResolver
 49
 50	// Diagnostic change callback
 51	onDiagnosticsChanged func(name string, count int)
 52
 53	// Diagnostic cache
 54	diagnostics *csync.VersionedMap[protocol.DocumentURI, []protocol.Diagnostic]
 55
 56	// Cached diagnostic counts to avoid map copy on every UI render.
 57	diagCountsCache   DiagnosticCounts
 58	diagCountsVersion uint64
 59	diagCountsMu      sync.Mutex
 60
 61	// Files are currently opened by the LSP
 62	openFiles *csync.Map[string, *OpenFileInfo]
 63
 64	// Server state
 65	serverState atomic.Value
 66}
 67
 68// New creates a new LSP client using the powernap implementation.
 69func New(ctx context.Context, name string, cfg config.LSPConfig, resolver config.VariableResolver, debug bool) (*Client, error) {
 70	client := &Client{
 71		name:        name,
 72		fileTypes:   cfg.FileTypes,
 73		diagnostics: csync.NewVersionedMap[protocol.DocumentURI, []protocol.Diagnostic](),
 74		openFiles:   csync.NewMap[string, *OpenFileInfo](),
 75		config:      cfg,
 76		ctx:         ctx,
 77		debug:       debug,
 78		resolver:    resolver,
 79	}
 80	client.serverState.Store(StateStarting)
 81
 82	if err := client.createPowernapClient(); err != nil {
 83		return nil, err
 84	}
 85
 86	return client, nil
 87}
 88
 89// Initialize initializes the LSP client and returns the server capabilities.
 90func (c *Client) Initialize(ctx context.Context, workspaceDir string) (*protocol.InitializeResult, error) {
 91	if err := c.client.Initialize(ctx, false); err != nil {
 92		return nil, fmt.Errorf("failed to initialize the lsp client: %w", err)
 93	}
 94
 95	// Convert powernap capabilities to protocol capabilities
 96	caps := c.client.GetCapabilities()
 97	protocolCaps := protocol.ServerCapabilities{
 98		TextDocumentSync: caps.TextDocumentSync,
 99		CompletionProvider: func() *protocol.CompletionOptions {
100			if caps.CompletionProvider != nil {
101				return &protocol.CompletionOptions{
102					TriggerCharacters:   caps.CompletionProvider.TriggerCharacters,
103					AllCommitCharacters: caps.CompletionProvider.AllCommitCharacters,
104					ResolveProvider:     caps.CompletionProvider.ResolveProvider,
105				}
106			}
107			return nil
108		}(),
109	}
110
111	result := &protocol.InitializeResult{
112		Capabilities: protocolCaps,
113	}
114
115	c.registerHandlers()
116
117	return result, nil
118}
119
120// Close closes the LSP client.
121func (c *Client) Close(ctx context.Context) error {
122	c.CloseAllFiles(ctx)
123
124	// Shutdown and exit the client
125	if err := c.client.Shutdown(ctx); err != nil {
126		slog.Warn("Failed to shutdown LSP client", "error", err)
127	}
128
129	return c.client.Exit()
130}
131
132// createPowernapClient creates a new powernap client with the current configuration.
133func (c *Client) createPowernapClient() error {
134	workDir, err := os.Getwd()
135	if err != nil {
136		return fmt.Errorf("failed to get working directory: %w", err)
137	}
138
139	rootURI := string(protocol.URIFromPath(workDir))
140	c.workDir = workDir
141
142	command, err := c.resolver.ResolveValue(c.config.Command)
143	if err != nil {
144		return fmt.Errorf("invalid lsp command: %w", err)
145	}
146
147	clientConfig := powernap.ClientConfig{
148		Command:     home.Long(command),
149		Args:        c.config.Args,
150		RootURI:     rootURI,
151		Environment: maps.Clone(c.config.Env),
152		Settings:    c.config.Options,
153		InitOptions: c.config.InitOptions,
154		WorkspaceFolders: []protocol.WorkspaceFolder{
155			{
156				URI:  rootURI,
157				Name: filepath.Base(workDir),
158			},
159		},
160	}
161
162	powernapClient, err := powernap.NewClient(clientConfig)
163	if err != nil {
164		return fmt.Errorf("failed to create lsp client: %w", err)
165	}
166
167	c.client = powernapClient
168	return nil
169}
170
171// registerHandlers registers the standard LSP notification and request handlers.
172func (c *Client) registerHandlers() {
173	c.RegisterServerRequestHandler("workspace/applyEdit", HandleApplyEdit)
174	c.RegisterServerRequestHandler("workspace/configuration", HandleWorkspaceConfiguration)
175	c.RegisterServerRequestHandler("client/registerCapability", HandleRegisterCapability)
176	c.RegisterNotificationHandler("window/showMessage", func(ctx context.Context, method string, params json.RawMessage) {
177		if c.debug {
178			HandleServerMessage(ctx, method, params)
179		}
180	})
181	c.RegisterNotificationHandler("textDocument/publishDiagnostics", func(_ context.Context, _ string, params json.RawMessage) {
182		HandleDiagnostics(c, params)
183	})
184}
185
186// Restart closes the current LSP client and creates a new one with the same configuration.
187func (c *Client) Restart() error {
188	var openFiles []string
189	for uri := range c.openFiles.Seq2() {
190		openFiles = append(openFiles, string(uri))
191	}
192
193	closeCtx, cancel := context.WithTimeout(c.ctx, 10*time.Second)
194	defer cancel()
195
196	if err := c.Close(closeCtx); err != nil {
197		slog.Warn("Error closing client during restart", "name", c.name, "error", err)
198	}
199
200	c.SetServerState(StateStopped)
201
202	c.diagCountsCache = DiagnosticCounts{}
203	c.diagCountsVersion = 0
204
205	if err := c.createPowernapClient(); err != nil {
206		return err
207	}
208
209	initCtx, cancel := context.WithTimeout(c.ctx, 30*time.Second)
210	defer cancel()
211
212	c.SetServerState(StateStarting)
213
214	if err := c.client.Initialize(initCtx, false); err != nil {
215		c.SetServerState(StateError)
216		return fmt.Errorf("failed to initialize lsp client: %w", err)
217	}
218
219	c.registerHandlers()
220
221	if err := c.WaitForServerReady(initCtx); err != nil {
222		slog.Error("Server failed to become ready after restart", "name", c.name, "error", err)
223		c.SetServerState(StateError)
224		return err
225	}
226
227	for _, uri := range openFiles {
228		if err := c.OpenFile(initCtx, uri); err != nil {
229			slog.Warn("Failed to reopen file after restart", "file", uri, "error", err)
230		}
231	}
232	return nil
233}
234
235// ServerState represents the state of an LSP server
236type ServerState int
237
238const (
239	StateStopped ServerState = iota
240	StateStarting
241	StateReady
242	StateError
243	StateDisabled
244)
245
246// GetServerState returns the current state of the LSP server
247func (c *Client) GetServerState() ServerState {
248	if val := c.serverState.Load(); val != nil {
249		return val.(ServerState)
250	}
251	return StateStarting
252}
253
254// SetServerState sets the current state of the LSP server
255func (c *Client) SetServerState(state ServerState) {
256	c.serverState.Store(state)
257}
258
259// GetName returns the name of the LSP client
260func (c *Client) GetName() string {
261	return c.name
262}
263
264// SetDiagnosticsCallback sets the callback function for diagnostic changes
265func (c *Client) SetDiagnosticsCallback(callback func(name string, count int)) {
266	c.onDiagnosticsChanged = callback
267}
268
269// WaitForServerReady waits for the server to be ready
270func (c *Client) WaitForServerReady(ctx context.Context) error {
271	// Set initial state
272	c.SetServerState(StateStarting)
273
274	// Create a context with timeout
275	ctx, cancel := context.WithTimeout(ctx, 30*time.Second)
276	defer cancel()
277
278	// Try to ping the server with a simple request
279	ticker := time.NewTicker(500 * time.Millisecond)
280	defer ticker.Stop()
281
282	if c.debug {
283		slog.Debug("Waiting for LSP server to be ready...")
284	}
285
286	c.openKeyConfigFiles(ctx)
287
288	for {
289		select {
290		case <-ctx.Done():
291			c.SetServerState(StateError)
292			return fmt.Errorf("timeout waiting for LSP server to be ready")
293		case <-ticker.C:
294			// Check if client is running
295			if !c.client.IsRunning() {
296				if c.debug {
297					slog.Debug("LSP server not ready yet", "server", c.name)
298				}
299				continue
300			}
301
302			// Server is ready
303			c.SetServerState(StateReady)
304			if c.debug {
305				slog.Debug("LSP server is ready")
306			}
307			return nil
308		}
309	}
310}
311
312// OpenFileInfo contains information about an open file
313type OpenFileInfo struct {
314	Version int32
315	URI     protocol.DocumentURI
316}
317
318// HandlesFile checks if this LSP client handles the given file based on its
319// extension and whether it's within the working directory.
320func (c *Client) HandlesFile(path string) bool {
321	// Check if file is within working directory.
322	absPath, err := filepath.Abs(path)
323	if err != nil {
324		slog.Debug("Cannot resolve path", "name", c.name, "file", path, "error", err)
325		return false
326	}
327	relPath, err := filepath.Rel(c.workDir, absPath)
328	if err != nil || strings.HasPrefix(relPath, "..") {
329		slog.Debug("File outside workspace", "name", c.name, "file", path, "workDir", c.workDir)
330		return false
331	}
332
333	// If no file types are specified, handle all files (backward compatibility).
334	if len(c.fileTypes) == 0 {
335		return true
336	}
337
338	kind := powernap.DetectLanguage(path)
339	name := strings.ToLower(filepath.Base(path))
340	for _, filetype := range c.fileTypes {
341		suffix := strings.ToLower(filetype)
342		if !strings.HasPrefix(suffix, ".") {
343			suffix = "." + suffix
344		}
345		if strings.HasSuffix(name, suffix) || filetype == string(kind) {
346			slog.Debug("Handles file", "name", c.name, "file", name, "filetype", filetype, "kind", kind)
347			return true
348		}
349	}
350	slog.Debug("Doesn't handle file", "name", c.name, "file", name)
351	return false
352}
353
354// OpenFile opens a file in the LSP server.
355func (c *Client) OpenFile(ctx context.Context, filepath string) error {
356	if !c.HandlesFile(filepath) {
357		return nil
358	}
359
360	uri := string(protocol.URIFromPath(filepath))
361
362	if _, exists := c.openFiles.Get(uri); exists {
363		return nil // Already open
364	}
365
366	// Skip files that do not exist or cannot be read
367	content, err := os.ReadFile(filepath)
368	if err != nil {
369		return fmt.Errorf("error reading file: %w", err)
370	}
371
372	// Notify the server about the opened document
373	if err = c.client.NotifyDidOpenTextDocument(ctx, uri, string(powernap.DetectLanguage(filepath)), 1, string(content)); err != nil {
374		return err
375	}
376
377	c.openFiles.Set(uri, &OpenFileInfo{
378		Version: 1,
379		URI:     protocol.DocumentURI(uri),
380	})
381
382	return nil
383}
384
385// NotifyChange notifies the server about a file change.
386func (c *Client) NotifyChange(ctx context.Context, filepath string) error {
387	uri := string(protocol.URIFromPath(filepath))
388
389	content, err := os.ReadFile(filepath)
390	if err != nil {
391		return fmt.Errorf("error reading file: %w", err)
392	}
393
394	fileInfo, isOpen := c.openFiles.Get(uri)
395	if !isOpen {
396		return fmt.Errorf("cannot notify change for unopened file: %s", filepath)
397	}
398
399	// Increment version
400	fileInfo.Version++
401
402	// Create change event
403	changes := []protocol.TextDocumentContentChangeEvent{
404		{
405			Value: protocol.TextDocumentContentChangeWholeDocument{
406				Text: string(content),
407			},
408		},
409	}
410
411	return c.client.NotifyDidChangeTextDocument(ctx, uri, int(fileInfo.Version), changes)
412}
413
414// IsFileOpen checks if a file is currently open.
415func (c *Client) IsFileOpen(filepath string) bool {
416	uri := string(protocol.URIFromPath(filepath))
417	_, exists := c.openFiles.Get(uri)
418	return exists
419}
420
421// CloseAllFiles closes all currently open files.
422func (c *Client) CloseAllFiles(ctx context.Context) {
423	for uri := range c.openFiles.Seq2() {
424		if c.debug {
425			slog.Debug("Closing file", "file", uri)
426		}
427		if err := c.client.NotifyDidCloseTextDocument(ctx, uri); err != nil {
428			slog.Warn("Error closing file", "uri", uri, "error", err)
429			continue
430		}
431		c.openFiles.Del(uri)
432	}
433}
434
435// GetFileDiagnostics returns diagnostics for a specific file.
436func (c *Client) GetFileDiagnostics(uri protocol.DocumentURI) []protocol.Diagnostic {
437	diags, _ := c.diagnostics.Get(uri)
438	return diags
439}
440
441// GetDiagnostics returns all diagnostics for all files.
442func (c *Client) GetDiagnostics() map[protocol.DocumentURI][]protocol.Diagnostic {
443	return c.diagnostics.Copy()
444}
445
446// GetDiagnosticCounts returns cached diagnostic counts by severity.
447// Uses the VersionedMap version to avoid recomputing on every call.
448func (c *Client) GetDiagnosticCounts() DiagnosticCounts {
449	currentVersion := c.diagnostics.Version()
450
451	c.diagCountsMu.Lock()
452	defer c.diagCountsMu.Unlock()
453
454	if currentVersion == c.diagCountsVersion {
455		return c.diagCountsCache
456	}
457
458	// Recompute counts.
459	counts := DiagnosticCounts{}
460	for _, diags := range c.diagnostics.Seq2() {
461		for _, diag := range diags {
462			switch diag.Severity {
463			case protocol.SeverityError:
464				counts.Error++
465			case protocol.SeverityWarning:
466				counts.Warning++
467			case protocol.SeverityInformation:
468				counts.Information++
469			case protocol.SeverityHint:
470				counts.Hint++
471			}
472		}
473	}
474
475	c.diagCountsCache = counts
476	c.diagCountsVersion = currentVersion
477	return counts
478}
479
480// OpenFileOnDemand opens a file only if it's not already open.
481func (c *Client) OpenFileOnDemand(ctx context.Context, filepath string) error {
482	// Check if the file is already open
483	if c.IsFileOpen(filepath) {
484		return nil
485	}
486
487	// Open the file
488	return c.OpenFile(ctx, filepath)
489}
490
491// GetDiagnosticsForFile ensures a file is open and returns its diagnostics.
492func (c *Client) GetDiagnosticsForFile(ctx context.Context, filepath string) ([]protocol.Diagnostic, error) {
493	documentURI := protocol.URIFromPath(filepath)
494
495	// Make sure the file is open
496	if !c.IsFileOpen(filepath) {
497		if err := c.OpenFile(ctx, filepath); err != nil {
498			return nil, fmt.Errorf("failed to open file for diagnostics: %w", err)
499		}
500
501		// Give the LSP server a moment to process the file
502		time.Sleep(100 * time.Millisecond)
503	}
504
505	// Get diagnostics
506	diagnostics, _ := c.diagnostics.Get(documentURI)
507
508	return diagnostics, nil
509}
510
511// ClearDiagnosticsForURI removes diagnostics for a specific URI from the cache.
512func (c *Client) ClearDiagnosticsForURI(uri protocol.DocumentURI) {
513	c.diagnostics.Del(uri)
514}
515
516// RegisterNotificationHandler registers a notification handler.
517func (c *Client) RegisterNotificationHandler(method string, handler transport.NotificationHandler) {
518	c.client.RegisterNotificationHandler(method, handler)
519}
520
521// RegisterServerRequestHandler handles server requests.
522func (c *Client) RegisterServerRequestHandler(method string, handler transport.Handler) {
523	c.client.RegisterHandler(method, handler)
524}
525
526// DidChangeWatchedFiles sends a workspace/didChangeWatchedFiles notification to the server.
527func (c *Client) DidChangeWatchedFiles(ctx context.Context, params protocol.DidChangeWatchedFilesParams) error {
528	return c.client.NotifyDidChangeWatchedFiles(ctx, params.Changes)
529}
530
531// openKeyConfigFiles opens important configuration files that help initialize the server.
532func (c *Client) openKeyConfigFiles(ctx context.Context) {
533	wd, err := os.Getwd()
534	if err != nil {
535		return
536	}
537
538	// Try to open each file, ignoring errors if they don't exist
539	for _, file := range c.config.RootMarkers {
540		file = filepath.Join(wd, file)
541		if _, err := os.Stat(file); err == nil {
542			// File exists, try to open it
543			if err := c.OpenFile(ctx, file); err != nil {
544				slog.Error("Failed to open key config file", "file", file, "error", err)
545			} else {
546				slog.Debug("Opened key config file for initialization", "file", file)
547			}
548		}
549	}
550}
551
552// WaitForDiagnostics waits until diagnostics change or the timeout is reached.
553func (c *Client) WaitForDiagnostics(ctx context.Context, d time.Duration) {
554	ticker := time.NewTicker(200 * time.Millisecond)
555	defer ticker.Stop()
556	timeout := time.After(d)
557	pv := c.diagnostics.Version()
558	for {
559		select {
560		case <-ctx.Done():
561			return
562		case <-timeout:
563			return
564		case <-ticker.C:
565			if pv != c.diagnostics.Version() {
566				return
567			}
568		}
569	}
570}
571
572// FindReferences finds all references to the symbol at the given position.
573func (c *Client) FindReferences(ctx context.Context, filepath string, line, character int, includeDeclaration bool) ([]protocol.Location, error) {
574	if err := c.OpenFileOnDemand(ctx, filepath); err != nil {
575		return nil, err
576	}
577	// NOTE: line and character should be 0-based.
578	// See: https://microsoft.github.io/language-server-protocol/specifications/lsp/3.17/specification/#position
579	return c.client.FindReferences(ctx, filepath, line-1, character-1, includeDeclaration)
580}