Expose system prompt through ContextRouter interface. Add basic BaseContextRouter

Author: cmsparks

packages/agents/src/mcp/client-prompt.ts

@@ -17,7 +17,7 @@ import type { RequestOptions } from "@modelcontextprotocol/sdk/shared/protocol.j
 import type { AgentsOAuthProvider } from "./do-oauth-client-provider";
 import { jsonSchema, type ToolSet } from "ai";
 import { nanoid } from "nanoid";
-import { unstable_getMcpPrompt } from "./client-prompt";
+import { BaseContextRouter, type ContextRouter } from "./context-router";
 
 /**
  * Utility class that aggregates multiple MCP clients into one
@@ -33,7 +33,8 @@ export class MCPClientManager {
    */
   constructor(
     private _name: string,
-    private _version: string
+    private _version: string,
+    public unstable_contextRouter = new BaseContextRouter(),
   ) {}
 
   /**
@@ -177,42 +178,6 @@ export class MCPClientManager {
     return { serverId };
   }
 
-  /**
-   * @returns namespaced list of tools
-   */
-  listTools(): NamespacedData["tools"] {
-    return getNamespacedData(this.mcpConnections, "tools");
-  }
-
-  /**
-   * @returns a set of tools that you can use with the AI SDK
-   */
-  unstable_getAITools(): ToolSet {
-    return Object.fromEntries(
-      getNamespacedData(this.mcpConnections, "tools").map((tool) => {
-        return [
-          `${tool.serverId}_${tool.name}`,
-          {
-            parameters: jsonSchema(tool.inputSchema),
-            description: tool.description,
-            execute: async (args) => {
-              const result = await this.callTool({
-                name: tool.name,
-                arguments: args,
-                serverId: tool.serverId,
-              });
-              if (result.isError) {
-                // @ts-expect-error TODO we should fix this
-                throw new Error(result.content[0].text);
-              }
-              return result;
-            },
-          },
-        ];
-      })
-    );
-  }
-
   /**
    * Closes all connections to MCP servers
    */
@@ -236,17 +201,31 @@ export class MCPClientManager {
   }
 
   /**
-   * @returns namespaced list of prompts
+   * @returns namespaced list of tools
    */
-  listPrompts(): NamespacedData["prompts"] {
-    return getNamespacedData(this.mcpConnections, "prompts");
+  listTools(): NamespacedData["tools"] {
+    return this.unstable_contextRouter.listTools(this)
+  }
+
+  /**
+   * @returns a set of tools that you can use with the AI SDK
+   */
+  unstable_getAITools(): ToolSet {
+    return this.unstable_contextRouter.getAITools(this)
   }
 
   /**
    * @returns namespaced list of tools
    */
   listResources(): NamespacedData["resources"] {
-    return getNamespacedData(this.mcpConnections, "resources");
+    return this.unstable_contextRouter.listResources(this);
+  }
+
+  /**
+   * @returns namespaced list of prompts
+   */
+  listPrompts(): NamespacedData["prompts"] {
+    return getNamespacedData(this.mcpConnections, "prompts");
   }
 
   /**
@@ -308,12 +287,12 @@ export class MCPClientManager {
    *
    * @param includeResources Whether to include resources in the prompt. This may be useful if you include a tool to fetch resources OR if the servers you connect to interact with resources.
    */
-  unstable_getMcpPrompt(includeResources = true): string {
-    return unstable_getMcpPrompt(this.mcpConnections, includeResources);
+  unstable_systemPrompt(): string {
+    return this.unstable_contextRouter.systemPrompt(this)
   }
 }
 
-type NamespacedData = {
+export type NamespacedData = {
   tools: (Tool & { serverId: string })[];
   prompts: (Prompt & { serverId: string })[];
   resources: (Resource & { serverId: string })[];

packages/agents/src/mcp/client.ts

@@ -0,0 +1,89 @@
+import { jsonSchema, type CoreMessage, type ToolSet, type LanguageModelV1, type UserContent } from "ai";
+import  { getNamespacedData, type MCPClientManager, type NamespacedData } from "./client";
+import type { MCPClientConnection } from "./client-connection";
+import type { Resource } from "@modelcontextprotocol/sdk/types.js"
+
+/**
+ * A context router provides an interface for:
+ * - Filtering tools and resources
+ * - Exposes a system prompt based on the MCP state
+ */
+export interface ContextRouter {
+  systemPrompt(connManager: MCPClientManager): string
+  
+  listTools(connManager: MCPClientManager): NamespacedData["tools"]
+  getAITools(connManager: MCPClientManager): ToolSet
+  listResources(connManager: MCPClientManager): NamespacedData["resources"]
+}
+
+/**
+ * The base ContextRouter:
+ * - Does not filter tools or resources
+ * - Provides a system prompt
+ */
+export class BaseContextRouter implements ContextRouter {
+  constructor(public includeResources = true) {}
+
+  listTools(connManager: MCPClientManager) {
+    return getNamespacedData(connManager.mcpConnections, "tools");
+  }
+
+  getAITools(connManager: MCPClientManager): ToolSet {
+    return Object.fromEntries(
+      this.listTools(connManager).map((tool) => {
+        return [
+          `${tool.serverId}_${tool.name}`,
+          {
+            parameters: jsonSchema(tool.inputSchema),
+            description: tool.description,
+            execute: async (args) => {
+              const result = await connManager.callTool({
+                name: tool.name,
+                arguments: args,
+                serverId: tool.serverId,
+              });
+              if (result.isError) {
+                // @ts-expect-error TODO we should fix this
+                throw new Error(result.content[0].text);
+              }
+              return result;
+            },
+          },
+        ];
+      })
+    )
+  }
+
+  listResources(connManager: MCPClientManager) {
+    return getNamespacedData(connManager.mcpConnections, "resources");
+  }
+
+  systemPrompt(connManager: MCPClientManager): string {
+    return `<integrations_list>
+  You have access to multiple integrations via Model Context Protocol (MCP). These integrations provide you with tools which you can use to execute to complete tasks or retrieive information.
+
+  ${this.includeResources && "Each integration, provides a list of resources, which are included in the list of integrations below."}
+
+  Here is a list of all of the integrations you have access to, with instructions if necessary:
+  
+  ${Object.entries(connManager.mcpConnections).map(([_id, conn]) => BaseContextRouter.serverContext(conn, this.includeResources))}
+<integrations_list>`
+  }
+
+  static serverContext(conn: MCPClientConnection, includeResources: boolean) {
+    return `<integration>
+  ${conn.serverInfo && `<integration_name>${conn.serverInfo.name}</integration_name>`}
+  ${conn.instructions && `<integration_instructions>${conn.instructions}</integration_instructions>`}
+  ${includeResources && `<resources_list>${conn.resources.map((resource) => BaseContextRouter.resourceContext(resource))}</resources_list>`}
+<integration>`;
+  }
+  
+  static resourceContext(resource: Resource) {
+    return `<resource>
+  <name>${resource.name}</name>
+  <uri>${resource.uri}</uri>
+  <description>${resource.description}</description>
+  <mimeType>${resource.mimeType}</mimeType>
+</resource>`;
+  }
+}