Add support for in-memory MCP servers

- Add InMemoryTransport support for local servers running in the same process - Update MCPConnectionManager to pass full connection object in callbacks for better context - Filter out local servers from persistence (they auto-connect on startup) - Add 'In-Memory' badge in UI to distinguish local servers - Fix connection state management to properly update connectionsRef - Auto-connect to available in-memory servers on startup This enables testing MCP functionality without external servers by running servers directly in the browser using the InMemoryTransport from the SDK.
2025-06-19 16:55:47 +01:00 · 2025-06-19 16:55:47 +01:00 · e53be9ec42
commit e53be9ec42
parent c9feacd763
6 changed files with 136 additions and 28 deletions
--- a/src/components/MCPTab.tsx
+++ b/src/components/MCPTab.tsx
@ -222,6 +222,11 @@ export function MCPTab() {
                      <h4 className="text-sm font-medium text-gray-900 dark:text-white truncate">
                        {connection.name}
                      </h4>
+                      {connection.url === 'local' && (
+                        <span className="px-1.5 py-0.5 text-xs bg-purple-100 dark:bg-purple-900/50 text-purple-800 dark:text-purple-200 rounded">
+                          In-Memory
+                        </span>
+                      )}
                    </div>
                    <p className={`text-xs mt-1 ${getStatusColor(connection.status)}`}>
                      {connection.status}
--- a/src/contexts/MCPContext.tsx
+++ b/src/contexts/MCPContext.tsx
@ -13,6 +13,7 @@ import type {
 } from '@/types/mcp';
 import type { Tool } from '@/types/inference';
 import { MCPConnectionManager } from '@/mcp/connection';
+import { availableServers } from '@/mcp/servers';

 const MCPContext = createContext<MCPContextValue | null>(null);

@ -56,24 +57,18 @@ export function MCPProvider({ children }: MCPProviderProps) {

  // Function to broadcast messages to all callbacks (stable reference)
  const broadcastMessage = useCallback((
-    connectionId: string, 
+    connection: MCPConnection, 
    client: any, 
    message: any, 
    direction: 'sent' | 'received', 
    extra?: any
  ) => {
-    // console.log(`MCP Message [${direction}] from ${connectionId}:`, message); // Debug log
-    
-    // Find connection name using ref to get latest state
-    const connection = connectionsRef.current.find(conn => conn.id === connectionId);
-    const serverName = connection?.name || 'Unknown Server';
-    
    // Store message
    const storedMessage: MCPMonitorMessage = {
      id: uuidv4(),
      timestamp: new Date(),
-      connectionId,
-      serverName,
+      connectionId: connection.id,
+      serverName: connection?.name || 'Unknown Server',
      message,
      direction,
      extra,
@ -82,7 +77,7 @@ export function MCPProvider({ children }: MCPProviderProps) {
    
    messageCallbacksRef.current.forEach(callback => {
      try {
-        callback(connectionId, client, message, direction, extra);
+        callback(connection.id, client, message, direction, extra);
      } catch (error) {
        console.error('Error in MCP message callback:', error);
      }
@ -118,11 +113,12 @@ export function MCPProvider({ children }: MCPProviderProps) {
              
              // Set up callback for connection state updates
              manager.setConnectionUpdateCallback(() => {
-                setConnections(prev => 
-                  prev.map(conn => 
+                setConnections(prev => {
+                  connectionsRef.current = prev.map(conn => 
                    conn.id === connectionId ? manager.getConnection() : conn
                  )
-                );
+                  return connectionsRef.current
+                });
              });
              
              // Set up message callback
@ -149,6 +145,50 @@ export function MCPProvider({ children }: MCPProviderProps) {
      } catch (error) {
        console.error('Failed to load persisted MCP connections:', error);
      }
+      
+      // After loading persisted connections, add local servers
+      await addLocalServers();
+    };
+
+    const addLocalServers = async () => {
+      // Check if local servers are already added
+      const hasLocalServers = connectionsRef.current.some(conn => conn.url === 'local');
+      if (hasLocalServers) {
+        return;
+      }
+
+      // Add each available local server
+      for (const serverConfig of availableServers) {
+        try {
+          const connectionId = uuidv4();
+          const manager = new MCPConnectionManager(connectionId, serverConfig);
+          
+          // Set up callback for connection state updates
+          manager.setConnectionUpdateCallback(() => {
+            setConnections(prev => {
+               connectionsRef.current = prev.map(conn => 
+                  conn.id === connectionId ? manager.getConnection() : conn
+                );
+               return connectionsRef.current;
+            });
+          });
+          
+          // Set up message callback
+          manager.setMessageCallback(broadcastMessage);
+          
+          // Add to managers map
+          setManagers(prev => new Map(prev).set(connectionId, manager));
+          
+          // Add initial connection state
+          setConnections(prev => [...prev, manager.getConnection()]);
+          
+          // Auto-connect local servers
+          await manager.connect();
+          console.log(`Connected to local server: ${serverConfig.name}`);
+        } catch (error) {
+          console.error(`Failed to connect to local server ${serverConfig.name}:`, error);
+        }
+      }
    };

    loadPersistedConnections();
@ -157,7 +197,7 @@ export function MCPProvider({ children }: MCPProviderProps) {
  const persistConnections = useCallback(() => {
    try {
      // Store both config and connection ID to maintain OAuth token association
-      const connectionData = connections.map(conn => ({
+      const connectionData = connections.filter(conn => conn.url !== 'local').map(conn => ({
        id: conn.id,
        config: conn.config
      }));
--- a/src/mcp/connection.ts
+++ b/src/mcp/connection.ts
@ -3,6 +3,7 @@
 import { Client } from '@modelcontextprotocol/sdk/client/index.js';
 import { SSEClientTransport } from '@modelcontextprotocol/sdk/client/sse.js';
 import { StreamableHTTPClientTransport } from '@modelcontextprotocol/sdk/client/streamableHttp.js';
+import { InMemoryTransport } from '@modelcontextprotocol/sdk/inMemory.js';
 import type { Transport } from '@modelcontextprotocol/sdk/shared/transport.js';
 import {
  auth,
@ -235,7 +236,7 @@ export class MCPConnectionManager {
  private healthCheckInterval?: NodeJS.Timeout;
  private oauthProvider?: MCPOAuthProvider;
  private onConnectionUpdate?: () => void;
-  private onMessage?: (connectionId: string, client: any, message: any, direction: 'sent' | 'received', extra?: any) => void;
+  private onMessage?: (connection: MCPConnection, client: any, message: any, direction: 'sent' | 'received', extra?: any) => void;

  constructor(id: string, config: MCPServerConfig) {
    this.connection = {
@ -265,7 +266,7 @@ export class MCPConnectionManager {
  }

  // Set callback for message monitoring
-  setMessageCallback(callback: (connectionId: string, client: any, message: any, direction: 'sent' | 'received', extra?: any) => void): void {
+  setMessageCallback(callback: (connection: MCPConnection, client: any, message: any, direction: 'sent' | 'received', extra?: any) => void): void {
    this.onMessage = callback;
  }

@ -286,14 +287,20 @@ export class MCPConnectionManager {
      // Clear any existing connections
      await this.disconnect();

-      try {
-        await this.tryStreamableHttp();
-        this.connection.transport = 'streamable-http';
-      } catch (error) {
-        // TODO: jerome - if this is a TypeError: failed to fetch, then there is likely a CORS (or 
-        // Access-Control-Expose-Headers) issue with the server.
-        await this.trySSE();
-        this.connection.transport = 'sse';
+      // Check if this is a local server
+      if (this.connection.url === 'local') {
+        await this.tryInMemory();
+        this.connection.transport = 'inmemory';
+      } else {
+        try {
+          await this.tryStreamableHttp();
+          this.connection.transport = 'streamable-http';
+        } catch (error) {
+          // TODO: jerome - if this is a TypeError: failed to fetch, then there is likely a CORS (or 
+          // Access-Control-Expose-Headers) issue with the server.
+          await this.trySSE();
+          this.connection.transport = 'sse';
+        }
      }

      // Initialize client capabilities
@ -435,6 +442,27 @@ export class MCPConnectionManager {
    }
  }

+  private async tryInMemory(): Promise<void> {
+    try {
+      if (!this.connection.config.localServer) {
+        throw new Error('Local server function not provided');
+      }
+
+      // Create linked transport pair
+      const [clientTransport, serverTransport] = InMemoryTransport.createLinkedPair();
+      
+      // Create and connect the server
+      const server = this.connection.config.localServer();
+      await server.connect(serverTransport);
+      
+      // Initialize the client with its transport
+      await this.initializeClient(clientTransport);
+    } catch (error) {
+      console.log('InMemory connection failed:', error);
+      throw error;
+    }
+  }
+
  private async initializeClient(transport: Transport): Promise<void> {
    try {
      const debugTransport = new DebugTransport(transport);
@ -453,13 +481,13 @@ export class MCPConnectionManager {
      // Set up message callbacks to broadcast to UI after client is created
      debugTransport.onsendmessage_ = async (message, options) => {
        if (this.onMessage && this.client) {
-          this.onMessage(this.connection.id, this.client, message, 'sent', { options });
+          this.onMessage(this.connection, this.client, message, 'sent', { options });
        }
      };

      debugTransport.onreceivemessage_ = (message, extra) => {
        if (this.onMessage && this.client) {
-          this.onMessage(this.connection.id, this.client, message, 'received', extra);
+          this.onMessage(this.connection, this.client, message, 'received', extra);
        }
      };

--- a/src/mcp/servers/index.ts
+++ b/src/mcp/servers/index.ts
@ -0,0 +1,10 @@
+import { createServer as createTestServer } from "./test";
+import type { MCPServerConfig } from "@/types/mcp";
+
+export const availableServers: MCPServerConfig[] = [
+    {
+        name: "In-Memory Test Server",
+        url: "local",
+        localServer: createTestServer
+    }
+];
--- a/src/mcp/servers/test/index.ts
+++ b/src/mcp/servers/test/index.ts
@ -0,0 +1,24 @@
+import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js";
+import { z } from "zod";
+
+
+export function createServer() {
+    // Create an MCP server
+    const server = new McpServer({
+        name: "demo-server",
+        version: "1.0.0"
+    });
+
+    // Add an addition tool
+    server.registerTool("add",
+        {
+            title: "Addition Tool",
+            description: "Add two numbers",
+            inputSchema: { a: z.number(), b: z.number() },
+        },
+        async ({ a, b }) => ({
+            content: [{ type: "text", text: String(a + b) }]
+        })
+    );
+    return server;
+}
--- a/src/types/mcp.ts
+++ b/src/types/mcp.ts
@ -8,7 +8,7 @@ import type { TransportSendOptions } from '@modelcontextprotocol/sdk/shared/tran

 export interface MCPServerConfig {
  name: string;                 // User-provided display name
-  url: string;                  // Server endpoint URL
+  url: string;                  // Server endpoint URL (or 'local' for in-memory servers)
  authType?: 'none' | 'oauth'; // Default: none
  oauthConfig?: {
    clientId?: string;
@ -18,6 +18,7 @@ export interface MCPServerConfig {
    redirectUri?: string;       // Override default redirect URI
  };
  maxReconnectAttempts?: number; // Default: 5
+  localServer?: () => any;      // Function to create the local server instance (when url === 'local')
 }

 export interface MCPMessage {
@ -73,7 +74,7 @@ export interface MCPConnection {
  url: string;                   // Server URL
  status: 'connecting' | 'connected' | 'failed' | 'disconnected';
  client?: Client;               // MCP SDK client instance
-  transport?: 'sse' | 'streamable-http';
+  transport?: 'sse' | 'streamable-http' | 'inmemory';
  authType?: 'none' | 'oauth';
  
  // Available capabilities