run on local to double check

Author: xzrderek

eval_protocol/mcp/client/connection.py

@@ -9,6 +9,7 @@
 import hashlib
 import json
 import logging
+import time
 from contextlib import AsyncExitStack
 from typing import Any, Dict, List, Optional, Tuple
 
@@ -27,9 +28,6 @@ class MCPConnectionManager:
     def __init__(self):
         self._tools_cache: Dict[str, List[Dict]] = {}
         self._tools_cache_lock = asyncio.Lock()
-        # Shared HTTP client for control plane requests with high connection limits
-        self._shared_client: Optional[httpx.AsyncClient] = None
-        self._client_lock = asyncio.Lock()
 
     async def initialize_session(self, session: MCPSession) -> None:
         """
@@ -147,8 +145,6 @@ 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"
 
@@ -177,16 +173,23 @@ async def discover_tools(self, session: MCPSession) -> List[Dict]:
 
         cache_key = session.base_url
 
-        # Check cache first (should be pre-warmed during initialization)
+        # Fast path: Check cache first without lock (safe for reads)
+        if cache_key in self._tools_cache:
+            cached_tools = self._tools_cache[cache_key]
+            logger.debug(f"Using cached tools for session {session.session_id} ({len(cached_tools)} tools)")
+            return cached_tools
+
+        # Slow path: Cache miss - use lock only for writing
         async with self._tools_cache_lock:
+            # Double-check pattern: another task might have cached it while we waited
             if cache_key in self._tools_cache:
                 cached_tools = self._tools_cache[cache_key]
                 logger.debug(f"Using cached tools for session {session.session_id} ({len(cached_tools)} tools)")
                 return cached_tools
 
-        # Fallback: if cache miss (shouldn't happen with pre-warming), fetch directly
-        logger.warning(f"Cache miss for {cache_key} - this shouldn't happen with pre-warming")
-        mcp_session = session._mcp_session
+            # Fallback: if cache miss (shouldn't happen with pre-warming), fetch directly
+            logger.warning(f"Cache miss for {cache_key} - this shouldn't happen with pre-warming")
+            mcp_session = session._mcp_session
 
         tools_response = await mcp_session.list_tools()
         tools = tools_response.tools if hasattr(tools_response, "tools") else []
@@ -233,74 +236,118 @@ async def get_initial_state(self, session: MCPSession) -> Any:
         Returns:
             Initial observation/state
         """
+        method_start = time.time()
+        session_id_short = session.session_id[:8] if len(session.session_id) > 8 else session.session_id
+        logger.info(f"### 🌟 GET_INITIAL_STATE_START: timestamp: {method_start}, session_id: {session_id_short}...")
+
         if not session._mcp_session:
+            logger.error(f"### ❌ SESSION_NOT_INITIALIZED: session_id: {session_id_short}")
             raise RuntimeError("Session not initialized")
 
         # Try to get initial state from control plane endpoint first
         initial_observation = None
 
         try:
-            import httpx
-
             # Extract base URL and session ID from the MCP session
+            url_extract_start = time.time()
+            logger.info(
+                f"### 🔍 URL_EXTRACT_START: timestamp: {url_extract_start}, elapsed: {url_extract_start - method_start:.6f}s, session_id: {session_id_short}..."
+            )
+
             base_url = session.base_url.rstrip("/").removesuffix("/mcp")
             session_id = session.session_id
 
+            url_extract_end = time.time()
+            logger.info(
+                f"### 🔍 URL_EXTRACT_END: timestamp: {url_extract_end}, elapsed: {url_extract_end - method_start:.6f}s, duration: {url_extract_end - url_extract_start:.6f}s, base_url: {base_url}, session_id: {session_id_short}..."
+            )
+
             if session_id:
+                headers_start = time.time()
+                logger.info(
+                    f"### 🔍 HEADERS_CREATE_START: timestamp: {headers_start}, elapsed: {headers_start - method_start:.6f}s, session_id: {session_id_short}..."
+                )
+
                 headers = {"mcp-session-id": session_id}
 
