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