e2e_test for typescript

Author: clareliguori

e2e_tests/typescript/package-lock.json

@@ -0,0 +1,23 @@
+{
+  "name": "mcp-e2e-test-typescript",
+  "version": "0.1.0",
+  "main": "build/main.js",
+  "types": "build/main.d.ts",
+  "scripts": {
+    "build": "tsc",
+    "test": "node build/main.js",
+    "lint": "eslint . --ext .ts"
+  },
+  "dependencies": {
+    "@aws-sdk/client-bedrock-runtime": "^3.750.0",
+    "@modelcontextprotocol/sdk": "^1.5.0",
+    "winston": "^3.17.0"
+  },
+  "devDependencies": {
+    "@eslint/js": "^9.8.0",
+    "@types/node": "^22.13.5",
+    "eslint": "^9.8.0",
+    "typescript": "^5.7.3",
+    "typescript-eslint": "^8.0.0"
+  }
+}

e2e_tests/typescript/package.json

@@ -0,0 +1,18 @@
+{
+  "stdioServers": {
+    "fetch": {
+      "command": "uvx",
+      "args": ["mcp-server-fetch"]
+    }
+  },
+  "lambdaFunctionServers": {
+    "time": {
+      "functionName": "mcp-server-time",
+      "region": "us-east-2"
+    },
+    "weatherAlerts": {
+      "functionName": "mcp-server-weather-alerts",
+      "region": "us-east-2"
+    }
+  }
+}

e2e_tests/typescript/servers_config.json

@@ -0,0 +1,134 @@
+import { Server } from "./server_clients/server.js";
+import { Servers } from "./server_clients/servers.js";
+import { LLMClient } from "./llm_client.js";
+import logger from "./logger.js";
+import {
+  ContentBlock,
+  ConverseCommandOutput,
+  Message,
+} from "@aws-sdk/client-bedrock-runtime";
+
+/**
+ * Orchestrates the interaction between user, LLM, and tools.
+ */
+export class ChatSession {
+  private servers: Server[];
+  private llmClient: LLMClient;
+  private userUtterances: string[] = [];
+
+  constructor(
+    servers: Server[],
+    llmClient: LLMClient,
+    userUtterances: string[]
+  ) {
+    this.servers = servers;
+    this.llmClient = llmClient;
+    this.userUtterances = userUtterances;
+  }
+
+  /**
+   * Process the LLM response and execute tools if needed.
+   * @param serversManager The servers manager.
+   * @param llmResponse The response from the Bedrock Converse API.
+   * @returns The result of tool execution, if any.
+   */
+  async executeRequestedTools(
+    serversManager: Servers,
+    llmResponse: ConverseCommandOutput
+  ): Promise<Message | null> {
+    const stopReason = llmResponse.stopReason;
+
+    if (stopReason === "tool_use") {
+      try {
+        const toolResponses: ContentBlock[] = [];
+        for (const contentItem of llmResponse.output!.message!.content!) {
+          if ("toolUse" in contentItem) {
+            logger.info(`Executing tool: ${contentItem.toolUse!.name}`);
+            logger.info(
+              `With arguments: ${JSON.stringify(contentItem.toolUse!.input)}`
+            );
+            const response = await serversManager.executeTool(
+              contentItem.toolUse!.name!,
+              contentItem.toolUse!.toolUseId!,
+              contentItem.toolUse!.input! as Record<string, any>
+            );
+            toolResponses.push(response);
+          }
+        }
+        return { role: "user", content: toolResponses };
+      } catch (e) {
+        throw new Error(`Failed to execute tool: ${e}`);
+      }
+    } else {
+      // Assume this catches stop reasons "end_turn", "stop_sequence", and "max_tokens"
+      return null;
+    }
+  }
+
+  /**
+   * Main chat session handler.
+   */
+  async start(): Promise<void> {
+    const serverManager = new Servers(this.servers);
+
+    try {
+      await serverManager.initialize();
+
+      const allTools = await serverManager.listTools();
+      const toolsDescription = allTools.map((tool) => tool.formatForLLM());
+
+      const systemPrompt = "You are a helpful assistant.";
+
+      const messages: Message[] = [];
+
+      for (const [i, userInput] of this.userUtterances.entries()) {
+        if (i != 0) {
+          console.log("\n**Pausing 5 seconds to avoid Bedrock throttling**");
+          await new Promise((resolve) => setTimeout(resolve, 5000));
+        }
+
+        console.log(`\You: ${userInput}`);
+
+        messages.push({ role: "user", content: [{ text: userInput }] });
+
+        const llmResponse = await this.llmClient.getResponse(
+          messages,
+          systemPrompt,
+          toolsDescription
+        );
+
+        logger.debug("\nAssistant: " + JSON.stringify(llmResponse, null, 2));
+        console.log(
+          `\nAssistant: ${llmResponse.output!.message!.content![0].text}`
+        );
+        messages.push(llmResponse.output!.message!);
+
+        const toolResults = await this.executeRequestedTools(
+          serverManager,
+          llmResponse
+        );
+
+        if (toolResults) {
+          logger.debug(
+            "\nTool Results: " + JSON.stringify(toolResults, null, 2)
+          );
+          messages.push(toolResults);
+          const finalResponse = await this.llmClient.getResponse(
+            messages,
+            systemPrompt,
+            toolsDescription
+          );
+          logger.debug(
+            "\nFinal response: " + JSON.stringify(finalResponse, null, 2)
+          );
+          console.log(
+            `\nAssistant: ${finalResponse.output!.message!.content![0].text}`
+          );
+          messages.push(finalResponse.output!.message!);
+        }
+      }
+    } finally {
+      await serverManager.close();
+    }
+  }
+}