+                headers_end = time.time()
+                logger.info(
+                    f"### 🔍 HEADERS_CREATE_END: timestamp: {headers_end}, elapsed: {headers_end - method_start:.6f}s, duration: {headers_end - headers_start:.6f}s, session_id: {session_id_short}..."
+                )
+
                 # Query initial state endpoint
                 try:
+                    timeout_start = time.time()
+                    logger.info(
+                        f"### 🔍 TIMEOUT_CONFIG_START: timestamp: {timeout_start}, elapsed: {timeout_start - method_start:.6f}s, session_id: {session_id_short}..."
+                    )
+
                     # Use shorter timeout for playback mode, longer timeout for high-concurrency initialization
                     # (50+ concurrent sessions need more time for initial state setup)
                     timeout = 3.0 if hasattr(session, "_is_playback_mode") and session._is_playback_mode else 15.0
 
-                    # TIMING: Get shared client
-                    client_start = __import__("time").time()
-                    client = await self._get_shared_client(timeout)
-                    client_time = __import__("time").time() - client_start
+                    timeout_end = time.time()
                     logger.info(
-                        f"DEBUG_CLIENT: Getting shared client took {client_time:.3f}s for {session.session_id}"
+                        f"### 🔍 TIMEOUT_CONFIG_END: timestamp: {timeout_end}, elapsed: {timeout_end - method_start:.6f}s, duration: {timeout_end - timeout_start:.6f}s, timeout: {timeout}s, session_id: {session_id_short}..."
                     )
 
+                    # TIMING: Get shared client
+                    # client = await self._get_shared_client(timeout)
+
                     # TIMING: HTTP request with shared client
