fix(adapters): Toolathlon spawn-time token resolution + exit-code verifier scoring

src/benchflow/adapters/_toolathlon.py

@@ -3,19 +3,22 @@
 from __future__ import annotations
 
 import json
+import logging
 import re
 import shlex
 import shutil
 from dataclasses import dataclass
+from importlib import resources
 from pathlib import Path
 from typing import Any
 
 import tomli_w
 import yaml
 
+logger = logging.getLogger(__name__)
+
 _NOOP_EXCLUDE_TAG = "__benchflow_exclude_no_tools__"
 _SAFE_NAME_RE = re.compile(r"[^A-Za-z0-9._-]+")
-_TOKEN_PLACEHOLDER_RE = re.compile(r"\$\{token\.([A-Za-z0-9_]+)\}")
 _TOOLATHLON_UVX_PACKAGE_PINS = {
     "office-word-mcp-server": "office-word-mcp-server==1.1.11",
 }
@@ -27,7 +30,6 @@ class _ToolathlonCredentialSpec:
     env: str
     b64_env: str
     trigger_servers: tuple[str, ...] = ()
-    token_defaults: tuple[tuple[str, str], ...] = ()
     copy_paths: tuple[str, ...] = ()
 
 
@@ -37,28 +39,12 @@ class _ToolathlonCredentialSpec:
         env="TOOLATHLON_GCP_SERVICE_ACCOUNT_JSON",
         b64_env="TOOLATHLON_GCP_SERVICE_ACCOUNT_JSON_B64",
         trigger_servers=("google-cloud",),
