add tests to LSP and fix an issue with requestDiagnostics, and add support for textDocument/definition, and fix mapper

Author: miparnisari

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,
@@ -102,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
 }
 
@@ -172,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,
@@ -284,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 {
@@ -337,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")
 
@@ -355,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).

internal/lsp/lsp_test.go

@@ -292,8 +292,10 @@ func TestDiagnosticsRefreshSupport(t *testing.T) {
 	// Initialize with diagnostic refresh support enabled
 	resp, serverState := sendAndReceive[lsp.InitializeResult](tester, "initialize", InitializeParams{
 		Capabilities: ClientCapabilities{
-			Diagnostics: DiagnosticWorkspaceClientCapabilities{
-				RefreshSupport: true,
+			Workspace: WorkspaceClientCapabilities{
+				Diagnostics: DiagnosticWorkspaceClientCapabilities{
+					RefreshSupport: true,
+				},
 			},
 		},
 	})
@@ -305,8 +307,10 @@ func TestDiagnosticsRefreshSupport(t *testing.T) {
 	tester2 := newLSPTester(t)
 	resp2, serverState2 := sendAndReceive[lsp.InitializeResult](tester2, "initialize", InitializeParams{
 		Capabilities: ClientCapabilities{
-			Diagnostics: DiagnosticWorkspaceClientCapabilities{
-				RefreshSupport: false,
+			Workspace: WorkspaceClientCapabilities{
+				Diagnostics: DiagnosticWorkspaceClientCapabilities{
+					RefreshSupport: false,
+				},
 			},
 		},
 	})
@@ -357,6 +361,219 @@ func TestUnmarshalParamsErrors(t *testing.T) {
 	}().Code)
 }
 
+func TestMultiFileNoDiagnostics(t *testing.T) {
+	tester := newLSPTester(t)
+	tester.initialize()
+
+	tester.setFileContents("file:///testdir/users.zed", "definition user {}")
+	tester.setFileContents("file:///testdir/root.zed", `use import
+
+import "users.zed"
+
+definition resource {
+	relation viewer: user
+	permission view = viewer
+}
+`)
+
+	resp, _ := sendAndReceive[FullDocumentDiagnosticReport](tester, "textDocument/diagnostic",
+		TextDocumentDiagnosticParams{
+			TextDocument: TextDocument{URI: "file:///testdir/root.zed"},
+		})
+	require.Equal(t, "full", resp.Kind)
+	require.Empty(t, resp.Items)
+}
+
+func TestMultiFileUndefinedDefinitionDiagnostics(t *testing.T) {
+	tester := newLSPTester(t)
+	tester.initialize()
+
+	tester.setFileContents("file:///testdir/broken.zed", `
+definition resource {
+	relation viewer: organization
+	permission view = viewer
+}`)
+	tester.setFileContents("file:///testdir/root.zed", `use import
+
+import "broken.zed"
+`)
+
+	resp, _ := sendAndReceive[FullDocumentDiagnosticReport](tester, "textDocument/diagnostic",
+		TextDocumentDiagnosticParams{
+			TextDocument: TextDocument{URI: "file:///testdir/root.zed"},
+		})
+	require.Equal(t, "full", resp.Kind)
+	require.Len(t, resp.Items, 1)
+	require.Equal(t, lsp.Error, resp.Items[0].Severity)
+	t.Log(resp.Items[0].Message)
+	require.Contains(t, resp.Items[0].Message, "could not lookup definition `organization` for relation `viewer`: object definition `organization` not found")
+}
+
+func TestMultiFileBrokenImportDiagnostics(t *testing.T) {
+	tester := newLSPTester(t)
+	tester.initialize()
+
+	tester.setFileContents("file:///testdir/root.zed", `use import
+import "unknown.zed"
+`)
+
+	resp, _ := sendAndReceive[FullDocumentDiagnosticReport](tester, "textDocument/diagnostic",
+		TextDocumentDiagnosticParams{
+			TextDocument: TextDocument{URI: "file:///testdir/root.zed"},
+		})
+	require.Equal(t, "full", resp.Kind)
+	require.Len(t, resp.Items, 1)
+	require.Equal(t, lsp.Error, resp.Items[0].Severity)
+	require.Contains(t, resp.Items[0].Message, "failed to read import \"unknown.zed\": open unknown.zed: no such file or director")
+}
+
+func TestDefinitionSameFileTypeReference(t *testing.T) {
+	tester := newLSPTester(t)
+	tester.initialize()
+
+	tester.setFileContents("file:///test", `definition user {}
+
+definition resource {
+	relation viewer: user
+	permission view = viewer
+}
+`)
+
+	// Click on "user" in "relation viewer: user" (line 3, character 18)
+	resp, _ := sendAndReceive[lsp.Location](tester, "textDocument/definition",
+		lsp.TextDocumentPositionParams{
+			TextDocument: lsp.TextDocumentIdentifier{URI: "file:///test"},
+			Position:     lsp.Position{Line: 3, Character: 18},
+		})
+	require.Equal(t, lsp.DocumentURI("file:///test"), resp.URI)
+	require.Equal(t, 0, resp.Range.Start.Line)
+	require.Equal(t, len("definition "), resp.Range.Start.Character)
+}
+
+func TestDefinitionSameFileRelationReference(t *testing.T) {
+	tester := newLSPTester(t)
+	tester.initialize()
+
+	tester.setFileContents("file:///test", `definition user {}
+
+definition resource {
+	relation viewer: user
+	permission view = viewer
+}
+`)
+
+	// Click on "viewer" in "permission view = viewer" (line 4, character 19)
+	resp, _ := sendAndReceive[lsp.Location](tester, "textDocument/definition",
+		lsp.TextDocumentPositionParams{
+			TextDocument: lsp.TextDocumentIdentifier{URI: "file:///test"},
+			Position:     lsp.Position{Line: 4, Character: 19},
+		})
+	require.Equal(t, lsp.DocumentURI("file:///test"), resp.URI)
+	require.Equal(t, 3, resp.Range.Start.Line)
+	require.Equal(t, len("\trelation "), resp.Range.Start.Character)
+}
+
+func TestDefinitionCrossFileTypeReference(t *testing.T) {
+	tester := newLSPTester(t)
+	tester.initialize()
+
+	tester.setFileContents("file:///testdir/users.zed", "definition user {}")
+	tester.setFileContents("file:///testdir/root.zed", `use import
+
+import "users.zed"
+
+definition resource {
+	relation viewer: user
+	permission view = viewer
+}
+`)
+
+	// Click on "user" in "relation viewer: user" (line 5, character 18)
+	// It should point to "users.zed"
+	resp, _ := sendAndReceive[lsp.Location](tester, "textDocument/definition",
+		lsp.TextDocumentPositionParams{
+			TextDocument: lsp.TextDocumentIdentifier{URI: "file:///testdir/root.zed"},
+			Position:     lsp.Position{Line: 5, Character: 18},
+		})
+	require.Equal(t, lsp.DocumentURI("file:///testdir/users.zed"), resp.URI)
+	require.Equal(t, 0, resp.Range.Start.Line)
+	require.Equal(t, len("definition "), resp.Range.Start.Character)
+}
+
+func TestDefinitionImportReference(t *testing.T) {
+	tester := newLSPTester(t)
+	tester.initialize()
+
+	tester.setFileContents("file:///testdir/users.zed", "definition user {}")
+	tester.setFileContents("file:///testdir/root.zed", `use import
+
+import "users.zed"
+
+definition resource {
+	relation viewer: user
+	permission view = viewer
+}
+`)
+
+	// Click on import "users.zed" (line 2, character 10)
+	// It should point on the very begginning of "users.zed"
+	resp, _ := sendAndReceive[lsp.Location](tester, "textDocument/definition",
+		lsp.TextDocumentPositionParams{
+			TextDocument: lsp.TextDocumentIdentifier{URI: "file:///testdir/root.zed"},
+			Position:     lsp.Position{Line: 2, Character: 10},
+		})
+	require.Equal(t, lsp.DocumentURI("file:///testdir/users.zed"), resp.URI)
+	require.Equal(t, 0, resp.Range.Start.Line)
+	require.Equal(t, 0, resp.Range.Start.Character)
+}
+
+func TestDefinitionCrossFileCaveatReference(t *testing.T) {
+	tester := newLSPTester(t)
+	tester.initialize()
+
+	tester.setFileContents("file:///testdir/caveats.zed", `caveat some_caveat(some_param int) {
+	some_param < 100
+}`)
+	tester.setFileContents("file:///testdir/root.zed", `use import
+
+import "caveats.zed"
+
+definition user {}
+
+definition resource {
+	relation viewer: user with some_caveat
+}
+`)
+
+	// Click on "some_caveat" in "relation viewer: user with some_caveat" (line 7, character 30)
+	// "\trelation viewer: user with some_caveat"
+	//  0         1         2         3
+	//  0123456789012345678901234567890123456789
+	resp, _ := sendAndReceive[lsp.Location](tester, "textDocument/definition",
+		lsp.TextDocumentPositionParams{
+			TextDocument: lsp.TextDocumentIdentifier{URI: "file:///testdir/root.zed"},
+			Position:     lsp.Position{Line: 7, Character: 30},
+		})
+	require.Equal(t, lsp.DocumentURI("file:///testdir/caveats.zed"), resp.URI)
+	require.Equal(t, 0, resp.Range.Start.Line)
+	require.Equal(t, len("caveat "), resp.Range.Start.Character)
+}
+
+func TestDefinitionNoReference(t *testing.T) {
+	tester := newLSPTester(t)
+	tester.initialize()
+
+	tester.setFileContents("file:///test", "definition user {}")
+
+	// Click on whitespace / keyword where no reference exists
+	resp, _ := sendAndReceive[*lsp.Location](tester, "textDocument/definition",
+		lsp.TextDocumentPositionParams{
+			TextDocument: lsp.TextDocumentIdentifier{URI: "file:///test"},
+			Position:     lsp.Position{Line: 0, Character: 0},
+		})
+	require.Nil(t, resp)
+}
+
 func TestInvalidParams(t *testing.T) {
 	err := invalidParams(errors.New("test error"))
 	require.Equal(t, int64(jsonrpc2.CodeInvalidParams), err.Code)