client.go

  1package lsp
  2
  3import (
  4	"context"
  5	"encoding/json"
  6	"fmt"
  7	"log/slog"
  8	"maps"
  9	"os"
 10	"path/filepath"
 11	"slices"
 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/fsext"
 19	"github.com/charmbracelet/crush/internal/home"
 20	powernap "github.com/charmbracelet/x/powernap/pkg/lsp"
 21	"github.com/charmbracelet/x/powernap/pkg/lsp/protocol"
 22	"github.com/charmbracelet/x/powernap/pkg/transport"
 23)
 24
 25// DiagnosticCounts holds the count of diagnostics by severity.
 26type DiagnosticCounts struct {
 27	Error       int
 28	Warning     int
 29	Information int
 30	Hint        int
 31}
 32
 33type Client struct {
 34	client *powernap.Client
 35	name   string
 36	debug  bool
 37
 38	// Working directory this LSP is scoped to.
 39	cwd string
 40
 41	// File types this LSP server handles (e.g., .go, .rs, .py)
 42	fileTypes []string
 43
 44	// Configuration for this LSP client
 45	config config.LSPConfig
 46
 47	// Original context and resolver for recreating the client
 48	ctx      context.Context
 49	resolver config.VariableResolver
 50
 51	// Diagnostic change callback
 52	onDiagnosticsChanged func(name string, count int)
 53
 54	// Diagnostic cache
 55	diagnostics *csync.VersionedMap[protocol.DocumentURI, []protocol.Diagnostic]
 56
 57	// Cached diagnostic counts to avoid map copy on every UI render.
 58	diagCountsCache   DiagnosticCounts
 59	diagCountsVersion uint64
 60	diagCountsMu      sync.Mutex
 61
 62	// Files are currently opened by the LSP
 63	openFiles *csync.Map[string, *OpenFileInfo]
 64
 65	// Server state
 66	serverState atomic.Value
 67}
 68
 69// New creates a new LSP client using the powernap implementation.
 70func New(
 71	ctx context.Context,
 72	name string,
 73	cfg config.LSPConfig,
 74	resolver config.VariableResolver,
 75	cwd string,
 76	debug bool,
 77) (*Client, error) {
 78	client := &Client{
 79		name:        name,
 80		fileTypes:   cfg.FileTypes,
 81		diagnostics: csync.NewVersionedMap[protocol.DocumentURI, []protocol.Diagnostic](),
 82		openFiles:   csync.NewMap[string, *OpenFileInfo](),
 83		config:      cfg,
 84		ctx:         ctx,
 85		debug:       debug,
 86		resolver:    resolver,
 87		cwd:         cwd,
 88	}
 89	client.serverState.Store(StateStopped)
 90
 91	if err := client.createPowernapClient(); err != nil {
 92		return nil, err
 93	}
 94
 95	return client, nil
 96}
 97
 98// Initialize initializes the LSP client and returns the server capabilities.
 99func (c *Client) Initialize(ctx context.Context, workspaceDir string) (*protocol.InitializeResult, error) {
100	if err := c.client.Initialize(ctx, false); err != nil {
101		return nil, fmt.Errorf("failed to initialize the lsp client: %w", err)
102	}
103
104	// Convert powernap capabilities to protocol capabilities
105	caps := c.client.GetCapabilities()
106	protocolCaps := protocol.ServerCapabilities{
107		TextDocumentSync: caps.TextDocumentSync,
108		CompletionProvider: func() *protocol.CompletionOptions {
109			if caps.CompletionProvider != nil {
110				return &protocol.CompletionOptions{
111					TriggerCharacters:   caps.CompletionProvider.TriggerCharacters,
112					AllCommitCharacters: caps.CompletionProvider.AllCommitCharacters,
113					ResolveProvider:     caps.CompletionProvider.ResolveProvider,
114				}
115			}
116			return nil
117		}(),
118	}
119
120	result := &protocol.InitializeResult{
121		Capabilities: protocolCaps,
122	}
123
124	c.registerHandlers()
125
126	return result, nil
127}
128
129// closeTimeout is the maximum time to wait for a graceful LSP shutdown.
130const closeTimeout = 5 * time.Second
131
132// Kill kills the client without doing anything else.
133func (c *Client) Kill() { c.client.Kill() }
134
135// Close closes all open files in the client, then shuts down gracefully.
136// If shutdown takes longer than closeTimeout, it falls back to Kill().
137func (c *Client) Close(ctx context.Context) error {
138	c.CloseAllFiles(ctx)
139
140	// Use a timeout to prevent hanging on unresponsive LSP servers.
141	// jsonrpc2's send lock doesn't respect context cancellation, so we
142	// need to fall back to Kill() which closes the underlying connection.
143	closeCtx, cancel := context.WithTimeout(ctx, closeTimeout)
144	defer cancel()
145
146	done := make(chan error, 1)
147	go func() {
148		if err := c.client.Shutdown(closeCtx); err != nil {
149			slog.Warn("Failed to shutdown LSP client", "error", err)
150		}
151		done <- c.client.Exit()
152	}()
153
154	select {
155	case err := <-done:
156		return err
157	case <-closeCtx.Done():
158		c.client.Kill()
159		return closeCtx.Err()
160	}
161}
162
163// createPowernapClient creates a new powernap client with the current configuration.
164func (c *Client) createPowernapClient() error {
165	rootURI := string(protocol.URIFromPath(c.cwd))
166
167	command, err := c.resolver.ResolveValue(c.config.Command)
168	if err != nil {
169		return fmt.Errorf("invalid lsp command: %w", err)
170	}
171
172	clientConfig := powernap.ClientConfig{
173		Command:     home.Long(command),
174		Args:        c.config.Args,
175		RootURI:     rootURI,
176		Environment: maps.Clone(c.config.Env),
177		Settings:    c.config.Options,
178		InitOptions: c.config.InitOptions,
179		WorkspaceFolders: []protocol.WorkspaceFolder{
180			{
181				URI:  rootURI,
182				Name: filepath.Base(c.cwd),
183			},
184		},
185	}
186
187	powernapClient, err := powernap.NewClient(clientConfig)
188	if err != nil {
189		return fmt.Errorf("failed to create lsp client: %w", err)
190	}
191
192	c.client = powernapClient
193	return nil
194}
195
196// registerHandlers registers the standard LSP notification and request handlers.
197func (c *Client) registerHandlers() {
198	c.RegisterServerRequestHandler("workspace/applyEdit", HandleApplyEdit(c.client.GetOffsetEncoding()))
199	c.RegisterServerRequestHandler("workspace/configuration", HandleWorkspaceConfiguration)
200	c.RegisterServerRequestHandler("client/registerCapability", HandleRegisterCapability)
201	c.RegisterNotificationHandler("window/showMessage", func(ctx context.Context, method string, params json.RawMessage) {
202		if c.debug {
203			HandleServerMessage(ctx, method, params)
204		}
205	})
206	c.RegisterNotificationHandler("textDocument/publishDiagnostics", func(_ context.Context, _ string, params json.RawMessage) {
207		HandleDiagnostics(c, params)
208	})
209}
210
211// Restart closes the current LSP client and creates a new one with the same configuration.
212func (c *Client) Restart() error {
213	var openFiles []string
214	for uri := range c.openFiles.Seq2() {
215		openFiles = append(openFiles, string(uri))
216	}
217
218	closeCtx, cancel := context.WithTimeout(c.ctx, 10*time.Second)
219	defer cancel()
220
221	if err := c.Close(closeCtx); err != nil {
222		slog.Warn("Error closing client during restart", "name", c.name, "error", err)
223	}
224
225	c.SetServerState(StateStopped)
226
227	c.diagCountsCache = DiagnosticCounts{}
228	c.diagCountsVersion = 0
229
230	if err := c.createPowernapClient(); err != nil {
231		return err
232	}
233
234	initCtx, cancel := context.WithTimeout(c.ctx, 30*time.Second)
235	defer cancel()
236
237	c.SetServerState(StateStarting)
238
239	if err := c.client.Initialize(initCtx, false); err != nil {
240		c.SetServerState(StateError)
241		return fmt.Errorf("failed to initialize lsp client: %w", err)
242	}
243
244	c.registerHandlers()
245
246	if err := c.WaitForServerReady(initCtx); err != nil {
247		slog.Error("Server failed to become ready after restart", "name", c.name, "error", err)
248		c.SetServerState(StateError)
249		return err
250	}
251
252	for _, uri := range openFiles {
253		if err := c.OpenFile(initCtx, uri); err != nil {
254			slog.Warn("Failed to reopen file after restart", "file", uri, "error", err)
255		}
256	}
257	return nil
258}
259
260// ServerState represents the state of an LSP server
261type ServerState int
262
263const (
264	StateUnstarted ServerState = iota
265	StateStarting
266	StateReady
267	StateError
268	StateStopped
269	StateDisabled
270)
271
272// GetServerState returns the current state of the LSP server
273func (c *Client) GetServerState() ServerState {
274	if val := c.serverState.Load(); val != nil {
275		return val.(ServerState)
276	}
277	return StateStarting
278}
279
280// SetServerState sets the current state of the LSP server
281func (c *Client) SetServerState(state ServerState) {
282	c.serverState.Store(state)
283}
284
285// GetName returns the name of the LSP client
286func (c *Client) GetName() string {
287	return c.name
288}
289
290// FileTypes returns the file types this LSP client handles
291func (c *Client) FileTypes() []string {
292	return slices.Clone(c.fileTypes)
293}
294
295// SetDiagnosticsCallback sets the callback function for diagnostic changes
296func (c *Client) SetDiagnosticsCallback(callback func(name string, count int)) {
297	c.onDiagnosticsChanged = callback
298}
299
300// WaitForServerReady waits for the server to be ready
301func (c *Client) WaitForServerReady(ctx context.Context) error {
302	// Set initial state
303	c.SetServerState(StateStarting)
304
305	// Try to ping the server with a simple request
306	ticker := time.NewTicker(500 * time.Millisecond)
307	defer ticker.Stop()
308
309	if c.debug {
310		slog.Debug("Waiting for LSP server to be ready...")
311	}
312
313	c.openKeyConfigFiles(ctx)
314
315	for {
316		select {
317		case <-ctx.Done():
318			c.SetServerState(StateError)
319			return fmt.Errorf("timeout waiting for LSP server to be ready")
320		case <-ticker.C:
321			// Check if client is running
322			if !c.client.IsRunning() {
323				if c.debug {
324					slog.Debug("LSP server not ready yet", "server", c.name)
325				}
326				continue
327			}
328
329			// Server is ready
330			c.SetServerState(StateReady)
331			if c.debug {
332				slog.Debug("LSP server is ready")
333			}
334			return nil
335		}
336	}
337}
338
339// OpenFileInfo contains information about an open file
340type OpenFileInfo struct {
341	Version int32
342	URI     protocol.DocumentURI
343}
344
345// HandlesFile checks if this LSP client handles the given file based on its
346// extension and whether it's within the working directory.
347func (c *Client) HandlesFile(path string) bool {
348	if c == nil {
349		return false
350	}
351	if !fsext.HasPrefix(path, c.cwd) {
352		slog.Debug("File outside workspace", "name", c.name, "file", path, "workDir", c.cwd)
353		return false
354	}
355	return handlesFiletype(c.name, c.fileTypes, path)
356}
357
358// OpenFile opens a file in the LSP server.
359func (c *Client) OpenFile(ctx context.Context, filepath string) error {
360	if !c.HandlesFile(filepath) {
361		return nil
362	}
363
364	uri := string(protocol.URIFromPath(filepath))
365
366	if _, exists := c.openFiles.Get(uri); exists {
367		return nil // Already open
368	}
369
370	// Skip files that do not exist or cannot be read
371	content, err := os.ReadFile(filepath)
372	if err != nil {
373		return fmt.Errorf("error reading file: %w", err)
374	}
375
376	// Notify the server about the opened document
377	if err = c.client.NotifyDidOpenTextDocument(ctx, uri, string(powernap.DetectLanguage(filepath)), 1, string(content)); err != nil {
378		return err
379	}
380
381	c.openFiles.Set(uri, &OpenFileInfo{
382		Version: 1,
383		URI:     protocol.DocumentURI(uri),
384	})
385
386	return nil
387}
388
389// NotifyChange notifies the server about a file change.
390func (c *Client) NotifyChange(ctx context.Context, filepath string) error {
391	if c == nil {
392		return nil
393	}
394	uri := string(protocol.URIFromPath(filepath))
395
396	content, err := os.ReadFile(filepath)
397	if err != nil {
398		return fmt.Errorf("error reading file: %w", err)
399	}
400
401	fileInfo, isOpen := c.openFiles.Get(uri)
402	if !isOpen {
403		return fmt.Errorf("cannot notify change for unopened file: %s", filepath)
404	}
405
406	// Increment version
407	fileInfo.Version++
408
409	// Create change event
410	changes := []protocol.TextDocumentContentChangeEvent{
411		{
412			Value: protocol.TextDocumentContentChangeWholeDocument{
413				Text: string(content),
414			},
415		},
416	}
417
418	return c.client.NotifyDidChangeTextDocument(ctx, uri, int(fileInfo.Version), changes)
419}
420
421// IsFileOpen checks if a file is currently open.
422func (c *Client) IsFileOpen(filepath string) bool {
423	uri := string(protocol.URIFromPath(filepath))
424	_, exists := c.openFiles.Get(uri)
425	return exists
426}
427
428// CloseAllFiles closes all currently open files.
429func (c *Client) CloseAllFiles(ctx context.Context) {
430	for uri := range c.openFiles.Seq2() {
431		if c.debug {
432			slog.Debug("Closing file", "file", uri)
433		}
434		if err := c.client.NotifyDidCloseTextDocument(ctx, uri); err != nil {
435			slog.Warn("Error closing file", "uri", uri, "error", err)
436			continue
437		}
438		c.openFiles.Del(uri)
439	}
440}
441
442// GetFileDiagnostics returns diagnostics for a specific file.
443func (c *Client) GetFileDiagnostics(uri protocol.DocumentURI) []protocol.Diagnostic {
444	diags, _ := c.diagnostics.Get(uri)
445	return diags
446}
447
448// GetDiagnostics returns all diagnostics for all files.
449func (c *Client) GetDiagnostics() map[protocol.DocumentURI][]protocol.Diagnostic {
450	if c == nil {
451		return nil
452	}
453	return c.diagnostics.Copy()
454}
455
456// GetDiagnosticCounts returns cached diagnostic counts by severity.
457// Uses the VersionedMap version to avoid recomputing on every call.
458func (c *Client) GetDiagnosticCounts() DiagnosticCounts {
459	if c == nil {
460		return DiagnosticCounts{}
461	}
462	currentVersion := c.diagnostics.Version()
463
464	c.diagCountsMu.Lock()
465	defer c.diagCountsMu.Unlock()
466
467	if currentVersion == c.diagCountsVersion {
468		return c.diagCountsCache
469	}
470
471	// Recompute counts.
472	counts := DiagnosticCounts{}
473	for _, diags := range c.diagnostics.Seq2() {
474		for _, diag := range diags {
475			switch diag.Severity {
476			case protocol.SeverityError:
477				counts.Error++
478			case protocol.SeverityWarning:
479				counts.Warning++
480			case protocol.SeverityInformation:
481				counts.Information++
482			case protocol.SeverityHint:
483				counts.Hint++
484			}
485		}
486	}
487
488	c.diagCountsCache = counts
489	c.diagCountsVersion = currentVersion
490	return counts
491}
492
493// OpenFileOnDemand opens a file only if it's not already open.
494func (c *Client) OpenFileOnDemand(ctx context.Context, filepath string) error {
495	if c == nil {
496		return nil
497	}
498	// Check if the file is already open
499	if c.IsFileOpen(filepath) {
500		return nil
501	}
502
503	// Open the file
504	return c.OpenFile(ctx, filepath)
505}
506
507// RegisterNotificationHandler registers a notification handler.
508func (c *Client) RegisterNotificationHandler(method string, handler transport.NotificationHandler) {
509	c.client.RegisterNotificationHandler(method, handler)
510}
511
512// RegisterServerRequestHandler handles server requests.
513func (c *Client) RegisterServerRequestHandler(method string, handler transport.Handler) {
514	c.client.RegisterHandler(method, handler)
515}
516
517// openKeyConfigFiles opens important configuration files that help initialize the server.
518func (c *Client) openKeyConfigFiles(ctx context.Context) {
519	// Try to open each file, ignoring errors if they don't exist
520	for _, file := range c.config.RootMarkers {
521		file = filepath.Join(c.cwd, file)
522		if _, err := os.Stat(file); err == nil {
523			// File exists, try to open it
524			if err := c.OpenFile(ctx, file); err != nil {
525				slog.Error("Failed to open key config file", "file", file, "error", err)
526			} else {
527				slog.Debug("Opened key config file for initialization", "file", file)
528			}
529		}
530	}
531}
532
533// NotifyWorkspaceChange sends a workspace-level file change notification to
534// trigger re-analysis of all files. This is useful when the overall project
535// state may have changed (e.g., after a project-wide refactoring) and
536// diagnostics for files not currently being edited may be stale.
537func (c *Client) NotifyWorkspaceChange(ctx context.Context) error {
538	if c == nil {
539		return nil
540	}
541	return c.client.NotifyDidChangeWatchedFiles(ctx, []protocol.FileEvent{
542		{URI: protocol.DocumentURI(protocol.URIFromPath(c.cwd)), Type: protocol.Changed},
543	})
544}
545
546// RefreshOpenFiles re-notifies the LSP server about all currently open files,
547// which triggers re-analysis and fresh diagnostics for the entire project.
548func (c *Client) RefreshOpenFiles(ctx context.Context) {
549	if c == nil {
550		return
551	}
552	for uri, info := range c.openFiles.Seq2() {
553		path, err := protocol.DocumentURI(uri).Path()
554		if err != nil {
555			slog.Warn("Failed to convert URI to path", "uri", uri, "error", err)
556			continue
557		}
558		content, err := os.ReadFile(path)
559		if err != nil {
560			slog.Warn("Failed to read file for refresh", "path", path, "error", err)
561			continue
562		}
563		info.Version++
564		changes := []protocol.TextDocumentContentChangeEvent{
565			{
566				Value: protocol.TextDocumentContentChangeWholeDocument{
567					Text: string(content),
568				},
569			},
570		}
571		if err := c.client.NotifyDidChangeTextDocument(ctx, uri, int(info.Version), changes); err != nil {
572			slog.Warn("Failed to notify file change", "uri", uri, "error", err)
573		}
574	}
575}
576
577// WaitForDiagnostics waits until diagnostics stop changing for a settling
578// period, indicating the LSP server has finished processing. If no
579// diagnostics change within firstChangeDuration, it returns early since the
580// server likely isn't going to republish.
581func (c *Client) WaitForDiagnostics(ctx context.Context, timeout time.Duration) {
582	if c == nil {
583		return
584	}
585
586	const (
587		firstChangeDuration = 1 * time.Second
588		settleDuration      = 300 * time.Millisecond
589	)
590
591	deadline := time.NewTimer(timeout)
592	defer deadline.Stop()
593	firstChangeTimer := time.NewTimer(min(timeout, firstChangeDuration))
594	defer firstChangeTimer.Stop()
595	previousVersion := c.diagnostics.Version()
596	ticker := time.NewTicker(100 * time.Millisecond)
597	defer ticker.Stop()
598
599	for {
600		select {
601		case <-ctx.Done():
602			return
603		case <-deadline.C:
604			return
605		case <-firstChangeTimer.C:
606			// No change arrived quickly — server isn't republishing.
607			return
608		case <-ticker.C:
609			currentVersion := c.diagnostics.Version()
610			if currentVersion != previousVersion {
611				// Diagnostics changed — now wait for them to settle.
612				c.waitForDiagnosticsToSettle(ctx, deadline.C, settleDuration)
613				return
614			}
615		}
616	}
617}
618
619// waitForDiagnosticsToSettle waits until diagnostics version stays the same
620// for settleDuration, indicating the LSP server has finished publishing.
621func (c *Client) waitForDiagnosticsToSettle(ctx context.Context, deadline <-chan time.Time, settleDuration time.Duration) {
622	lastVersion := c.diagnostics.Version()
623	settleTicker := time.NewTicker(50 * time.Millisecond)
624	defer settleTicker.Stop()
625
626	// Track how long the version has been stable.
627	stableStart := time.Now()
628
629	for {
630		select {
631		case <-ctx.Done():
632			return
633		case <-deadline:
634			return
635		case <-settleTicker.C:
636			currentVersion := c.diagnostics.Version()
637			if currentVersion != lastVersion {
638				// New change detected — reset the stable timer.
639				lastVersion = currentVersion
640				stableStart = time.Now()
641			} else if time.Since(stableStart) >= settleDuration {
642				// Diagnostics have been stable for the settle duration.
643				return
644			}
645		}
646	}
647}
648
649// FindReferences finds all references to the symbol at the given position.
650func (c *Client) FindReferences(ctx context.Context, filepath string, line, character int, includeDeclaration bool) ([]protocol.Location, error) {
651	if err := c.OpenFileOnDemand(ctx, filepath); err != nil {
652		return nil, err
653	}
654
655	// Add timeout to prevent hanging on slow LSP servers.
656	ctx, cancel := context.WithTimeout(ctx, 5*time.Second)
657	defer cancel()
658
659	// NOTE: line and character should be 0-based.
660	// See: https://microsoft.github.io/language-server-protocol/specifications/lsp/3.17/specification/#position
661	return c.client.FindReferences(ctx, filepath, line-1, character-1, includeDeclaration)
662}