fix: For integ tests, add retries for server initialization

Author: clareliguori

e2e_tests/python/server_clients/server.py

@@ -66,7 +66,7 @@ async def execute_tool(
         tool_name: str,
         tool_use_id: str,
         arguments: dict[str, Any],
-        retries: int = 2,
+        retries: int = 3,
         delay: float = 1.0,
     ) -> Any:
         """Execute a tool with retry mechanism.

e2e_tests/python/server_clients/servers.py

@@ -1,3 +1,4 @@
+import asyncio
 import logging
 
 from server_clients.server import Server
@@ -17,11 +18,44 @@ async def __aenter__(self):
         initialized_servers = []
         for server in self.servers:
             logging.info(f"Starting server: {server.name}")
-            server = await server.__aenter__()
-            initialized_servers.append(server)
+            initialized_server = await self._initialize_server_with_retry(server)
+            initialized_servers.append(initialized_server)
+
         self.servers = initialized_servers
         return self
 
+    async def _initialize_server_with_retry(
+        self, server: Server, retries: int = 3, delay: float = 1.0
+    ) -> Server:
+        """Initialize a server with retry mechanism"""
+        attempt = 0
+        while attempt < retries:
+            try:
+                return await server.__aenter__()
+            except Exception as e:
+                attempt += 1
+                logging.warning(
+                    f"Error initializing server {server.name}: {e}. Attempt {attempt} of {retries}."
+                )
+
+                try:
+                    await server.__aexit__(None, None, None)
+                except Exception as exit_error:
+                    logging.warning(
+                        f"Error closing server {server.name} during retry: {exit_error}"
+                    )
+
+                if attempt < retries:
+                    logging.info(f"Retrying in {delay} seconds...")
+                    await asyncio.sleep(delay)
+                else:
+                    logging.error(
+                        f"Max retries reached for server {server.name}. Failing."
+                    )
+                    raise Exception(
+                        f"Error initializing server {server.name}: {str(e)}"
+                    )
+
     async def __aexit__(self, exc_type, exc_val, exc_tb):
         logging.info("Stopping servers")
         for server in reversed(self.servers):

e2e_tests/typescript/src/server_clients/server.ts

@@ -78,7 +78,7 @@ export abstract class Server {
     toolName: string,
     toolUseId: string,
     args: Record<string, any>,
-    retries: number = 2,
+    retries: number = 3,
     delay: number = 1.0
   ): Promise<ToolResultBlock> {
     let attempt = 0;

e2e_tests/typescript/src/server_clients/servers.ts

@@ -20,7 +20,49 @@ export class Servers {
     logger.info("Starting servers");
     for (const server of this.servers) {
       logger.info(`Starting server: ${server.name}`);
-      await server.initialize();
+      await this.initializeServerWithRetry(server);
+    }
+  }
+
+  /**
+   * Initialize a server with retry mechanism
+   */
+  private async initializeServerWithRetry(
+    server: Server,
+    retries: number = 3,
+    delay: number = 1.0
+  ): Promise<void> {
+    let attempt = 0;
+    while (attempt < retries) {
+      try {
+        await server.initialize();
+        return;
+      } catch (e) {
+        attempt += 1;
+        logger.warn(
+          `Error initializing server ${server.name}: ${e}. Attempt ${attempt} of ${retries}.`
+        );
+
+        try {
+          await server.close();
+        } catch (closeError) {
+          logger.warn(
+            `Error closing server ${server.name} during retry: ${closeError}`
+          );
+        }
+
+        if (attempt < retries) {
+          logger.info(`Retrying in ${delay} seconds...`);
+          await new Promise((resolve) => setTimeout(resolve, delay * 1000));
+        } else {
+          logger.error(
+            `Max retries reached for server ${server.name}. Failing.`
+          );
+          throw new Error(
+            `Error initializing server ${server.name}: ${String(e)}`
+          );
+        }
+      }
     }
   }
 

examples/chatbots/python/server_clients/server.py

@@ -66,7 +66,7 @@ async def execute_tool(
         tool_name: str,
         tool_use_id: str,
         arguments: dict[str, Any],
-        retries: int = 2,
+        retries: int = 3,
         delay: float = 1.0,
     ) -> Any:
         """Execute a tool with retry mechanism.

examples/chatbots/typescript/src/server_clients/server.ts

@@ -79,7 +79,7 @@ export abstract class Server {
     toolName: string,
     toolUseId: string,
     args: Record<string, any>,
-    retries: number = 2,
+    retries: number = 3,
     delay: number = 1.0
   ): Promise<ToolResultBlock> {
     let attempt = 0;