Add pprof & mcp metrics in api-service and request retry in aenv sdk

aenv/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "aenvironment"
-version = "0.1.5"
+version = "0.1.6"
 description = "AEnvironment Python SDK - Production-grade environment for AI agent tools"
 readme = "README.md"
 requires-python = ">=3.10"

aenv/src/aenv/__init__.py

@@ -20,7 +20,12 @@
 """
 
 from aenv.core.environment import Environment
-from aenv.core.exceptions import AEnvError, EnvironmentError, ToolError
+from aenv.core.exceptions import (
+    AEnvError,
+    EnvironmentError,
+    ToolError,
+    UnrecoverableEnvironmentError,
+)
 from aenv.core.function_registry import (
     register_function,
     register_health,
@@ -30,7 +35,7 @@
 from aenv.core.models import EnvInstance, EnvStatus
 from aenv.core.tool import Tool, get_registry, register_tool
 
-__version__ = "0.1.5"
+__version__ = "0.1.6"
 __all__ = [
     "Tool",
     "register_tool",
@@ -42,6 +47,7 @@
     "AEnvError",
     "ToolError",
     "EnvironmentError",
+    "UnrecoverableEnvironmentError",
     "EnvInstance",
     "EnvStatus",
     "setup_logging",

aenv/src/aenv/core/__init__.py

@@ -20,6 +20,7 @@
     NetworkError,
     ToolError,
     ToolTimeoutError,
+    UnrecoverableEnvironmentError,
 )
 from aenv.core.function_registry import (
     get_function_registry,
@@ -39,4 +40,5 @@
     "ToolError",
     "ToolTimeoutError",
     "NetworkError",
+    "UnrecoverableEnvironmentError",
 ]

aenv/src/aenv/core/environment.py

@@ -20,6 +20,7 @@
 import asyncio
 import json
 import os
+import random
 import traceback
 from datetime import datetime, timezone
 from typing import Any, Dict, List, Optional, Tuple
@@ -33,7 +34,11 @@
 from fastmcp.client.transports import StreamableHttpTransport
 
 from aenv.client.scheduler_client import AEnvSchedulerClient
-from aenv.core.exceptions import EnvironmentError, ToolError
+from aenv.core.exceptions import (
+    EnvironmentError,
+    ToolError,
+    UnrecoverableEnvironmentError,
+)
 from aenv.core.logging import getLogger
 from aenv.core.models import EnvInstance, EnvStatus
 from aenv.core.tool import Tool
@@ -166,21 +171,59 @@ def __init__(
         self._client: Optional[AEnvSchedulerClient] = None
         self._mcp_client: Optional[Client] = None
         self._mcp_session_active: bool = False
+        self._consecutive_tool_errors: int = 0
+        self._CIRCUIT_BREAKER_THRESHOLD = 5
 
     def _log_prefix(self) -> str:
-        """Get log prefix with instance ID."""
+        """Get log prefix with instance ID and SDK version."""
+        from aenv import __version__
+
         instance_id = (
             getattr(self._instance, "id", "None") if self._instance else "None"
         )
-        return f"[ENV:{instance_id}]"
+        return f"[ENV:{instance_id}][sdk:v{__version__}]"
+
+    async def _backoff(self, attempt: int, base: float = 2.0) -> None:
+        """Exponential backoff with jitter."""
+        wait = base**attempt + random.uniform(0, 1)
+        logger.debug(f"{self._log_prefix()} Backoff {wait:.1f}s (attempt {attempt})")
+        await asyncio.sleep(wait)
+
+    async def _rebuild_mcp_client(self) -> None:
+        """Destroy and recreate MCP client for session recovery."""
+        logger.warning(f"{self._log_prefix()} Full MCP client rebuild")
+        if self._mcp_client is not None:
+            try:
+                await asyncio.wait_for(self._mcp_client.close(), timeout=2.0)
+            except Exception as e:
+                logger.debug(f"{self._log_prefix()} Error closing MCP client: {e}")
+            self._mcp_client = None
+        self._mcp_session_active = False
+        await asyncio.sleep(0.5)
 
     async def __aenter__(self):
         """Async context manager entry."""
         await self.initialize()
+        max_attempts = 3
+        for attempt in range(max_attempts):
+            try:
+                await self._ensure_mcp_session()
+                return self
+            except (Exception, asyncio.CancelledError) as e:
+                if attempt < max_attempts - 1:
+                    logger.warning(
+                        f"{self._log_prefix()} Initial session failed "
+                        f"({attempt+1}/{max_attempts}), retrying: {type(e).__name__}: {e}"
+                    )
+                    await self._rebuild_mcp_client()
+                    await asyncio.sleep(1.0)
+                else:
+                    raise
         return self
 
     async def __aexit__(self, exc_type, exc_val, exc_tb):
         """Async context manager exit."""
+        await self._close_mcp_session()
         await self.release()
 
     async def initialize(self) -> bool:
@@ -244,6 +287,17 @@ async def initialize(self) -> bool:
                 f"Failed to initialize environment '{self.env_name}': {str(e)}"
             ) from e
 
+    async def _close_mcp_session(self):
+        """Close the persistent MCP session if one is active."""
+        if self._mcp_client and self._mcp_session_active:
+            logger.info(f"{self._log_prefix()} Closing persistent MCP session")
+            try:
+                await self._mcp_client.__aexit__(None, None, None)
+            except Exception as e:
+                logger.warning(f"{self._log_prefix()} Error closing MCP session: {e}")
+            finally:
+                self._mcp_session_active = False
+
     async def release(self):
         """Release environment resources."""
         logger.info(
@@ -254,7 +308,7 @@ async def release(self):
             try:
                 await self._mcp_client.close()
                 logger.debug(f"{self._log_prefix()} MCP client closed")
-            except Exception as e:
+            except (Exception, asyncio.CancelledError) as e:
                 logger.warning(
                     f"{self._log_prefix()} Failed to close MCP client: {str(e)}"
                 )
@@ -621,6 +675,12 @@ async def call_tool(
         """
         Execute a tool with given arguments using MCP client.
 
