109 lines
3.6 KiB
TypeScript
109 lines
3.6 KiB
TypeScript
/**
|
|
* Copyright (c) Microsoft Corporation.
|
|
*
|
|
* Licensed under the Apache License, Version 2.0 (the "License");
|
|
* you may not use this file except in compliance with the License.
|
|
* You may obtain a copy of the License at
|
|
*
|
|
* http://www.apache.org/licenses/LICENSE-2.0
|
|
*
|
|
* Unless required by applicable law or agreed to in writing, software
|
|
* distributed under the License is distributed on an "AS IS" BASIS,
|
|
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
* See the License for the specific language governing permissions and
|
|
* limitations under the License.
|
|
*/
|
|
|
|
import debug from 'debug';
|
|
import { Server } from '@modelcontextprotocol/sdk/server/index.js';
|
|
import { CallToolRequestSchema, ListToolsRequestSchema } from '@modelcontextprotocol/sdk/types.js';
|
|
import { ManualPromise } from '../utils/manualPromise.js';
|
|
import { logUnhandledError } from '../utils/log.js';
|
|
|
|
import type { Tool, CallToolResult, CallToolRequest } from '@modelcontextprotocol/sdk/types.js';
|
|
import type { Transport } from '@modelcontextprotocol/sdk/shared/transport.js';
|
|
export type { Server } from '@modelcontextprotocol/sdk/server/index.js';
|
|
export type { Tool, CallToolResult, CallToolRequest } from '@modelcontextprotocol/sdk/types.js';
|
|
|
|
const serverDebug = debug('pw:mcp:server');
|
|
|
|
export interface ServerBackend {
|
|
name: string;
|
|
version: string;
|
|
initialize?(server: Server): Promise<void>;
|
|
listTools(): Promise<Tool[]>;
|
|
callTool(name: string, args: CallToolRequest['params']['arguments']): Promise<CallToolResult>;
|
|
serverClosed?(): void;
|
|
}
|
|
|
|
export type ServerBackendFactory = () => ServerBackend;
|
|
|
|
export async function connect(serverBackendFactory: ServerBackendFactory, transport: Transport, runHeartbeat: boolean) {
|
|
const backend = serverBackendFactory();
|
|
const server = createServer(backend, runHeartbeat);
|
|
await server.connect(transport);
|
|
}
|
|
|
|
export function createServer(backend: ServerBackend, runHeartbeat: boolean): Server {
|
|
const initializedPromise = new ManualPromise<void>();
|
|
const server = new Server({ name: backend.name, version: backend.version }, {
|
|
capabilities: {
|
|
tools: {},
|
|
}
|
|
});
|
|
|
|
server.setRequestHandler(ListToolsRequestSchema, async () => {
|
|
serverDebug('listTools');
|
|
const tools = await backend.listTools();
|
|
return { tools };
|
|
});
|
|
|
|
let heartbeatRunning = false;
|
|
server.setRequestHandler(CallToolRequestSchema, async request => {
|
|
serverDebug('callTool', request);
|
|
await initializedPromise;
|
|
|
|
if (runHeartbeat && !heartbeatRunning) {
|
|
heartbeatRunning = true;
|
|
startHeartbeat(server);
|
|
}
|
|
|
|
try {
|
|
return await backend.callTool(request.params.name, request.params.arguments || {});
|
|
} catch (error) {
|
|
return {
|
|
content: [{ type: 'text', text: '### Result\n' + String(error) }],
|
|
isError: true,
|
|
};
|
|
}
|
|
});
|
|
addServerListener(server, 'initialized', () => {
|
|
backend.initialize?.(server).then(() => initializedPromise.resolve()).catch(logUnhandledError);
|
|
});
|
|
addServerListener(server, 'close', () => backend.serverClosed?.());
|
|
return server;
|
|
}
|
|
|
|
const startHeartbeat = (server: Server) => {
|
|
const beat = () => {
|
|
Promise.race([
|
|
server.ping(),
|
|
new Promise((_, reject) => setTimeout(() => reject(new Error('ping timeout')), 5000)),
|
|
]).then(() => {
|
|
setTimeout(beat, 3000);
|
|
}).catch(() => {
|
|
void server.close();
|
|
});
|
|
};
|
|
|
|
beat();
|
|
};
|
|
|
|
function addServerListener(server: Server, event: 'close' | 'initialized', listener: () => void) {
|
|
const oldListener = server[`on${event}`];
|
|
server[`on${event}`] = () => {
|
|
oldListener?.();
|
|
listener();
|
|
};
|
|
}
|