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..5b2ffdd --- /dev/null +++ b/extension/lsp/proxy.test.default.ts @@ -0,0 +1,184 @@ +//===----------------------------------------------------------------------===// +// 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'; + +// 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: { + 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..d58fd7a --- /dev/null +++ b/extension/test/fake-lsp-server.js @@ -0,0 +1,61 @@ +/* global Buffer, process */ + +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; } /**