+        Retry strategy (idempotent-safe):
+        - Session establishment failures are retried (tool was never sent to server)
+        - Once call_tool_mcp() is invoked, the tool MAY have executed on the server.
+          On failure after invocation, we rebuild the session but do NOT re-invoke the
+          tool, because the tool may be non-idempotent (e.g. bash "echo x >> file").
+
         Args:
             tool_name: Name of the tool (format: "env_name/tool_name" or just "tool_name")
             arguments: Tool arguments
@@ -630,11 +690,18 @@ async def call_tool(
             Tool execution result
 
         Raises:
-            ToolError: If tool execution fails
-            ToolTimeoutError: If execution times out
+            ToolError: If tool execution fails after invocation
+            EnvironmentError: If session cannot be established
         """
         await self._ensure_initialized()
 
+        # Circuit breaker: fail fast if too many consecutive tool errors
+        if self._consecutive_tool_errors >= self._CIRCUIT_BREAKER_THRESHOLD:
+            raise UnrecoverableEnvironmentError(
+                f"Circuit breaker open: {self._consecutive_tool_errors} consecutive failures",
+                env_name=self.env_name,
+            )
+
         # Parse tool name
         if "/" in tool_name:
             env_name, actual_tool_name = tool_name.split("/", 1)
@@ -649,8 +716,39 @@ async def call_tool(
             f"{self._log_prefix()} Executing tool: {actual_tool_name} in environment {self.env_name}, arguments={arguments}"
         )
 
+        # Establish session (safe to retry)
+        max_session_attempts = 3
+        client = None
+
+        for attempt in range(max_session_attempts):
+            try:
+                client = await self._ensure_mcp_session()
+                break
+            except (Exception, asyncio.CancelledError) as e:
+                if (
+                    isinstance(e, httpx.HTTPStatusError)
+                    and 400 <= e.response.status_code < 500
+                ):
+                    raise EnvironmentError(
+                        f"Non-retryable error establishing session: HTTP {e.response.status_code}"
+                    ) from e
+
+                if attempt < max_session_attempts - 1:
+                    logger.warning(
+                        f"{self._log_prefix()} Session establish failed ({attempt+1}/{max_session_attempts}), "
+                        f"retrying: {type(e).__name__}: {e}"
+                    )
+                    await self._backoff(attempt, base=1.5)
+                    await self._rebuild_mcp_client()
+                    continue
+
+                raise EnvironmentError(
+                    f"Failed to establish session after {max_session_attempts} attempts: {e}",
+                    env_name=self.env_name,
+                ) from e
+
+        # Execute tool (not safe to retry — may be non-idempotent)
         try:
-            client = await self._ensure_mcp_session()
             result = await client.call_tool_mcp(
                 name=actual_tool_name, arguments=arguments, timeout=timeout
             )
@@ -666,20 +764,28 @@ async def call_tool(
                     else:
                         content.append({"type": "text", "text": str(item)})
 
+            self._consecutive_tool_errors = 0
             return ToolResult(content=content, is_error=result.isError)
 
-        except Exception as e:
+        except (Exception, asyncio.CancelledError) as e:
+            self._mcp_session_active = False
+            self._consecutive_tool_errors += 1
+            # Rebuild client to break ClosedResourceError cascade
+            await self._rebuild_mcp_client()
+
+            # Tool may have executed — do NOT retry
             logger.error(
-                f"{self._log_prefix()} Tool execution encountered an issue: {str(e)} | "
-                f"Type: {type(e).__name__} | "
+                f"{self._log_prefix()} Tool '{actual_tool_name}' failed after invocation "
+                f"({type(e).__name__}: {e}). Tool may have executed on server — "
+                f"NOT retrying (non-idempotent safety). | "
                 f"Environment: {self.env_name} | "
-                f"Tool: {actual_tool_name} | "
                 f"Arguments: {arguments} | "
                 f"Timeout: {timeout or self.timeout}s | "
                 f"MCP URL: {self.aenv_data_url}"
             )
             raise ToolError(
-                f"Tool '{actual_tool_name}' execution encountered an issue: {str(e)}"
+                f"Tool '{actual_tool_name}' execution failed: {e}. "
+                f"Tool may have already executed on server — not retried for safety."
             )
 
     async def get_env_info(self) -> Dict[str, Any]:
@@ -953,7 +1059,7 @@ async def _ensure_mcp_session(self) -> Client:
             if self._mcp_client is not None:
                 try:
                     await self._mcp_client.close()
-                except Exception as e:
+                except (Exception, asyncio.CancelledError) as e:
                     logger.debug(
                         f"{self._log_prefix()} Error closing stale MCP client: {e}"
                     )
@@ -968,10 +1074,7 @@ async def _ensure_mcp_session(self) -> Client:
                     f"{self._log_prefix()} MCP session established and will be reused"
                 )
                 return client
-            except Exception as e:
+            except (Exception, asyncio.CancelledError) as e:
                 self._mcp_client = None
                 self._mcp_session_active = False
-                logger.error(
-                    f"{self._log_prefix()} Failed to establish MCP session: {e}"
-                )
                 raise EnvironmentError(f"Failed to establish MCP session: {e}") from e

aenv/src/aenv/core/exceptions.py

@@ -82,6 +82,12 @@ def __init__(
         self.status_code = status_code
 
 
+class UnrecoverableEnvironmentError(EnvironmentError):
+    """Environment is in an unrecoverable state (circuit breaker open)."""
+
+    pass
+
+
 class NetworkError(AEnvError):
     """Exception raised for network-related errors."""