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