From 694574f21ca468d79775bc8ade129369989200a4 Mon Sep 17 00:00:00 2001 From: Omar Habra Date: Mon, 30 Mar 2026 17:42:08 -0700 Subject: [PATCH 1/2] [Mojo] Serialize LSP proxy packet writes BEGIN_PUBLIC [Mojo] Serialize LSP proxy packet writes Auto-triggered completion could race with didChange because the proxy wrote outgoing LSP packets independently. Serialize proxy writes through a single queue and register pending requests before sending request packets so completion requests no longer overtake earlier document-change notifications. Add regression coverage for the proxy write ordering and a completion test that exercises trigger-character completion after typing '.' in a Mojo file. END_PUBLIC --- extension/lsp/completion.test.default.ts | 166 +++++++++++++++++++++ extension/lsp/proxy.test.default.ts | 181 +++++++++++++++++++++++ extension/test/fake-lsp-server.js | 59 ++++++++ lsp-proxy/src/MojoLSPServer.ts | 44 ++++-- 4 files changed, 434 insertions(+), 16 deletions(-) create mode 100644 extension/lsp/completion.test.default.ts create mode 100644 extension/lsp/proxy.test.default.ts create mode 100644 extension/test/fake-lsp-server.js diff --git a/extension/lsp/completion.test.default.ts b/extension/lsp/completion.test.default.ts new file mode 100644 index 0000000..5853336 --- /dev/null +++ b/extension/lsp/completion.test.default.ts @@ -0,0 +1,166 @@ +//===----------------------------------------------------------------------===// +// Copyright (c) 2025, Modular Inc. All rights reserved. +// +// Licensed under the Apache License v2.0 with LLVM Exceptions: +// https://llvm.org/LICENSE.txt +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +//===----------------------------------------------------------------------===// + +import * as assert from 'assert'; +import * as fs from 'fs/promises'; +import * as path from 'path'; +import * as vscode from 'vscode'; +import { extension } from '../extension'; + +const modularHome = path.join(process.env.HOME!, '.modular'); +const modularConfig = path.join(modularHome, 'modular.cfg'); +const modularPackageRoot = path.join( + modularHome, + 'pkg', + 'packages.modular.com_mojo', +); +const modularLsp = path.join(modularPackageRoot, 'bin', 'mojo-lsp-server'); + +async function pathExists(target: string): Promise { + try { + await fs.lstat(target); + return true; + } catch { + return false; + } +} + +function createDerivedConfig(): string { + return [ + '[max]', + 'version = 0.0.0-test', + '', + '[mojo-max]', + `lsp_server_path = ${path.join(modularPackageRoot, 'bin', 'mojo-lsp-server')}`, + `mblack_path = ${path.join(modularPackageRoot, 'lib', 'mblack', 'mblack')}`, + `lldb_plugin_path = ${path.join(modularPackageRoot, 'lib', 'libMojoLLDB.dylib')}`, + `lldb_vscode_path = ${path.join(modularPackageRoot, 'bin', 'lldb-dap')}`, + `driver_path = ${path.join(modularPackageRoot, 'bin', 'mojo')}`, + `lldb_visualizers_path = ${path.join(modularPackageRoot, 'lib', 'lldb-visualizers')}`, + `lldb_path = ${path.join(modularPackageRoot, 'bin', 'lldb')}`, + '', + ].join('\n'); +} + +async function withTimeout( + promise: PromiseLike, + timeoutMs: number, + label: string, +): Promise { + let timeoutHandle: NodeJS.Timeout | undefined; + const timeout = new Promise((_, reject) => { + timeoutHandle = setTimeout(() => { + reject(new Error(`${label} timed out after ${timeoutMs}ms`)); + }, timeoutMs); + }); + + try { + return await Promise.race([promise, timeout]); + } finally { + if (timeoutHandle) { + clearTimeout(timeoutHandle); + } + } +} + +suite('Completion', function () { + test('trigger-character completion should work immediately after typing "."', async function () { + if (!(await pathExists(modularConfig)) || !(await pathExists(modularLsp))) { + this.skip(); + return; + } + + const workspaceFolder = vscode.workspace.workspaceFolders?.[0]; + assert.ok(workspaceFolder, 'expected a workspace folder'); + + const workspaceRoot = workspaceFolder.uri.fsPath; + const derivedPath = path.join(workspaceRoot, '.derived'); + let derivedExisted = false; + try { + const derivedStat = await fs.lstat(derivedPath); + derivedExisted = true; + if (!derivedStat.isDirectory()) { + await fs.rm(derivedPath, { recursive: true, force: true }); + derivedExisted = false; + } + } catch { + derivedExisted = false; + } + + const createdDerivedDir = !derivedExisted; + if (createdDerivedDir) { + await fs.mkdir(derivedPath, { recursive: true }); + } + await fs.writeFile( + path.join(derivedPath, path.basename(modularConfig)), + createDerivedConfig(), + ); + + const testFile = vscode.Uri.file( + path.join(workspaceRoot, 'completion-trigger-race.test.mojo'), + ); + + try { + await vscode.workspace.fs.writeFile( + testFile, + Buffer.from('import math\n\nfn main():\n math'), + ); + + await vscode.commands.executeCommand('mojo.extension.restart'); + + const document = await vscode.workspace.openTextDocument(testFile); + const editor = await vscode.window.showTextDocument(document); + assert.strictEqual(document.languageId, 'mojo'); + await withTimeout( + extension.lspManager!.tryStartLanguageClient(document), + 30000, + 'language client startup', + ); + assert.strictEqual( + extension.lspManager!.lspClient?.name, + 'Mojo Language Client', + ); + + const line = document.lineAt(document.lineCount - 1); + const cursor = line.range.end; + editor.selection = new vscode.Selection(cursor, cursor); + + await vscode.commands.executeCommand('default:type', { text: '.' }); + + const completion = await withTimeout( + vscode.commands.executeCommand( + 'vscode.executeCompletionItemProvider', + document.uri, + editor.selection.active, + '.', + ), + 30000, + 'trigger-character completion request', + ); + + assert.ok(completion, 'expected a completion result'); + assert.ok( + completion.items.length > 0, + 'expected trigger-character completion items after typing "."', + ); + } finally { + await vscode.commands.executeCommand( + 'workbench.action.closeActiveEditor', + ); + await vscode.workspace.fs.delete(testFile, { useTrash: false }); + if (createdDerivedDir) { + await fs.rm(derivedPath, { recursive: true, force: true }); + } + } + }); +}); diff --git a/extension/lsp/proxy.test.default.ts b/extension/lsp/proxy.test.default.ts new file mode 100644 index 0000000..9fc1b50 --- /dev/null +++ b/extension/lsp/proxy.test.default.ts @@ -0,0 +1,181 @@ +//===----------------------------------------------------------------------===// +// Copyright (c) 2025, Modular Inc. All rights reserved. +// +// Licensed under the Apache License v2.0 with LLVM Exceptions: +// https://llvm.org/LICENSE.txt +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +//===----------------------------------------------------------------------===// + +import * as assert from 'assert'; +import * as path from 'path'; + +const { MojoLSPServer } = require('../../lsp-proxy/out/MojoLSPServer') as { + MojoLSPServer: new (args: { + initializationOptions: { + serverArgs: string[]; + serverEnv: NodeJS.ProcessEnv; + serverPath: string; + }; + logger: (message: string) => void; + onExit: (status: { + code: number | null; + signal: NodeJS.Signals | null; + }) => void; + onNotification: (method: string, params: unknown) => void; + onOutgoingRequest: (id: unknown, method: string, params: unknown) => void; + }) => { + dispose(): void; + sendNotification(params: unknown, method: string): void; + sendRequest(params: unknown, method: string): Promise; + }; +}; + +async function withTimeout( + promise: PromiseLike, + timeoutMs: number, + label: string, +): Promise { + let timeoutHandle: NodeJS.Timeout | undefined; + const timeout = new Promise((_, reject) => { + timeoutHandle = setTimeout(() => { + reject(new Error(`${label} timed out after ${timeoutMs}ms`)); + }, timeoutMs); + }); + + try { + return await Promise.race([promise, timeout]); + } finally { + if (timeoutHandle) { + clearTimeout(timeoutHandle); + } + } +} + +function createServer() { + const fakeServerScript = path.join( + __dirname, + '..', + '..', + 'extension', + 'test', + 'fake-lsp-server.js', + ); + + return new MojoLSPServer({ + initializationOptions: { + serverArgs: [fakeServerScript], + serverEnv: process.env, + serverPath: process.execPath, + }, + logger: () => {}, + onExit: () => {}, + onNotification: () => {}, + onOutgoingRequest: () => {}, + }) as any; +} + +suite('LSP Proxy', function () { + test('sendRequest should handle immediate responses', async function () { + const server = createServer(); + + try { + const result = (await withTimeout( + server.sendRequest({ value: 1 }, 'test/immediate'), + 5000, + 'immediate response request', + )) as { method: string; params: { value: number } }; + + assert.strictEqual(result.method, 'test/immediate'); + assert.deepStrictEqual(result.params, { value: 1 }); + } finally { + server.dispose(); + } + }); + + test('sendNotification should block later writes until its write callback finishes', async function () { + const server = createServer(); + const stdin = server.serverProcess.stdin; + const originalWrite = stdin.write.bind(stdin); + const queuedWrites: Array<{ + chunk: string | Uint8Array; + callback?: () => void; + }> = []; + + let writeCount = 0; + let secondWriteBeforeRelease = false; + let firstWriteReleased = false; + let releaseFirstWrite!: () => void; + + stdin.write = ( + chunk: string | Uint8Array, + callback?: (() => void) | BufferEncoding, + maybeCallback?: () => void, + ) => { + const resolvedCallback = + typeof callback === 'function' ? callback : maybeCallback; + + writeCount += 1; + if (writeCount === 1) { + releaseFirstWrite = () => { + firstWriteReleased = true; + originalWrite(chunk, () => { + resolvedCallback?.(); + for (const queuedWrite of queuedWrites) { + originalWrite(queuedWrite.chunk, queuedWrite.callback); + } + queuedWrites.length = 0; + }); + }; + return true; + } + + if (firstWriteReleased) { + return originalWrite(chunk, resolvedCallback); + } + + secondWriteBeforeRelease = true; + queuedWrites.push({ chunk, callback: resolvedCallback }); + return true; + }; + + try { + server.sendNotification( + { + textDocument: { uri: 'file:///test.mojo', version: 1 }, + contentChanges: [], + }, + 'textDocument/didChange', + ); + + const pendingCompletion = withTimeout( + server.sendRequest( + { + textDocument: { uri: 'file:///test.mojo' }, + position: { line: 0, character: 0 }, + }, + 'textDocument/completion', + ), + 5000, + 'completion request', + ); + + await new Promise((resolve) => setTimeout(resolve, 50)); + assert.strictEqual( + secondWriteBeforeRelease, + false, + 'request write started before the earlier notification write finished', + ); + + releaseFirstWrite(); + await pendingCompletion; + } finally { + stdin.write = originalWrite; + server.dispose(); + } + }); +}); diff --git a/extension/test/fake-lsp-server.js b/extension/test/fake-lsp-server.js new file mode 100644 index 0000000..0824d3f --- /dev/null +++ b/extension/test/fake-lsp-server.js @@ -0,0 +1,59 @@ +const header = 'Content-Length: '; +const separator = '\r\n\r\n'; + +let buffer = Buffer.alloc(0); + +function sendPacket(packet) { + const payload = Buffer.from(JSON.stringify(packet)); + process.stdout.write( + `${header}${payload.length}${separator}${payload.toString()}`, + ); +} + +function tryReadPacket() { + const asString = buffer.toString(); + if (!asString.startsWith(header)) { + return undefined; + } + + let index = header.length; + let contentLength = 0; + for (; index < asString.length; index += 1) { + const char = asString[index]; + if (char < '0' || char > '9') { + break; + } + contentLength = contentLength * 10 + Number.parseInt(char, 10); + } + + if (!asString.slice(index).startsWith(separator)) { + return undefined; + } + + const payloadStart = index + separator.length; + const payload = buffer.subarray(payloadStart, payloadStart + contentLength); + if (payload.length < contentLength) { + return undefined; + } + + buffer = buffer.subarray(payloadStart + contentLength); + return JSON.parse(payload.toString()); +} + +process.stdin.on('data', (chunk) => { + buffer = Buffer.concat([buffer, chunk]); + + let packet; + while ((packet = tryReadPacket()) !== undefined) { + if (packet.id !== undefined) { + sendPacket({ + jsonrpc: '2.0', + id: packet.id, + result: { + method: packet.method, + params: packet.params, + }, + }); + } + } +}); diff --git a/lsp-proxy/src/MojoLSPServer.ts b/lsp-proxy/src/MojoLSPServer.ts index 510b43c..4de6de2 100644 --- a/lsp-proxy/src/MojoLSPServer.ts +++ b/lsp-proxy/src/MojoLSPServer.ts @@ -45,6 +45,7 @@ export class MojoLSPServer extends DisposableContext { private serverProcess: ChildProcess; private lastSentRequestId: RequestId = -1; private pendingRequests = new Map(); + private pendingPacketWrites: Promise = Promise.resolve(); /** * @param initializationOptions The options needed to spawn the * mojo-lsp-server. @@ -82,8 +83,12 @@ export class MojoLSPServer extends DisposableContext { this.pushSubscription( new JSONRPCStream( this.serverProcess.stdout!, - (response: JSONObject) => - this.pendingRequests.get(response.id)!.responseStream.next(response), + (response: JSONObject) => { + const pendingRequest = this.pendingRequests.get(response.id); + if (pendingRequest !== undefined) { + pendingRequest.responseStream.next(response); + } + }, (notification: JSONObject) => onNotification(notification.method, notification.params), (request: JSONObject) => @@ -116,13 +121,15 @@ export class MojoLSPServer extends DisposableContext { ): Promise { const request = this.wrapRequest(params, method); const id = request.id; - await this.sendPacket(request); - const subject = new Subject(); this.pendingRequests.set(id, { params: params, responseStream: subject }); - const result = (await firstValueFrom(subject)).result; - this.pendingRequests.delete(id); - return result; + try { + await this.sendPacket(request); + const result = (await firstValueFrom(subject)).result; + return result; + } finally { + this.pendingRequests.delete(id); + } } /** @@ -131,7 +138,7 @@ export class MojoLSPServer extends DisposableContext { */ public sendNotification(params: T, method: string): void { const notification = this.wrapNotification(params, method); - this.sendPacket(notification); + void this.sendPacket(notification); } /** @@ -140,12 +147,12 @@ export class MojoLSPServer extends DisposableContext { */ public sendResponse(id: any, result: unknown): void { const response = this.wrapResponse(id, result); - this.sendPacket(response); + void this.sendPacket(response); } public sendError(id: any, error: unknown): void { const response = this.wrapResponse(id, undefined, error); - this.sendPacket(response); + void this.sendPacket(response); } /** @@ -163,12 +170,17 @@ export class MojoLSPServer extends DisposableContext { */ private async sendPacket(packet: T): Promise { const payload = Buffer.from(JSON.stringify(packet)); - return new Promise((resolve, _reject) => { - return this.serverProcess.stdin?.write( - `${protocolHeader}${payload.length}${protocolLineSeparator}${payload}`, - () => resolve(), - ); - }); + const queuedWrite = this.pendingPacketWrites.then( + () => + new Promise((resolve) => { + this.serverProcess.stdin?.write( + `${protocolHeader}${payload.length}${protocolLineSeparator}${payload}`, + () => resolve(), + ); + }), + ); + this.pendingPacketWrites = queuedWrite.catch(() => {}); + return queuedWrite; } /** From f19972f5a008f057fda374ecb6e306b4ecb2d2e2 Mon Sep 17 00:00:00 2001 From: Omar Habra Date: Tue, 31 Mar 2026 14:54:15 -0700 Subject: [PATCH 2/2] [Mojo] Fix lint for proxy regression tests BEGIN_PUBLIC [Mojo] Fix lint for proxy regression tests Address the GitHub lint failures in the new proxy regression tests by documenting the compiled proxy require in the TypeScript test and declaring the Node globals used by the fake LSP server helper. END_PUBLIC Signed-off-by: Omar Habra --- extension/lsp/proxy.test.default.ts | 3 +++ extension/test/fake-lsp-server.js | 2 ++ 2 files changed, 5 insertions(+) diff --git a/extension/lsp/proxy.test.default.ts b/extension/lsp/proxy.test.default.ts index 9fc1b50..5b2ffdd 100644 --- a/extension/lsp/proxy.test.default.ts +++ b/extension/lsp/proxy.test.default.ts @@ -14,6 +14,9 @@ import * as assert from 'assert'; import * as path from 'path'; +// We load the compiled proxy implementation because the extension test build +// cannot import the proxy sources directly from the referenced TS project. +// eslint-disable-next-line @typescript-eslint/no-require-imports const { MojoLSPServer } = require('../../lsp-proxy/out/MojoLSPServer') as { MojoLSPServer: new (args: { initializationOptions: { diff --git a/extension/test/fake-lsp-server.js b/extension/test/fake-lsp-server.js index 0824d3f..d58fd7a 100644 --- a/extension/test/fake-lsp-server.js +++ b/extension/test/fake-lsp-server.js @@ -1,3 +1,5 @@ +/* global Buffer, process */ + const header = 'Content-Length: '; const separator = '\r\n\r\n';