1package lsp
2
3import (
4 "bufio"
5 "context"
6 "encoding/json"
7 "fmt"
8 "io"
9 "log/slog"
10 "maps"
11 "os"
12 "os/exec"
13 "path/filepath"
14 "slices"
15 "strings"
16 "sync"
17 "sync/atomic"
18 "time"
19
20 "github.com/charmbracelet/crush/internal/config"
21 "github.com/charmbracelet/crush/internal/log"
22 "github.com/charmbracelet/crush/internal/lsp/protocol"
23)
24
25type Client struct {
26 Cmd *exec.Cmd
27 stdin io.WriteCloser
28 stdout *bufio.Reader
29 stderr io.ReadCloser
30
31 // Client name for identification
32 name string
33
34 // File types this LSP server handles (e.g., .go, .rs, .py)
35 fileTypes []string
36
37 // Diagnostic change callback
38 onDiagnosticsChanged func(name string, count int)
39
40 // Request ID counter
41 nextID atomic.Int32
42
43 // Response handlers
44 handlers map[int32]chan *Message
45 handlersMu sync.RWMutex
46
47 // Server request handlers
48 serverRequestHandlers map[string]ServerRequestHandler
49 serverHandlersMu sync.RWMutex
50
51 // Notification handlers
52 notificationHandlers map[string]NotificationHandler
53 notificationMu sync.RWMutex
54
55 // Diagnostic cache
56 diagnostics map[protocol.DocumentURI][]protocol.Diagnostic
57 diagnosticsMu sync.RWMutex
58
59 // Files are currently opened by the LSP
60 openFiles map[string]*OpenFileInfo
61 openFilesMu sync.RWMutex
62
63 // Server state
64 serverState atomic.Value
65
66 // Server capabilities as returned by initialize
67 caps protocol.ServerCapabilities
68 capsMu sync.RWMutex
69 capsSet atomic.Bool
70}
71
72// NewClient creates a new LSP client.
73func NewClient(ctx context.Context, name string, config config.LSPConfig) (*Client, error) {
74 cmd := exec.CommandContext(ctx, config.Command, config.Args...)
75
76 // Copy env
77 cmd.Env = slices.Concat(os.Environ(), config.ResolvedEnv())
78
79 stdin, err := cmd.StdinPipe()
80 if err != nil {
81 return nil, fmt.Errorf("failed to create stdin pipe: %w", err)
82 }
83
84 stdout, err := cmd.StdoutPipe()
85 if err != nil {
86 return nil, fmt.Errorf("failed to create stdout pipe: %w", err)
87 }
88
89 stderr, err := cmd.StderrPipe()
90 if err != nil {
91 return nil, fmt.Errorf("failed to create stderr pipe: %w", err)
92 }
93
94 client := &Client{
95 Cmd: cmd,
96 name: name,
97 fileTypes: config.FileTypes,
98 stdin: stdin,
99 stdout: bufio.NewReader(stdout),
100 stderr: stderr,
101 handlers: make(map[int32]chan *Message),
102 notificationHandlers: make(map[string]NotificationHandler),
103 serverRequestHandlers: make(map[string]ServerRequestHandler),
104 diagnostics: make(map[protocol.DocumentURI][]protocol.Diagnostic),
105 openFiles: make(map[string]*OpenFileInfo),
106 }
107
108 // Initialize server state
109 client.serverState.Store(StateStarting)
110
111 // Start the LSP server process
112 if err := cmd.Start(); err != nil {
113 return nil, fmt.Errorf("failed to start LSP server: %w", err)
114 }
115
116 // Handle stderr in a separate goroutine
117 go func() {
118 scanner := bufio.NewScanner(stderr)
119 for scanner.Scan() {
120 slog.Error("LSP Server", "err", scanner.Text())
121 }
122 if err := scanner.Err(); err != nil {
123 slog.Error("Error reading", "err", err)
124 }
125 }()
126
127 // Start message handling loop
128 go func() {
129 defer log.RecoverPanic("LSP-message-handler", func() {
130 slog.Error("LSP message handler crashed, LSP functionality may be impaired")
131 })
132 client.handleMessages()
133 }()
134
135 return client, nil
136}
137
138func (c *Client) RegisterNotificationHandler(method string, handler NotificationHandler) {
139 c.notificationMu.Lock()
140 defer c.notificationMu.Unlock()
141 c.notificationHandlers[method] = handler
142}
143
144func (c *Client) RegisterServerRequestHandler(method string, handler ServerRequestHandler) {
145 c.serverHandlersMu.Lock()
146 defer c.serverHandlersMu.Unlock()
147 c.serverRequestHandlers[method] = handler
148}
149
150func (c *Client) InitializeLSPClient(ctx context.Context, workspaceDir string) (*protocol.InitializeResult, error) {
151 initParams := protocol.ParamInitialize{
152 WorkspaceFoldersInitializeParams: protocol.WorkspaceFoldersInitializeParams{
153 WorkspaceFolders: []protocol.WorkspaceFolder{
154 {
155 URI: protocol.URI(protocol.URIFromPath(workspaceDir)),
156 Name: workspaceDir,
157 },
158 },
159 },
160
161 XInitializeParams: protocol.XInitializeParams{
162 ProcessID: int32(os.Getpid()),
163 ClientInfo: &protocol.ClientInfo{
164 Name: "mcp-language-server",
165 Version: "0.1.0",
166 },
167 RootPath: workspaceDir,
168 RootURI: protocol.URIFromPath(workspaceDir),
169 Capabilities: protocol.ClientCapabilities{
170 Workspace: protocol.WorkspaceClientCapabilities{
171 Configuration: true,
172 DidChangeConfiguration: protocol.DidChangeConfigurationClientCapabilities{
173 DynamicRegistration: true,
174 },
175 DidChangeWatchedFiles: protocol.DidChangeWatchedFilesClientCapabilities{
176 DynamicRegistration: true,
177 RelativePatternSupport: true,
178 },
179 },
180 TextDocument: protocol.TextDocumentClientCapabilities{
181 Synchronization: &protocol.TextDocumentSyncClientCapabilities{
182 DynamicRegistration: true,
183 DidSave: true,
184 },
185 Completion: protocol.CompletionClientCapabilities{
186 CompletionItem: protocol.ClientCompletionItemOptions{},
187 },
188 CodeLens: &protocol.CodeLensClientCapabilities{
189 DynamicRegistration: true,
190 },
191 DocumentSymbol: protocol.DocumentSymbolClientCapabilities{},
192 CodeAction: protocol.CodeActionClientCapabilities{
193 CodeActionLiteralSupport: protocol.ClientCodeActionLiteralOptions{
194 CodeActionKind: protocol.ClientCodeActionKindOptions{
195 ValueSet: []protocol.CodeActionKind{},
196 },
197 },
198 },
199 PublishDiagnostics: protocol.PublishDiagnosticsClientCapabilities{
200 VersionSupport: true,
201 },
202 SemanticTokens: protocol.SemanticTokensClientCapabilities{
203 Requests: protocol.ClientSemanticTokensRequestOptions{
204 Range: &protocol.Or_ClientSemanticTokensRequestOptions_range{},
205 Full: &protocol.Or_ClientSemanticTokensRequestOptions_full{},
206 },
207 TokenTypes: []string{},
208 TokenModifiers: []string{},
209 Formats: []protocol.TokenFormat{},
210 },
211 },
212 Window: protocol.WindowClientCapabilities{},
213 },
214 InitializationOptions: map[string]any{
215 "codelenses": map[string]bool{
216 "generate": true,
217 "regenerate_cgo": true,
218 "test": true,
219 "tidy": true,
220 "upgrade_dependency": true,
221 "vendor": true,
222 "vulncheck": false,
223 },
224 },
225 },
226 }
227
228 result, err := c.Initialize(ctx, initParams)
229 if err != nil {
230 return nil, fmt.Errorf("initialize failed: %w", err)
231 }
232
233 c.setCapabilities(result.Capabilities)
234
235 if err := c.Initialized(ctx, protocol.InitializedParams{}); err != nil {
236 return nil, fmt.Errorf("initialized notification failed: %w", err)
237 }
238
239 // Register handlers
240 c.RegisterServerRequestHandler("workspace/applyEdit", HandleApplyEdit)
241 c.RegisterServerRequestHandler("workspace/configuration", HandleWorkspaceConfiguration)
242 c.RegisterServerRequestHandler("client/registerCapability", HandleRegisterCapability)
243 c.RegisterNotificationHandler("window/showMessage", HandleServerMessage)
244 c.RegisterNotificationHandler("textDocument/publishDiagnostics", func(params json.RawMessage) {
245 HandleDiagnostics(c, params)
246 })
247
248 return &result, nil
249}
250
251func (c *Client) Close() error {
252 // Try to close all open files first
253 ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
254 defer cancel()
255
256 // Attempt to close files but continue shutdown regardless
257 c.CloseAllFiles(ctx)
258
259 // Close stdin to signal the server
260 if err := c.stdin.Close(); err != nil {
261 return fmt.Errorf("failed to close stdin: %w", err)
262 }
263
264 // Use a channel to handle the Wait with timeout
265 done := make(chan error, 1)
266 go func() {
267 done <- c.Cmd.Wait()
268 }()
269
270 // Wait for process to exit with timeout
271 select {
272 case err := <-done:
273 return err
274 case <-time.After(2 * time.Second):
275 // If we timeout, try to kill the process
276 if err := c.Cmd.Process.Kill(); err != nil {
277 return fmt.Errorf("failed to kill process: %w", err)
278 }
279 return fmt.Errorf("process killed after timeout")
280 }
281}
282
283type ServerState int
284
285const (
286 StateStarting ServerState = iota
287 StateReady
288 StateError
289)
290
291// GetServerState returns the current state of the LSP server
292func (c *Client) GetServerState() ServerState {
293 if val := c.serverState.Load(); val != nil {
294 return val.(ServerState)
295 }
296 return StateStarting
297}
298
299// SetServerState sets the current state of the LSP server
300func (c *Client) SetServerState(state ServerState) {
301 c.serverState.Store(state)
302}
303
304// GetName returns the name of the LSP client
305func (c *Client) GetName() string {
306 return c.name
307}
308
309// SetDiagnosticsCallback sets the callback function for diagnostic changes
310func (c *Client) SetDiagnosticsCallback(callback func(name string, count int)) {
311 c.onDiagnosticsChanged = callback
312}
313
314// WaitForServerReady waits for the server to be ready by polling the server
315// with a simple request until it responds successfully or times out
316func (c *Client) WaitForServerReady(ctx context.Context) error {
317 cfg := config.Get()
318
319 // Set initial state
320 c.SetServerState(StateStarting)
321
322 // Create a context with timeout
323 ctx, cancel := context.WithTimeout(ctx, 30*time.Second)
324 defer cancel()
325
326 // Try to ping the server with a simple request
327 ticker := time.NewTicker(500 * time.Millisecond)
328 defer ticker.Stop()
329
330 if cfg.Options.DebugLSP {
331 slog.Debug("Waiting for LSP server to be ready...")
332 }
333
334 c.openKeyConfigFiles(ctx)
335
336 for {
337 select {
338 case <-ctx.Done():
339 c.SetServerState(StateError)
340 return fmt.Errorf("timeout waiting for LSP server to be ready")
341 case <-ticker.C:
342 // Try a ping method appropriate for this server type
343 if err := c.ping(ctx); err != nil {
344 if cfg.Options.DebugLSP {
345 slog.Debug("LSP server not ready yet", "error", err, "server", c.name)
346 }
347 continue
348 }
349
350 // Server responded successfully
351 c.SetServerState(StateReady)
352 if cfg.Options.DebugLSP {
353 slog.Debug("LSP server is ready")
354 }
355 return nil
356 }
357 }
358}
359
360// ServerType represents the type of LSP server
361type ServerType int
362
363const (
364 ServerTypeUnknown ServerType = iota
365 ServerTypeGo
366 ServerTypeTypeScript
367 ServerTypeRust
368 ServerTypePython
369 ServerTypeGeneric
370)
371
372// detectServerType tries to determine what type of LSP server we're dealing with
373func (c *Client) detectServerType() ServerType {
374 if c.Cmd == nil {
375 return ServerTypeUnknown
376 }
377
378 cmdPath := strings.ToLower(c.Cmd.Path)
379
380 switch {
381 case strings.Contains(cmdPath, "gopls"):
382 return ServerTypeGo
383 case strings.Contains(cmdPath, "typescript") || strings.Contains(cmdPath, "vtsls") || strings.Contains(cmdPath, "tsserver"):
384 return ServerTypeTypeScript
385 case strings.Contains(cmdPath, "rust-analyzer"):
386 return ServerTypeRust
387 case strings.Contains(cmdPath, "pyright") || strings.Contains(cmdPath, "pylsp") || strings.Contains(cmdPath, "python"):
388 return ServerTypePython
389 default:
390 return ServerTypeGeneric
391 }
392}
393
394// openKeyConfigFiles opens important configuration files that help initialize the server
395func (c *Client) openKeyConfigFiles(ctx context.Context) {
396 workDir := config.Get().WorkingDir()
397 serverType := c.detectServerType()
398
399 var filesToOpen []string
400
401 switch serverType {
402 case ServerTypeTypeScript:
403 // TypeScript servers need these config files to properly initialize
404 filesToOpen = []string{
405 filepath.Join(workDir, "tsconfig.json"),
406 filepath.Join(workDir, "package.json"),
407 filepath.Join(workDir, "jsconfig.json"),
408 }
409
410 // Also find and open a few TypeScript files to help the server initialize
411 c.openTypeScriptFiles(ctx, workDir)
412 case ServerTypeGo:
413 filesToOpen = []string{
414 filepath.Join(workDir, "go.mod"),
415 filepath.Join(workDir, "go.sum"),
416 }
417 case ServerTypeRust:
418 filesToOpen = []string{
419 filepath.Join(workDir, "Cargo.toml"),
420 filepath.Join(workDir, "Cargo.lock"),
421 }
422 }
423
424 // Try to open each file, ignoring errors if they don't exist
425 for _, file := range filesToOpen {
426 if _, err := os.Stat(file); err == nil {
427 // File exists, try to open it
428 if err := c.OpenFile(ctx, file); err != nil {
429 slog.Debug("Failed to open key config file", "file", file, "error", err)
430 } else {
431 slog.Debug("Opened key config file for initialization", "file", file)
432 }
433 }
434 }
435}
436
437// ping sends a ping request...
438func (c *Client) ping(ctx context.Context) error {
439 if _, err := c.Symbol(ctx, protocol.WorkspaceSymbolParams{}); err == nil {
440 return nil
441 }
442 // This is a very lightweight request that should work for most servers
443 return c.Notify(ctx, "$/cancelRequest", protocol.CancelParams{ID: "1"})
444}
445
446// openTypeScriptFiles finds and opens TypeScript files to help initialize the server
447func (c *Client) openTypeScriptFiles(ctx context.Context, workDir string) {
448 cfg := config.Get()
449 filesOpened := 0
450 maxFilesToOpen := 5 // Limit to a reasonable number of files
451
452 // Find and open TypeScript files
453 err := filepath.WalkDir(workDir, func(path string, d os.DirEntry, err error) error {
454 if err != nil {
455 return err
456 }
457
458 // Skip directories and non-TypeScript files
459 if d.IsDir() {
460 // Skip common directories to avoid wasting time
461 if shouldSkipDir(path) {
462 return filepath.SkipDir
463 }
464 return nil
465 }
466
467 // Check if we've opened enough files
468 if filesOpened >= maxFilesToOpen {
469 return filepath.SkipAll
470 }
471
472 // Check file extension
473 ext := filepath.Ext(path)
474 if ext == ".ts" || ext == ".tsx" || ext == ".js" || ext == ".jsx" {
475 // Try to open the file
476 if err := c.OpenFile(ctx, path); err == nil {
477 filesOpened++
478 if cfg.Options.DebugLSP {
479 slog.Debug("Opened TypeScript file for initialization", "file", path)
480 }
481 }
482 }
483
484 return nil
485 })
486
487 if err != nil && cfg.Options.DebugLSP {
488 slog.Debug("Error walking directory for TypeScript files", "error", err)
489 }
490
491 if cfg.Options.DebugLSP {
492 slog.Debug("Opened TypeScript files for initialization", "count", filesOpened)
493 }
494}
495
496// shouldSkipDir returns true if the directory should be skipped during file search
497func shouldSkipDir(path string) bool {
498 dirName := filepath.Base(path)
499
500 // Skip hidden directories
501 if strings.HasPrefix(dirName, ".") {
502 return true
503 }
504
505 // Skip common directories that won't contain relevant source files
506 skipDirs := map[string]bool{
507 "node_modules": true,
508 "dist": true,
509 "build": true,
510 "coverage": true,
511 "vendor": true,
512 "target": true,
513 }
514
515 return skipDirs[dirName]
516}
517
518type OpenFileInfo struct {
519 Version int32
520 URI protocol.DocumentURI
521}
522
523// HandlesFile checks if this LSP client handles the given file based on its
524// extension.
525func (c *Client) HandlesFile(path string) bool {
526 // If no file types are specified, handle all files (backward compatibility)
527 if len(c.fileTypes) == 0 {
528 return true
529 }
530
531 name := strings.ToLower(filepath.Base(path))
532 for _, filetpe := range c.fileTypes {
533 suffix := strings.ToLower(filetpe)
534 if !strings.HasPrefix(suffix, ".") {
535 suffix = "." + suffix
536 }
537 if strings.HasSuffix(name, suffix) {
538 slog.Debug("handles file", "name", c.name, "file", name, "filetype", filetpe)
539 return true
540 }
541 }
542 slog.Debug("doesn't handle file", "name", c.name, "file", name)
543 return false
544}
545
546func (c *Client) OpenFile(ctx context.Context, filepath string) error {
547 if !c.HandlesFile(filepath) {
548 return nil
549 }
550
551 uri := string(protocol.URIFromPath(filepath))
552
553 c.openFilesMu.Lock()
554 if _, exists := c.openFiles[uri]; exists {
555 c.openFilesMu.Unlock()
556 return nil // Already open
557 }
558 c.openFilesMu.Unlock()
559
560 // Skip files that do not exist or cannot be read
561 content, err := os.ReadFile(filepath)
562 if err != nil {
563 return fmt.Errorf("error reading file: %w", err)
564 }
565
566 params := protocol.DidOpenTextDocumentParams{
567 TextDocument: protocol.TextDocumentItem{
568 URI: protocol.DocumentURI(uri),
569 LanguageID: DetectLanguageID(uri),
570 Version: 1,
571 Text: string(content),
572 },
573 }
574
575 if err := c.DidOpen(ctx, params); err != nil {
576 return err
577 }
578
579 c.openFilesMu.Lock()
580 c.openFiles[uri] = &OpenFileInfo{
581 Version: 1,
582 URI: protocol.DocumentURI(uri),
583 }
584 c.openFilesMu.Unlock()
585
586 return nil
587}
588
589func (c *Client) NotifyChange(ctx context.Context, filepath string) error {
590 uri := string(protocol.URIFromPath(filepath))
591
592 content, err := os.ReadFile(filepath)
593 if err != nil {
594 return fmt.Errorf("error reading file: %w", err)
595 }
596
597 c.openFilesMu.Lock()
598 fileInfo, isOpen := c.openFiles[uri]
599 if !isOpen {
600 c.openFilesMu.Unlock()
601 return fmt.Errorf("cannot notify change for unopened file: %s", filepath)
602 }
603
604 // Increment version
605 fileInfo.Version++
606 version := fileInfo.Version
607 c.openFilesMu.Unlock()
608
609 params := protocol.DidChangeTextDocumentParams{
610 TextDocument: protocol.VersionedTextDocumentIdentifier{
611 TextDocumentIdentifier: protocol.TextDocumentIdentifier{
612 URI: protocol.DocumentURI(uri),
613 },
614 Version: version,
615 },
616 ContentChanges: []protocol.TextDocumentContentChangeEvent{
617 {
618 Value: protocol.TextDocumentContentChangeWholeDocument{
619 Text: string(content),
620 },
621 },
622 },
623 }
624
625 return c.DidChange(ctx, params)
626}
627
628func (c *Client) CloseFile(ctx context.Context, filepath string) error {
629 cfg := config.Get()
630 uri := string(protocol.URIFromPath(filepath))
631
632 c.openFilesMu.Lock()
633 if _, exists := c.openFiles[uri]; !exists {
634 c.openFilesMu.Unlock()
635 return nil // Already closed
636 }
637 c.openFilesMu.Unlock()
638
639 params := protocol.DidCloseTextDocumentParams{
640 TextDocument: protocol.TextDocumentIdentifier{
641 URI: protocol.DocumentURI(uri),
642 },
643 }
644
645 if cfg.Options.DebugLSP {
646 slog.Debug("Closing file", "file", filepath)
647 }
648 if err := c.DidClose(ctx, params); err != nil {
649 return err
650 }
651
652 c.openFilesMu.Lock()
653 delete(c.openFiles, uri)
654 c.openFilesMu.Unlock()
655
656 return nil
657}
658
659func (c *Client) IsFileOpen(filepath string) bool {
660 uri := string(protocol.URIFromPath(filepath))
661 c.openFilesMu.RLock()
662 defer c.openFilesMu.RUnlock()
663 _, exists := c.openFiles[uri]
664 return exists
665}
666
667// CloseAllFiles closes all currently open files
668func (c *Client) CloseAllFiles(ctx context.Context) {
669 cfg := config.Get()
670 c.openFilesMu.Lock()
671 filesToClose := make([]string, 0, len(c.openFiles))
672
673 // First collect all URIs that need to be closed
674 for uri := range c.openFiles {
675 // Convert URI back to file path using proper URI handling
676 filePath, err := protocol.DocumentURI(uri).Path()
677 if err != nil {
678 slog.Error("Failed to convert URI to path for file closing", "uri", uri, "error", err)
679 continue
680 }
681 filesToClose = append(filesToClose, filePath)
682 }
683 c.openFilesMu.Unlock()
684
685 // Then close them all
686 for _, filePath := range filesToClose {
687 err := c.CloseFile(ctx, filePath)
688 if err != nil && cfg.Options.DebugLSP {
689 slog.Warn("Error closing file", "file", filePath, "error", err)
690 }
691 }
692
693 if cfg.Options.DebugLSP {
694 slog.Debug("Closed all files", "files", filesToClose)
695 }
696}
697
698func (c *Client) GetFileDiagnostics(uri protocol.DocumentURI) []protocol.Diagnostic {
699 c.diagnosticsMu.RLock()
700 defer c.diagnosticsMu.RUnlock()
701
702 return c.diagnostics[uri]
703}
704
705// GetDiagnostics returns all diagnostics for all files
706func (c *Client) GetDiagnostics() map[protocol.DocumentURI][]protocol.Diagnostic {
707 c.diagnosticsMu.RLock()
708 defer c.diagnosticsMu.RUnlock()
709
710 return maps.Clone(c.diagnostics)
711}
712
713// OpenFileOnDemand opens a file only if it's not already open
714// This is used for lazy-loading files when they're actually needed
715func (c *Client) OpenFileOnDemand(ctx context.Context, filepath string) error {
716 // Check if the file is already open
717 if c.IsFileOpen(filepath) {
718 return nil
719 }
720
721 // Open the file
722 return c.OpenFile(ctx, filepath)
723}
724
725// GetDiagnosticsForFile ensures a file is open and returns its diagnostics
726// This is useful for on-demand diagnostics when using lazy loading
727func (c *Client) GetDiagnosticsForFile(ctx context.Context, filepath string) ([]protocol.Diagnostic, error) {
728 documentURI := protocol.URIFromPath(filepath)
729
730 // Make sure the file is open
731 if !c.IsFileOpen(filepath) {
732 if err := c.OpenFile(ctx, filepath); err != nil {
733 return nil, fmt.Errorf("failed to open file for diagnostics: %w", err)
734 }
735
736 // Give the LSP server a moment to process the file
737 time.Sleep(100 * time.Millisecond)
738 }
739
740 // Get diagnostics
741 c.diagnosticsMu.RLock()
742 diagnostics := c.diagnostics[documentURI]
743 c.diagnosticsMu.RUnlock()
744
745 return diagnostics, nil
746}
747
748// ClearDiagnosticsForURI removes diagnostics for a specific URI from the cache
749func (c *Client) ClearDiagnosticsForURI(uri protocol.DocumentURI) {
750 c.diagnosticsMu.Lock()
751 defer c.diagnosticsMu.Unlock()
752 delete(c.diagnostics, uri)
753}