support reset mcp session

Author: mayinghan

eval_protocol/mcp/client/connection.py

@@ -12,8 +12,10 @@
 from contextlib import AsyncExitStack
 from typing import Any, Dict, List, Optional, Tuple
 
+from mcp.types import EmptyResult
 from mcp.client.session import ClientSession
 from mcp.client.streamable_http import streamablehttp_client
+from pydantic import BaseModel
 
 from ...types import MCPSession
 
@@ -101,7 +103,7 @@ async def initialize_session(self, session: MCPSession) -> None:
 
                 # Update the session ID to match what the server generated
                 session.session_id = server_session_id
-                logger.debug(f"Updated session ID to match server: {server_session_id}")
+                logger.info(f"Updated session ID to match server: {server_session_id}")
 
         # PRE-WARM: Discover and cache tools immediately after session initialization
         # This prevents concurrent list_tools() calls later
@@ -133,6 +135,24 @@ async def _prewarm_tools_cache(self, session: MCPSession) -> None:
                 self._tools_cache[cache_key] = tool_schemas
                 logger.debug(f"✅ PRE-WARMED {len(tool_schemas)} tools for{cache_key}")
 
+    async def reset_session(self, session: MCPSession) -> None:
+        """
+        Clean session data in remote mcp server for the given session
+        """
+        import httpx
+
+        base_url = session.base_url.rstrip("/").removesuffix("/mcp")
+        url = f"{base_url}/control/reset_session"
+
+        headers = {"mcp-session-id": session.session_id}
+        body = {"seed": session.seed}
+
+        timeout = httpx.Timeout(3.0)
+        async with httpx.AsyncClient(timeout=timeout) as client:
+            resp = await client.post(url, headers=headers, json=body)
+            resp.raise_for_status()
+            logger.debug(f"Session {session.session_id}: reset_session -> {resp.json()}")
+
     async def discover_tools(self, session: MCPSession) -> List[Dict]:
         """
         Discover available tools from an MCP session.
@@ -422,6 +442,8 @@ async def call_tool(self, session: MCPSession, tool_name: str, arguments: Dict)
         mcp_session = session._mcp_session
 
         # 1. Execute the tool call via MCP protocol (DATA PLANE)
+        print("session.session_id", session._mcp_session)
+        print(session._mcp_session._write_stream._closed)
         tool_result = await mcp_session.call_tool(tool_name, arguments)
 
         # Extract data plane results (observation only)

eval_protocol/mcp/mcpgym.py

@@ -230,14 +230,23 @@ def _get_or_create_session(self, ctx: Context) -> Dict[str, Any]:
     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:
+        async def reset_session_endpoint(request: Request) -> JSONResponse:
             session_id = request.headers.get("mcp-session-id")
+            body = await request.json()
+            seed = body.get("seed", None)
+            print(f"🔍 _register_session_reset_endpoint: Resetting session, session_id: {session_id}, seed: {seed}")
             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)
+                    env, obs, _ = self._new_env(seed=seed)
+                    self.sessions[session_id] = {
+                        "env": env,
+                        "obs": obs,
+                        "session_data": {},
+                        "session_id": session_id,
+                    }
+                    print(f"🔍 _register_session_reset_endpoint: Finished reset session, session_id: {session_id}")
             return JSONResponse({"message": "Session reset successfully"})
 
     def _discover_and_register_control_plane_endpoints(self):
@@ -336,7 +345,7 @@ def _update_control_plane(self, reward: float, terminated: bool, truncated: bool
 
         # Log control plane update (for debugging)
         print(
-            f"🎛️  Control plane updated: reward={reward}, terminated={terminated}, step={self.control_plane_state['step_count']}"
+            f"🎛️  Control plane updated: reward={reward}, terminated={terminated}, step={self.control_plane_state['step_count']}, total_reward={self.control_plane_state['total_reward']}"
         )
 
     def _get_or_create_session_control_plane(self, session_id: str) -> Dict[str, Any]:
@@ -378,7 +387,7 @@ def _update_session_control_plane(
 
         # Log control plane update
         print(
-            f"🎛️  Session {session_id[:16]}... control plane: reward={reward}, terminated={terminated}, step={control_plane['step_count']}"
+            f"🎛️  Session {session_id[:16]}... control plane: reward={reward}, terminated={terminated}, step={control_plane['step_count']}, total_reward={control_plane['total_reward']}"
         )
 
     def get_control_plane_state(self, session_id: str) -> Optional[Dict[str, Any]]:

eval_protocol/mcp_env.py

@@ -17,7 +17,7 @@
     policy = ep.FireworksPolicy(model_id="accounts/fireworks/models/qwen3-235b-a22b")
 
     # Create environments with evaluation_rows configuration
-    envs = ep.make("http://localhost:8000/mcp", evaluation_rows=evaluation_rows)
+    envs = await await ep.make("http://localhost:8000/mcp", evaluation_rows=evaluation_rows)
 
     # Execute tool-calling rollouts
     evaluation_rows = await ep.rollout(envs, policy=policy, steps=512)
@@ -56,14 +56,23 @@
 logger = logging.getLogger(__name__)
 
 
-def make(
+async def reset_mcp_sessions(envs: GeneralMCPVectorEnv):
+    """
+    Reset mcp server sessions
+    """
+    tasks = [envs.connection_manager.reset_session(session) for session in envs.sessions]
+    await asyncio.gather(*tasks)
+
+
+async def make(
     env_spec: str,
     evaluation_rows: Optional[List[EvaluationRow]] = None,
     dataset: Optional[List[Dict]] = None,
     n: Optional[int] = None,
     seeds: Optional[List[int]] = None,
     model_id: str = "unknown",
     user_prompt_formatter: Optional[Callable] = None,
+    reset_sessions: bool = False,
 ) -> GeneralMCPVectorEnv:
     """
     Create general MCP environments driven by evaluation_rows configuration.
