UNPKG

adk-typescript

Version:

TypeScript port of Google's Agent Development Kit (ADK)

197 lines (196 loc) 7.7 kB
"use strict"; Object.defineProperty(exports, "__esModule", { value: true }); exports.MCPSessionManager = exports.SseServerParams = exports.AsyncExitStack = exports.ClosedResourceError = void 0; exports.retryOnClosedResource = retryOnClosedResource; // Import real MCP SDK types with correct paths const index_js_1 = require("@modelcontextprotocol/sdk/client/index.js"); const stdio_js_1 = require("@modelcontextprotocol/sdk/client/stdio.js"); const sse_js_1 = require("@modelcontextprotocol/sdk/client/sse.js"); // Define the closed resource error that MCP might throw class ClosedResourceError extends Error { constructor(message) { super(message); this.name = 'ClosedResourceError'; } } exports.ClosedResourceError = ClosedResourceError; // AsyncExitStack implementation for resource management class AsyncExitStack { constructor() { this.callbacks = []; } /** * Enter an async context and register its cleanup function * @param context An object with an aclose method or a cleanup function * @returns The context that was entered */ async enterAsyncContext(context) { if (typeof context === 'function' && context.length === 0) { // Add type assertion to ensure the function conforms to CleanupCallback this.callbacks.push(context); } else if (context && typeof context.aclose === 'function') { const closeable = context; this.callbacks.push(() => closeable.aclose()); } else { throw new Error('Context must have an aclose method or be a cleanup function'); } return context; } /** * Close all registered contexts in reverse order */ async aclose() { const errors = []; // Close contexts in reverse order (LIFO) while (this.callbacks.length) { const callback = this.callbacks.pop(); if (callback) { try { await callback(); } catch (error) { errors.push(error instanceof Error ? error : new Error(String(error))); } } } // If any errors occurred, throw the first one if (errors.length > 0) { throw errors[0]; } } } exports.AsyncExitStack = AsyncExitStack; // SseServerParams class for SSE connections class SseServerParams { constructor({ url, headers = {}, timeout = 5, sseReadTimeout = 300 }) { this.url = url; this.headers = headers; this.timeout = timeout; this.sseReadTimeout = sseReadTimeout; } } exports.SseServerParams = SseServerParams; /** * A decorator factory that creates a function to retry operations when resources are closed. * @param asyncReinitFuncName Name of the method to call for reinitialization */ function retryOnClosedResource(asyncReinitFuncName) { return function (_target, _propertyKey, descriptor) { const originalMethod = descriptor.value; descriptor.value = async function (...args) { try { return await originalMethod.apply(this, args); } catch (error) { // Check if error is a ClosedResourceError if (error instanceof Error && (error.name === 'ClosedResourceError' || error.message.includes('closed'))) { try { // @ts-ignore if (typeof this[asyncReinitFuncName] === 'function') { // @ts-ignore await this[asyncReinitFuncName](); } else { throw new Error(`Function ${asyncReinitFuncName} does not exist in decorated class.`); } } catch (reinitError) { throw new Error(`Error reinitializing: ${reinitError}`); } return await originalMethod.apply(this, args); } throw error; } }; return descriptor; }; } /** * Manages MCP client sessions. * This class provides methods for creating and initializing MCP client sessions, * handling different connection parameters (Stdio and SSE). */ class MCPSessionManager { /** * Initializes the MCP session manager. * @param connectionParams Parameters for the MCP connection (Stdio or SSE) * @param exitStack AsyncExitStack to manage the session lifecycle * @param errlog Optional error logging stream */ constructor(connectionParams, exitStack, errlog) { this.connectionParams = connectionParams; this.exitStack = exitStack; this.errlog = errlog; } /** * Creates a new MCP client session * @returns A promise that resolves to a new ClientSession */ async createSession() { return MCPSessionManager.initializeSession({ connectionParams: this.connectionParams, exitStack: this.exitStack, errlog: this.errlog }); } /** * Initializes an MCP client session * @param params Session initialization parameters * @returns A promise that resolves to the initialized ClientSession */ static async initializeSession({ connectionParams, exitStack, errlog }) { // Create the MCP client with real MCP SDK const client = new index_js_1.Client({ name: 'adk-mcp-client', version: '1.0.0' }); // Create proper transport based on connection type if ('command' in connectionParams) { // This is StdioServerParameters const transport = new stdio_js_1.StdioClientTransport({ command: connectionParams.command, args: connectionParams.args // Note: StdioClientTransport doesn't accept errlog parameter }); // Connect to the transport await client.connect(transport); // Add to exit stack for proper cleanup await exitStack.enterAsyncContext(async () => { try { // Use transport.close() instead of client.disconnect() await transport.close(); } catch (error) { console.error('Error closing transport:', error); } }); } else if ('url' in connectionParams) { // This is SseServerParams // Create a proper URL object for SSEClientTransport const url = new URL(connectionParams.url); // Based on GitHub examples, SSEClientTransport takes just the URL // without options, or we need to use the correct option format const transport = new sse_js_1.SSEClientTransport(url); // Connect to the transport await client.connect(transport); // Add to exit stack for proper cleanup await exitStack.enterAsyncContext(async () => { try { // Use transport.close() instead of client.disconnect() await transport.close(); } catch (error) { console.error('Error closing transport:', error); } }); } else { throw new Error('Unable to initialize connection. Connection should be StdioServerParameters or SseServerParams'); } return client; } } exports.MCPSessionManager = MCPSessionManager;