client.go

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