From e53be9ec42ed6e66e8be08680ebda7d680714032 Mon Sep 17 00:00:00 2001 From: Jerome Date: Thu, 19 Jun 2025 16:55:47 +0100 Subject: [PATCH] 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. --- src/components/MCPTab.tsx | 5 +++ src/contexts/MCPContext.tsx | 68 +++++++++++++++++++++++++++-------- src/mcp/connection.ts | 52 ++++++++++++++++++++------- src/mcp/servers/index.ts | 10 ++++++ src/mcp/servers/test/index.ts | 24 +++++++++++++ src/types/mcp.ts | 5 +-- 6 files changed, 136 insertions(+), 28 deletions(-) create mode 100644 src/mcp/servers/index.ts create mode 100644 src/mcp/servers/test/index.ts diff --git a/src/components/MCPTab.tsx b/src/components/MCPTab.tsx index c8f136d..64dbd43 100644 --- a/src/components/MCPTab.tsx +++ b/src/components/MCPTab.tsx @@ -222,6 +222,11 @@ export function MCPTab() {

{connection.name}

+ {connection.url === 'local' && ( + + In-Memory + + )}

{connection.status} diff --git a/src/contexts/MCPContext.tsx b/src/contexts/MCPContext.tsx index ede9045..d35fce6 100644 --- 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(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 })); diff --git a/src/mcp/connection.ts b/src/mcp/connection.ts index 704c612..56c356d 100644 --- 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 { + 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 { 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); } }; diff --git a/src/mcp/servers/index.ts b/src/mcp/servers/index.ts new file mode 100644 index 0000000..4976b9b --- /dev/null +++ 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 + } +]; \ No newline at end of file diff --git a/src/mcp/servers/test/index.ts b/src/mcp/servers/test/index.ts new file mode 100644 index 0000000..6c6d4ce --- /dev/null +++ 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; +} \ No newline at end of file diff --git a/src/types/mcp.ts b/src/types/mcp.ts index 1eab0d3..fc28b92 100644 --- 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