Add MCP server registry and support server references in tasks

examples/cli/src/commands/mcp.ts

@@ -5,6 +5,7 @@
  */
 
 import type { DataPortSchemaObject } from "@workglow/util";
+import type { McpServerRecord } from "@workglow/tasks";
 import type { Command } from "commander";
 import { loadConfig } from "../config";
 import {
@@ -13,7 +14,8 @@ import {
   resolveInput,
   validateInput,
 } from "../input";
-import { createMcpStorage, McpServerRecordSchema } from "../storage";
+import { McpServerRecordSchema } from "@workglow/tasks";
+import { createMcpServerRepository } from "../storage";
 import { formatTable } from "../util";
 import type { SearchPage, SearchSelectItem } from "../ui/render";
 
@@ -100,19 +102,22 @@ async function searchMcpRegistry(
 }
 
 function mapMcpRegistryResult(server: McpRegistryServer): Record<string, unknown> {
-  const name = server.name.split("/").pop() ?? server.name;
+  const serverId = server.name.split("/").pop() ?? server.name;
+  const title = server.title ?? serverId;
 
   if (server.remotes && server.remotes.length > 0) {
     const remote = server.remotes[0];
     return {
-      name,
+      server_id: serverId,
+      title,
+      description: server.description,
       transport: remote.type,
       server_url: remote.url,
     };
   }
 
   const pkg = server.packages?.[0];
-  if (!pkg) return { name };
+  if (!pkg) return { server_id: serverId, title, description: server.description };
 
   let command: string;
   let args: string[];
@@ -143,7 +148,9 @@ function mapMcpRegistryResult(server: McpRegistryServer): Record<string, unknown
   }
 
   const result: Record<string, unknown> = {
-    name,
+    server_id: serverId,
+    title,
+    description: server.description,
     transport: "stdio",
     command,
     args,
@@ -168,55 +175,56 @@ export function registerMcpCommand(program: Command): void {
     .description("List all configured MCP servers")
     .action(async () => {
       const config = await loadConfig();
-      const storage = createMcpStorage(config);
-      await storage.setupDirectory();
+      const repo = createMcpServerRepository(config);
+      await repo.setupDatabase();
 
-      const all = await storage.getAll();
+      const all = await repo.enumerateAllServers();
       if (!all || all.length === 0) {
         console.log("No MCP servers found.");
         return;
       }
 
       const rows = all.map((entry) => ({
-        name: String(entry.name ?? ""),
+        server_id: String(entry.server_id ?? ""),
+        title: String(entry.title ?? ""),
         transport: String(entry.transport ?? ""),
         server_url: String(entry.server_url ?? ""),
         command: String(entry.command ?? ""),
       }));
-      console.log(formatTable(rows, ["name", "transport", "server_url", "command"]));
+      console.log(formatTable(rows, ["server_id", "title", "transport", "server_url", "command"]));
     });
 
   mcp
     .command("detail")
-    .argument("[id]", "MCP server name to show")
+    .argument("[id]", "MCP server id to show")
     .description("Show full details of an MCP server")
     .action(async (id: string | undefined) => {
       const config = await loadConfig();
-      const storage = createMcpStorage(config);
-      await storage.setupDirectory();
+      const repo = createMcpServerRepository(config);
+      await repo.setupDatabase();
 
       let targetId = id;
       if (!targetId) {
         if (!process.stdin.isTTY) {
           console.error("Error: specify an id or run interactively.");
           process.exit(1);
         }
-        const all = await storage.getAll();
+        const all = await repo.enumerateAllServers();
         if (!all || all.length === 0) {
           console.log("No MCP servers found.");
           return;
         }
         const { renderSelectPrompt } = await import("../ui/render");
         const options = all.map((e) => ({
-          label: `${e.name}  ${e.transport ?? ""}  ${e.server_url ?? e.command ?? ""}`,
-          value: String(e.name),
+          label: `${e.server_id}  ${e.transport ?? ""}  ${e.server_url ?? e.command ?? ""}`,
+          value: String(e.server_id),
         }));
         const selected = await renderSelectPrompt(options, "Select MCP server:");
         if (!selected) return;
         targetId = selected;
       }
 
-      const entry = await storage.get({ name: targetId });
+      const entry = await repo.findByName(targetId);
       if (!entry) {
         console.error(`MCP server "${targetId}" not found.`);
         process.exit(1);
@@ -227,35 +235,35 @@ export function registerMcpCommand(program: Command): void {
 
   mcp
     .command("remove")
-    .argument("[id]", "MCP server name to remove")
-    .description("Remove an MCP server by name")
+    .argument("[id]", "MCP server id to remove")
+    .description("Remove an MCP server by id")
     .action(async (id: string | undefined) => {
       const config = await loadConfig();
-      const storage = createMcpStorage(config);
-      await storage.setupDirectory();
+      const repo = createMcpServerRepository(config);
+      await repo.setupDatabase();
 
       let targetId = id;
       if (!targetId) {
         if (!process.stdin.isTTY) {
           console.error("Error: specify an id or run interactively.");
           process.exit(1);
         }
-        const all = await storage.getAll();
+        const all = await repo.enumerateAllServers();
         if (!all || all.length === 0) {
           console.log("No MCP servers to remove.");
           return;
         }
         const { renderSelectPrompt } = await import("../ui/render");
         const options = all.map((e) => ({
-          label: `${e.name}  ${e.transport ?? ""}  ${e.server_url ?? e.command ?? ""}`,
-          value: String(e.name),
+          label: `${e.server_id}  ${e.transport ?? ""}  ${e.server_url ?? e.command ?? ""}`,
+          value: String(e.server_id),
         }));
         const selected = await renderSelectPrompt(options, "Select MCP server to remove:");
         if (!selected) return;
         targetId = selected;
       }
 
-      await storage.delete({ name: targetId });
+      await repo.removeServer(targetId);
       console.log(`MCP server "${targetId}" removed.`);
     });
 
@@ -316,11 +324,11 @@ export function registerMcpCommand(program: Command): void {
       }
 
       const config = await loadConfig();
-      const storage = createMcpStorage(config);
-      await storage.setupDirectory();
+      const repo = createMcpServerRepository(config);
+      await repo.setupDatabase();
 
-      await storage.put(input as Record<string, unknown>);
-      console.log(`MCP server "${input.name}" added.`);
+      await repo.addServer(input as McpServerRecord);
+      console.log(`MCP server "${input.server_id}" added.`);
     });
 
   mcp
@@ -375,10 +383,10 @@ export function registerMcpCommand(program: Command): void {
       }
 
       const config = await loadConfig();
-      const storage = createMcpStorage(config);
-      await storage.setupDirectory();
+      const repo = createMcpServerRepository(config);
+      await repo.setupDatabase();
 
-      await storage.put(input as Record<string, unknown>);
-      console.log(`MCP server "${input.name}" added.`);
+      await repo.addServer(input as McpServerRecord);
+      console.log(`MCP server "${input.server_id}" added.`);
     });
 }

examples/cli/src/lib.ts

@@ -9,5 +9,5 @@ export {
   createModelRepository,
   createWorkflowRepository,
   createAgentRepository,
-  createMcpStorage,
+  createMcpServerRepository,
 } from "./storage";

examples/cli/src/storage.ts

@@ -11,34 +11,13 @@ import {
   TaskGraphSchema,
   TaskGraphPrimaryKeyNames,
 } from "@workglow/task-graph";
+import {
+  McpServerRepository,
+  McpServerRecordSchema,
+  McpServerPrimaryKeyNames,
+} from "@workglow/tasks";
 import type { CliConfig } from "./config";
 
-export const McpServerRecordSchema = {
-  type: "object",
-  properties: {
-    name: { type: "string", "x-auto-generated": false },
-    transport: { type: "string", enum: ["stdio", "sse", "streamable-http"] },
-    server_url: { type: "string" },
-    command: { type: "string" },
-    args: { type: "array", items: { type: "string" } },
-    env: { type: "object", additionalProperties: { type: "string" } },
-    auth_type: { type: "string" },
-    auth_token: { type: "string" },
-    auth_client_id: { type: "string" },
-    auth_client_secret: { type: "string" },
-    auth_private_key: { type: "string" },
-    auth_algorithm: { type: "string" },
-    auth_jwt_bearer_assertion: { type: "string" },
-    auth_redirect_url: { type: "string" },
-    auth_scope: { type: "string" },
-    auth_client_name: { type: "string" },
-    auth_jwt_lifetime_seconds: { type: "number" },
-  },
-  required: ["name", "transport"],
-} as const;
-
-export const McpServerPrimaryKeyNames = ["name"] as const;
-
 export function createModelRepository(config: CliConfig): ModelRepository {
   const storage = new FsFolderTabularStorage(
     config.directories.models,
@@ -68,13 +47,12 @@ export function createAgentRepository(config: CliConfig): TaskGraphTabularReposi
   });
 }
 
-export function createMcpStorage(config: CliConfig): FsFolderTabularStorage<
-  typeof McpServerRecordSchema,
-  typeof McpServerPrimaryKeyNames
-> {
-  return new FsFolderTabularStorage(
-    config.directories.mcps,
-    McpServerRecordSchema,
-    McpServerPrimaryKeyNames
+export function createMcpServerRepository(config: CliConfig): McpServerRepository {
+  return new McpServerRepository(
+    new FsFolderTabularStorage(
+      config.directories.mcps,
+      McpServerRecordSchema,
+      McpServerPrimaryKeyNames
+    )
   );
 }

packages/task-graph/src/task/TaskRunner.ts

@@ -133,6 +133,17 @@ export class TaskRunner<
     await this.handleStart(config);
 
     try {
+      // Resolve schema-annotated config properties (e.g., mcp-server references)
+      const configSchemaResult = (this.task.constructor as typeof Task).configSchema();
+      if (configSchemaResult) {
+        const resolvedConfig = await resolveSchemaInputs(
+          this.task.config as Record<string, unknown>,
+          configSchemaResult,
+          { registry: this.registry }
+        );
+        Object.assign(this.task.config, resolvedConfig);
+      }
+
       this.task.setInput(overrides);
 
       // Resolve schema-annotated inputs (models, repositories) before validation
@@ -185,6 +196,16 @@ export class TaskRunner<
         this.task.runOutputData = outputs ?? ({} as Output);
       }
 
+      // Resolve schema-annotated output properties
+      const outSchema = this.task.outputSchema();
+      if (outSchema && this.task.runOutputData) {
+        this.task.runOutputData = (await resolveSchemaInputs(
+          this.task.runOutputData as Record<string, unknown>,
+          outSchema,
+          { registry: this.registry }
+        )) as Output;
+      }
+
       await this.handleComplete();
 
       return this.task.runOutputData as Output;
@@ -205,6 +226,17 @@ export class TaskRunner<
     if (this.task.status === TaskStatus.PROCESSING) {
       return this.task.runOutputData as Output;
     }
+    // Resolve schema-annotated config properties (e.g., mcp-server references)
+    const configSchemaResult = (this.task.constructor as typeof Task).configSchema();
+    if (configSchemaResult) {
+      const resolvedConfig = await resolveSchemaInputs(
+        this.task.config as Record<string, unknown>,
+        configSchemaResult,
+        { registry: this.registry }
+      );
+      Object.assign(this.task.config, resolvedConfig);
+    }
+
     this.task.setInput(overrides);
 
     // Resolve schema-annotated inputs (models, repositories) before validation
@@ -230,6 +262,16 @@ export class TaskRunner<
 
       this.task.runOutputData = resultReactive;
 
+      // Resolve schema-annotated output properties
+      const outSchema = this.task.outputSchema();
+      if (outSchema && this.task.runOutputData) {
+        this.task.runOutputData = (await resolveSchemaInputs(
+          this.task.runOutputData as Record<string, unknown>,
+          outSchema,
+          { registry: this.registry }
+        )) as Output;
+      }
+
       await this.handleCompleteReactive();
     } catch (err: any) {
       await this.handleErrorReactive();

packages/tasks/src/common.ts

@@ -27,6 +27,10 @@ export * from "./task/mcp/McpListTask";
 export * from "./task/mcp/McpPromptGetTask";
 export * from "./task/mcp/McpResourceReadTask";
 export * from "./task/mcp/McpToolCallTask";
+export * from "./mcp-server/McpServerSchema";
+export * from "./mcp-server/McpServerRepository";
+export * from "./mcp-server/InMemoryMcpServerRepository";
+export * from "./mcp-server/McpServerRegistry";
 export * from "./task/string/StringConcatTask";
 export * from "./task/string/StringIncludesTask";
 export * from "./task/string/StringJoinTask";

packages/tasks/src/mcp-server/InMemoryMcpServerRepository.ts

@@ -0,0 +1,19 @@
+/**
+ * @license
+ * Copyright 2025 Steven Roussey <sroussey@gmail.com>
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import { InMemoryTabularStorage } from "@workglow/storage";
+import { McpServerRepository } from "./McpServerRepository";
+import { McpServerPrimaryKeyNames, McpServerRecordSchema } from "./McpServerSchema";
+
+/**
+ * In-memory implementation of an MCP server repository.
+ * Provides storage and retrieval for MCP server configurations.
+ */
+export class InMemoryMcpServerRepository extends McpServerRepository {
+  constructor() {
+    super(new InMemoryTabularStorage(McpServerRecordSchema, McpServerPrimaryKeyNames));
+  }
+}