client.go

  1package lsp
  2
  3import (
  4	"context"
  5	"encoding/json"
  6	"fmt"
  7	"log/slog"
  8	"os"
  9	"path/filepath"
 10	"slices"
 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	args, err := c.config.ResolvedArgs(c.resolver)
172	if err != nil {
173		return fmt.Errorf("invalid lsp args: %w", err)
174	}
175
176	envs, err := c.config.ResolvedEnv(c.resolver)
177	if err != nil {
178		return fmt.Errorf("invalid lsp env: %w", err)
179	}
180
181	clientConfig := powernap.ClientConfig{
182		Command:     home.Long(command),
183		Args:        args,
184		RootURI:     rootURI,
185		Environment: envs,
186		Settings:    c.config.Options,
187		InitOptions: c.config.InitOptions,
188		WorkspaceFolders: []protocol.WorkspaceFolder{
189			{
190				URI:  rootURI,
191				Name: filepath.Base(c.cwd),
192			},
193		},
194	}
195
196	powernapClient, err := powernap.NewClient(clientConfig)
197	if err != nil {
198		return fmt.Errorf("failed to create lsp client: %w", err)
199	}
200
201	c.client = powernapClient
202	return nil
203}
204
205// registerHandlers registers the standard LSP notification and request handlers.
206func (c *Client) registerHandlers() {
207	c.RegisterServerRequestHandler("workspace/applyEdit", HandleApplyEdit(c.client.GetOffsetEncoding()))
208	c.RegisterServerRequestHandler("workspace/configuration", HandleWorkspaceConfiguration)
209	c.RegisterServerRequestHandler("client/registerCapability", HandleRegisterCapability)
210	c.RegisterNotificationHandler("window/showMessage", func(ctx context.Context, method string, params json.RawMessage) {
211		if c.debug {
212			HandleServerMessage(ctx, method, params)
213		}
214	})
215	c.RegisterNotificationHandler("textDocument/publishDiagnostics", func(_ context.Context, _ string, params json.RawMessage) {
216		HandleDiagnostics(c, params)
217	})
218}
219
220// Restart closes the current LSP client and creates a new one with the same configuration.
221func (c *Client) Restart() error {
222	var openFiles []string
223	for uri := range c.openFiles.Seq2() {
224		openFiles = append(openFiles, string(uri))
225	}
226
227	closeCtx, cancel := context.WithTimeout(c.ctx, 10*time.Second)
228	defer cancel()
229
230	if err := c.Close(closeCtx); err != nil {
231		slog.Warn("Error closing client during restart", "name", c.name, "error", err)
232	}
233
234	c.SetServerState(StateStopped)
235
236	c.diagCountsCache = DiagnosticCounts{}
237	c.diagCountsVersion = 0
238
239	if err := c.createPowernapClient(); err != nil {
240		return err
241	}
242
243	initCtx, cancel := context.WithTimeout(c.ctx, 30*time.Second)
244	defer cancel()
245
246	c.SetServerState(StateStarting)
247
248	if err := c.client.Initialize(initCtx, false); err != nil {
249		c.SetServerState(StateError)
250		return fmt.Errorf("failed to initialize lsp client: %w", err)
251	}
252
253	c.registerHandlers()
254
255	if err := c.WaitForServerReady(initCtx); err != nil {
256		slog.Error("Server failed to become ready after restart", "name", c.name, "error", err)
257		c.SetServerState(StateError)
258		return err
259	}
260
261	for _, uri := range openFiles {
262		if err := c.OpenFile(initCtx, uri); err != nil {
263			slog.Warn("Failed to reopen file after restart", "file", uri, "error", err)
264		}
265	}
266	return nil
267}
268
269// ServerState represents the state of an LSP server
270type ServerState int
271
272const (
273	StateUnstarted ServerState = iota
274	StateStarting
275	StateReady
276	StateError
277	StateStopped
278	StateDisabled
279)
280
281// GetServerState returns the current state of the LSP server
282func (c *Client) GetServerState() ServerState {
283	if val := c.serverState.Load(); val != nil {
284		return val.(ServerState)
285	}
286	return StateStarting
287}
288
289// SetServerState sets the current state of the LSP server
290func (c *Client) SetServerState(state ServerState) {
291	c.serverState.Store(state)
292}
293
294// GetName returns the name of the LSP client
295func (c *Client) GetName() string {
296	return c.name
297}
298
299// FileTypes returns the file types this LSP client handles
300func (c *Client) FileTypes() []string {
301	return slices.Clone(c.fileTypes)
302}
303
304// SetDiagnosticsCallback sets the callback function for diagnostic changes
305func (c *Client) SetDiagnosticsCallback(callback func(name string, count int)) {
306	c.onDiagnosticsChanged = callback
307}
308
309// WaitForServerReady waits for the server to be ready
310func (c *Client) WaitForServerReady(ctx context.Context) error {
311	// Set initial state
312	c.SetServerState(StateStarting)
313
314	// Try to ping the server with a simple request
315	ticker := time.NewTicker(500 * time.Millisecond)
316	defer ticker.Stop()
317
318	if c.debug {
319		slog.Debug("Waiting for LSP server to be ready...")
320	}
321
322	c.openKeyConfigFiles(ctx)
323
324	for {
325		select {
326		case <-ctx.Done():
327			c.SetServerState(StateError)
328			return fmt.Errorf("timeout waiting for LSP server to be ready")
329		case <-ticker.C:
330			// Check if client is running
331			if !c.client.IsRunning() {
332				if c.debug {
333					slog.Debug("LSP server not ready yet", "server", c.name)
334				}
335				continue
336			}
337
338			// Server is ready
339			c.SetServerState(StateReady)
340			if c.debug {
341				slog.Debug("LSP server is ready")
342			}
343			return nil
344		}
345	}
346}
347
348// OpenFileInfo contains information about an open file
349type OpenFileInfo struct {
350	Version int32
351	URI     protocol.DocumentURI
352}
353
354// HandlesFile checks if this LSP client handles the given file based on its
355// extension and whether it's within the working directory.
356func (c *Client) HandlesFile(path string) bool {
357	if c == nil {
358		return false
359	}
360	if !fsext.HasPrefix(path, c.cwd) {
361		slog.Debug("File outside workspace", "name", c.name, "file", path, "workDir", c.cwd)
362		return false
363	}
364	return handlesFiletype(c.name, c.fileTypes, path)
365}
366
367// OpenFile opens a file in the LSP server.
368func (c *Client) OpenFile(ctx context.Context, filepath string) error {
369	if !c.HandlesFile(filepath) {
370		return nil
371	}
372
373	uri := string(protocol.URIFromPath(filepath))
374
375	if _, exists := c.openFiles.Get(uri); exists {
376		return nil // Already open
377	}
378
379	// Skip files that do not exist or cannot be read
380	content, err := os.ReadFile(filepath)
381	if err != nil {
382		return fmt.Errorf("error reading file: %w", err)
383	}
384
385	// Notify the server about the opened document
386	if err = c.client.NotifyDidOpenTextDocument(ctx, uri, string(powernap.DetectLanguage(filepath)), 1, string(content)); err != nil {
387		return err
388	}
389
390	c.openFiles.Set(uri, &OpenFileInfo{
391		Version: 1,
392		URI:     protocol.DocumentURI(uri),
393	})
394
395	return nil
396}
397
398// NotifyChange notifies the server about a file change.
399func (c *Client) NotifyChange(ctx context.Context, filepath string) error {
400	if c == nil {
401		return nil
402	}
403	uri := string(protocol.URIFromPath(filepath))
404
405	content, err := os.ReadFile(filepath)
406	if err != nil {
407		return fmt.Errorf("error reading file: %w", err)
408	}
409
410	fileInfo, isOpen := c.openFiles.Get(uri)
411	if !isOpen {
412		return fmt.Errorf("cannot notify change for unopened file: %s", filepath)
413	}
414
415	// Increment version
416	fileInfo.Version++
417
418	// Create change event
419	changes := []protocol.TextDocumentContentChangeEvent{
420		{
421			Value: protocol.TextDocumentContentChangeWholeDocument{
422				Text: string(content),
423			},
424		},
425	}
426
427	return c.client.NotifyDidChangeTextDocument(ctx, uri, int(fileInfo.Version), changes)
428}
429
430// IsFileOpen checks if a file is currently open.
431func (c *Client) IsFileOpen(filepath string) bool {
432	uri := string(protocol.URIFromPath(filepath))
433	_, exists := c.openFiles.Get(uri)
434	return exists
435}
436
437// CloseAllFiles closes all currently open files.
438func (c *Client) CloseAllFiles(ctx context.Context) {
439	for uri := range c.openFiles.Seq2() {
440		if c.debug {
441			slog.Debug("Closing file", "file", uri)
442		}
443		if err := c.client.NotifyDidCloseTextDocument(ctx, uri); err != nil {
444			slog.Warn("Error closing file", "uri", uri, "error", err)
445			continue
446		}
447		c.openFiles.Del(uri)
448	}
449}
450
451// GetFileDiagnostics returns diagnostics for a specific file.
452func (c *Client) GetFileDiagnostics(uri protocol.DocumentURI) []protocol.Diagnostic {
453	diags, _ := c.diagnostics.Get(uri)
454	return diags
455}
456
457// GetDiagnostics returns all diagnostics for all files.
458func (c *Client) GetDiagnostics() map[protocol.DocumentURI][]protocol.Diagnostic {
459	if c == nil {
460		return nil
461	}
462	return c.diagnostics.Copy()
463}
464
465// GetDiagnosticCounts returns cached diagnostic counts by severity.
466// Uses the VersionedMap version to avoid recomputing on every call.
467func (c *Client) GetDiagnosticCounts() DiagnosticCounts {
468	if c == nil {
469		return DiagnosticCounts{}
470	}
471	currentVersion := c.diagnostics.Version()
472
473	c.diagCountsMu.Lock()
474	defer c.diagCountsMu.Unlock()
475
476	if currentVersion == c.diagCountsVersion {
477		return c.diagCountsCache
478	}
479
480	// Recompute counts.
481	counts := DiagnosticCounts{}
482	for _, diags := range c.diagnostics.Seq2() {
483		for _, diag := range diags {
484			switch diag.Severity {
485			case protocol.SeverityError:
486				counts.Error++
487			case protocol.SeverityWarning:
488				counts.Warning++
489			case protocol.SeverityInformation:
490				counts.Information++
491			case protocol.SeverityHint:
492				counts.Hint++
493			}
494		}
495	}
496
497	c.diagCountsCache = counts
498	c.diagCountsVersion = currentVersion
499	return counts
500}
501
502// OpenFileOnDemand opens a file only if it's not already open.
503func (c *Client) OpenFileOnDemand(ctx context.Context, filepath string) error {
504	if c == nil {
505		return nil
506	}
507	// Check if the file is already open
508	if c.IsFileOpen(filepath) {
509		return nil
510	}
511
512	// Open the file
513	return c.OpenFile(ctx, filepath)
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// openKeyConfigFiles opens important configuration files that help initialize the server.
527func (c *Client) openKeyConfigFiles(ctx context.Context) {
528	// Try to open each file, ignoring errors if they don't exist
529	for _, file := range c.config.RootMarkers {
530		file = filepath.Join(c.cwd, file)
531		if _, err := os.Stat(file); err == nil {
532			// File exists, try to open it
533			if err := c.OpenFile(ctx, file); err != nil {
534				slog.Error("Failed to open key config file", "file", file, "error", err)
535			} else {
536				slog.Debug("Opened key config file for initialization", "file", file)
537			}
538		}
539	}
540}
541
542// NotifyWorkspaceChange sends a workspace-level file change notification to
543// trigger re-analysis of all files. This is useful when the overall project
544// state may have changed (e.g., after a project-wide refactoring) and
545// diagnostics for files not currently being edited may be stale.
546func (c *Client) NotifyWorkspaceChange(ctx context.Context) error {
547	if c == nil {
548		return nil
549	}
550	return c.client.NotifyDidChangeWatchedFiles(ctx, []protocol.FileEvent{
551		{URI: protocol.DocumentURI(protocol.URIFromPath(c.cwd)), Type: protocol.Changed},
552	})
553}
554
555// RefreshOpenFiles re-notifies the LSP server about all currently open files,
556// which triggers re-analysis and fresh diagnostics for the entire project.
557func (c *Client) RefreshOpenFiles(ctx context.Context) {
558	if c == nil {
559		return
560	}
561	for uri, info := range c.openFiles.Seq2() {
562		path, err := protocol.DocumentURI(uri).Path()
563		if err != nil {
564			slog.Warn("Failed to convert URI to path", "uri", uri, "error", err)
565			continue
566		}
567		content, err := os.ReadFile(path)
568		if err != nil {
569			slog.Warn("Failed to read file for refresh", "path", path, "error", err)
570			continue
571		}
572		info.Version++
573		changes := []protocol.TextDocumentContentChangeEvent{
574			{
575				Value: protocol.TextDocumentContentChangeWholeDocument{
576					Text: string(content),
577				},
578			},
579		}
580		if err := c.client.NotifyDidChangeTextDocument(ctx, uri, int(info.Version), changes); err != nil {
581			slog.Warn("Failed to notify file change", "uri", uri, "error", err)
582		}
583	}
584}
585
586// WaitForDiagnostics waits until diagnostics stop changing for a settling
587// period, indicating the LSP server has finished processing. If no
588// diagnostics change within firstChangeDuration, it returns early since the
589// server likely isn't going to republish.
590func (c *Client) WaitForDiagnostics(ctx context.Context, timeout time.Duration) {
591	if c == nil {
592		return
593	}
594
595	const (
596		firstChangeDuration = 1 * time.Second
597		settleDuration      = 300 * time.Millisecond
598	)
599
600	deadline := time.NewTimer(timeout)
601	defer deadline.Stop()
602	firstChangeTimer := time.NewTimer(min(timeout, firstChangeDuration))
603	defer firstChangeTimer.Stop()
604	previousVersion := c.diagnostics.Version()
605	ticker := time.NewTicker(100 * time.Millisecond)
606	defer ticker.Stop()
607
608	for {
609		select {
610		case <-ctx.Done():
611			return
612		case <-deadline.C:
613			return
614		case <-firstChangeTimer.C:
615			// No change arrived quickly — server isn't republishing.
616			return
617		case <-ticker.C:
618			currentVersion := c.diagnostics.Version()
619			if currentVersion != previousVersion {
620				// Diagnostics changed — now wait for them to settle.
621				c.waitForDiagnosticsToSettle(ctx, deadline.C, settleDuration)
622				return
623			}
624		}
625	}
626}
627
628// waitForDiagnosticsToSettle waits until diagnostics version stays the same
629// for settleDuration, indicating the LSP server has finished publishing.
630func (c *Client) waitForDiagnosticsToSettle(ctx context.Context, deadline <-chan time.Time, settleDuration time.Duration) {
631	lastVersion := c.diagnostics.Version()
632	settleTicker := time.NewTicker(50 * time.Millisecond)
633	defer settleTicker.Stop()
634
635	// Track how long the version has been stable.
636	stableStart := time.Now()
637
638	for {
639		select {
640		case <-ctx.Done():
641			return
642		case <-deadline:
643			return
644		case <-settleTicker.C:
645			currentVersion := c.diagnostics.Version()
646			if currentVersion != lastVersion {
647				// New change detected — reset the stable timer.
648				lastVersion = currentVersion
649				stableStart = time.Now()
650			} else if time.Since(stableStart) >= settleDuration {
651				// Diagnostics have been stable for the settle duration.
652				return
653			}
654		}
655	}
656}
657
658// FindReferences finds all references to the symbol at the given position.
659func (c *Client) FindReferences(ctx context.Context, filepath string, line, character int, includeDeclaration bool) ([]protocol.Location, error) {
660	if err := c.OpenFileOnDemand(ctx, filepath); err != nil {
661		return nil, err
662	}
663
664	// Add timeout to prevent hanging on slow LSP servers.
665	ctx, cancel := context.WithTimeout(ctx, 5*time.Second)
666	defer cancel()
667
668	// NOTE: line and character should be 0-based.
669	// See: https://microsoft.github.io/language-server-protocol/specifications/lsp/3.17/specification/#position
670	return c.client.FindReferences(ctx, filepath, line-1, character-1, includeDeclaration)
671}