-        token_defaults=(
-            (
-                "gcp_service_account_path",
-                "/workspace/configs/gcp-service_account.keys.json",
-            ),
-        ),
     ),
     _ToolathlonCredentialSpec(
         path="configs/google_credentials.json",
         env="TOOLATHLON_GOOGLE_CREDENTIALS_JSON",
         b64_env="TOOLATHLON_GOOGLE_CREDENTIALS_JSON_B64",
         trigger_servers=("google_calendar", "google_sheet"),
-        token_defaults=(
-            (
-                "google_oauth2_credentials_path",
-                "/workspace/configs/google_credentials.json",
-            ),
-            (
-                "google_oauth2_token_path",
-                "/workspace/agent_workspace/.toolathlon/google_credentials.json",
-            ),
-        ),
         copy_paths=(
             "agent_workspace/.toolathlon/google_credentials.json",
             "agent_workspace/.toolathlon/calendar_credentials.json",
@@ -77,11 +63,59 @@ class _ToolathlonCredentialSpec:
 _TOOLATHLON_CREDENTIAL_REF_RE = re.compile(
     "|".join(re.escape(spec.path) for spec in _TOOLATHLON_CREDENTIAL_SPECS)
 )
-_TOOLATHLON_TOKEN_DEFAULTS = {
-    token: value
-    for spec in _TOOLATHLON_CREDENTIAL_SPECS
-    for token, value in spec.token_defaults
-}
+
+# Secret ``token.X`` values baked into the container-side global
+# ``token_key_session.py`` (see ``_toolathlon_container._GLOBAL_TOKENS``),
+# resolved host-side from BenchFlow's environment/.env at setup time.
+_TOOLATHLON_TOKEN_SECRET_ENVS = (
+    "TOOLATHLON_GCP_PROJECT_ID",
+    "TOOLATHLON_MAPS_API_KEY",
+    "TOOLATHLON_SERPER_API_KEY",
+    "TOOLATHLON_GITHUB_TOKEN",
+    "TOOLATHLON_HF_TOKEN",
+    "TOOLATHLON_WANDB_API_KEY",
+    "TOOLATHLON_NOTION_KEY",
+    "TOOLATHLON_NOTION_KEY_EVAL",
+    "TOOLATHLON_GOOGLE_CLIENT_ID",
+    "TOOLATHLON_GOOGLE_CLIENT_SECRET",
+    "TOOLATHLON_GOOGLE_REFRESH_TOKEN",
+    "TOOLATHLON_SNOWFLAKE_ACCOUNT",
+    "TOOLATHLON_SNOWFLAKE_USER",
+    "TOOLATHLON_SNOWFLAKE_ROLE",
+    "TOOLATHLON_SNOWFLAKE_WAREHOUSE",
+    "TOOLATHLON_SNOWFLAKE_DATABASE",
+    "TOOLATHLON_SNOWFLAKE_SCHEMA",
+)
+
+_TOOLATHLON_CONTAINER_MODULE_PATH = "/workspace/.toolathlon/toolathlon_container.py"
+
+# MCP server flags whose value is a working DIRECTORY the server (and the task's
+# evaluator) expect to exist. Upstream servers create these lazily, so an agent
+# that never exercises the server leaves the directory absent and evaluators
+# that ``os.listdir`` it crash. The launcher ``mkdir -p``s these. Explicit
+# allow-list, not a heuristic: every entry takes a directory (never a file).
+_TOOLATHLON_MCP_DIR_ARG_FLAGS = frozenset(
+    {
+        "--storage-path",
+        "--attachment_download_path",
+        "--attachment_upload_path",
+        "--email_export_path",
+    }
+)
+
+
+def _toolathlon_container_python(variant: str) -> str:
+    # Absolute so FastMCP can spawn the launcher even without PATH in the MCP
+    # server subprocess environment.
+    return "/opt/venv/bin/python3" if variant == "gym" else "/usr/bin/python3"
+
+
+def _toolathlon_container_module_source() -> str:
+    return (
+        resources.files("benchflow.adapters")
+        .joinpath("_toolathlon_container.py")
+        .read_text()
+    )
 
 
 def _safe_name(value: str) -> str:
@@ -131,10 +165,24 @@ def materialize_toolathlon(ctx: Any, output_dir: Path, *, variant: str) -> None:
         credential_refs = _toolathlon_credential_refs_for_task(
             upstream_task, task_config, variant=variant
         )
-        mcp_servers = [
-            _toolathlon_mcp_server(ctx, server_name, variant=variant)
-            for server_name in task_config.get("needed_mcp_servers", [])
-        ]
+        mcp_servers = []
+        for server_name in task_config.get("needed_mcp_servers", []):
+            try:
+                mcp_servers.append(
+                    _toolathlon_mcp_server(
+                        ctx, server_name, variant=variant, task_name=task_name
+                    )
+                )
+            except FileNotFoundError:
+                # ``task_config`` may name a local harness tool (e.g.
+                # ``web_search``) rather than an MCP server — upstream resolves
+                # those from its agent-tool registry, so no configs/mcp_servers
+                # yaml exists. Drop it instead of failing the adaptation.
+                logger.warning(
+                    "Toolathlon task %s: no MCP config for %r; dropping server",
+                    task_name,
+                    server_name,
+                )
 
         _write_text(task_dir / "instruction.md", _toolathlon_instruction(upstream_task))
         _write_toml(
@@ -218,6 +266,14 @@ def _toolathlon_task_toml(
 ) -> dict[str, Any]:
     benchmark_name = "toolathlon-gym" if variant == "gym" else "toolathlon"
     setup_commands: list[dict[str, Any]] = []
+    setup_commands.append(
+        {
+            "command": _toolathlon_token_setup_command(variant),
+            "cwd": "/workspace",
+            "timeout_sec": 60.0,
+            "env": _toolathlon_token_setup_env(),
+        }
+    )
     credential_command = _toolathlon_credential_setup_command(credential_refs)
     if credential_command:
         setup_commands.append(
@@ -370,6 +426,50 @@ def _toolathlon_credential_setup_command(credential_refs: set[str]) -> str | Non
     )
 
 
+def _toolathlon_token_setup_env() -> dict[str, str]:
+    """Env for the token-setup command: the container helper (base64) plus the
+    secret token values, resolved host-side from BenchFlow's environment/.env."""
+    import base64
+
+    module_b64 = base64.b64encode(
+        _toolathlon_container_module_source().encode()
+    ).decode()
+    env: dict[str, str] = {"TOOLATHLON_CONTAINER_MODULE_B64": module_b64}
+    for name in _TOOLATHLON_TOKEN_SECRET_ENVS:
+        env[name] = f"${{{name}:-}}"
+    return env
+
+
+def _toolathlon_token_setup_command(variant: str) -> str:
+    """Stage the container helper and generate the global token_key_session.py.
+
+    The helper is decoded from ``TOOLATHLON_CONTAINER_MODULE_B64`` and its
+    ``write-config`` mode bakes the injected secret tokens into
+    ``/workspace/configs/token_key_session.py`` (gitignored upstream, so absent
+    from the image). MCP servers later resolve ``${token.X}`` against it via the
+    same helper's ``launch`` mode.
+    """
+    python_cmd = _toolathlon_container_python(variant)
+    module_path = shlex.quote(_TOOLATHLON_CONTAINER_MODULE_PATH)
+    return "\n".join(
+        [
+            "set -e",
+            "mkdir -p /workspace/.toolathlon",
+            f"{python_cmd} - <<'PY'",
+            "import base64, os, pathlib",
+            f"target = pathlib.Path({_TOOLATHLON_CONTAINER_MODULE_PATH!r})",
+            "target.write_bytes(",
+            "    base64.b64decode(os.environ['TOOLATHLON_CONTAINER_MODULE_B64'])",
+            ")",
+            "target.chmod(0o755)",
+            "PY",
+            "chmod -R a+rX /workspace/.toolathlon",
+            f"{python_cmd} {module_path} write-config",
+            "chmod 0644 /workspace/configs/token_key_session.py",
+        ]
+    )
+
+
 def _toolathlon_setup_command(*, task_name: str, variant: str) -> str:
     python_cmd = (
         "/opt/venv/bin/python3" if variant == "gym" else "/usr/local/bin/uv run python"
@@ -455,8 +555,16 @@ def _toolathlon_test_sh(*, task_name: str, variant: str) -> str:
             'if [ "$status" -eq 0 ]; then',
             "  printf '{\"reward\": 1.0}\\n' > /logs/verifier/reward.json",
             "  printf '1.0\\n' > /logs/verifier/reward.txt",
-            "elif grep -q 'Traceback (most recent call last):' \"$EVAL_LOG\"; then",
-            '  echo "BenchFlow Toolathlon verifier setup error: evaluator crashed" >&2',
+            # Upstream evaluators score by exit code: a non-zero exit is a task
+            # FAIL (reward 0), and several signal failure by *raising* (e.g.
+            # ``raise ValueError('Some tests FAILED')``) rather than returning
+            # False. So a traceback alone is not an infra error. Only escalate on
+            # signatures that can't be an agent failure — a broken evaluator
+            # ENVIRONMENT (missing module/dependency, permission denied) — which
+            # a rerun cannot fix and which should surface, not silently score 0.
+            "elif grep -qE "
+            "'ModuleNotFoundError|ImportError|PermissionError: ' \"$EVAL_LOG\"; then",
+            '  echo "BenchFlow Toolathlon verifier setup error: evaluator environment failure" >&2',
             '  exit "$status"',
             "else",
             "  printf '{\"reward\": 0.0}\\n' > /logs/verifier/reward.json",
@@ -477,11 +585,15 @@ def _toolathlon_test_sh(*, task_name: str, variant: str) -> str:
 
 
 def _toolathlon_mcp_server(
-    ctx: Any, server_name: str, *, variant: str
+    ctx: Any, server_name: str, *, variant: str, task_name: str
 ) -> dict[str, Any]:
     config_path = _toolathlon_mcp_config_path(ctx, server_name)
     data = yaml.safe_load(config_path.read_text())
     params = data.get("params") or {}
+    # Static path placeholders (``${agent_workspace}`` …) resolve to fixed
+    # in-container paths now; ``${token.X}`` are left intact and resolved at
+    # spawn time by the container launcher, the only layer that sees per-task
+    # and runtime-computed token values.
     command = _replace_toolathlon_placeholders(
         str(params.get("command", "")), variant=variant
     )
@@ -502,21 +614,34 @@ def _toolathlon_mcp_server(
         )
     cwd = params.get("cwd")
     cwd = _replace_toolathlon_placeholders(str(cwd), variant=variant) if cwd else None
-    if variant == "official" and command in {"uv", "uvx"}:
+    if command in {"uv", "uvx"}:
         command = f"/usr/local/bin/{command}"
         if args:
             args = [_TOOLATHLON_UVX_PACKAGE_PINS.get(arg, arg) for arg in args]
+    # Wrap the real server in the container launcher so ``${token.X}`` in argv
+    # and env resolve at spawn time. ``TOOLATHLON_TASK_DIR`` tells the launcher
+    # which task's token_key_session.py overrides the global one.
+    env = dict(env)
+    env["TOOLATHLON_TASK_DIR"] = f"/workspace/tasks/finalpool/{task_name}"
+    ensure_dirs = [
+        args[i + 1]
+        for i, arg in enumerate(args[:-1])
+        if arg in _TOOLATHLON_MCP_DIR_ARG_FLAGS
+    ]
+    if ensure_dirs:
+        # ``:``-joined for the Linux container launcher (not os.pathsep, which
+        # would follow the host running the adapter).
+        env["TOOLATHLON_ENSURE_DIRS"] = ":".join(ensure_dirs)
     payload: dict[str, Any] = {
         "name": str(data.get("name") or server_name),
         "transport": "stdio",
-        "command": command,
-        "args": args,
+        "command": _toolathlon_container_python(variant),
+        "args": [_TOOLATHLON_CONTAINER_MODULE_PATH, "launch", command, *args],
+        "env": env,
         "exclude_tags": [_NOOP_EXCLUDE_TAG],
     }
     if cwd:
         payload["cwd"] = cwd
-    if env:
-        payload["env"] = env
     return payload
 
 
@@ -577,6 +702,13 @@ def _toolathlon_mcp_config_path(ctx: Any, server_name: str) -> Path:
 
 
 def _replace_toolathlon_placeholders(value: str, *, variant: str) -> str:
+    """Resolve static path placeholders to fixed in-container paths.
+
+    ``${token.X}`` placeholders are intentionally left untouched: they are
+    resolved at MCP spawn time by the container launcher (see
+    ``_toolathlon_container.py``), the only layer that can see per-task and
+    runtime-computed token values.
+    """
     local_servers = (
         "/opt/local_servers" if variant == "gym" else "/workspace/local_servers"
     )
@@ -588,15 +720,7 @@ def _replace_toolathlon_placeholders(value: str, *, variant: str) -> str:
     }
     for old, new in replacements.items():
         value = value.replace(old, new)
-
-    def token_replacement(match: re.Match[str]) -> str:
-        token_name = match.group(1)
-        if variant == "official" and token_name in _TOOLATHLON_TOKEN_DEFAULTS:
-            return _TOOLATHLON_TOKEN_DEFAULTS[token_name]
-        name = token_name.upper()
-        return f"${{TOOLATHLON_{name}}}"
-
-    return _TOKEN_PLACEHOLDER_RE.sub(token_replacement, value)
+    return value
 
 
 def _toolathlon_dockerfile(ctx: Any, *, variant: str) -> str: