mcp sampling (#1722)

* Add support for MCP model provider and enhance debugging capabilities

* Add MCP Client Sampling configuration to language model providers

* Add model specification to emojifier script

* Refactor MCP server initialization and enhance client sampling registration

* Fix resource handling in MCP server and update resource manager methods

* Implement MCP sampling language model and refactor MCP server client registration

* Add parent language model support to MCP server and worker

* Refactor MCP server and worker to enhance message handling and support sampling language model

* Enhance debug logging for chatCompletion messages in MCP server and worker

* Refactor message handling in createWorkerLanguageModel for improved clarity and maintainability

Author: pelikhan

.vscode/settings.json

@@ -71,6 +71,7 @@
         "dbgi",
         "dbgp",
         "dbgr",
+        "dbgs",
         "dbgt",
         "ddir",
         "debugify",
@@ -235,6 +236,7 @@
         "oldsrc",
         "ollama",
         "olmo",
+        "oninitialized",
         "onnx",
         "onvsc",
         "openai",

packages/api/src/api.ts

@@ -55,11 +55,19 @@ export async function run(
     /**
      * Handles messages
      */
-    onMessage?: (data: { type: "resourceChange" } & Resource) => Awaitable<void>;
+    onMessage?: (
+      data: { type: "resourceChange" } & Resource,
+      postMessage: (data: any) => void,
+    ) => Awaitable<void>;
+    /**
+     * Enable client language model as parent.
+     */
+    parentLanguageModel?: boolean;
   },
 ): Promise<GenerationResult> {
   if (!scriptId) throw new Error("scriptId is required");
   dbg(`run ${scriptId}`);
+  // eslint-disable-next-line no-param-reassign
   if (typeof files === "string") files = [files];
 
   const { signal, onMessage, ...rest } = options || {};
@@ -77,7 +85,7 @@ export async function run(
   dbg(`start ${workerJs}`);
   const worker = new Worker(workerJs, { workerData, name: options?.label });
   return new Promise((resolve, reject) => {
-    const abort = () => {
+    const abort = (): void => {
       if (worker) {
         dbg(`abort`);
         reject(new Error("aborted")); // fail early
@@ -92,12 +100,15 @@ export async function run(
         signal?.removeEventListener("abort", abort);
         resolve(res.result);
       } else if (onMessage) {
-        await onMessage(res);
+        await onMessage(res, (data) => {
+          dbg(`postMessage %O`, data);
+          worker.postMessage(data);
+        });
       } else {
         dbg(`unknown message type ${type}`);
       }
     });
-    worker.on("error", (reason) => {
+    worker.on("error", (reason: string) => {
       dbg(`error ${reason}`);
       signal?.removeEventListener("abort", abort);
       reject(reason);

packages/api/src/worker.ts

@@ -6,12 +6,15 @@ import { delay } from "es-toolkit";
 import { NodeHost } from "@genaiscript/runtime";
 import {
   RESOURCE_CHANGE,
+  genaiscriptDebug,
   installGlobals,
   overrideStdoutWithStdErr,
   runtimeHost,
+  createWorkerLanguageModel,
 } from "@genaiscript/core";
 import type { Resource } from "@genaiscript/core";
 import { runScriptInternal } from "./run.js";
+const dbg = genaiscriptDebug("worker");
 
 /**
  * Handles worker thread execution based on the provided data type.
@@ -28,12 +31,13 @@ import { runScriptInternal } from "./run.js";
  *     - Handles resource change events and communicates them to the parent thread.
  *     - Ensures compatibility with Windows by setting the SystemRoot environment variable.
  */
-export async function worker() {
+export async function worker(): Promise<void> {
   overrideStdoutWithStdErr();
   installGlobals();
   const { type, ...data } = workerData as {
     type: string;
   };
+  dbg(`worker data: %O`, data);
   await NodeHost.install(undefined, undefined); // Install NodeHost with environment options
   if (process.platform === "win32") {
     // https://github.com/Azure/azure-sdk-for-js/issues/32374
@@ -55,9 +59,13 @@ export async function worker() {
       const { scriptId, files, options } = data as {
         scriptId: string;
         files: string[];
-        options: object;
+        options: { parentLanguageModel?: boolean };
       };
-      const { result } = await runScriptInternal(scriptId, files, options);
+      if (options.parentLanguageModel) {
+        dbg(`using parent language model`);
+        runtimeHost.clientLanguageModel = createWorkerLanguageModel();
+      }
+      const { result } = await runScriptInternal(scriptId, files, options as any);
       await delay(0); // flush streams
       parentPort.postMessage({ type: "run", result });
       break;

packages/cli/src/mcpserver.ts

@@ -9,12 +9,14 @@ import {
   deleteUndefinedValues,
   ensureDotGenaiscriptPath,
   errorMessage,
+  genaiscriptDebug,
   logVerbose,
   logWarn,
   runtimeHost,
   setConsoleColors,
   splitMarkdownTextImageParts,
   toStrictJSONSchema,
+  mcpRequestSample,
 } from "@genaiscript/core";
 import type {
   GenerationResult,
@@ -39,8 +41,8 @@ import type {
 import { applyRemoteOptions } from "./remote.js";
 import type { RemoteOptions } from "./remote.js";
 import { startProjectWatcher } from "./watch.js";
-import debug from "debug";
-const dbg = debug("genaiscript:mcp:server");
+import { workerData } from "worker_threads";
+const dbg = genaiscriptDebug("mcp:server");
 
 /**
  * Starts the MCP server.
@@ -57,13 +59,14 @@ export async function startMcpServer(
     RemoteOptions & {
       startup?: string;
     },
-) {
+): Promise<void> {
   setConsoleColors(false);
   logVerbose(`mcp server: starting...`);
 
   await ensureDotGenaiscriptPath();
   await applyRemoteOptions(options);
   const { startup } = options || {};
+  let samplingSupported = false;
 
   const watcher = await startProjectWatcher(options);
   logVerbose(`mcp server: watching ${watcher.cwd}`);
@@ -89,10 +92,30 @@ export async function startMcpServer(
       },
     },
   );
-  watcher.addEventListener("change", async () => {
-    logVerbose(`mcp server: tools changed`);
-    await server.sendToolListChanged();
-  });
+  watcher.addEventListener(
+    "change",
+    async () => {
+      logVerbose(`mcp server: tools changed`);
+      await server.sendToolListChanged();
+    },
+    false,
+  );
+  const onMessage = async (data: any, postMessage: (data: any) => void) => {
+    if (data.type === RESOURCE_CHANGE) {
+      await runtimeHost.resources.upsertResource(data.reference, data.content);
+    } else if (data.type === "chatCompletion") {
+      if (!samplingSupported) throw new Error("Sampling not supported by client");
+      // Handle chat completion messages if needed
+      dbg(`chatCompletion message received: %O`, data);
+      const { request, ...rest } = data;
+      const response = await mcpRequestSample(server, data.request);
+      const msg = { ...rest, response };
+      dbg(`chatCompletion response: %O`, msg);
+      postMessage(msg);
+    } else {
+      dbg(`unknown message type: ${data.type}`);
+    }
+  };
   server.setRequestHandler(ListToolsRequestSchema, async () => {
     dbg(`fetching scripts from watcher`);
     const scripts = await watcher.scripts();
@@ -112,14 +135,15 @@ export async function startMcpServer(
           properties: {},
         };
         const outputSchema = responseSchema ? toStrictJSONSchema(responseSchema) : undefined;
-        if (accept !== "none")
+        if (accept !== "none") {
           scriptSchema.properties.files = {
             type: "array",
             items: {
               type: "string",
               description: `Filename or globs relative to the workspace used by the script.${accept ? ` Accepts: ${accept}` : ""}`,
             },
           };
+        }
         if (!description) logWarn(`script ${id} has no description`);
         return deleteUndefinedValues({
           name: id,
@@ -146,14 +170,16 @@ export async function startMcpServer(
         vars: vars as Record<string, string | number | boolean | object>,
         runTrace: false,
         outputTrace: false,
+        parentLanguageModel: samplingSupported,
+        onMessage,
       })) || { status: "error", error: { message: "run failed" } };
       dbg(`res: %s`, res.status);
       if (res.error) dbg(`error: %O`, res.error);
       const isError = res.status !== "success" || !!res.error;
       const text = res?.error?.message || (res.json ? JSON.stringify(res.json) : res.text) || "";
       dbg(`inlining images`);
       const parts = await splitMarkdownTextImageParts(text, {
-        dir: res.env.runDir,
+        dir: res.env?.runDir,
         convertToDataUri: true,
       });
       dbg(`parts: %O`, parts);
@@ -193,31 +219,42 @@ export async function startMcpServer(
     if (!resource) dbg(`resource not found: ${uri}`);
     return resource as ReadResourceResult;
   });
-  runtimeHost.resources.addEventListener(CHANGE, async () => {
-    await server.sendResourceListChanged();
-  });
-  runtimeHost.resources.addEventListener(RESOURCE_CHANGE, async (e) => {
-    const ev = e as CustomEvent<Resource>;
-    await server.sendResourceUpdated({
-      uri: ev.detail.reference.uri,
-    });
-  });
+  runtimeHost.resources.addEventListener(
+    CHANGE,
+    async () => {
+      await server.sendResourceListChanged();
+    },
+    false,
+  );
+  runtimeHost.resources.addEventListener(
+    RESOURCE_CHANGE,
+    async (e) => {
+      const ev = e as CustomEvent<Resource>;
+      await server.sendResourceUpdated({
+        uri: ev.detail.reference.uri,
+      });
+    },
+    false,
+  );
+
+  server.oninitialized = async () => {
+    dbg(`server/client connection initialized`);
+    // Check if client supports sampling
+    const clientCapabilities = server.getClientCapabilities();
+    dbg(`client capabilities: %O`, clientCapabilities);
+    samplingSupported = !!clientCapabilities?.sampling;
+
+    if (startup) {
+      logVerbose(`startup script: ${startup}`);
+      await run(startup, [], {
+        vars: {},
+        parentLanguageModel: samplingSupported,
+        onMessage,
+      });
+    }
+  };
 
   const transport = new StdioServerTransport();
   dbg(`connecting server with transport`);
   await server.connect(transport);
-
-  if (startup) {
-    logVerbose(`startup script: ${startup}`);
-    await run(startup, [], {
-      vars: {},
-      onMessage: async (data) => {
-        if (data.type === RESOURCE_CHANGE) {
-          await runtimeHost.resources.upsetResource(data.reference, data.content);
-        } else {
-          dbg(`unknown message type: ${data.type}`);
-        }
-      },
-    });
-  }
 }

packages/cli/src/watch.ts

@@ -1,7 +1,8 @@
 // Copyright (c) Microsoft Corporation.
 // Licensed under the MIT License.
 
-import { FSWatcher, watch } from "chokidar";
+import { watch } from "chokidar";
+import type { FSWatcher } from "chokidar";
 import { basename, resolve } from "node:path";
 import {
   CHANGE,
@@ -40,11 +41,11 @@ export class ProjectWatcher extends EventTarget {
     signal?.addEventListener("abort", this.close.bind(this));
   }
 
-  get cwd() {
+  get cwd(): string {
     return this.options.cwd;
   }
 
-  async open() {
+  async open(): Promise<void> {
     if (this._watcher) return;
 
     dbg(`starting`);
@@ -77,7 +78,7 @@ export class ProjectWatcher extends EventTarget {
       depth: 30,
       cwd,
     });
-    const changed = () => {
+    const changed = (): void => {
       dbg(`changed`);
       this.dispatchEvent(new Event(CHANGE));
     };
@@ -90,27 +91,27 @@ export class ProjectWatcher extends EventTarget {
     this.dispatchEvent(new Event(OPEN));
   }
 
-  private async refresh() {
+  private async refresh(): Promise<void> {
     this._project = undefined;
   }
 
-  async project() {
+  async project(): Promise<Project> {
     if (!this._project) {
       dbg(`building project`);
       this._project = await buildProject();
     }
     return this._project;
   }
 
-  async scripts() {
+  async scripts(): Promise<PromptScript[]> {
     if (!this._scripts) {
       const project = await this.project();
       this._scripts = filterScripts(project.scripts, this.options);
     }
     return this._scripts?.slice(0);
   }
 
-  async close() {
+  async close(): Promise<void> {
     dbg(`closing`);
     await this._watcher?.close();
     this._watcher = undefined;
@@ -133,7 +134,7 @@ export async function startProjectWatcher(
     paths?: ElementOrArray<string>;
     cwd?: string;
   } & CancellationOptions,
-) {
+): Promise<ProjectWatcher> {
   const { paths = ".", cwd = resolve("."), ...rest } = options || {};
   const watcher = new ProjectWatcher({ paths, cwd, ...rest });
   await watcher.open();

packages/core/src/ast.ts

@@ -85,7 +85,7 @@ export function collectFolders(
  * @returns The script with the matching ID, or undefined if no match is found.
  */
 export function resolveScript(prj: Project, system: SystemPromptInstance) {
-  return prj?.scripts?.find((t) => t.id == system.id); // Find and return the template with the matching ID
+  return prj?.scripts?.find((t) => t.id === system.id); // Find and return the template with the matching ID
 }
 
 export interface ScriptFilterOptions {

packages/core/src/constants.ts

@@ -205,6 +205,7 @@ export const MODEL_PROVIDER_WINDOWS_AI = "windows";
 export const MODEL_PROVIDER_DOCKER_MODEL_RUNNER = "docker";
 export const MODEL_PROVIDER_ECHO = "echo";
 export const MODEL_PROVIDER_NONE = "none";
+export const MODEL_PROVIDER_MCP = "mcp";
 
 export const MODEL_GITHUB_COPILOT_CHAT_CURRENT = MODEL_PROVIDER_GITHUB_COPILOT_CHAT + ":current";