e2e_tests/typescript/src/chat_session.ts

@@ -0,0 +1,48 @@
+import { BedrockRuntimeClient } from '@aws-sdk/client-bedrock-runtime';
+import * as fs from 'fs';
+
+/**
+ * Manages configuration for the MCP client and the Bedrock client.
+ */
+export class Configuration {
+  modelId: string;
+  region: string;
+
+  /**
+   * Initialize configuration.
+   */
+  constructor(
+    modelId: string = 'anthropic.claude-3-5-sonnet-20241022-v2:0',
+    region: string = 'us-west-2'
+  ) {
+    this.modelId = modelId;
+    this.region = region;
+  }
+
+  /**
+   * Load server configuration from JSON file.
+   * @param filePath Path to the JSON configuration file.
+   * @returns Dict containing server configuration.
+   * @throws Error if configuration file doesn't exist or is invalid JSON.
+   */
+  static loadConfig(filePath: string): Record<string, any> {
+    try {
+      const fileContent = fs.readFileSync(filePath, 'utf8');
+      return JSON.parse(fileContent);
+    } catch (e) {
+      if ((e as NodeJS.ErrnoException).code === 'ENOENT') {
+        throw new Error(`Configuration file not found: ${filePath}`);
+      } else {
+        throw new Error(`Error parsing configuration file: ${e}`);
+      }
+    }
+  }
+
+  /**
+   * Get a Bedrock runtime client.
+   * @returns The Bedrock client.
+   */
+  get bedrockClient(): BedrockRuntimeClient {
+    return new BedrockRuntimeClient({ region: this.region });
+  }
+}

e2e_tests/typescript/src/configuration.ts

@@ -0,0 +1,47 @@
+import {
+  BedrockRuntimeClient,
+  ConverseCommand,
+  ConverseCommandOutput,
+  Message,
+  Tool,
+} from "@aws-sdk/client-bedrock-runtime";
+
+/**
+ * Manages communication with Bedrock.
+ */
+export class LLMClient {
+  private bedrockClient: BedrockRuntimeClient;
+  private modelId: string;
+
+  constructor(bedrockClient: BedrockRuntimeClient, modelId: string) {
+    this.bedrockClient = bedrockClient;
+    this.modelId = modelId;
+  }
+
+  /**
+   * Get a response from the LLM, using the Bedrock Converse API.
+   * @param messages A list of message for the model.
+   * @param systemPrompt The system prompt to use.
+   * @param tools The list of tools to make available to the model.
+   * @returns The LLM's full response.
+   */
+  async getResponse(
+    messages: Message[],
+    systemPrompt: string,
+    tools: Tool[]
+  ): Promise<ConverseCommandOutput> {
+    const command = new ConverseCommand({
+      modelId: this.modelId,
+      messages,
+      system: [{ text: systemPrompt }],
+      inferenceConfig: {
+        maxTokens: 4096,
+        temperature: 0.7,
+        topP: 1,
+      },
+      toolConfig: { tools },
+    });
+
+    return await this.bedrockClient.send(command);
+  }
+}

e2e_tests/typescript/src/llm_client.ts

@@ -0,0 +1,25 @@
+import winston from 'winston';
+
+// Configure the logger
+const logger = winston.createLogger({
+  level: process.env.LOG_LEVEL || 'info',
+  format: winston.format.combine(
+    winston.format.timestamp(),
+    winston.format.printf(({ level, message, timestamp }) => {
+      return `${timestamp} - ${level.toUpperCase()}: ${message}`;
+    })
+  ),
+  transports: [
+    new winston.transports.Console({
+      format: winston.format.combine(
+        winston.format.colorize(),
+        winston.format.timestamp(),
+        winston.format.printf(({ level, message, timestamp }) => {
+          return `${timestamp} - ${level.toUpperCase()}: ${message}`;
+        })
+      ),
+    }),
+  ],
+});
+
+export default logger;

e2e_tests/typescript/src/logger.ts

@@ -0,0 +1,55 @@
+import { Configuration } from "./configuration.js";
+import { ChatSession } from "./chat_session.js";
+import { LLMClient } from "./llm_client.js";
+import { StdioServer } from "./server_clients/stdio_server.js";
+import { LambdaFunctionClient } from "./server_clients/lambda_function.js";
+import { Server } from "./server_clients/server.js";
+import logger from "./logger.js";
+
+/**
+ * Initialize and run the chat session.
+ */
+async function main(): Promise<void> {
+  const config = new Configuration();
+  const serverConfig = Configuration.loadConfig("./servers_config.json");
+
+  const servers: Server[] = [];
+
+  // Initialize stdio servers
+  for (const [name, srvConfig] of Object.entries(serverConfig.stdioServers)) {
+    servers.push(new StdioServer(name, srvConfig as Record<string, any>));
+  }
+
+  // Initialize Lambda function servers
+  for (const [name, srvConfig] of Object.entries(
+    serverConfig.lambdaFunctionServers
+  )) {
+    servers.push(
+      new LambdaFunctionClient(name, srvConfig as Record<string, any>)
+    );
+  }
+
+  const userUtterances = [
+    "Hello!",
+    "What is the current time in Seattle?",
+    "Are there any weather alerts right now?",
+    "Who is Tom Cruise?",
+  ];
+
+  const llmClient = new LLMClient(config.bedrockClient, config.modelId);
+  const chatSession = new ChatSession(servers, llmClient, userUtterances);
+
+  await chatSession.start();
+}
+
+// Handle errors
+process.on("unhandledRejection", (reason, promise) => {
+  logger.error("Unhandled Rejection at:", promise, "reason:", reason);
+  process.exit(1);
+});
+
+// Run the main function
+main().catch((error) => {
+  logger.error("Error in main:", error);
+  process.exit(1);
+});

e2e_tests/typescript/src/main.ts

@@ -0,0 +1,38 @@
+import {
+  LambdaFunctionParameters,
+  LambdaFunctionClientTransport,
+} from "mcp-lambda";
+import { Server } from "./server.js";
+import logger from "../logger.js";
+
+/**
+ * Manages MCP server connections and tool execution for servers running in Lambda functions.
+ */
+export class LambdaFunctionClient extends Server {
+  /**
+   * Initialize the server connection.
+   * @throws ValueError if initialization parameters are invalid
+   * @throws RuntimeError if server fails to initialize
+   */
+  async initialize(): Promise<void> {
+    if (!this.config.functionName) {
+      throw new Error(
+        "The functionName must be a valid string and cannot be None."
+      );
+    }
+
+    const serverParams: LambdaFunctionParameters = {
+      functionName: this.config.functionName,
+      regionName: this.config.region,
+    };
+
+    const transport = new LambdaFunctionClientTransport(serverParams);
+
+    try {
+      await this.client.connect(transport);
+    } catch (e) {
+      logger.error(`Error initializing server ${this.name}: ${e}`);
+      throw e;
+    }
+  }
+}