Add LSP DocumentSymbolProvider support (#4096)

Author: stefanvanburen

CHANGELOG.md

@@ -2,7 +2,7 @@
 
 ## [Unreleased]
 
-- No changes yet.
+- Add `textDocument/documentSymbol` support for `buf lsp serve`.
 
 ## [v1.59.0] - 2025-10-20
 

private/buf/buflsp/file.go

@@ -23,6 +23,7 @@ import (
 	"fmt"
 	"io"
 	"io/fs"
+	"iter"
 	"log/slog"
 	"os"
 	"slices"
@@ -1162,3 +1163,39 @@ func readAllAsString(reader io.Reader) (string, error) {
 	}
 	return builder.String(), nil
 }
+
+// GetSymbols retrieves symbols for the file. If a query is passed, matches only symbols matching
+// the case-insensitive substring match to the symbol.
+//
+// This operation requires [IndexSymbols].
+func (f *file) GetSymbols(query string) iter.Seq[protocol.SymbolInformation] {
+	return func(yield func(protocol.SymbolInformation) bool) {
+		if f.ir.IsZero() {
+			return
+		}
+		// Search through all symbols in this file.
+		for _, sym := range f.symbols {
+			if sym.ir.IsZero() {
+				continue
+			}
+			// Only include definitions: static and referenceable symbols.
+			// Skip references, imports, builtins, and tags
+			_, isStatic := sym.kind.(*static)
+			_, isReferenceable := sym.kind.(*referenceable)
+			if !isStatic && !isReferenceable {
+				continue
+			}
+			symbolInfo := sym.GetSymbolInformation()
+			if symbolInfo.Name == "" {
+				continue // Symbol information not supported for this symbol.
+			}
+			// Filter by query (case-insensitive substring match)
+			if query != "" && !strings.Contains(strings.ToLower(symbolInfo.Name), query) {
+				continue
+			}
+			if !yield(symbolInfo) {
+				return
+			}
+		}
+	}
+}

private/buf/buflsp/server.go

@@ -19,7 +19,6 @@ import (
 	"errors"
 	"fmt"
 	"runtime/debug"
-	"slices"
 	"strings"
 	"unicode/utf16"
 
@@ -32,6 +31,8 @@ import (
 
 const (
 	serverName = "buf-lsp"
+
+	maxSymbolResults = 1000
 )
 
 const (
@@ -141,6 +142,7 @@ func (s *server) Initialize(
 				Full: true,
 			},
 			WorkspaceSymbolProvider: true,
+			DocumentSymbolProvider:  true,
 		},
 		ServerInfo: info,
 	}, nil
@@ -529,45 +531,34 @@ func (s *server) Symbols(
 	ctx context.Context,
 	params *protocol.WorkspaceSymbolParams,
 ) ([]protocol.SymbolInformation, error) {
-	const maxResults = 1000 // Limit results to avoid overwhelming clients
 	query := strings.ToLower(params.Query)
-
 	var results []protocol.SymbolInformation
 	for _, file := range s.fileManager.uriToFile.Range {
-		if file.ir.IsZero() {
-			continue
-		}
-		// Search through all symbols in this file.
-		for _, sym := range file.symbols {
-			if sym.ir.IsZero() {
-				continue
-			}
-			// Only include definitions: static and referenceable symbols.
-			// Skip references, imports, builtins, and tags
-			_, isStatic := sym.kind.(*static)
-			_, isReferenceable := sym.kind.(*referenceable)
-			if !isStatic && !isReferenceable {
-				continue
-			}
-			symbolInfo := sym.GetSymbolInformation()
-			if symbolInfo.Name == "" {
-				continue // Symbol information not supported for this symbol.
-			}
-			// Filter by query (case-insensitive substring match)
-			if query != "" && !strings.Contains(strings.ToLower(symbolInfo.Name), query) {
-				continue
-			}
-			results = append(results, symbolInfo)
-			if len(results) >= maxResults {
+		for symbol := range file.GetSymbols(query) {
+			results = append(results, symbol)
+			if len(results) > maxSymbolResults {
 				break
 			}
 		}
 	}
-	slices.SortFunc(results, func(a, b protocol.SymbolInformation) int {
-		if a.Name != b.Name {
-			return strings.Compare(a.Name, b.Name)
-		}
-		return strings.Compare(a.ContainerName, b.ContainerName)
-	})
 	return results, nil
 }
+
+// DocumentSymbol is the entry point for document symbol search.
+func (s *server) DocumentSymbol(ctx context.Context, params *protocol.DocumentSymbolParams) (
+	result []any, // []protocol.SymbolInformation
+	err error,
+) {
+	file := s.fileManager.Get(params.TextDocument.URI)
+	if file == nil {
+		return nil, nil
+	}
+	anyResults := []any{}
+	for symbol := range file.GetSymbols("") {
+		anyResults = append(anyResults, symbol)
+		if len(anyResults) > maxSymbolResults {
+			break
+		}
+	}
+	return anyResults, nil
+}