-                    request_start = __import__("time").time()
-                    initial_state_response = await client.get(
-                        f"{base_url}/control/initial_state",
-                        headers=headers,
-                        timeout=timeout,
+                    request_start = time.time()
+                    logger.info(
+                        f"### 🌐 HTTP_REQUEST_START: timestamp: {request_start}, elapsed: {request_start - method_start:.6f}s, url: {base_url}/control/initial_state, session_id: {session_id_short}..."
                     )
-                    request_time = __import__("time").time() - request_start
-                    logger.info(f"DEBUG_REQUEST: HTTP request took {request_time:.3f}s for {session.session_id}")
-                    if initial_state_response.status_code == 200:
-                        initial_observation = initial_state_response.json()
-                        logger.info(
-                            f"Session {session.session_id}: ✅ Successfully fetched session-aware initial state from control plane endpoint"
+
+                    timeout = 3.0 if hasattr(session, "_is_playback_mode") and session._is_playback_mode else 15.0
+
+                    async with httpx.AsyncClient(timeout=timeout) as client:
+                        initial_state_response = await client.get(
+                            f"{base_url}/control/initial_state",
+                            headers=headers,
+                            timeout=timeout,
                         )
-                    else:
-                        logger.warning(
-                            f"Control plane initial state endpoint returned {initial_state_response.status_code}"
+                        request_time = time.time() - request_start
+
+                        request_end = time.time()
+                        logger.info(
+                            f"### 🌐 HTTP_REQUEST_END: timestamp: {request_end}, elapsed: {request_end - method_start:.6f}s, duration: {request_time:.6f}s, status_code: {initial_state_response.status_code}, session_id: {session_id_short}..."
                         )
+
+                        if initial_state_response.status_code == 200:
+                            initial_observation = initial_state_response.json()
+                            success_end = time.time()
+                            logger.info(
+                                f"### ✅ RETURN: timestamp: {success_end}, total_duration: {success_end - method_start:.6f}s, session_id: {session_id_short}..."
+                            )
+                            # return initial_observation
+                        else:
+                            error_time = time.time()
+                            logger.warning(
+                                f"### ⚠️ HTTP_ERROR_RESPONSE: timestamp: {error_time}, elapsed: {error_time - method_start:.6f}s, status_code: {initial_state_response.status_code}, session_id: {session_id_short}"
+                            )
                 except httpx.TimeoutException:
-                    logger.warning(f"Control plane initial state endpoint timed out after {timeout}s")
+                    timeout_error_time = time.time()
+                    logger.warning(
+                        f"### ⏰ HTTP_TIMEOUT: timestamp: {timeout_error_time}, elapsed: {timeout_error_time - method_start:.6f}s, timeout: {timeout}s, session_id: {session_id_short}"
+                    )
                 except Exception as e:
-                    logger.warning(f"Failed to query initial state endpoint: {e}")
+                    http_error_time = time.time()
+                    logger.warning(
+                        f"### ❌ HTTP_ERROR: timestamp: {http_error_time}, elapsed: {http_error_time - method_start:.6f}s, error: {str(e)}, session_id: {session_id_short}"
+                    )
 
         except Exception as e:
-            logger.warning(f"Failed to query control plane initial state endpoint: {e}")
-
-        # Fallback to MCP resource if control plane endpoint fails (backward compatibility)
-        if initial_observation is None:
-            logger.debug(f"Session {session.session_id}: Falling back to MCP resource for initial state")
-            initial_observation = await self._get_initial_state_from_mcp_resource(session)
-
-        # Ensure we have some observation
-        if initial_observation is None:
-            logger.debug(f"Session {session.session_id}: Using default initial state")
-            initial_observation = {
-                "observation": "default_initial_state",
-                "session_id": session.session_id,
-            }
+            general_error_time = time.time()
+            logger.warning(
+                f"### ❌ GENERAL_ERROR: timestamp: {general_error_time}, elapsed: {general_error_time - method_start:.6f}s, error: {str(e)}, session_id: {session_id_short}"
+            )
+
+        method_end = time.time()
+        logger.info(
+            f"### 🔴 GET_INITIAL_STATE_END: timestamp: {method_end}, total_duration: {method_end - method_start:.6f}s, session_id: {session_id_short}..."
+        )
 
         return initial_observation
 
@@ -509,9 +556,6 @@ async def call_tool(self, session: MCPSession, tool_name: str, arguments: Dict)
         control_plane_info = {}
 
         try:
-            # Query control plane endpoints following the new architecture
-            import httpx
-
             # Extract base URL and session ID from the MCP session
             base_url = session.base_url.rstrip("/").removesuffix("/mcp")
             # Use the session ID from the established MCP session
@@ -601,47 +645,3 @@ async def close_session(self, session: MCPSession) -> None:
             finally:
                 session._exit_stack = None
                 session._mcp_session = None
-
-    async def _get_shared_client(self, timeout: float) -> httpx.AsyncClient:
-        """
-        Get or create a shared HTTP client with high connection limits for concurrent requests.
-
-        Args:
-            timeout: Timeout for requests
-
-        Returns:
-            Shared httpx.AsyncClient instance
-        """
-        # Fast path: if client exists and is not closed, return it immediately
-        if self._shared_client is not None and not self._shared_client.is_closed:
-            return self._shared_client
-
-        # Slow path: need to create client (use lock only for creation)
-        async with self._client_lock:
-            # Double-check pattern: another task might have created it while we waited
-            if self._shared_client is None or self._shared_client.is_closed:
-                # Create HTTP client with high connection limits for concurrent initial state requests
-                limits = httpx.Limits(
-                    max_keepalive_connections=None,  # Unlimited keep-alive connections
-                    max_connections=None,  # Unlimited total connection pool size
-                    keepalive_expiry=30.0,  # Keep connections alive for 30s
-                )
-
-                self._shared_client = httpx.AsyncClient(
-                    timeout=timeout,
-                    limits=limits,
-                    # Enable connection pooling and keep-alive
-                    http2=False,  # Disable HTTP/2 for better connection pooling with many concurrent requests
-                )
-                logger.info(
-                    "Created shared HTTP client with unlimited connection limits for MCP control plane requests"
-                )
-
-        return self._shared_client
-
-    async def close_shared_client(self):
-        """Close the shared HTTP client when shutting down."""
-        async with self._client_lock:
-            if self._shared_client and not self._shared_client.is_closed:
-                await self._shared_client.aclose()
-                self._shared_client = None

eval_protocol/mcp/execution/manager.py

@@ -191,9 +191,7 @@ async def _execute_rollout(
         dataset_row = envs.dataset_rows[rollout_idx]
         rollout_start = time.time()
         elapsed_from_main_start = rollout_start - start_time
-        logger.info(
-            f"DEBUG4. Starting rollout {dataset_row.id} at {datetime.fromtimestamp(rollout_start).strftime('%H:%M:%S.%f')[:-3]} (+{elapsed_from_main_start:.3f}s from start)"
-        )
+        logger.info(f"DEBUG4. Starting rollout {dataset_row.id} at {rollout_start}")
 
         # Initialize trajectory
         trajectory = Trajectory(
@@ -219,7 +217,7 @@ async def _execute_rollout(
         temp_start = time.time()
         current_observation, tool_schema = await envs.reset(session)
         logger.info(
-            f"DEBUG6: User simulator get_init_state took {time.time() - temp_start:.3f}s for {session.session_id}"
+            f"DEBUG6: User simulator get_init_state took {time.time() - temp_start:.3f}s for {session.session_id}, started at {temp_start}"
         )
         system_prompt = dataset_row.system_prompt
 
@@ -240,7 +238,7 @@ async def _execute_rollout(
 
             # Get initial messages in tau2-bench format for user simulator
             user_simulator_state = user_simulator.get_init_state()
-            user_message, user_simulator_state = user_simulator.generate_next_message(
+            user_message, user_simulator_state = await user_simulator.generate_next_message(
                 AssistantMessage(role="assistant", content="Hi! How can I help you today?"),
                 user_simulator_state,
             )
@@ -277,11 +275,11 @@ async def _execute_rollout(
                 if user_simulator_messages and isinstance(user_simulator_messages[-1], AssistantMessage):
                     # Generate user response using the simulator
                     temp_start1 = time.time()
-                    user_message, user_simulator_state = user_simulator.generate_next_message(
+                    user_message, user_simulator_state = await user_simulator.generate_next_message(
                         user_simulator_messages[-1], user_simulator_state
                     )
                     logger.info(
-                        f"DEBUG8: User simulator generate_next_message took {time.time() - temp_start1:.3f}s for {dataset_row.id}"
+                        f"DEBUG8: User simulator generate_next_message took {time.time() - temp_start1:.3f}s for {dataset_row.id}, started at {temp_start1}"
                     )
                     user_content = user_message.content if user_message.content else ""
 
@@ -297,7 +295,9 @@ async def _execute_rollout(
             while not turn_completed and not trajectory.terminated:
                 temp_start2 = time.time()
                 tool_calls, usage_stats = await policy(tool_schema, rollout_idx, conversation_history)
-                logger.info(f"DEBUG9: Policy took {time.time() - temp_start2:.3f}s for {dataset_row.id}")
+                logger.info(
+                    f"DEBUG9: Policy took {time.time() - temp_start2:.3f}s for {dataset_row.id}, started at {temp_start2}"
+                )
 
                 # If no tool call is generated, turn is finished
                 if len(tool_calls) == 1:
@@ -316,7 +316,9 @@ async def _execute_rollout(
                     # Execute tool call for this environment
                     temp_start3 = time.time()
                     observation, reward, rollout_end, info = await envs.step(rollout_idx, tool_call)
-                    logger.info(f"DEBUG10: Env step took {time.time() - temp_start3:.3f}s for {dataset_row.id}")
+                    logger.info(
+                        f"DEBUG10: Env step took {time.time() - temp_start3:.3f}s for {dataset_row.id}, started at {temp_start3}"
+                    )
 
                     tool_response = envs.format_tool_response(observation)
 
@@ -464,9 +466,7 @@ async def _execute_rollout(
         logger.info(
             f"✅ Rollout {rollout_idx} completed: {trajectory.steps} steps, reward: {trajectory.total_reward:.2f}, termination: {trajectory.termination_reason}, in thread {threading.current_thread().name}"
         )
-        logger.info(
-            f"DEBUG11: Rollout {dataset_row.id} completed at {datetime.fromtimestamp(time.time()).strftime('%H:%M:%S.%f')[:-3]} (+{time.time() - rollout_start:.3f}s from start)"
-        )
+        logger.info(f"DEBUG11: Rollout {dataset_row.id} completed at {time.time()}, started at {rollout_start}")
         return trajectory
 
     async def _get_control_plane_status(self, session) -> Optional[Dict[str, Any]]: