fix(deps): pin patched LiteLLM and repair integration CI

eval_protocol/rewards/code_execution.py

@@ -833,7 +833,10 @@ def execute_code_with_e2b(
                 "error": "API key is required for E2B execution. Set it using the api_key parameter or E2B_API_KEY environment variable.",
             }
 
-        with Sandbox(api_key=api_key) as sandbox:
+        # E2B v2's factory carries the sandbox access token required by envd.
+        # Retain constructor support for older optional code-interpreter SDKs.
+        sandbox_factory = getattr(Sandbox, "create", Sandbox)
+        with sandbox_factory(api_key=api_key) as sandbox:
             stdout = []
             stderr = []
 
@@ -849,8 +852,6 @@ def capture_stderr(output):
                 else:
                     stderr.append(str(output))
 
-            sandbox.on_exit = lambda *args: None  # type: ignore[method-assign, assignment]
-
             if language.lower() in ["python", "py"]:
                 file_path = "/code/script.py"
                 cmd = "python3 /code/script.py"
@@ -866,7 +867,9 @@ def capture_stderr(output):
 
             try:
                 fs_handler = None
-                if _E2B_SOURCE == "e2b_code_interpreter":
+                if hasattr(sandbox, "files"):
+                    fs_handler = sandbox.files
+                elif _E2B_SOURCE == "e2b_code_interpreter":
                     if hasattr(sandbox, "filesystem"):
                         fs_handler = sandbox.filesystem
                 elif _E2B_SOURCE == "e2b":

pyproject.toml

