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// SetDiagnosticsCallback sets the callback function for diagnostic changes
290func (c *Client) SetDiagnosticsCallback(callback func(name string, count int)) {
291 c.onDiagnosticsChanged = callback
292}
293
294// WaitForServerReady waits for the server to be ready
295func (c *Client) WaitForServerReady(ctx context.Context) error {
296 // Set initial state
297 c.SetServerState(StateStarting)
298
299 // Try to ping the server with a simple request
300 ticker := time.NewTicker(500 * time.Millisecond)
301 defer ticker.Stop()
302
303 if c.debug {
304 slog.Debug("Waiting for LSP server to be ready...")
305 }
306
307 c.openKeyConfigFiles(ctx)
308
309 for {
310 select {
311 case <-ctx.Done():
312 c.SetServerState(StateError)
313 return fmt.Errorf("timeout waiting for LSP server to be ready")
314 case <-ticker.C:
315 // Check if client is running
316 if !c.client.IsRunning() {
317 if c.debug {
318 slog.Debug("LSP server not ready yet", "server", c.name)
319 }
320 continue
321 }
322
323 // Server is ready
324 c.SetServerState(StateReady)
325 if c.debug {
326 slog.Debug("LSP server is ready")
327 }
328 return nil
329 }
330 }
331}
332
333// OpenFileInfo contains information about an open file
334type OpenFileInfo struct {
335 Version int32
336 URI protocol.DocumentURI
337}
338
339// HandlesFile checks if this LSP client handles the given file based on its
340// extension and whether it's within the working directory.
341func (c *Client) HandlesFile(path string) bool {
342 if c == nil {
343 return false
344 }
345 if !fsext.HasPrefix(path, c.cwd) {
346 slog.Debug("File outside workspace", "name", c.name, "file", path, "workDir", c.cwd)
347 return false
348 }
349 return handlesFiletype(c.name, c.fileTypes, path)
350}
351
352// OpenFile opens a file in the LSP server.
353func (c *Client) OpenFile(ctx context.Context, filepath string) error {
354 if !c.HandlesFile(filepath) {
355 return nil
356 }
357
358 uri := string(protocol.URIFromPath(filepath))
359
360 if _, exists := c.openFiles.Get(uri); exists {
361 return nil // Already open
362 }
363
364 // Skip files that do not exist or cannot be read
365 content, err := os.ReadFile(filepath)
366 if err != nil {
367 return fmt.Errorf("error reading file: %w", err)
368 }
369
370 // Notify the server about the opened document
371 if err = c.client.NotifyDidOpenTextDocument(ctx, uri, string(powernap.DetectLanguage(filepath)), 1, string(content)); err != nil {
372 return err
373 }
374
375 c.openFiles.Set(uri, &OpenFileInfo{
376 Version: 1,
377 URI: protocol.DocumentURI(uri),
378 })
379
380 return nil
381}
382
383// NotifyChange notifies the server about a file change.
384func (c *Client) NotifyChange(ctx context.Context, filepath string) error {
385 if c == nil {
386 return nil
387 }
388 uri := string(protocol.URIFromPath(filepath))
389
390 content, err := os.ReadFile(filepath)
391 if err != nil {
392 return fmt.Errorf("error reading file: %w", err)
393 }
394
395 fileInfo, isOpen := c.openFiles.Get(uri)
396 if !isOpen {
397 return fmt.Errorf("cannot notify change for unopened file: %s", filepath)
398 }
399
400 // Increment version
401 fileInfo.Version++
402
403 // Create change event
404 changes := []protocol.TextDocumentContentChangeEvent{
405 {
406 Value: protocol.TextDocumentContentChangeWholeDocument{
407 Text: string(content),
408 },
409 },
410 }
411
412 return c.client.NotifyDidChangeTextDocument(ctx, uri, int(fileInfo.Version), changes)
413}
414
415// IsFileOpen checks if a file is currently open.
416func (c *Client) IsFileOpen(filepath string) bool {
417 uri := string(protocol.URIFromPath(filepath))
418 _, exists := c.openFiles.Get(uri)
419 return exists
420}
421
422// CloseAllFiles closes all currently open files.
423func (c *Client) CloseAllFiles(ctx context.Context) {
424 for uri := range c.openFiles.Seq2() {
425 if c.debug {
426 slog.Debug("Closing file", "file", uri)
427 }
428 if err := c.client.NotifyDidCloseTextDocument(ctx, uri); err != nil {
429 slog.Warn("Error closing file", "uri", uri, "error", err)
430 continue
431 }
432 c.openFiles.Del(uri)
433 }
434}
435
436// GetFileDiagnostics returns diagnostics for a specific file.
437func (c *Client) GetFileDiagnostics(uri protocol.DocumentURI) []protocol.Diagnostic {
438 diags, _ := c.diagnostics.Get(uri)
439 return diags
440}
441
442// GetDiagnostics returns all diagnostics for all files.
443func (c *Client) GetDiagnostics() map[protocol.DocumentURI][]protocol.Diagnostic {
444 if c == nil {
445 return nil
446 }
447 return c.diagnostics.Copy()
448}
449
450// GetDiagnosticCounts returns cached diagnostic counts by severity.
451// Uses the VersionedMap version to avoid recomputing on every call.
452func (c *Client) GetDiagnosticCounts() DiagnosticCounts {
453 if c == nil {
454 return DiagnosticCounts{}
455 }
456 currentVersion := c.diagnostics.Version()
457
458 c.diagCountsMu.Lock()
459 defer c.diagCountsMu.Unlock()
460
461 if currentVersion == c.diagCountsVersion {
462 return c.diagCountsCache
463 }
464
465 // Recompute counts.
466 counts := DiagnosticCounts{}
467 for _, diags := range c.diagnostics.Seq2() {
468 for _, diag := range diags {
469 switch diag.Severity {
470 case protocol.SeverityError:
471 counts.Error++
472 case protocol.SeverityWarning:
473 counts.Warning++
474 case protocol.SeverityInformation:
475 counts.Information++
476 case protocol.SeverityHint:
477 counts.Hint++
478 }
479 }
480 }
481
482 c.diagCountsCache = counts
483 c.diagCountsVersion = currentVersion
484 return counts
485}
486
487// OpenFileOnDemand opens a file only if it's not already open.
488func (c *Client) OpenFileOnDemand(ctx context.Context, filepath string) error {
489 if c == nil {
490 return nil
491 }
492 // Check if the file is already open
493 if c.IsFileOpen(filepath) {
494 return nil
495 }
496
497 // Open the file
498 return c.OpenFile(ctx, filepath)
499}
500
501// RegisterNotificationHandler registers a notification handler.
502func (c *Client) RegisterNotificationHandler(method string, handler transport.NotificationHandler) {
503 c.client.RegisterNotificationHandler(method, handler)
504}
505
506// RegisterServerRequestHandler handles server requests.
507func (c *Client) RegisterServerRequestHandler(method string, handler transport.Handler) {
508 c.client.RegisterHandler(method, handler)
509}
510
511// openKeyConfigFiles opens important configuration files that help initialize the server.
512func (c *Client) openKeyConfigFiles(ctx context.Context) {
513 // Try to open each file, ignoring errors if they don't exist
514 for _, file := range c.config.RootMarkers {
515 file = filepath.Join(c.cwd, file)
516 if _, err := os.Stat(file); err == nil {
517 // File exists, try to open it
518 if err := c.OpenFile(ctx, file); err != nil {
519 slog.Error("Failed to open key config file", "file", file, "error", err)
520 } else {
521 slog.Debug("Opened key config file for initialization", "file", file)
522 }
523 }
524 }
525}
526
527// WaitForDiagnostics waits until diagnostics change or the timeout is reached.
528func (c *Client) WaitForDiagnostics(ctx context.Context, d time.Duration) {
529 if c == nil {
530 return
531 }
532 ticker := time.NewTicker(200 * time.Millisecond)
533 defer ticker.Stop()
534 timeout := time.After(d)
535 pv := c.diagnostics.Version()
536 for {
537 select {
538 case <-ctx.Done():
539 return
540 case <-timeout:
541 return
542 case <-ticker.C:
543 if pv != c.diagnostics.Version() {
544 return
545 }
546 }
547 }
548}
549
550// FindReferences finds all references to the symbol at the given position.
551func (c *Client) FindReferences(ctx context.Context, filepath string, line, character int, includeDeclaration bool) ([]protocol.Location, error) {
552 if err := c.OpenFileOnDemand(ctx, filepath); err != nil {
553 return nil, err
554 }
555
556 // Add timeout to prevent hanging on slow LSP servers.
557 ctx, cancel := context.WithTimeout(ctx, 5*time.Second)
558 defer cancel()
559
560 // NOTE: line and character should be 0-based.
561 // See: https://microsoft.github.io/language-server-protocol/specifications/lsp/3.17/specification/#position
562 return c.client.FindReferences(ctx, filepath, line-1, character-1, includeDeclaration)
563}