Skip to content

Commit e53be9e

Browse files
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.
1 parent c9feacd commit e53be9e

File tree

6 files changed

+136
-28
lines changed

6 files changed

+136
-28
lines changed

src/components/MCPTab.tsx

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -222,6 +222,11 @@ export function MCPTab() {
222222
<h4 className="text-sm font-medium text-gray-900 dark:text-white truncate">
223223
{connection.name}
224224
</h4>
225+
{connection.url === 'local' && (
226+
<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">
227+
In-Memory
228+
</span>
229+
)}
225230
</div>
226231
<p className={`text-xs mt-1 ${getStatusColor(connection.status)}`}>
227232
{connection.status}

src/contexts/MCPContext.tsx

Lines changed: 54 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@ import type {
1313
} from '@/types/mcp';
1414
import type { Tool } from '@/types/inference';
1515
import { MCPConnectionManager } from '@/mcp/connection';
16+
import { availableServers } from '@/mcp/servers';
1617

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

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

5758
// Function to broadcast messages to all callbacks (stable reference)
5859
const broadcastMessage = useCallback((
59-
connectionId: string,
60+
connection: MCPConnection,
6061
client: any,
6162
message: any,
6263
direction: 'sent' | 'received',
6364
extra?: any
6465
) => {
65-
// console.log(`MCP Message [${direction}] from ${connectionId}:`, message); // Debug log
66-
67-
// Find connection name using ref to get latest state
68-
const connection = connectionsRef.current.find(conn => conn.id === connectionId);
69-
const serverName = connection?.name || 'Unknown Server';
70-
7166
// Store message
7267
const storedMessage: MCPMonitorMessage = {
7368
id: uuidv4(),
7469
timestamp: new Date(),
75-
connectionId,
76-
serverName,
70+
connectionId: connection.id,
71+
serverName: connection?.name || 'Unknown Server',
7772
message,
7873
direction,
7974
extra,
@@ -82,7 +77,7 @@ export function MCPProvider({ children }: MCPProviderProps) {
8277

8378
messageCallbacksRef.current.forEach(callback => {
8479
try {
85-
callback(connectionId, client, message, direction, extra);
80+
callback(connection.id, client, message, direction, extra);
8681
} catch (error) {
8782
console.error('Error in MCP message callback:', error);
8883
}
@@ -118,11 +113,12 @@ export function MCPProvider({ children }: MCPProviderProps) {
118113

119114
// Set up callback for connection state updates
120115
manager.setConnectionUpdateCallback(() => {
121-
setConnections(prev =>
122-
prev.map(conn =>
116+
setConnections(prev => {
117+
connectionsRef.current = prev.map(conn =>
123118
conn.id === connectionId ? manager.getConnection() : conn
124119
)
125-
);
120+
return connectionsRef.current
121+
});
126122
});
127123

128124
// Set up message callback
@@ -149,6 +145,50 @@ export function MCPProvider({ children }: MCPProviderProps) {
149145
} catch (error) {
150146
console.error('Failed to load persisted MCP connections:', error);
151147
}
148+
149+
// After loading persisted connections, add local servers
150+
await addLocalServers();
151+
};
152+
153+
const addLocalServers = async () => {
154+
// Check if local servers are already added
155+
const hasLocalServers = connectionsRef.current.some(conn => conn.url === 'local');
156+
if (hasLocalServers) {
157+
return;
158+
}
159+
160+
// Add each available local server
161+
for (const serverConfig of availableServers) {
162+
try {
163+
const connectionId = uuidv4();
164+
const manager = new MCPConnectionManager(connectionId, serverConfig);
165+
166+
// Set up callback for connection state updates
167+
manager.setConnectionUpdateCallback(() => {
168+
setConnections(prev => {
169+
connectionsRef.current = prev.map(conn =>
170+
conn.id === connectionId ? manager.getConnection() : conn
171+
);
172+
return connectionsRef.current;
173+
});
174+
});
175+
176+
// Set up message callback
177+
manager.setMessageCallback(broadcastMessage);
178+
179+
// Add to managers map
180+
setManagers(prev => new Map(prev).set(connectionId, manager));
181+
182+
// Add initial connection state
183+
setConnections(prev => [...prev, manager.getConnection()]);
184+
185+
// Auto-connect local servers
186+
await manager.connect();
187+
console.log(`Connected to local server: ${serverConfig.name}`);
188+
} catch (error) {
189+
console.error(`Failed to connect to local server ${serverConfig.name}:`, error);
190+
}
191+
}
152192
};
153193

154194
loadPersistedConnections();
@@ -157,7 +197,7 @@ export function MCPProvider({ children }: MCPProviderProps) {
157197
const persistConnections = useCallback(() => {
158198
try {
159199
// Store both config and connection ID to maintain OAuth token association
160-
const connectionData = connections.map(conn => ({
200+
const connectionData = connections.filter(conn => conn.url !== 'local').map(conn => ({
161201
id: conn.id,
162202
config: conn.config
163203
}));

src/mcp/connection.ts

Lines changed: 40 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@
33
import { Client } from '@modelcontextprotocol/sdk/client/index.js';
44
import { SSEClientTransport } from '@modelcontextprotocol/sdk/client/sse.js';
55
import { StreamableHTTPClientTransport } from '@modelcontextprotocol/sdk/client/streamableHttp.js';
6+
import { InMemoryTransport } from '@modelcontextprotocol/sdk/inMemory.js';
67
import type { Transport } from '@modelcontextprotocol/sdk/shared/transport.js';
78
import {
89
auth,
@@ -235,7 +236,7 @@ export class MCPConnectionManager {
235236
private healthCheckInterval?: NodeJS.Timeout;
236237
private oauthProvider?: MCPOAuthProvider;
237238
private onConnectionUpdate?: () => void;
238-
private onMessage?: (connectionId: string, client: any, message: any, direction: 'sent' | 'received', extra?: any) => void;
239+
private onMessage?: (connection: MCPConnection, client: any, message: any, direction: 'sent' | 'received', extra?: any) => void;
239240

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

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

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

289-
try {
290-
await this.tryStreamableHttp();
291-
this.connection.transport = 'streamable-http';
292-
} catch (error) {
293-
// TODO: jerome - if this is a TypeError: failed to fetch, then there is likely a CORS (or
294-
// Access-Control-Expose-Headers) issue with the server.
295-
await this.trySSE();
296-
this.connection.transport = 'sse';
290+
// Check if this is a local server
291+
if (this.connection.url === 'local') {
292+
await this.tryInMemory();
293+
this.connection.transport = 'inmemory';
294+
} else {
295+
try {
296+
await this.tryStreamableHttp();
297+
this.connection.transport = 'streamable-http';
298+
} catch (error) {
299+
// TODO: jerome - if this is a TypeError: failed to fetch, then there is likely a CORS (or
300+
// Access-Control-Expose-Headers) issue with the server.
301+
await this.trySSE();
302+
this.connection.transport = 'sse';
303+
}
297304
}
298305

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

445+
private async tryInMemory(): Promise<void> {
446+
try {
447+
if (!this.connection.config.localServer) {
448+
throw new Error('Local server function not provided');
449+
}
450+
451+
// Create linked transport pair
452+
const [clientTransport, serverTransport] = InMemoryTransport.createLinkedPair();
453+
454+
// Create and connect the server
455+
const server = this.connection.config.localServer();
456+
await server.connect(serverTransport);
457+
458+
// Initialize the client with its transport
459+
await this.initializeClient(clientTransport);
460+
} catch (error) {
461+
console.log('InMemory connection failed:', error);
462+
throw error;
463+
}
464+
}
465+
438466
private async initializeClient(transport: Transport): Promise<void> {
439467
try {
440468
const debugTransport = new DebugTransport(transport);
@@ -453,13 +481,13 @@ export class MCPConnectionManager {
453481
// Set up message callbacks to broadcast to UI after client is created
454482
debugTransport.onsendmessage_ = async (message, options) => {
455483
if (this.onMessage && this.client) {
456-
this.onMessage(this.connection.id, this.client, message, 'sent', { options });
484+
this.onMessage(this.connection, this.client, message, 'sent', { options });
457485
}
458486
};
459487

460488
debugTransport.onreceivemessage_ = (message, extra) => {
461489
if (this.onMessage && this.client) {
462-
this.onMessage(this.connection.id, this.client, message, 'received', extra);
490+
this.onMessage(this.connection, this.client, message, 'received', extra);
463491
}
464492
};
465493

src/mcp/servers/index.ts

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,10 @@
1+
import { createServer as createTestServer } from "./test";
2+
import type { MCPServerConfig } from "@/types/mcp";
3+
4+
export const availableServers: MCPServerConfig[] = [
5+
{
6+
name: "In-Memory Test Server",
7+
url: "local",
8+
localServer: createTestServer
9+
}
10+
];

src/mcp/servers/test/index.ts

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,24 @@
1+
import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js";
2+
import { z } from "zod";
3+
4+
5+
export function createServer() {
6+
// Create an MCP server
7+
const server = new McpServer({
8+
name: "demo-server",
9+
version: "1.0.0"
10+
});
11+
12+
// Add an addition tool
13+
server.registerTool("add",
14+
{
15+
title: "Addition Tool",
16+
description: "Add two numbers",
17+
inputSchema: { a: z.number(), b: z.number() },
18+
},
19+
async ({ a, b }) => ({
20+
content: [{ type: "text", text: String(a + b) }]
21+
})
22+
);
23+
return server;
24+
}

src/types/mcp.ts

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,7 @@ import type { TransportSendOptions } from '@modelcontextprotocol/sdk/shared/tran
88

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

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

7980
// Available capabilities

0 commit comments

Comments
 (0)