Merge branch 'main' into copilot/fix-weakened-error-glitch

Author: jakebailey

internal/fourslash/_scripts/convertFourslash.mts

@@ -2941,7 +2941,8 @@ func Test${testName}(t *testing.T) {
     ${failingTests.has(testName) ? "t.Skip()" : ""}
     defer testutil.RecoverAndFail(t, "Panic on fourslash test")
 	const content = ${content}
-    f := fourslash.NewFourslash(t, nil /*capabilities*/, content)
+    f, done := fourslash.NewFourslash(t, nil /*capabilities*/, content)
+    defer done()
     ${isServer ? `f.MarkTestAsStradaServer()\n` : ""}${commands}
 }`;
     return template;

internal/fourslash/fourslash.go

@@ -2,13 +2,15 @@ package fourslash
 
 import (
 	"context"
+	"errors"
 	"fmt"
 	"io"
 	"maps"
 	"runtime"
 	"slices"
 	"strconv"
 	"strings"
+	"sync"
 	"testing"
 	"unicode/utf8"
 
@@ -36,6 +38,7 @@ import (
 	"github.com/microsoft/typescript-go/internal/vfs"
 	"github.com/microsoft/typescript-go/internal/vfs/iovfs"
 	"github.com/microsoft/typescript-go/internal/vfs/vfstest"
+	"golang.org/x/sync/errgroup"
 	"gotest.tools/v3/assert"
 )
 
@@ -62,6 +65,10 @@ type FourslashTest struct {
 	selectionEnd         *lsproto.Position
 
 	isStradaServer bool // Whether this is a fourslash server test in Strada. !!! Remove once we don't need to diff baselines.
+
+	// Async message handling
+	pendingRequests   map[lsproto.ID]chan *lsproto.ResponseMessage
+	pendingRequestsMu sync.Mutex
 }
 
 type scriptInfo struct {
@@ -137,13 +144,14 @@ var parseCache = project.ParseCache{
 	},
 }
 
-func NewFourslash(t *testing.T, capabilities *lsproto.ClientCapabilities, content string) *FourslashTest {
+func NewFourslash(t *testing.T, capabilities *lsproto.ClientCapabilities, content string) (*FourslashTest, func()) {
 	repo.SkipIfNoTypeScriptSubmodule(t)
 	if !bundled.Embedded {
 		// Without embedding, we'd need to read all of the lib files out from disk into the MapFS.
 		// Just skip this for now.
 		t.Skip("bundled files are not embedded")
 	}
+
 	fileName := getBaseFileNameFromTest(t) + tspath.ExtensionTs
 	testfs := make(map[string]any)
 	scriptInfos := make(map[string]*scriptInfo)
@@ -211,16 +219,6 @@ func NewFourslash(t *testing.T, capabilities *lsproto.ClientCapabilities, conten
 		ParseCache: &parseCache,
 	})
 
-	go func() {
-		defer func() {
-			outputWriter.Close()
-		}()
-		err := server.Run(context.TODO())
-		if err != nil {
-			t.Error("server error:", err)
-		}
-	}()
-
 	converters := lsconv.NewConverters(lsproto.PositionEncodingKindUTF8, func(fileName string) *lsconv.LSPLineMap {
 		scriptInfo, ok := scriptInfos[fileName]
 		if !ok {
@@ -240,11 +238,26 @@ func NewFourslash(t *testing.T, capabilities *lsproto.ClientCapabilities, conten
 		converters:      converters,
 		baselines:       make(map[baselineCommand]*strings.Builder),
 		openFiles:       make(map[string]struct{}),
+		pendingRequests: make(map[lsproto.ID]chan *lsproto.ResponseMessage),
 	}
 
+	ctx, cancel := context.WithCancel(t.Context())
+	g, ctx := errgroup.WithContext(ctx)
+
+	// Start server goroutine
+	g.Go(func() error {
+		defer outputWriter.Close()
+		return server.Run(ctx)
+	})
+
+	// Start async message router
+	g.Go(func() error {
+		return f.messageRouter(ctx)
+	})
+
 	// !!! temporary; remove when we have `handleDidChangeConfiguration`/implicit project config support
 	// !!! replace with a proper request *after initialize*
-	f.server.SetCompilerOptionsForInferredProjects(t.Context(), compilerOptions)
+	f.server.SetCompilerOptionsForInferredProjects(ctx, compilerOptions)
 	f.initialize(t, capabilities)
 
 	if testData.isStateBaseliningEnabled() {
@@ -258,11 +271,132 @@ func NewFourslash(t *testing.T, capabilities *lsproto.ClientCapabilities, conten
 	}
 
 	_, testPath, _, _ := runtime.Caller(1)
-	t.Cleanup(func() {
+	return f, func() {
+		t.Helper()
+		cancel()
 		inputWriter.Close()
+		if err := g.Wait(); err != nil && !errors.Is(err, context.Canceled) {
+			t.Errorf("goroutine error: %v", err)
+		}
 		f.verifyBaselines(t, testPath)
-	})
-	return f
+	}
+}
+
+// messageRouter runs in a goroutine and routes incoming messages from the server.
+// It handles responses to client requests and server-initiated requests.
+func (f *FourslashTest) messageRouter(ctx context.Context) error {
+	for {
+		if ctx.Err() != nil {
+			return nil
+		}
+
+		msg, err := f.out.Read()
+		if err != nil {
+			if errors.Is(err, io.EOF) || ctx.Err() != nil {
+				return nil
+			}
+			return fmt.Errorf("failed to read message: %w", err)
+		}
+
+		// Validate message can be marshaled
+		if err := json.MarshalWrite(io.Discard, msg); err != nil {
+			if ctx.Err() != nil {
+				return nil
+			}
+
+			return fmt.Errorf("failed to encode message as JSON: %w", err)
+		}
+
+		switch msg.Kind {
+		case lsproto.MessageKindResponse:
+			f.handleResponse(ctx, msg.AsResponse())
+		case lsproto.MessageKindRequest:
+			if err := f.handleServerRequest(ctx, msg.AsRequest()); err != nil {
+				return err
+			}
+		case lsproto.MessageKindNotification:
+			// Server-initiated notifications (e.g., publishDiagnostics) are currently ignored
+			// in fourslash tests
+		}
+	}
+}
+
+// handleResponse routes a response message to the waiting request goroutine.
+func (f *FourslashTest) handleResponse(ctx context.Context, resp *lsproto.ResponseMessage) {
+	if resp.ID == nil {
+		return
+	}
+
+	f.pendingRequestsMu.Lock()
+	respChan, ok := f.pendingRequests[*resp.ID]
+	if ok {
+		delete(f.pendingRequests, *resp.ID)
+	}
+	f.pendingRequestsMu.Unlock()
+
+	if ok {
+		select {
+		case respChan <- resp:
+			// sent response
+		case <-ctx.Done():
+			// context cancelled
+		}
+	}
+}
+
+// handleServerRequest handles requests initiated by the server (e.g., workspace/configuration).
+func (f *FourslashTest) handleServerRequest(ctx context.Context, req *lsproto.RequestMessage) error {
+	var response *lsproto.ResponseMessage
+
+	switch req.Method {
+	case lsproto.MethodWorkspaceConfiguration:
+		// Return current user preferences
+		response = &lsproto.ResponseMessage{
+			ID:      req.ID,
+			JSONRPC: req.JSONRPC,
+			Result:  []any{f.userPreferences},
+		}
+
+	case lsproto.MethodClientRegisterCapability:
+		// Accept all capability registrations
+		response = &lsproto.ResponseMessage{
+			ID:      req.ID,
+			JSONRPC: req.JSONRPC,
+			Result:  lsproto.Null{},
+		}
+
+	case lsproto.MethodClientUnregisterCapability:
+		// Accept all capability unregistrations
+		response = &lsproto.ResponseMessage{
+			ID:      req.ID,
+			JSONRPC: req.JSONRPC,
+			Result:  lsproto.Null{},
+		}
+
+	default:
+		// Unknown server request
+		response = &lsproto.ResponseMessage{
+			ID:      req.ID,
+			JSONRPC: req.JSONRPC,
+			Error: &lsproto.ResponseError{
+				Code:    int32(lsproto.ErrorCodeMethodNotFound),
+				Message: fmt.Sprintf("Unknown method: %s", req.Method),
+			},
+		}
+	}
+
+	// Send response back to server
+	if ctx.Err() != nil {
+		return nil
+	}
+
+	if err := f.in.Write(response.Message()); err != nil {
+		if ctx.Err() != nil {
+			return nil
+		}
+		return fmt.Errorf("failed to write server request response: %w", err)
+	}
+	return nil
 }
 
 func getBaseFileNameFromTest(t *testing.T) string {
@@ -300,16 +434,22 @@ func (f *FourslashTest) initialize(t *testing.T, capabilities *lsproto.ClientCap
 	params := &lsproto.InitializeParams{
 		Locale: ptrTo("en-US"),
 		InitializationOptions: &lsproto.InitializationOptions{
-			// Hack: disable push diagnostics entirely, since the fourslash runner does not
-			// yet gracefully handle non-request messages.
-			DisablePushDiagnostics:           ptrTo(true),
 			CodeLensShowLocationsCommandName: ptrTo(showCodeLensLocationsCommandName),
 		},
 	}
 	params.Capabilities = getCapabilitiesWithDefaults(capabilities)
-	// !!! check for errors?
-	sendRequestWorker(t, f, lsproto.InitializeInfo, params)
+	resp, _, ok := sendRequestWorker(t, f, lsproto.InitializeInfo, params)
+	if !ok {
+		t.Fatalf("Initialize request failed")
+	}
+	if resp.AsResponse().Error != nil {
+		t.Fatalf("Initialize request returned error: %s", resp.AsResponse().Error.String())
+	}
 	sendNotificationWorker(t, f, lsproto.InitializedInfo, &lsproto.InitializedParams{})
+
+	// Wait for the initial configuration exchange to complete
+	// The server will send workspace/configuration as part of handleInitialized
+	<-f.server.InitComplete()
 }
 
 var (
@@ -412,45 +552,36 @@ func getCapabilitiesWithDefaults(capabilities *lsproto.ClientCapabilities) *lspr
 
 func sendRequestWorker[Params, Resp any](t *testing.T, f *FourslashTest, info lsproto.RequestInfo[Params, Resp], params Params) (*lsproto.Message, Resp, bool) {
 	id := f.nextID()
-	req := info.NewRequestMessage(
-		lsproto.NewID(lsproto.IntegerOrString{Integer: &id}),
-		params,
-	)
-	f.writeMsg(t, req.Message())
-	resp := f.readMsg(t)
-	if resp == nil {
-		return nil, *new(Resp), false
-	}
+	reqID := lsproto.NewID(lsproto.IntegerOrString{Integer: &id})
+	req := info.NewRequestMessage(reqID, params)
 
-	// currently, the only request that may be sent by the server during a client request is one `config` request
-	// !!! remove if `config` is handled in initialization and there are no other server-initiated requests
-	if resp.Kind == lsproto.MessageKindRequest {
-		req := resp.AsRequest()
+	// Create response channel and register it
+	responseChan := make(chan *lsproto.ResponseMessage, 1)
+	f.pendingRequestsMu.Lock()
+	f.pendingRequests[*reqID] = responseChan
+	f.pendingRequestsMu.Unlock()
 
-		assert.Equal(t, req.Method, lsproto.MethodWorkspaceConfiguration, "Unexpected request received: %s", req.Method)
-		res := lsproto.ResponseMessage{
-			ID:      req.ID,
-			JSONRPC: req.JSONRPC,
-			Result:  []any{f.userPreferences},
-		}
-		f.writeMsg(t, res.Message())
-		req = f.readMsg(t).AsRequest()
+	// Send the request
+	f.writeMsg(t, req.Message())
 
-		assert.Equal(t, req.Method, lsproto.MethodClientRegisterCapability, "Unexpected request received: %s", req.Method)
-		res = lsproto.ResponseMessage{
-			ID:      req.ID,
-			JSONRPC: req.JSONRPC,
-			Result:  lsproto.Null{},
+	// Wait for response with context
+	ctx := t.Context()
+	var resp *lsproto.ResponseMessage
+	select {
+	case <-ctx.Done():
+		f.pendingRequestsMu.Lock()
+		delete(f.pendingRequests, *reqID)
+		f.pendingRequestsMu.Unlock()
+		t.Fatalf("Request cancelled: %v", ctx.Err())
+		return nil, *new(Resp), false
+	case resp = <-responseChan:
+		if resp == nil {
+			return nil, *new(Resp), false
 		}
-		f.writeMsg(t, res.Message())
-		resp = f.readMsg(t)
 	}
 
-	if resp == nil {
-		return nil, *new(Resp), false
-	}
-	result, ok := resp.AsResponse().Result.(Resp)
-	return resp, result, ok
+	result, ok := resp.Result.(Resp)
+	return resp.Message(), result, ok
 }
 
 func sendNotificationWorker[Params any](t *testing.T, f *FourslashTest, info lsproto.NotificationInfo[Params], params Params) {
@@ -467,16 +598,6 @@ func (f *FourslashTest) writeMsg(t *testing.T, msg *lsproto.Message) {
 	}
 }
 
-func (f *FourslashTest) readMsg(t *testing.T) *lsproto.Message {
-	// !!! filter out response by id etc
-	msg, err := f.out.Read()
-	if err != nil {
-		t.Fatalf("failed to read message: %v", err)
-	}
-	assert.NilError(t, json.MarshalWrite(io.Discard, msg), "failed to encode message as JSON")
-	return msg
-}
-
 func sendRequest[Params, Resp any](t *testing.T, f *FourslashTest, info lsproto.RequestInfo[Params, Resp], params Params) Resp {
 	t.Helper()
 	prefix := f.getCurrentPositionPrefix()

internal/fourslash/tests/autoImportCompletion_test.go

@@ -24,7 +24,8 @@ import {someVar} from "./a.ts";
 someVar;
 a/**/
 `