@@ -76,19 +85,20 @@ def make(
         seeds: List of seeds (for backward compatibility)
         model_id: Model identifier
         user_prompt_formatter: Optional callback for formatting user prompts
+        reset_sessions: Whether to reset sessions before returning the environment
 
     Returns:
         General MCP environment that works with any MCP server
 
     Example:
         # EvaluationRow approach (preferred)
-        envs = ep.make("http://localhost:8000/mcp", evaluation_rows=evaluation_rows)
+        envs = await await ep.make("http://localhost:8000/mcp", evaluation_rows=evaluation_rows)
 
         # Dataset approach (backward compatibility)
-        envs = ep.make("http://localhost:8000/mcp", dataset=dataset)
+        envs = await await ep.make("http://localhost:8000/mcp", dataset=dataset)
 
         # Legacy approach (backward compatibility)
-        envs = ep.make("http://localhost:8000/mcp", n=10, seeds=seeds)
+        envs = await await ep.make("http://localhost:8000/mcp", n=10, seeds=seeds)
     """
     # Parse environment specification - make sure URL format is correct
     base_url = env_spec
@@ -161,8 +171,6 @@ def make(
             )
             sessions.append(session)
 
-        return GeneralMCPVectorEnv(sessions, dataset_rows, user_prompt_formatter)
-
     else:
         # Legacy approach for backward compatibility
         if n is None:
@@ -199,10 +207,14 @@ def make(
             )
             sessions.append(session)
 
-        mcp_envs = GeneralMCPVectorEnv(sessions, dataset_rows, user_prompt_formatter)
-        tasks = [mcp_envs.connection_manager.initialize_session(session) for session in sessions]
-        asyncio.run(asyncio.gather(*tasks))
-        return mcp_envs
+    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)
+
+    if reset_sessions:
+        await reset_mcp_sessions(mcp_envs)
+
+    return mcp_envs
 
 
 async def rollout(

eval_protocol/pytest/default_mcp_gym_rollout_processor.py

@@ -182,49 +182,47 @@ def __exit__(self, exc_type, exc_val, exc_tb):
         return False  # Don't suppress exceptions
 
 
-
-async def default_mcp_gym_rollout_processor(rows: List[EvaluationRow], config: RolloutProcessorConfig) -> List[EvaluationRow]:
+async def default_mcp_gym_rollout_processor(
+    rows: List[EvaluationRow], config: RolloutProcessorConfig
+) -> List[EvaluationRow]:
     """
     Rollout processor for tau bench environments.
-    
+
     This processor starts an MCP server, creates tau bench environments, and runs rollouts
     using the eval_protocol framework, following the pattern from test_tau2_e2e.py.
-    
+
     Args:
         rows: List of EvaluationRow objects containing messages and dataset info in input_metadata
         config: RolloutProcessorConfig with model and other parameters
-    
+
     Returns:
         List of EvaluationRow objects with completed conversations
     """
     server = MCPServerManager(config.server_script_path, port=9700)
-    
+
     try:
         server.start()
-        
+
         policy = ep.LiteLLMPolicy(
             model_id=config.model,
-            temperature=config.input_params.get('temperature', 0.0),
-            max_tokens=config.input_params.get('max_tokens', 4096),
+            temperature=config.input_params.get("temperature", 0.0),
+            max_tokens=config.input_params.get("max_tokens", 4096),
         )
-        
+
         # Create MCP environments directly from evaluation_rows
-        envs = ep.make(
-            'http://localhost:9700/mcp/',
+        envs = await ep.make(
+            "http://localhost:9700/mcp/",
             evaluation_rows=rows,
             model_id=policy.model_id,
         )
-        
+
         # Run rollout with environments and policy
         evaluation_rows = await ep.rollout(
-            envs, 
-            policy=policy, 
-            steps=config.steps, 
-            max_concurrent_rollouts=config.max_concurrent_rollouts
+            envs, policy=policy, steps=config.steps, max_concurrent_rollouts=config.max_concurrent_rollouts
         )
-        
+
         return evaluation_rows
-        
+
     finally:
         # Always clean up the server
         server.stop()

eval_protocol/types/types.py

@@ -1,6 +1,8 @@
 from dataclasses import dataclass, field
 from enum import Enum
 from typing import Any, Dict, List, Optional
+from mcp.client.session import ClientSession
+from contextlib import AsyncExitStack
 
 
 class TerminationReason(str, Enum):
@@ -50,8 +52,8 @@ class MCPSession:
     last_observation: Any = None
 
     # Persistent MCP connection components
-    _exit_stack: Optional[Any] = None
-    _mcp_session: Optional[Any] = None
+    _exit_stack: Optional[AsyncExitStack] = None
+    _mcp_session: Optional[ClientSession] = None
 
 
 @dataclass

examples/blackjack_mcp/tests/test_record_and_replay_e2e.py

@@ -215,7 +215,7 @@ async def test_production_server_record_and_replay(production_server, blackjack_
         assert playback_policy.is_playback_mode(), "Should be in playback mode in CI"
 
         # Create environments for playback
-        playback_envs = ep.make(
+        playback_envs = await ep.make(
             "http://localhost:9500/mcp/",
             dataset=blackjack_dataset,
             model_id=playback_policy.model_id,
@@ -250,7 +250,7 @@ async def test_production_server_record_and_replay(production_server, blackjack_
     assert not policy.is_playback_mode(), "Should be in recording mode initially"
 
     # Create environments
-    envs = ep.make(
+    envs = await ep.make(
         "http://localhost:9500/mcp/",
         dataset=blackjack_dataset,
         model_id=policy.model_id,
@@ -310,7 +310,7 @@ async def test_production_server_record_and_replay(production_server, blackjack_
     assert playback_policy.is_playback_mode(), "Should be in playback mode"
 
     # Create new environments for playback
-    playback_envs = ep.make(
+    playback_envs = await ep.make(
         "http://localhost:9500/mcp/",
         dataset=blackjack_dataset,
         model_id=playback_policy.model_id,
@@ -462,7 +462,7 @@ async def test_blackjack_step_by_step(conda_isolation_recording_file):
         ]
 
         # Create environment pointing to conda-isolated server
-        envs = ep.make(
+        envs = await ep.make(
             f"http://localhost:{port}/mcp/",
             dataset=test_dataset,
             model_id=policy.model_id,
@@ -570,7 +570,7 @@ async def test_multi_environment_sessions(multi_env_dataset, multi_env_recording
         policy = create_blackjack_static_policy(action_sequence=["HIT", "HIT", "STICK"])
 
         # Create multiple environments
-        envs = ep.make(
+        envs = await ep.make(
             f"http://localhost:{server.port}/mcp/",
             dataset=multi_env_dataset,
             model_id=policy.model_id,
@@ -992,7 +992,7 @@ async def test_fireworks_multi_environment_sessions(multi_env_dataset, fireworks
         assert playback_policy.is_playback_mode(), "Should be in playback mode in CI"
 
         # Create environments for playback
-        playback_envs = ep.make(
+        playback_envs = await ep.make(
             "http://localhost:9500/mcp/",
             dataset=multi_env_dataset,
             model_id=playback_policy.model_id,
@@ -1033,7 +1033,7 @@ async def test_fireworks_multi_environment_sessions(multi_env_dataset, fireworks
         assert not policy.is_playback_mode(), "Should be in recording mode initially"
 
         # Create multiple environments
-        envs = ep.make(
+        envs = await ep.make(
             f"http://localhost:{server.port}/mcp/",
             dataset=multi_env_dataset,
             model_id=policy.model_id,
@@ -1149,7 +1149,7 @@ async def test_control_plane_state_querying(multi_env_dataset):
         policy = create_blackjack_static_policy(action_sequence=["HIT", "STAND"])
 
         # Create environments
-        envs = ep.make(
+        envs = await ep.make(
             f"http://localhost:{server.port}/mcp/",
             dataset=multi_env_dataset[:2],  # Use only 2 environments for faster testing
             model_id=policy.model_id,