specialize duty for session and execution manager for further customization

Author: mayinghan

eval_protocol/mcp/execution/manager.py

@@ -22,7 +22,6 @@
 
 from ...models import CompletionParams, EvaluationRow, InputMetadata, Message
 from ...types import MCPSession, MCPToolCall, TerminationReason, Trajectory
-from ..client.connection import MCPConnectionManager
 
 if TYPE_CHECKING:
     from ..session.manager import GeneralMCPVectorEnv
@@ -33,43 +32,9 @@
 
 class ExecutionManager:
     """
-    Unified manager that handles both MCP session lifecycle and rollout execution.
-
-    Combines the functionality of SessionManager and RolloutManager for better
-    organization and reduced complexity.
+    Manage rollout for MCP environments.
     """
 
-    def __init__(self):
-        """Initialize the execution manager."""
-        self.connection_manager = MCPConnectionManager()
-
-    async def initialize_sessions(self, sessions: List[MCPSession]) -> None:
-        """
-        Initialize multiple MCP sessions in parallel.
-
-        Args:
-            sessions: List of MCPSessions to initialize
-        """
-        tasks = [self.connection_manager.initialize_session(session) for session in sessions]
-        await asyncio.gather(*tasks)
-
-    async def close_sessions(self, sessions: List[MCPSession]) -> None:
-        """
-        Close multiple MCP sessions in parallel.
-
-        Args:
-            sessions: List of MCPSessions to close
-        """
-        tasks = [asyncio.create_task(self.connection_manager.close_session(session)) for session in sessions]
-
-        if tasks:
-            try:
-                # Wait for all close operations to complete
-                await asyncio.gather(*tasks, return_exceptions=True)
-            except asyncio.CancelledError:
-                # Handle cancellation gracefully (especially important for Python 3.12)
-                logger.debug("Close operation was cancelled, but sessions are marked as closed")
-
     async def execute_rollouts(
         self,
         envs: "GeneralMCPVectorEnv",
@@ -178,7 +143,7 @@ async def _execute_with_semaphore(idx):
             for msg in trajectory.conversation_history:
                 # Create a copy to avoid modifying the original
                 msg_dict = dict(msg)
-                
+
                 # Handle multimodal content (list of content blocks) by extracting text
                 if isinstance(msg_dict.get("content"), list):
                     text_content = None
@@ -187,7 +152,7 @@ async def _execute_with_semaphore(idx):
                             text_content = content_block.get("text")
                             break
                     msg_dict["content"] = text_content or ""
-                
+
                 messages.append(Message.model_validate(msg_dict))
 
             input_metadata = InputMetadata(

eval_protocol/mcp/mcpgym.py

@@ -115,6 +115,7 @@ def __init__(self, server_name: str, adapter: EnvironmentAdapter, seed: Optional
         # Register tools and control plane endpoints
         self._register_tools()
         self._discover_and_register_control_plane_endpoints()
+        self._register_session_reset_endpoint()
 
     def _get_session_id(self, ctx: Context) -> str:
         """
@@ -226,6 +227,19 @@ def _get_or_create_session(self, ctx: Context) -> Dict[str, Any]:
 
             return self.sessions[session_id]
 
+    def _register_session_reset_endpoint(self):
+
+        @self.mcp.custom_route("/control/reset_session", methods=["POST"])
+        async def reset_session_endpoint(request: Request, ctx: Context) -> JSONResponse:
+            session_id = request.headers.get("mcp-session-id")
+            if not session_id:
+                return JSONResponse({"error": "Missing mcp-session-id header"}, status_code=400)
+            with self.session_lock:
+                if session_id in self.sessions:
+                    del self.sessions[session_id]
+            self.sessions[session_id] = self._get_or_create_session(ctx)
+            return JSONResponse({"message": "Session reset successfully"})
+
     def _discover_and_register_control_plane_endpoints(self):
         """
         Discover and register control plane endpoints on the subclass instance.

eval_protocol/mcp/session/manager.py

@@ -11,7 +11,7 @@
 from typing import Any, Callable, Dict, List, Optional, Tuple, Union
 
 from ...types import DatasetRow, MCPSession, MCPToolCall
-from ..execution.manager import ExecutionManager
+from ..client.connection import MCPConnectionManager
 
 logger = logging.getLogger(__name__)
 
@@ -44,7 +44,7 @@ def __init__(
         self.user_prompt_formatter = user_prompt_formatter or self._default_formatter
         self.n = len(sessions)
         self.tool_schemas = []  # Discovered from MCP servers
-        self.execution_manager = ExecutionManager()
+        self.connection_manager = MCPConnectionManager()
         self.usage_stats = {}  # llm usage stats for monitoring
 
         if len(sessions) != len(dataset_rows):
@@ -58,17 +58,14 @@ async def reset(self, session: MCPSession) -> Tuple[Any, List[Dict]]:
 
         This is thread-safe and can be called from worker threads.
         """
-        # Establish a persistent session for each environment.
-        await self.execution_manager.connection_manager.initialize_session(session)
-
         # Get available tools from MCP server
-        tool_schemas = await self.execution_manager.connection_manager.discover_tools(session)
+        tool_schemas = await self.connection_manager.discover_tools(session)
 
         if not self.tool_schemas:
             self.tool_schemas = tool_schemas
 
         # PROPER MCP PATTERN: Get initial state from resources during session establishment
-        initial_observation = await self.execution_manager.connection_manager.get_initial_state(session)
+        initial_observation = await self.connection_manager.get_initial_state(session)
 
         # Update session state
         session.terminated = False
@@ -119,7 +116,7 @@ async def step(self, env_index: int, tool_call: MCPToolCall) -> Tuple[Any, float
             )
 
         # Execute the tool call via MCP protocol
-        observation, reward, done, info = await self.execution_manager.connection_manager.call_tool(
+        observation, reward, done, info = await self.connection_manager.call_tool(
             session, tool_call.tool_name, tool_call.arguments
         )
 
@@ -223,5 +220,6 @@ def _default_formatter(self, template: str, obs: Any, context: Dict) -> Union[st
     async def close(self):
         """Closes all MCP sessions."""
         print(f"🧹 Closing {self.n} MCP sessions...")
-        await self.execution_manager.close_sessions(self.sessions)
+        tasks = [self.connection_manager.close_session(session) for session in self.sessions]
+        await asyncio.gather(*tasks)
         print(f"✅ All MCP sessions closed.")

eval_protocol/mcp_env.py

@@ -51,11 +51,12 @@
 from .mcp.session.manager import GeneralMCPVectorEnv
 from .models import EvaluationRow
 from .types import DatasetRow, MCPSession, MCPToolCall
+import asyncio
 
 logger = logging.getLogger(__name__)
 
 
-def make(
+async def make(
     env_spec: str,
     evaluation_rows: Optional[List[EvaluationRow]] = None,
     dataset: Optional[List[Dict]] = None,
@@ -104,17 +105,17 @@ def make(
     if evaluation_rows:
         for i, row in enumerate(evaluation_rows):
             dataset_info = row.input_metadata.dataset_info if row.input_metadata else {}
-            
+
             system_message = row.get_system_message()
             system_prompt = system_message.content or ""
-            
+
             dataset_entry = {
                 "id": row.input_metadata.row_id if row.input_metadata and row.input_metadata.row_id else f"task_{i}",
                 "system_prompt": system_prompt,
                 "user_prompt_template": dataset_info.get("user_prompt_template", ""),
                 "environment_context": dataset_info.get("environment_context", {}),
                 "user_simulation": dataset_info.get("user_simulation", {}),
-                "evaluation_criteria": dataset_info.get("evaluation_criteria", {})
+                "evaluation_criteria": dataset_info.get("evaluation_criteria", {}),
             }
             internal_dataset.append(dataset_entry)
     elif dataset:
@@ -198,7 +199,10 @@ def make(
             )
             sessions.append(session)
 
-        return GeneralMCPVectorEnv(sessions, dataset_rows, user_prompt_formatter)
+        mcp_envs = GeneralMCPVectorEnv(sessions, dataset_rows, user_prompt_formatter)
+        tasks = [mcp_envs.connection_manager.initialize_session(session) for session in sessions]
+        await asyncio.gather(*tasks)
+        return mcp_envs
 
 
 async def rollout(