Fix trigger-character completion race in LSP proxy

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<boolean> {
+  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<T>(
+  promise: PromiseLike<T>,
+  timeoutMs: number,
+  label: string,
+): Promise<T> {
+  let timeoutHandle: NodeJS.Timeout | undefined;
+  const timeout = new Promise<never>((_, 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.CompletionList>(
+          '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 });
+      }
+    }
+  });
+});

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<unknown>;
+  };
+};
+
+async function withTimeout<T>(
+  promise: PromiseLike<T>,
+  timeoutMs: number,
+  label: string,
+): Promise<T> {
+  let timeoutHandle: NodeJS.Timeout | undefined;
+  const timeout = new Promise<never>((_, 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();
+    }
+  });
+});