lsp: Wait for client to send initialized at boot (#1909)

Now the server will wait for the client to send the initialized (empty)
message before starting work. Previously we got away with sending diagnostics
but now we have a custom test location handler, we need to ensure that we wait
for that to be registered before sending test location notifications.

So we have moved the load workspace contents to the post initialized state, and
added a gate to block the test locations worker from running if the client does
not support it.

Signed-off-by: Charlie Egan <charlie_egan@apple.com>

Author: charlieegan3

cmd/languageserver.go

@@ -42,10 +42,16 @@ func init() {
 				verbose = true
 			}
 
+			logLevel := log.LevelMessage
+			if verbose {
+				logLevel = log.LevelDebug
+			}
+
 			opts := &lsp.LanguageServerOptions{
-				Logger:       log.NewLogger(log.LevelMessage, os.Stderr),
+				Logger:       log.NewLogger(logLevel, os.Stderr),
 				FeatureFlags: lsp.DefaultServerFeatureFlags(),
 			}
+
 			ls := lsp.NewLanguageServer(ctx, opts)
 
 			conf := connection.LoggingConfig{Logger: opts.Logger, LogInbound: verbose, LogOutbound: verbose}

internal/lsp/server.go

@@ -146,6 +146,9 @@ type LanguageServer struct {
 	regoRouter      *rego.RegoRouter
 	testingCompiler *ast.Compiler
 
+	// initializationGate blocks workers until the initialized notification is received
+	initializationGate chan struct{}
+
 	commandRequest       chan types.ExecuteCommandParams
 	lintWorkspaceJobs    chan lintWorkspaceJob
 	lintFileJobs         chan lintFileJob
@@ -213,6 +216,7 @@ func NewLanguageServerMinimal(ctx context.Context, opts *LanguageServerOptions,
 		regoStore:                   store,
 		log:                         opts.Logger,
 		featureFlags:                *featureFlags,
+		initializationGate:          make(chan struct{}),
 		lintFileJobs:                make(chan lintFileJob, 10),
 		lintWorkspaceJobs:           make(chan lintWorkspaceJob, 10),
 		builtinsPositionJobs:        make(chan lintFileJob, 10),
@@ -260,13 +264,13 @@ func (l *LanguageServer) Handle(ctx context.Context, _ *jsonrpc2.Conn, req *json
 	case "initialize":
 		return handler.WithContextAndParams(ctx, req, l.handleInitialize)
 	case "initialized":
-		return l.handleInitialized()
+		return l.handleInitialized(ctx)
 	case "textDocument/definition":
 		return handler.WithParams(req, l.handleTextDocumentDefinition)
 	case "textDocument/diagnostic":
 		return l.handleTextDocumentDiagnostic()
 	case "textDocument/didOpen":
-		return handler.WithContextAndParams(ctx, req, l.handleTextDocumentDidOpen)
+		return handler.WithParams(req, l.handleTextDocumentDidOpen)
 	case "textDocument/didClose":
 		return handler.WithParams(req, l.handleTextDocumentDidClose)
 	case "textDocument/didSave":
@@ -365,13 +369,15 @@ func (l *LanguageServer) StartDiagnosticsWorker(ctx context.Context) {
 					continue
 				}
 
-				// Send test locations update after parse completes
-				if parseSuccess {
-					l.testLocationJobs <- lintFileJob{Reason: job.Reason, URI: job.URI}
-				} else {
-					// Parse failed, send empty test locations
-					if err := l.sendTestLocations(ctx, job.URI, []any{}); err != nil {
-						l.log.Message("failed to send empty test locations after parse failure: %s", err)
+				// Send test locations update after parse completes (if client supports it)
+				if l.client.SupportsOPATestProvider() {
+					if parseSuccess {
+						l.testLocationJobs <- lintFileJob{Reason: job.Reason, URI: job.URI}
+					} else {
+						// Parse failed, send empty test locations
+						if err := l.sendTestLocations(ctx, job.URI, []any{}); err != nil {
+							l.log.Message("failed to send empty test locations after parse failure: %s", err)
+						}
 					}
 				}
 
@@ -1516,12 +1522,11 @@ func (l *LanguageServer) handleTextDocumentDefinition(params types.DefinitionPar
 }
 
 func (l *LanguageServer) handleTextDocumentDidOpen(
-	ctx context.Context,
 	params types.DidOpenTextDocumentParams,
 ) (any, error) {
 	// then we have started the server, and not yet received a suitable root to use.
 	if l.workspaceRootURI == "" {
-		err := l.updateRootURI(ctx,
+		err := l.updateRootURI(
 			// get the URI of the file's immediate parent
 			l.fromPath(filepath.Dir(uri.ToPath(params.TextDocument.URI))),
 		)
@@ -2034,7 +2039,7 @@ func (l *LanguageServer) handleInitialize(ctx context.Context, params types.Init
 	}
 
 	if params.RootURI != "" {
-		err := l.updateRootURI(ctx, params.RootURI)
+		err := l.updateRootURI(params.RootURI)
 		if err != nil {
 			l.log.Message("failed to set rootURI: %w", err)
 		}
@@ -2044,7 +2049,7 @@ func (l *LanguageServer) handleInitialize(ctx context.Context, params types.Init
 			l.log.Message("cannot operate with more than one workspace folder, using: %s", (*params.WorkspaceFolders)[0].URI)
 		}
 
-		err := l.updateRootURI(ctx, (*params.WorkspaceFolders)[0].URI)
+		err := l.updateRootURI((*params.WorkspaceFolders)[0].URI)
 		if err != nil {
 			l.log.Message("failed to set rootURI to workspace folder: %w", err)
 		}
@@ -2053,7 +2058,7 @@ func (l *LanguageServer) handleInitialize(ctx context.Context, params types.Init
 	return initializeResult, nil
 }
 
-func (l *LanguageServer) updateRootURI(ctx context.Context, rootURI string) error {
+func (l *LanguageServer) updateRootURI(rootURI string) error {
 	// rootURI not expected to have a trailing slash, remove if present for
 	// consistency
 	normalizedRootURI := strings.TrimSuffix(rootURI, string(os.PathSeparator))
@@ -2105,19 +2110,6 @@ func (l *LanguageServer) updateRootURI(ctx context.Context, rootURI string) erro
 		l.log.Message("no config file found for workspace")
 	}
 
-	_, failed, err := l.loadWorkspaceContents(ctx, false)
-	for _, f := range failed {
-		l.log.Message("failed to load file %s: %s", f.URI, f.Error)
-	}
-
-	if err != nil {
-		l.log.Message("failed to load workspace contents: %s", err)
-	}
-
-	// 'OverwriteAggregates' is set to populate the cache's initial aggregate state.
-	// Subsequent runs of lintWorkspaceJobs will not set this and use the cached state.
-	l.lintWorkspaceJobs <- lintWorkspaceJob{Reason: "server initialize", OverwriteAggregates: true}
-
 	return nil
 }
 
@@ -2161,13 +2153,15 @@ func (l *LanguageServer) loadWorkspaceContents(ctx context.Context, newOnly bool
 			return nil // continue processing other files
 		}
 
-		if parseSuccess {
-			l.testLocationJobs <- lintFileJob{Reason: "server initialized", URI: fileURI}
-		} else {
-			// this is covering the case where the client starts and the state
-			// is different (the client remembers test locations and state).
-			if err := l.sendTestLocations(ctx, fileURI, []any{}); err != nil {
-				l.log.Message("failed to send empty test locations after parse failure: %s", err)
+		if l.client.SupportsOPATestProvider() {
+			if parseSuccess {
+				l.testLocationJobs <- lintFileJob{Reason: "server initialized", URI: fileURI}
+			} else {
+				// this is covering the case where the client starts and the state
+				// is different (the client remembers test locations and state).
+				if err := l.sendTestLocations(ctx, fileURI, []any{}); err != nil {
+					l.log.Message("failed to send empty test locations after parse failure: %s", err)
+				}
 			}
 		}
 
@@ -2187,12 +2181,33 @@ func (l *LanguageServer) loadWorkspaceContents(ctx context.Context, newOnly bool
 	return changedOrNewURIs, failed, nil
 }
 
-func (l *LanguageServer) handleInitialized() (any, error) {
-	// if running without config, then we should send the diagnostic request now
-	// otherwise it'll happen when the config is loaded
-	if !l.configWatcher.IsWatching() {
-		l.lintWorkspaceJobs <- lintWorkspaceJob{Reason: "server initialized"}
-	}
+func (l *LanguageServer) handleInitialized(ctx context.Context) (any, error) {
+	// Signal workers that initialization handshake is complete
+	close(l.initializationGate)
+
+	// Load workspace contents and start jobs asynchronously
+	// This allows us to respond to the client immediately while workspace
+	// loading happens in the background
+	go func() {
+		_, failed, err := l.loadWorkspaceContents(ctx, false)
+		for _, f := range failed {
+			l.log.Message("failed to load file %s: %s", f.URI, f.Error)
+		}
+
+		if err != nil {
+			l.log.Message("failed to load workspace contents: %s", err)
+		}
+
+		// 'OverwriteAggregates' is set to populate the cache's initial aggregate state.
+		// Subsequent runs of lintWorkspaceJobs will not set this and use the cached state.
+		l.lintWorkspaceJobs <- lintWorkspaceJob{Reason: "server initialize", OverwriteAggregates: true}
+
+		// if running without config, then we should send the diagnostic request now
+		// otherwise it'll happen when the config is loaded
+		if !l.configWatcher.IsWatching() {
+			l.lintWorkspaceJobs <- lintWorkspaceJob{Reason: "server initialized"}
+		}
+	}()
 
 	return emptyStruct, nil
 }

internal/lsp/server_test.go

@@ -98,6 +98,7 @@ func createAndInitServer(t *testing.T, ctx context.Context, tempDir string, clie
 			EnableDebugCodelens:       new(true),
 			EnableExplorer:            new(true),
 			EvalCodelensDisplayInline: new(true),
+			EnableServerTesting:       new(true),
 		},
 	}
 

internal/lsp/testing.go

@@ -11,6 +11,17 @@ import (
 )
 
 func (l *LanguageServer) StartTestLocationsWorker(ctx context.Context) {
+	// Wait for initialization to complete before starting to check worker is needed
+	<-l.initializationGate
+
+	if !l.client.SupportsOPATestProvider() {
+		l.log.Debug("Test locations worker exiting - client does not support opaTestProvider")
+
+		return
+	}
+
+	l.log.Debug("Test locations worker starting")
+
 	for {
 		select {
 		case <-ctx.Done():
@@ -27,6 +38,12 @@ func (l *LanguageServer) StartTestLocationsWorker(ctx context.Context) {
 // This is called after the parse has completed in the lintFileJobs worker to
 // ensure we send the latest.
 func (l *LanguageServer) processTestLocationsUpdate(ctx context.Context, fileURI string) error {
+	if !l.client.SupportsOPATestProvider() {
+		l.log.Message("processTestLocationsUpdate called but client does not support opaTestProvider")
+
+		return nil
+	}
+
 	if l.ignoreURI(fileURI) {
 		return nil
 	}

internal/lsp/types/internal.go

@@ -85,6 +85,12 @@ func (c Client) SupportsEvalCodelensDisplayInline() bool {
 		*c.InitOptions.EvalCodelensDisplayInline
 }
 
+func (c Client) SupportsOPATestProvider() bool {
+	return c.InitOptions != nil &&
+		c.InitOptions.EnableServerTesting != nil &&
+		*c.InitOptions.EnableServerTesting
+}
+
 // ServerContext is a type which is used to contain things from the server's
 // state that is needed in RegalContext.
 type ServerContext struct {

internal/lsp/types/types.go

@@ -34,6 +34,9 @@ type (
 		// EnableExplorer, if set, will enable the regal.explorer command
 		// and related functionality.
 		EnableExplorer *bool `json:"enableExplorer,omitempty"`
+		// EnableServerTesting, if set, will enable test location notifications
+		// via the regal/testLocations and test running handler.
+		EnableServerTesting *bool `json:"enableServerTesting,omitempty"`
 	}
 
 	InitializeParams struct {