-	f := fourslash.NewFourslash(t, nil /*capabilities*/, content)
+	f, done := fourslash.NewFourslash(t, nil /*capabilities*/, content)
+	defer done()
 	f.VerifyCompletions(t, "", &fourslash.CompletionsExpectedList{
 		UserPreferences: &lsutil.UserPreferences{
 			IncludeCompletionsForModuleExports:    core.TSTrue,
@@ -69,7 +70,8 @@ import {someVar} from "./a.ts";
 someVar;
 a/**/
 `
-	f := fourslash.NewFourslash(t, nil /*capabilities*/, content)
+	f, done := fourslash.NewFourslash(t, nil /*capabilities*/, content)
+	defer done()
 	f.VerifyCompletions(t, "", &fourslash.CompletionsExpectedList{
 		UserPreferences: &lsutil.UserPreferences{
 			IncludeCompletionsForModuleExports:    core.TSTrue,
@@ -100,7 +102,8 @@ import { aa, someVar } from "./a.ts";
 someVar;
 b/**/
 `
-	f := fourslash.NewFourslash(t, nil /*capabilities*/, content)
+	f, done := fourslash.NewFourslash(t, nil /*capabilities*/, content)
+	defer done()
 	f.VerifyCompletions(t, "", &fourslash.CompletionsExpectedList{
 		UserPreferences: &lsutil.UserPreferences{
 			IncludeCompletionsForModuleExports:    core.TSTrue,
@@ -153,6 +156,7 @@ const m = import("./src//*6*/");
 import {} from "./src//*7*/";
 import mod = require("./src//*8*/");
 const m = import("./src//*9*/");`
-	f := fourslash.NewFourslash(t, nil /*capabilities*/, content)
+	f, done := fourslash.NewFourslash(t, nil /*capabilities*/, content)
+	defer done()
 	f.BaselineAutoImportsCompletions(t, []string{"1", "3", "6", "9", "2", "4", "5", "7", "8"})
 }