small fix

Author: xzrderek

eval_protocol/mcp/client/connection.py

@@ -214,14 +214,12 @@ async def get_initial_state(self, session: MCPSession) -> Any:
                     # 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
-
                     async with httpx.AsyncClient(timeout=timeout) as client:
                         initial_state_response = await client.get(
                             f"{base_url}/control/initial_state",
                             headers=headers,
                             timeout=timeout,
                         )
-
                         if initial_state_response.status_code == 200:
                             initial_observation = initial_state_response.json()
                             logger.info(
@@ -234,7 +232,8 @@ async def get_initial_state(self, session: MCPSession) -> Any:
                 except httpx.TimeoutException:
                     logger.warning(f"Control plane initial state endpoint timed out after {timeout}s")
                 except Exception as e:
-                    logger.warning(f"Failed to query control plane initial state endpoint: {e}")
+                    logger.warning(f"Failed to query initial state endpoint: {e}")
+
         except Exception as e:
             logger.warning(f"Failed to query control plane initial state endpoint: {e}")
 

eval_protocol/mcp/execution/manager.py

@@ -11,9 +11,7 @@
 import os
 import threading
 import time
-from concurrent.futures import ThreadPoolExecutor, as_completed
-from dataclasses import asdict, dataclass
-from datetime import datetime
+from dataclasses import asdict
 from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Union
 
 from openai.types import CompletionUsage
@@ -260,6 +258,7 @@ async def _execute_rollout(
                 {"role": "system", "content": system_prompt},
                 {"role": "user", "content": user_prompt},
             ]
+
             logger.info(f"🎯 Starting rollout {rollout_idx} in thread {threading.current_thread().name}")
 
             # Run rollout loop for this specific environment

eval_protocol/mcp/mcpgym.py

@@ -20,7 +20,6 @@
 import logging
 import os
 import threading
-import time
 from abc import ABC, abstractmethod
 from concurrent.futures import ThreadPoolExecutor
 from datetime import date, datetime

tests/pytest/test_tau_bench_airline.py

@@ -58,7 +58,7 @@ def tau_bench_airline_to_evaluation_row(data: List[Dict[str, Any]]) -> List[Eval
 
         rows.append(eval_row)
 
-    return rows[0:3]
+    return rows
 
 
 @evaluation_test(