@@ -25,13 +25,13 @@ dependencies = [
     "openai>=1.78.1",
     "aiosqlite",
     "aiohttp",
-    "mcp>=1.9.2",
+    "mcp>=1.30.0,<2",
     "PyYAML>=5.0",
     "hydra-core>=1.3.2",
     "omegaconf>=2.3.0",
     "httpx>=0.24.0",
     "anthropic>=0.59.0",
-    "litellm==1.84.3",
+    "litellm==1.105.0",
     "pytest>=6.0.0",
     "pytest-asyncio>=0.21.0",
     "peewee>=3.18.2",
@@ -70,7 +70,8 @@ dev = [
     "versioneer>=0.20",
     "openai>=1.78.1",
     "pre-commit",
-    "e2b",
+    # Reproducible v2 sandbox tests without changing the optional tracing stack.
+    "e2b==2.34.0",
     "pytest-cov",
     "pytest-xdist",
     "docker==7.1.0",

scripts/run_sharded_tests.sh

@@ -71,9 +71,11 @@ if [ "$DRY_RUN" = "--dry-run" ]; then
 fi
 
 # Run tests for this shard
+# Keep evaluation tests and their shared-processor verification in the same worker.
 # shellcheck disable=SC2086
 exec pytest \
 	-n auto \
+	--dist=loadfile \
 	--ignore=eval_protocol/benchmarks/ \
 	--ignore=eval_protocol/quickstart/ \
 	-v --durations=10 \

tests/klavis_helpers.py

@@ -0,0 +1,93 @@
+"""Helpers shared by live Klavis evaluations and offline tests."""
+
+import json
+import logging
+import os
+
+from openai import AsyncOpenAI
+from pydantic import BaseModel
+
+from eval_protocol.models import EvaluateResult, EvaluationRow, Message
+
+logger = logging.getLogger(__name__)
+MODEL = os.environ.get("EP_KLAVIS_MODEL", "accounts/fireworks/models/gpt-oss-120b")
+
+
+class ResponseFormat(BaseModel):
+    score: float
+
+
+async def judge_klavis_response(row: EvaluationRow) -> EvaluationRow:
+    """Judge a completed Klavis rollout; service errors must fail the live test."""
+    ground_truth = row.ground_truth
+
+    async with AsyncOpenAI(
+        api_key=os.environ["FIREWORKS_API_KEY"], base_url="https://api.fireworks.ai/inference/v1"
+    ) as client:
+        response = await client.chat.completions.create(
+            model=MODEL,
+            messages=[
+                {
+                    "role": "system",
+                    "content": "You are judging the output of the model versus the ground truth. Return score = 1 if the output contains the ground truth, 0 otherwise.",
+                },
+                {
+                    "role": "user",
+                    "content": f"Final model output: {row.messages[-1].content}\nGround truth: {ground_truth}",
+                },
+            ],
+            response_format={
+                "type": "json_schema",
+                "json_schema": {"name": "ResponseFormat", "schema": ResponseFormat.model_json_schema()},
+            },
+        )
+        response_text = response.choices[0].message.content
+        logger.info("response_text: %s", response_text)
+        try:
+            parsed = json.loads(response_text or "{}")
+            score = parsed.get("score", 0.0)
+        except (json.JSONDecodeError, TypeError):
+            logger.warning("Failed to parse response as JSON: %s", response_text)
+            score = 0.0
+
+        row.evaluation_result = EvaluateResult(score=score, reason=response_text)
+    return row
+
+
+def klavis_gmail_sandbox_dataset_adapter(rows: list[dict]) -> list[EvaluationRow]:
+    """Dataset adapter for sandbox JSONL rows.
+
+    Supports the new schema:
+      - initialize_data: dict (passed to Klavis sandbox initializer)
+      - messages: str (task instruction)
+      - ground_truth: dict (expected final sandbox state)
+
+    """
+    adapted: list[EvaluationRow] = []
+    system_prompt = (
+        "You are a helpful assistant with access to Gmail. "
+        "You can send emails, draft emails, and manage messages, etc."
+    )
+
+    for r in rows:
+        if isinstance(r.get("messages"), str) and "initialize_data" in r:
+            init_data = r.get("initialize_data") or {}
+            task = r.get("messages") or ""
+            ground_truth = r.get("ground_truth")
+
+            row = EvaluationRow(
+                messages=[
+                    Message(role="system", content=system_prompt),
+                    Message(role="user", content=task),
+                ],
+                ground_truth=ground_truth,
+            )
+            row.input_metadata.session_data = {
+                "initialize_data": init_data,
+                "task": task,
+            }
+            adapted.append(row)
+        else:
+            adapted.append(EvaluationRow(**r))
+
+    return adapted

tests/pytest/test_pytest_klavis_mcp.py

@@ -1,66 +1,25 @@
-from eval_protocol.models import EvaluateResult, EvaluationRow, Message
-from eval_protocol.pytest import AgentRolloutProcessor, evaluation_test
-from openai import AsyncOpenAI
-import json
-from pydantic import BaseModel
-import logging
-
-logger = logging.getLogger(__name__)
 import os
 
+import pytest
 
-class ResponseFormat(BaseModel):
-    score: float
-
-
-"""
-You should copy https://painted-tennis-ebc.notion.site/MCPMark-Source-Hub-23181626b6d7805fb3a7d59c63033819
-into your Notion for the notion test.
-"""
+from eval_protocol.models import EvaluationRow
+from eval_protocol.pytest import AgentRolloutProcessor
+from eval_protocol.pytest.evaluation_test import evaluation_test
+from tests.klavis_helpers import MODEL, judge_klavis_response
 
 
+# Live datasets require a Klavis connection to Gmail/Notion and their seeded test data.
+@pytest.mark.skipif(
+    not (os.environ.get("KLAVIS_AUTH_TOKEN") and os.environ.get("FIREWORKS_API_KEY")),
+    reason="Live Klavis MCP test requires KLAVIS_AUTH_TOKEN and FIREWORKS_API_KEY",
+)
 @evaluation_test(
     input_dataset=["tests/pytest/datasets/klavis_mcp_test.jsonl"],
     rollout_processor=AgentRolloutProcessor(),
-    completion_params=[{"model": "fireworks_ai/accounts/fireworks/models/kimi-k2p5"}],
+    completion_params=[{"model": f"fireworks_ai/{MODEL}"}],
     mode="pointwise",
     mcp_config_path="tests/pytest/mcp_configurations/klavis_strata_mcp.json",
 )
 async def test_pytest_klavis_mcp(row: EvaluationRow) -> EvaluationRow:
-    ground_truth = row.ground_truth
-    # check if the final messages contains the ground truth
-
-    async with AsyncOpenAI(
-        api_key=os.environ["FIREWORKS_API_KEY"], base_url="https://api.fireworks.ai/inference/v1"
-    ) as client:
-        response = await client.chat.completions.create(
-            model="accounts/fireworks/models/kimi-k2p5",
-            messages=[
-                {
-                    "role": "system",
-                    "content": "You are judging the output of the model versus the ground truth. Return score = 1 if the output contains the ground truth, 0 otherwise.",
-                },
-                {
-                    "role": "user",
-                    "content": f"Final model output: {row.messages[-1].content}\nGround truth: {ground_truth}",
-                },
-            ],
-            response_format={
-                "type": "json_schema",
-                "json_schema": {"name": "ResponseFormat", "schema": ResponseFormat.model_json_schema()},
-            },
-        )
-        response_text = response.choices[0].message.content
-        logger.info("response_text: %s", response_text)
-        try:
-            parsed = json.loads(response_text or "{}")
-            score = parsed.get("score", 0.0)
-        except (json.JSONDecodeError, TypeError):
-            logger.warning("Failed to parse response as JSON: %s", response_text)
-            score = 0.0
-
-        row.evaluation_result = EvaluateResult(
-            score=score,
-            reason=response_text,
-        )
-    return row
+    assert row.rollout_status.is_finished(), row.rollout_status
+    return await judge_klavis_response(row)

tests/pytest/test_pytest_klavis_sandbox.py

@@ -2,73 +2,43 @@
 import logging
 import os
 
-from eval_protocol.models import EvaluateResult, EvaluationRow, Message
-from eval_protocol.pytest import KlavisSandboxRolloutProcessor, evaluation_test
+import pytest
 from openai import AsyncOpenAI
 from pydantic import BaseModel
 
+from eval_protocol.models import EvaluateResult, EvaluationRow
+from eval_protocol.pytest import KlavisSandboxRolloutProcessor
+from eval_protocol.pytest.evaluation_test import evaluation_test
+from tests.klavis_helpers import MODEL, klavis_gmail_sandbox_dataset_adapter
+
 logger = logging.getLogger(__name__)
 
 
 class ResponseFormat(BaseModel):
     score: float
-
-
-def klavis_gmail_sandbox_dataset_adapter(rows: list[dict]) -> list[EvaluationRow]:
-    """Dataset adapter for sandbox JSONL rows.
-
-    Supports the new schema:
-      - initialize_data: dict (passed to Klavis sandbox initializer)
-      - messages: str (task instruction)
-      - ground_truth: dict (expected final sandbox state)
-
-    """
-    adapted: list[EvaluationRow] = []
-    system_prompt = (
-        "You are a helpful assistant with access to Gmail. "
-        "You can send emails, draft emails, and manage messages, etc."
-    )
-
-    for r in rows:
-        if isinstance(r.get("messages"), str) and "initialize_data" in r:
-            init_data = r.get("initialize_data") or {}
-            task = r.get("messages") or ""
-            ground_truth = r.get("ground_truth")
-
-            row = EvaluationRow(
-                messages=[
-                    Message(role="system", content=system_prompt),
-                    Message(role="user", content=task),
-                ],
-                ground_truth=ground_truth,
-            )
-            row.input_metadata.session_data = {
-                "initialize_data": init_data,
-                "task": task,
-            }
-            adapted.append(row)
-        else:
-            adapted.append(EvaluationRow(**r))
-
-    return adapted
 
 
+@pytest.mark.skipif(
+    not (os.environ.get("KLAVIS_API_KEY") and os.environ.get("FIREWORKS_API_KEY")),
+    reason="Live Klavis Gmail test requires KLAVIS_API_KEY and FIREWORKS_API_KEY",
+)
 @evaluation_test(
     input_dataset=["tests/pytest/datasets/klavis_gmail_sandbox_test.jsonl"],
     rollout_processor=KlavisSandboxRolloutProcessor(
         server_name="gmail",
     ),
-    completion_params=[{"model": "fireworks_ai/accounts/fireworks/models/kimi-k2p5"}],
+    completion_params=[{"model": f"fireworks_ai/{MODEL}"}],
     mode="pointwise",
     dataset_adapter=klavis_gmail_sandbox_dataset_adapter,
 )
 async def test_pytest_gmail_sandbox(row: EvaluationRow) -> EvaluationRow:
     """
     Evaluate Gmail sandbox results by comparing with ground truth using LLM judge.
-    
+
     The sandbox data is exported after agent execution and compared with expected output.
     Sandbox data is available in row.execution_metadata.extra["sandbox_data"].
     """
+    assert row.rollout_status.is_finished(), row.rollout_status
     ground_truth = row.ground_truth
     sandbox_data = row.execution_metadata.extra.get("sandbox_data", {}) if row.execution_metadata.extra else {}
     final_message = row.messages[-1].content if row.messages else ""
@@ -83,11 +53,10 @@ async def test_pytest_gmail_sandbox(row: EvaluationRow) -> EvaluationRow:
     async with AsyncOpenAI(
         api_key=os.environ["FIREWORKS_API_KEY"], base_url="https://api.fireworks.ai/inference/v1"
     ) as client:
-
         evaluation_prompt = f"""You are evaluating an AI agent's performance on a Gmail sandbox task.
 
 Task:
-{task or (row.messages[-1].content if row.messages else 'N/A')}
+{task or (row.messages[-1].content if row.messages else "N/A")}
 
 Initial Gmail Sandbox State (initialize_data):
 {json.dumps(initialize_data, indent=2, default=str)}
@@ -110,7 +79,7 @@ async def test_pytest_gmail_sandbox(row: EvaluationRow) -> EvaluationRow:
 
         try:
             response = await client.chat.completions.create(
-                model="accounts/fireworks/models/kimi-k2p5",
+                model=MODEL,
                 messages=[
                     {
                         "role": "system",
@@ -134,9 +103,6 @@ async def test_pytest_gmail_sandbox(row: EvaluationRow) -> EvaluationRow:
             row.evaluation_result = EvaluateResult(score=score)
         except Exception as e:
             logger.error(f"Error during LLM evaluation: {str(e)}", exc_info=True)
-            row.evaluation_result = EvaluateResult(
-                score=0.0,
-                reason=f"Evaluation error: {str(e)}",
-            )
+            raise
 
     return row