import { Transport } from '@modelcontextprotocol/sdk/shared/transport.js';
import { JSONRPCRequest, JSONRPCResponse, JSONRPCError } from '@modelcontextprotocol/sdk/types.js';
/**
* In-memory transport for testing MCP servers
*/
export class MemoryTransport implements Transport {
private requestHandlers = new Map<string, (params: any) => Promise<any>>();
private responseHandlers = new Map<number, (response: JSONRPCResponse) => void>();
private errorHandlers = new Map<number, (error: JSONRPCError) => void>();
private nextId = 1;
private closed = false;
constructor() {}
async start(): Promise<void> {
// No-op for memory transport
}
async close(): Promise<void> {
this.closed = true;
this.requestHandlers.clear();
this.responseHandlers.clear();
this.errorHandlers.clear();
}
send(message: JSONRPCRequest | JSONRPCResponse): Promise<void> {
if (this.closed) {
throw new Error('Transport is closed');
}
// Handle responses
if ('id' in message && 'result' in message || 'error' in message) {
const response = message as JSONRPCResponse;
if (response.id !== null && response.id !== undefined) {
const handler = this.responseHandlers.get(response.id as number);
if (handler) {
handler(response);
this.responseHandlers.delete(response.id as number);
}
}
return Promise.resolve();
}
// Handle requests
const request = message as JSONRPCRequest;
const handler = this.requestHandlers.get(request.method);
if (!handler) {
const errorResponse: JSONRPCResponse = {
jsonrpc: '2.0',
id: request.id,
error: {
code: -32601,
message: `Method not found: ${request.method}`
}
};
this.send(errorResponse);
return Promise.resolve();
}
// Execute handler asynchronously
handler(request.params).then(
result => {
const response: JSONRPCResponse = {
jsonrpc: '2.0',
id: request.id,
result
};
this.send(response);
},
error => {
const response: JSONRPCResponse = {
jsonrpc: '2.0',
id: request.id,
error: {
code: -32603,
message: error.message || 'Internal error',
data: error
}
};
this.send(response);
}
);
return Promise.resolve();
}
onMessage(handler: (message: JSONRPCRequest | JSONRPCResponse) => void): void {
// Not needed for memory transport
}
onClose(handler: () => void): void {
// Not needed for memory transport
}
/**
* Register a handler for a specific method
*/
registerHandler(method: string, handler: (params: any) => Promise<any>): void {
this.requestHandlers.set(method, handler);
}
/**
* Call a method and wait for response
*/
async callMethod(method: string, params: any = {}): Promise<any> {
const id = this.nextId++;
const request: JSONRPCRequest = {
jsonrpc: '2.0',
id,
method,
params
};
return new Promise((resolve, reject) => {
this.responseHandlers.set(id, (response) => {
if ('error' in response && response.error) {
reject(new Error(response.error.message));
} else {
resolve(response.result);
}
});
this.send(request);
});
}
/**
* Call a tool and wait for response
*/
async callTool(toolName: string, args: any): Promise<any> {
return this.callMethod('tools/call', {
name: toolName,
arguments: args
});
}
}
