chore: make development system filesystem-aware

internal/lsp/handlers.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"errors"
 	"fmt"
+	"path/filepath"
 	"strings"
 
 	"github.com/jzelinskie/persistent"
@@ -34,11 +35,6 @@ func (s *Server) textDocDiagnostic(ctx context.Context, r *jsonrpc2.Request) (Fu
 		return FullDocumentDiagnosticReport{}, err
 	}
 
-	log.Info().
-		Str("uri", string(params.TextDocument.URI)).
-		Int("diagnostics", len(diagnostics)).
-		Msg("diagnostics complete")
-
 	return FullDocumentDiagnosticReport{
 		Kind:  "full",
 		Items: diagnostics,
@@ -57,10 +53,11 @@ func (s *Server) computeDiagnostics(ctx context.Context, uri lsp.DocumentURI) ([
 			return &jsonrpc2.Error{Code: jsonrpc2.CodeInternalError, Message: "file not found"}
 		}
 
+		overlayFS := newLSPOverlayFS(uriToSourceDir(uri), files)
 		devCtx, devErrs, err := development.NewDevContext(ctx, &developerv1.RequestContext{
 			Schema:        file.contents,
 			Relationships: nil,
-		})
+		}, development.WithSourceFS(overlayFS))
 		if err != nil {
 			return err
 		}
@@ -101,7 +98,6 @@ func (s *Server) computeDiagnostics(ctx context.Context, uri lsp.DocumentURI) ([
 		return nil, err
 	}
 
-	log.Info().Int("diagnostics", len(diagnostics)).Str("uri", string(uri)).Msg("computed diagnostics")
 	return diagnostics, nil
 }
 
@@ -171,15 +167,16 @@ func (s *Server) publishDiagnosticsIfNecessary(ctx context.Context, conn *jsonrp
 		return nil
 	}
 
-	log.Debug().
-		Str("uri", string(uri)).
-		Msg("publishing diagnostics")
-
 	diagnostics, err := s.computeDiagnostics(ctx, uri)
 	if err != nil {
 		return fmt.Errorf("failed to compute diagnostics: %w", err)
 	}
 
+	log.Info().
+		Str("uri", string(uri)).
+		Int("diagnostics", len(diagnostics)).
+		Msg("publishing diagnostics")
+
 	return conn.Notify(ctx, "textDocument/publishDiagnostics", lsp.PublishDiagnosticsParams{
 		URI:         uri,
 		Diagnostics: diagnostics,
@@ -197,7 +194,8 @@ func (s *Server) getCompiledContents(path lsp.DocumentURI, files *persistent.Map
 		return compiled, nil
 	}
 
-	justCompiled, derr, err := development.CompileSchema(file.contents)
+	overlayFS := newLSPOverlayFS(uriToSourceDir(path), files)
+	justCompiled, derr, err := development.CompileSchema(file.contents, development.WithSourceFS(overlayFS))
 	if err != nil {
 		return nil, err
 	}
@@ -282,6 +280,63 @@ func (s *Server) textDocHover(_ context.Context, r *jsonrpc2.Request) (*Hover, e
 	return hoverContents, nil
 }
 
+func (s *Server) textDocDefinition(_ context.Context, r *jsonrpc2.Request) (*lsp.Location, error) {
+	params, err := unmarshalParams[lsp.TextDocumentPositionParams](r)
+	if err != nil {
+		return nil, err
+	}
+
+	var location *lsp.Location
+	err = s.withFiles(func(files *persistent.Map[lsp.DocumentURI, trackedFile]) error {
+		compiled, err := s.getCompiledContents(params.TextDocument.URI, files)
+		if err != nil {
+			return err
+		}
+
+		resolver, err := development.NewSchemaPositionMapper(compiled)
+		if err != nil {
+			return err
+		}
+
+		position := input.Position{
+			LineNumber:     params.Position.Line,
+			ColumnPosition: params.Position.Character,
+		}
+
+		resolved, err := resolver.ReferenceAtPosition(input.Source("schema"), position)
+		if err != nil {
+			return err
+		}
+
+		if resolved == nil || resolved.TargetPosition == nil {
+			return nil
+		}
+
+		// Determine the target file URI from TargetSource.
+		targetURI := params.TextDocument.URI
+		if resolved.TargetSource != nil && *resolved.TargetSource != "schema" {
+			sourceDir := uriToSourceDir(params.TextDocument.URI)
+			targetURI = lsp.DocumentURI("file://" + filepath.Join(sourceDir, string(*resolved.TargetSource)))
+		}
+
+		nameStart := resolved.TargetPosition.ColumnPosition + resolved.TargetNamePositionOffset
+		location = &lsp.Location{
+			URI: targetURI,
+			Range: lsp.Range{
+				Start: lsp.Position{Line: resolved.TargetPosition.LineNumber, Character: nameStart},
+				End:   lsp.Position{Line: resolved.TargetPosition.LineNumber, Character: nameStart + len(resolved.Text)},
+			},
+		}
+
+		return nil
+	})
+	if err != nil {
+		return nil, err
+	}
+
+	return location, nil
+}
+
 func (s *Server) textDocFormat(_ context.Context, r *jsonrpc2.Request) ([]lsp.TextEdit, error) {
 	params, err := unmarshalParams[lsp.DocumentFormattingParams](r)
 	if err != nil {
@@ -335,8 +390,8 @@ func (s *Server) initialize(_ context.Context, r *jsonrpc2.Request) (any, error)
 		return nil, err
 	}
 
-	s.requestsDiagnostics = ip.Capabilities.Diagnostics.RefreshSupport
-	log.Debug().
+	s.requestsDiagnostics = ip.Capabilities.Workspace.Diagnostics.RefreshSupport
+	log.Info().
 		Bool("requestsDiagnostics", s.requestsDiagnostics).
 		Msg("initialize")
 
@@ -353,6 +408,7 @@ func (s *Server) initialize(_ context.Context, r *jsonrpc2.Request) (any, error)
 			DocumentFormattingProvider: true,
 			DiagnosticProvider:         &DiagnosticOptions{Identifier: "spicedb", InterFileDependencies: false, WorkspaceDiagnostics: false},
 			HoverProvider:              true,
+			DefinitionProvider:         true,
 		},
 	}, nil
 }

internal/lsp/lsp.go

@@ -96,6 +96,8 @@ func (s *Server) handle(ctx context.Context, conn *jsonrpc2.Conn, r *jsonrpc2.Re
 		result, err = s.textDocFormat(ctx, r)
 	case "textDocument/hover":
 		result, err = s.textDocHover(ctx, r)
+	case "textDocument/definition":
+		result, err = s.textDocDefinition(ctx, r)
 	default:
 		log.Ctx(ctx).Warn().
 			Str("method", r.Method).