lsp: Fix server signature race in test

E.g. https://github.com/StyraInc/regal/actions/runs/16962142183/job/48077281888?pr=1653

Signed-off-by: Charlie Egan <charlie@styra.com>

Author: Charlie Egan

internal/lsp/server.go

@@ -260,7 +260,7 @@ func (l *LanguageServer) Handle(ctx context.Context, _ *jsonrpc2.Conn, req *json
 	case "textDocument/completion":
 		return handler.WithContextAndParams(ctx, req, l.handleTextDocumentCompletion)
 	case "textDocument/signatureHelp":
-		return handler.WithParams(req, l.handleTextDocumentSignatureHelp)
+		return handler.WithContextAndParams(ctx, req, l.handleTextDocumentSignatureHelp)
 	case "workspace/didChangeWatchedFiles":
 		return handler.WithParams(req, l.handleWorkspaceDidChangeWatchedFiles)
 	case "workspace/diagnostic":
@@ -1449,7 +1449,10 @@ func (l *LanguageServer) handleTextDocumentHover(params types.TextDocumentHoverP
 	return nil, nil
 }
 
-func (l *LanguageServer) handleTextDocumentSignatureHelp(params types.SignatureHelpParams) (any, error) {
+func (l *LanguageServer) handleTextDocumentSignatureHelp(
+	ctx context.Context,
+	params types.SignatureHelpParams,
+) (any, error) {
 	if l.ignoreURI(params.TextDocument.URI) {
 		return nil, nil
 	}
@@ -1465,9 +1468,6 @@ func (l *LanguageServer) handleTextDocumentSignatureHelp(params types.SignatureH
 	}
 	input := ast.MustInterfaceToValue(inputData)
 
-	ctx, cancel := context.WithTimeout(context.Background(), time.Second)
-	defer cancel()
-
 	var queryResult map[string]any
 
 	if err := rego.CachedQueryEval(ctx, `data.regal.lsp["signature-help"]`, input, &queryResult); err != nil {

internal/lsp/server_signature_help_test.go

@@ -16,15 +16,10 @@ func TestTextDocumentSignatureHelp(t *testing.T) {
 	t.Parallel()
 
 	ctx, cancel := context.WithCancel(t.Context())
-	defer cancel()
+	t.Cleanup(cancel)
 
 	tempDir := t.TempDir()
 
-	ls := NewLanguageServer(ctx, &LanguageServerOptions{
-		LogWriter: newTestLogger(t),
-		LogLevel:  log.LevelDebug,
-	})
-
 	mainRegoURI := fileURIScheme + filepath.Join(tempDir, "main.rego")
 
 	content := `package example
@@ -33,20 +28,12 @@ allow if regex.match(` + "`foo`" + `, "bar")
 allow if count([1,2,3]) == 2
 allow if concat(",", "a", "b") == "b,a"`
 
-	ls.cache.SetFileContents(mainRegoURI, content)
-
 	builtins := map[string]*ast.Builtin{
 		"count":       ast.Count,
 		"concat":      ast.Concat,
 		"regex.match": ast.RegexMatch,
 	}
 
-	// this is needed to prepare the server for signature help queries
-	err := ls.storeBuiltinsAndCacheSignatureHelp(ctx, builtins)
-	if err != nil {
-		t.Fatalf("failed to store builtins and cache signature help: %s", err)
-	}
-
 	testCases := map[string]struct {
 		position       types.Position
 		expectedLabel  string
@@ -77,14 +64,27 @@ allow if concat(",", "a", "b") == "b,a"`
 		t.Run(name, func(t *testing.T) {
 			t.Parallel()
 
+			ls := NewLanguageServer(ctx, &LanguageServerOptions{
+				LogWriter: newTestLogger(t),
+				LogLevel:  log.LevelDebug,
+			})
+
+			ls.cache.SetFileContents(mainRegoURI, content)
+
+			// this is needed to prepare the server for signature help queries
+			err := ls.storeBuiltinsAndCacheSignatureHelp(ctx, builtins)
+			if err != nil {
+				t.Fatalf("failed to store builtins and cache signature help: %s", err)
+			}
+
 			params := types.SignatureHelpParams{
 				TextDocumentPositionParams: types.TextDocumentPositionParams{
 					TextDocument: types.TextDocumentIdentifier{URI: mainRegoURI},
 					Position:     tc.position,
 				},
 			}
 
-			res, err := ls.handleTextDocumentSignatureHelp(params)
+			res, err := ls.handleTextDocumentSignatureHelp(ctx, params)
 			if err != nil {
 				t.Fatalf("signature help should work, got error: %s", err)
 			}