Upload folder using huggingface_hub

core/a2a_with_tools.py

@@ -11,6 +11,37 @@ from any_agent import AgentConfig, AnyAgent
 from any_agent.serving import A2AServingConfig
 from any_agent.tools import a2a_tool_async
 
+# ----------------- Reliability Helpers -----------------
+import socket
+
+
+async def wait_for_port(port: int, host: str = "127.0.0.1", timeout: float = 20.0):
+    """Wait until a TCP port is accepting connections.
+
+    This prevents race conditions where we try to reach an agent whose
+    Uvicorn server has not finished booting yet.
+    """
+    start_time = asyncio.get_event_loop().time()
+    while (asyncio.get_event_loop().time() - start_time) < timeout:
+        try:
+            with socket.create_connection((host, port), timeout=0.5):
+                return  # Port is open
+        except OSError:
+            await asyncio.sleep(0.25)
+    raise RuntimeError(f"Port {port} on {host} did not open within {timeout} seconds")
+
+
+async def create_tool_with_retries(url: str, retries: int = 3, delay: float = 1.0):
+    """Wrapper around a2a_tool_async that retries on network errors."""
+    last_err = None
+    for _ in range(retries):
+        try:
+            return await a2a_tool_async(url)
+        except Exception as err:  # Broad but intentional for transient startup errors
+            last_err = err
+            await asyncio.sleep(delay)
+    raise last_err
+
 
 # Define output types
 class PickDecision(BaseModel):
@@ -106,14 +137,17 @@ class SimpleA2ADraftCoordinator:
         self.agents['zero_rb'] = await create_zero_rb_agent()
         self.agents['robust_rb'] = await create_robust_rb_agent()
         
-        # Wait for servers to start
-        await asyncio.sleep(2)
+        # Wait for servers to start robustly
+        await asyncio.gather(
+            wait_for_port(5001),
+            wait_for_port(5003),
+        )
         
-        # Create a2a tools for each agent
-        self.agent_tools['zero_rb'] = await a2a_tool_async(
+        # Create a2a tools for each agent (with retries)
+        self.agent_tools['zero_rb'] = await create_tool_with_retries(
             "http://localhost:5001/zero_rb_agent"
         )
-        self.agent_tools['robust_rb'] = await a2a_tool_async(
+        self.agent_tools['robust_rb'] = await create_tool_with_retries(
             "http://localhost:5003/robust_rb_agent"
         )