Add 'custom/chatCompletion' event in vscode extension and 'lmProvider' in lsp consuming it

.chronus/changes/lm-from-ide-2025-11-15-20-51-26.md

@@ -0,0 +1,7 @@
+---
+changeKind: internal
+packages:
+  - "@typespec/compiler"
+---
+
+Expose LSP connection through globalThis.

.chronus/changes/lm-from-ide-2025-11-15-21-3-28.md

@@ -0,0 +1,7 @@
+---
+changeKind: internal
+packages:
+  - typespec-vscode
+---
+
+Add 'custom/chatComplete' custom event to expose Language Model chatComplete capability to LSP

.vscode/launch.json

@@ -143,6 +143,7 @@
 
         //"ENABLE_SERVER_COMPILE_LOGGING": "true",
         //"ENABLE_UPDATE_MANAGER_LOGGING": "true",
+        //"ENABLE_LM_LOGGING": "true",
 
         "TYPESPEC_SERVER_NODE_OPTIONS": "--nolazy --inspect-brk=4242",
         "TYPESPEC_DEVELOPMENT_MODE": "true"

packages/compiler/src/server/server-compile-manager.ts

@@ -8,6 +8,7 @@ import {
   Program,
   ServerLog,
 } from "../index.js";
+import { getEnvironmentVariable } from "../utils/misc.js";
 import { ENABLE_SERVER_COMPILE_LOGGING } from "./constants.js";
 import { trackActionFunc } from "./server-track-action-task.js";
 import { UpdateManager } from "./update-manager.js";
@@ -45,8 +46,7 @@ export class ServerCompileManager {
     private log: (log: ServerLog) => void,
   ) {
     this.logDebug =
-      typeof process !== "undefined" &&
-      process?.env?.[ENABLE_SERVER_COMPILE_LOGGING]?.toLowerCase() === "true"
+      getEnvironmentVariable(ENABLE_SERVER_COMPILE_LOGGING)?.toLowerCase() === "true"
         ? (msg) => this.log({ level: "debug", message: msg })
         : () => {};
   }

packages/compiler/src/server/server.ts

@@ -153,6 +153,8 @@ function main() {
   documents.onDidClose(profile(s.documentClosed));
   documents.onDidOpen(profile(s.documentOpened));
 
+  (globalThis as any).lspConnection = connection;
+
   documents.listen(connection);
   connection.listen();
 }

packages/compiler/src/server/update-manager.ts

@@ -1,5 +1,6 @@
 import { TextDocumentIdentifier } from "vscode-languageserver";
 import { TextDocument } from "vscode-languageserver-textdocument";
+import { getEnvironmentVariable } from "../utils/misc.js";
 import { ENABLE_UPDATE_MANAGER_LOGGING } from "./constants.js";
 import { ServerLog } from "./types.js";
 
@@ -43,8 +44,7 @@ export class UpdateManager<T = void> {
     getDebounceDelay?: () => number,
   ) {
     this._log =
-      typeof process !== "undefined" &&
-      process?.env?.[ENABLE_UPDATE_MANAGER_LOGGING]?.toLowerCase() === "true"
+      getEnvironmentVariable(ENABLE_UPDATE_MANAGER_LOGGING)?.toLowerCase() === "true"
         ? (sl: ServerLog) => {
             log({ ...sl, message: `#FromUpdateManager(${this.name}): ${sl.message}` });
           }

packages/compiler/src/utils/misc.ts

@@ -488,3 +488,14 @@ class RekeyableMapImpl<K, V> implements RekeyableMap<K, V> {
 export function isPromise(value: unknown): value is Promise<unknown> {
   return !!value && typeof (value as any).then === "function";
 }
+
+export function getEnvironmentVariable(
+  envVarName: string,
+  defaultWhenNotAvailable?: string,
+): string | undefined {
+  // make sure we are fine in both node and browser environments
+  if (typeof process !== "undefined") {
+    return process?.env?.[envVarName] ?? defaultWhenNotAvailable;
+  }
+  return defaultWhenNotAvailable;
+}

packages/typespec-vscode/src/const.ts

@@ -1,3 +1,5 @@
 export const StartFileName = "main.tsp";
 export const TspConfigFileName = "tspconfig.yaml";
 export const EmptyGuid = "00000000-0000-0000-0000-000000000000";
+
+export const ENABLE_LM_LOGGING = "ENABLE_LM_LOGGING";

packages/typespec-vscode/src/lm/language-model.ts

@@ -0,0 +1,122 @@
+import { inspect } from "util";
+import { LanguageModelChat, LanguageModelChatMessage, lm } from "vscode";
+import { ENABLE_LM_LOGGING } from "../const";
+import logger, { LogItem } from "../log/logger";
+import { RetryResult, runWithRetry, runWithTimingLog } from "../utils";
+
+const lmModelCache = new Map<string, Thenable<LanguageModelChat[]>>();
+let lmParallelRequestCount = 0;
+
+export interface LmChatMesage {
+  role: "user" | "assist";
+  message: string;
+}
+
+export interface LmChatRequestOptions {
+  modelOptions?: { [name: string]: any };
+}
+
+export async function sendLmChatRequest(
+  messages: LmChatMesage[],
+  modelFamily: string,
+  options?: LmChatRequestOptions,
+  /** Only for logging purpose */
+  id?: string,
+): Promise<string | undefined> {
+  const logEnabled = process.env[ENABLE_LM_LOGGING] === "true";
+  const lmLog = (item: LogItem) => {
+    if (logEnabled || item.level === "error" || item.level === "warning") {
+      logger.log(
+        item.level,
+        `[ChatComplete][#${id ?? "N/A"}] ${item.message}`,
+        item.details,
+        item.options,
+      );
+    }
+  };
+  if (modelFamily === undefined || modelFamily.trim() === "") {
+    lmLog({ level: "warning", message: `No model family provided for chat completion.` });
+    return undefined;
+  }
+  if (messages === undefined || messages.length === 0) {
+    lmLog({ level: "warning", message: `No messages provided for chat completion.` });
+    return undefined;
+  }
+
+  let models: LanguageModelChat[] | undefined;
+  try {
+    models = await runWithRetry(
+      "Model Selection",
+      async () =>
+        await runWithTimingLog(
+          "Model Selection",
+          async () => {
+            let mp = lmModelCache.get(modelFamily);
+            if (!mp) {
+              mp = lm.selectChatModels({ family: modelFamily });
+              lmModelCache.set(modelFamily, mp);
+            }
+            const m = await mp;
+            if (!m || m.length === 0) {
+              lmModelCache.delete(modelFamily);
+              return RetryResult.Failed;
+            }
+            return m;
+          },
+          lmLog,
+        ),
+      lmLog,
+    );
+  } catch (e) {
+    lmLog({
+      level: "error",
+      message: `Error when selecting chat models for family ${modelFamily}`,
+      details: [e],
+    });
+    return undefined;
+  }
+
+  try {
+    lmParallelRequestCount++;
+    return await runWithTimingLog(
+      "Send request to LM",
+      async () => {
+        lmLog({ level: "debug", message: `LM parallelism increased to ${lmParallelRequestCount}` });
+        const selectedModel = models[0];
+        lmLog({
+          level: "debug",
+          message: `Requested model family: ${modelFamily}, selected: ${selectedModel.family}`,
+        });
+
+        const response = await selectedModel.sendRequest(
+          messages.map((m) => {
+            if (m.role === "user") {
+              return LanguageModelChatMessage.User(m.message);
+            } else if (m.role === "assist") {
+              return LanguageModelChatMessage.Assistant(m.message);
+            } else {
+              logger.debug(`Unknown chat message role: ${m.role}. Default to use User role.`);
+              return LanguageModelChatMessage.User(m.message);
+            }
+          }),
+          options,
+        );
+        let fullResponse = "";
+        for await (const chunk of response.text) {
+          fullResponse += chunk;
+        }
+        return fullResponse;
+      },
+      lmLog,
+    );
+  } catch (e) {
+    lmLog({
+      level: "error",
+      message: `Error when sending chat completion request to model ${models[0].name}. error: ${inspect(e)}`,
+    });
+    return undefined;
+  } finally {
+    lmParallelRequestCount--;
+    lmLog({ level: "debug", message: `LM parallelism reduced to ${lmParallelRequestCount}` });
+  }
+}

packages/typespec-vscode/src/tsp-language-client.ts

@@ -16,9 +16,14 @@ import {
   TextDocumentIdentifier,
 } from "vscode-languageclient/node.js";
 import { TspConfigFileName } from "./const.js";
+import { sendLmChatRequest } from "./lm/language-model.js";
 import logger from "./log/logger.js";
 import telemetryClient from "./telemetry/telemetry-client.js";
 import { resolveTypeSpecServer } from "./tsp-executable-resolver.js";
+import {
+  LspClientCustomRequest_ChatComplete_Name,
+  LspClientCustomRequest_ChatCompletion_Params,
+} from "./types.js";
 import {
   ExecOutput,
   isWhitespaceStringOrUndefined,
@@ -269,6 +274,15 @@ export class TspLanguageClient {
     const name = "TypeSpec";
     const id = "typespec";
     const lc = new LanguageClient(id, name, { run: exe, debug: exe }, options);
+
+    const sendLmChatRequestRequestName: LspClientCustomRequest_ChatComplete_Name =
+      "custom/chatCompletion";
+    lc.onRequest(
+      sendLmChatRequestRequestName,
+      (params: LspClientCustomRequest_ChatCompletion_Params) =>
+        sendLmChatRequest(params.messages, params.modelFamily, params.options, params.id),
+    );
+
     return new TspLanguageClient(lc, exe);
   }
 

packages/typespec-vscode/src/types.ts

@@ -1,3 +1,4 @@
+import { LmChatMesage, LmChatRequestOptions } from "./lm/language-model.js";
 import { TspLanguageClient } from "./tsp-language-client.js";
 import { InitTemplatesUrlSetting } from "./vscode-cmd/create-tsp-project.js";
 
@@ -74,3 +75,12 @@ export interface TypeSpecExtensionApi {
   /** Register more InitTemplateUrls which will be included in the Create TypeSpec Project scenario */
   registerInitTemplateUrls(items: InitTemplatesUrlSetting[]): void;
 }
+
+export type LspClientCustomRequest_ChatComplete_Name = "custom/chatCompletion";
+export interface LspClientCustomRequest_ChatCompletion_Params {
+  messages: LmChatMesage[];
+  modelFamily: string;
+  options?: LmChatRequestOptions;
+  /** Only for logging purpose */
+  id?: string;
+}

packages/typespec-vscode/src/utils.ts

@@ -7,7 +7,7 @@ import vscode, { CancellationToken } from "vscode";
 import { Executable } from "vscode-languageclient/node.js";
 import which from "which";
 import { parseDocument } from "yaml";
-import logger from "./log/logger.js";
+import logger, { LogItem } from "./log/logger.js";
 import { getDirectoryPath, isUrl, joinPaths } from "./path-utils.js";
 import { ResultCode } from "./types.js";
 
@@ -487,3 +487,70 @@ export function distinctArray<T>(arr: T[], compare: (a: T, b: T) => boolean): T[
   }
   return result;
 }
+
+/**
+ * the fn will be wrapped with following log:
+ * operationName started at startTime
+ *  fn()
+ * operationName finished at endTime in duration ms
+ *
+ * @param operationName for logging purpose only
+ */
+export async function runWithTimingLog<T>(
+  operationName: string,
+  fn: () => Promise<T>,
+  log: (item: LogItem) => void,
+): Promise<T> {
+  const start = new Date();
+  try {
+    log({ level: "debug", message: `${operationName} started at ${start.toISOString()}` });
+    return await fn();
+  } finally {
+    const end = new Date();
+    log({
+      level: "debug",
+      message: `${operationName} finished at ${end.toISOString()} in ${end.getTime() - start.getTime()} ms`,
+    });
+  }
+}
+
+export enum RetryResult {
+  Failed,
+}
+/**
+ *
+ * @param fn return RetryResult.Failed or throw to retry
+ * @param operationName for logging purpose only
+ */
+export async function runWithRetry<T>(
+  operationName: string,
+  fn: () => Promise<T | RetryResult.Failed>,
+  log: (item: LogItem) => void,
+  retryCount: number = 3,
+  retryIntervalInMs: number = 200,
+): Promise<T> {
+  for (let i = 0; i < retryCount; i++) {
+    try {
+      const result = await fn();
+      if (result === RetryResult.Failed) {
+        log({
+          level: "debug",
+          message: `${operationName} returned failed on attempt (${i}/${retryCount})`,
+        });
+      } else {
+        return result;
+      }
+    } catch (e) {
+      log({
+        level: "debug",
+        message: `${operationName} threw exception on attempt (${i}/${retryCount}): ${JSON.stringify(e)}`,
+      });
+    }
+    if (i < retryCount - 1) {
+      await new Promise((res) => setTimeout(res, retryIntervalInMs));
+    }
+  }
+  throw new Error(
+    `'${operationName}' failed after ${retryCount} retries, please check previous logs for details`,
+  );
+}