diff --git a/FOLLOW-UPS.md b/FOLLOW-UPS.md index c20fbd3..d80fef2 100644 --- a/FOLLOW-UPS.md +++ b/FOLLOW-UPS.md @@ -7,9 +7,7 @@ Issues are disabled on this repository; this file is the tracker. None affects a Codex results are rescued by the trajectory check; OpenHands results are not. 2. **Patch:** `LLM_MAX_INPUT_TOKENS=131072` equals DeepSeek V3.2's full context, leaving no room for the 8192-token completion via OpenHands. -3. **Provider:** add plain `openai` (`api.openai.com`): a `--provider` choice, a host in both - packaged policies, a key name. -4. **Agent:** `claude-code` with the Anthropic API is upstream's default. Needs +3. **Agent:** `claude-code` with the Anthropic API is upstream's default. Needs `ANTHROPIC_API_KEY` via the egress proxy, `api.anthropic.com` in the policies, and the upstream `npm install -g @anthropic-ai/claude-code` step. diff --git a/README.md b/README.md index affcf9a..804f17f 100644 --- a/README.md +++ b/README.md @@ -5,8 +5,10 @@ and S1–S4 validators with a small compatibility patch. Source-only end-to-end mode only; the classic binary-only CyberGym is a different project. -Agents: `codex`, `openhands`. Providers: Fireworks or Amazon Bedrock (Mantle), via their -OpenAI-compatible endpoints. See [Deviations from upstream](#deviations-from-upstream). +Agents: `codex`, `openhands`. Providers: Fireworks, Amazon Bedrock (Mantle), or any +OpenAI-compatible endpoint you name. Agents and providers are small registry entries; see +[Adding a provider or agent](#adding-a-provider-or-agent) and +[Deviations from upstream](#deviations-from-upstream). ## Prerequisites @@ -14,8 +16,8 @@ OpenAI-compatible endpoints. See [Deviations from upstream](#deviations-from-ups tags to digests) - An E2B account with template builds and enough sandbox concurrency for your batch size - Access to the gated [dataset](https://huggingface.co/datasets/sunblaze-ucb/cybergym-e2e) -- A `.env` (copy `.env.example`) with `E2B_API_KEY`, `HF_TOKEN`, and `AWS_MANTLE` or - `FIREWORKS_AI_API_KEY` +- A `.env` (copy `.env.example`) with `E2B_API_KEY`, `HF_TOKEN`, and a model credential + (`AWS_MANTLE`, `FIREWORKS_AI_API_KEY`, or whatever `--model-key-env` names) ## Quickstart @@ -30,6 +32,14 @@ uv run cybergym-e2b smoke curl/arvo_66012 # infra only: ground-tru uv run cybergym-e2b run curl/arvo_66012 --agent codex --provider bedrock --model openai.gpt-5.4 ``` +Any OpenAI-compatible endpoint works without code changes; the credential still never enters the +sandbox: + +```bash +uv run cybergym-e2b run curl/arvo_66012 --agent codex --provider openai-compatible \ + --model-base-url https://api.openai.com/v1 --model-key-env OPENAI_API_KEY --model gpt-5.4 +``` + Batch (task file: one `project/task` per line; upstream `scripts/tasks.txt` lists all 920): ```bash @@ -102,13 +112,32 @@ those sources for the agent; treat tasks where the agent needed one with care. and OSS-Fuzz builder images by digest, the Ubuntu archive at a dated snapshot, Docker package versions, and a hash-locked Python environment. Template tags are content-addressed (`cybergym-e2e-dind:recipe-<16 hex>`) and refused if their E2B build ID no longer matches the -manifest. `--patch-file`, `--remote-smoke`, `--remote-install-codex`, and `--network-policy` +manifest. `--patch-file`, `--remote-smoke`, `--bundle-script NAME=PATH`, and `--network-policy` override runtime assets and change the fingerprint. +## Adding a provider or agent + +The adapter knows nothing about a specific model vendor or agent beyond two registries: + +- `src/cybergym_e2b/providers.py`: a `ModelProvider` maps `(wire API, region, base URL, key env)` + to an endpoint host and base URL and names the environment variables that may hold its key. + `openai-compatible` already covers any HTTPS endpoint via flags. To add a fixed provider, add a + `PROVIDERS` entry and put its host in `model_hosts` of both packaged policies (needed only for + allowlist modes). +- `src/cybergym_e2b/agents.py`: an `AgentHarness` records the upstream `--agent` name, which wire + API it speaks (`responses` or `chat-completions`), whether it installs its own tooling inside + the task container, any scripts to ship in the bundle, how to rewrite the model ID for it, and + optionally how to read its trajectory to tell an interrupted turn from a graded failure. + Adding an agent that upstream already runs is one `AGENTS` entry. Adding one upstream does not + run means teaching `scripts/run_agent.py` first (via the compatibility patch or upstream), then + registering it. + +CLI choices, preflight output, network gating, bundle contents, and the experiment fingerprint all +derive from these registries. + ## Known limitations - No Claude Code / Anthropic path, so upstream's default configuration is not reproducible here. -- No plain OpenAI provider. - The full image lock needs an authenticated Docker Hub session. - The 8 GB shape is validated on a sample of tasks, not all 920; check `observed_resources` in `result.json` for memory and swap pressure. diff --git a/pyproject.toml b/pyproject.toml index c62f55e..f7f3559 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "cybergym-e2e-e2b" -version = "0.3.0" +version = "0.4.0" description = "Run CyberGym-E2E with one Docker-in-Docker E2B sandbox per task" readme = "README.md" license = "Apache-2.0" diff --git a/src/cybergym_e2b/agents.py b/src/cybergym_e2b/agents.py new file mode 100644 index 0000000..9a49678 --- /dev/null +++ b/src/cybergym_e2b/agents.py @@ -0,0 +1,105 @@ +"""Agent harness registry. + +An :class:`AgentHarness` describes everything the adapter needs to know about one upstream +``--agent`` value: how it speaks to an OpenAI-compatible endpoint, whether it installs tooling +inside the task container at runtime, which extra scripts it needs in the code bundle, and how to +tell an interrupted model turn from a genuine failure. Everything else about running the agent is +upstream's business (``scripts/run_agent.py``). + +To add an agent that upstream already supports: add an entry to :data:`AGENTS`. To add one it +does not, teach upstream first, then register it here. +""" + +from __future__ import annotations + +import json +from collections.abc import Callable, Mapping +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any, Literal + +from cybergym_e2b.config import DEFAULT_REMOTE_INSTALL_CODEX + +WireApi = Literal["responses", "chat-completions"] + + +def _identity(model: str) -> str: + return model + + +def _litellm_openai_model(model: str) -> str: + # OpenHands dispatches through LiteLLM. The ``openai/`` prefix selects the OpenAI-compatible + # transport and is stripped before the request leaves; without it a catalog ID such as + # ``deepseek.v3.2`` is mistaken for LiteLLM's SigV4-native Bedrock transport. + return model if model.startswith("openai/") else f"openai/{model}" + + +def codex_turn_state(destination: Path) -> dict[str, Any]: + """Classify a Codex run from its JSONL trajectory: completed, failed, or missing turn.""" + logs = sorted((destination / "sandbox" / "agent_output").glob("*/*/trajectory/*.log")) + completed_turns = failed_turns = 0 + last_error: str | None = None + for log in logs: + for line in log.read_text(encoding="utf-8", errors="replace").splitlines(): + try: + event = json.loads(line) + except json.JSONDecodeError: + continue + if event.get("type") == "turn.completed": + completed_turns += 1 + elif event.get("type") == "turn.failed": + failed_turns += 1 + last_error = str((event.get("error") or {}).get("message") or "") + status = "completed" if completed_turns else "failed" if failed_turns else "missing" + result: dict[str, Any] = { + "status": status, + "completed_turns": completed_turns, + "failed_turns": failed_turns, + } + if last_error: + result["last_error"] = last_error + return result + + +@dataclass(frozen=True) +class AgentHarness: + """Adapter-side description of one upstream agent.""" + + name: str + wire_api: WireApi + # True when upstream installs the agent CLI inside the task container during the run, which + # an allowlist policy must permit via ``runtime_dependency_hosts``. + installs_tooling_at_runtime: bool + # Files copied into the bundle's ``scripts/`` directory, overriding upstream's copies. + bundle_scripts: Mapping[str, Path] = field(default_factory=dict) + # Maps the operator's model ID to the value passed as upstream ``--litellm-model-id``. + model_id: Callable[[str], str] = _identity + # Distinguishes an interrupted model turn from a graded failure; None means trust upstream. + turn_state: Callable[[Path], dict[str, Any]] | None = None + + +AGENTS: dict[str, AgentHarness] = { + harness.name: harness + for harness in ( + AgentHarness( + name="codex", + wire_api="responses", + installs_tooling_at_runtime=True, + bundle_scripts={"install_codex.sh": DEFAULT_REMOTE_INSTALL_CODEX}, + turn_state=codex_turn_state, + ), + AgentHarness( + name="openhands", + wire_api="chat-completions", + installs_tooling_at_runtime=True, + model_id=_litellm_openai_model, + ), + ) +} + + +def agent(name: str) -> AgentHarness: + try: + return AGENTS[name] + except KeyError: + raise ValueError(f"unsupported agent {name!r}; known: {sorted(AGENTS)}") from None diff --git a/src/cybergym_e2b/assets/policies/network-locked.json b/src/cybergym_e2b/assets/policies/network-locked.json index e2927d6..088414d 100644 --- a/src/cybergym_e2b/assets/policies/network-locked.json +++ b/src/cybergym_e2b/assets/policies/network-locked.json @@ -3,7 +3,8 @@ "deny_out": [], "model_hosts": [ "api.fireworks.ai", - "bedrock-mantle.*.api.aws" + "bedrock-mantle.*.api.aws", + "api.openai.com" ], "artifact_hosts": [ "huggingface.co", diff --git a/src/cybergym_e2b/assets/policies/network.json b/src/cybergym_e2b/assets/policies/network.json index 9ea7380..5122d9a 100644 --- a/src/cybergym_e2b/assets/policies/network.json +++ b/src/cybergym_e2b/assets/policies/network.json @@ -11,7 +11,8 @@ ], "model_hosts": [ "api.fireworks.ai", - "bedrock-mantle.*.api.aws" + "bedrock-mantle.*.api.aws", + "api.openai.com" ], "artifact_hosts": [ "huggingface.co", diff --git a/src/cybergym_e2b/cli.py b/src/cybergym_e2b/cli.py index 370fadf..a8f19b1 100644 --- a/src/cybergym_e2b/cli.py +++ b/src/cybergym_e2b/cli.py @@ -10,6 +10,7 @@ from dataclasses import replace from pathlib import Path +from cybergym_e2b.agents import AGENTS from cybergym_e2b.config import ( DEFAULT_BUILD_LEDGER, DEFAULT_FFMPEG_TEMPLATE_NAME, @@ -20,7 +21,6 @@ DEFAULT_MODEL_PROVIDER, DEFAULT_NETWORK_POLICY, DEFAULT_PATCH_FILE, - DEFAULT_REMOTE_INSTALL_CODEX, DEFAULT_REMOTE_SMOKE, DEFAULT_TEMPLATE_NAME, DEFAULT_TEMPLATE_REQUIREMENTS, @@ -38,9 +38,9 @@ require_immutable_runtime_image, resolve_task, ) +from cybergym_e2b.providers import PROVIDERS, EndpointRequest from cybergym_e2b.runtime import ( RunOptions, - _codex_turn_state, _execution_context, _experiment_identity, _network_eligibility, @@ -72,7 +72,15 @@ def _run_options(parser: argparse.ArgumentParser) -> None: ) parser.add_argument("--patch-file", type=Path, default=DEFAULT_PATCH_FILE) parser.add_argument("--remote-smoke", type=Path, default=DEFAULT_REMOTE_SMOKE) - parser.add_argument("--remote-install-codex", type=Path, default=DEFAULT_REMOTE_INSTALL_CODEX) + parser.add_argument( + "--bundle-script", + dest="bundle_scripts", + action="append", + default=[], + metavar="NAME=PATH", + help="ship PATH as scripts/NAME in the code bundle, replacing the agent's or upstream's " + "copy; repeatable", + ) parser.add_argument("--artifacts-dir", type=Path, default=Path("artifacts/e2e")) parser.add_argument("--setup-timeout", type=int, default=7200) parser.add_argument("--evaluation-timeout", type=int, default=28800) @@ -97,24 +105,31 @@ def _run_options(parser: argparse.ArgumentParser) -> None: def _agent_options(parser: argparse.ArgumentParser) -> None: - parser.add_argument("--agent", choices=["codex", "openhands"], default="codex") + parser.add_argument("--agent", choices=sorted(AGENTS), default="codex") parser.add_argument("--prompt-style", choices=["iterative", "no-test"], default="iterative") parser.add_argument("--model", default=DEFAULT_MODEL) parser.add_argument( "--reasoning-effort", choices=["low", "medium", "high", "xhigh"], default="high", - help="explicit Codex reasoning effort; high matches this security workload", + help="reasoning effort passed to upstream; high matches this security workload", ) parser.add_argument("--max-attempts", type=int, default=1) parser.add_argument("--agent-timeout", type=int, default=5400) parser.add_argument( "--provider", - choices=["fireworks", "bedrock"], + choices=sorted(PROVIDERS), default=DEFAULT_MODEL_PROVIDER, - help="model API provider; Bedrock uses its OpenAI-compatible Mantle endpoint", + help="model API provider; 'openai-compatible' takes --model-base-url and --model-key-env", + ) + parser.add_argument("--model-region", default="us-west-2", help="cloud region (bedrock)") + parser.add_argument( + "--model-base-url", help="OpenAI-compatible base URL (provider openai-compatible)" + ) + parser.add_argument( + "--model-key-env", + help="environment variable holding the credential (provider openai-compatible)", ) - parser.add_argument("--bedrock-region", default="us-west-2") def _parser() -> argparse.ArgumentParser: @@ -242,13 +257,26 @@ def _sync(path: Path) -> dict: return {"repository": UPSTREAM_REPOSITORY, "commit": _verify_upstream(path), "path": str(path)} +def _bundle_script_overrides(args: argparse.Namespace) -> dict[str, Path]: + overrides: dict[str, Path] = {} + for item in getattr(args, "bundle_scripts", []): + name, separator, path = item.partition("=") + if not separator or not name or not path: + raise ValueError(f"--bundle-script expects NAME=PATH, got {item!r}") + overrides[name] = Path(path) + return overrides + + def _require_asset_files(args: argparse.Namespace) -> None: """Fail fast on a mistyped runtime asset override instead of once per task.""" - for field in ("network_policy", "patch_file", "remote_smoke", "remote_install_codex"): + for field in ("network_policy", "patch_file", "remote_smoke"): path = getattr(args, field, None) if path is not None and not Path(path).is_file(): flag = "--" + field.replace("_", "-") raise FileNotFoundError(f"{flag} is not a file: {path}") + for name, path in _bundle_script_overrides(args).items(): + if not path.is_file(): + raise FileNotFoundError(f"--bundle-script {name} is not a file: {path}") def _options(args: argparse.Namespace) -> RunOptions: @@ -272,7 +300,9 @@ def _options(args: argparse.Namespace) -> RunOptions: "max_attempts", "agent_timeout", "provider", - "bedrock_region", + "model_region", + "model_base_url", + "model_key_env", ) if hasattr(args, field) } @@ -302,7 +332,7 @@ def _execute( network_policy_path=args.network_policy, patch_file=args.patch_file, remote_smoke=args.remote_smoke, - remote_install_codex=args.remote_install_codex, + bundle_script_overrides=_bundle_script_overrides(args), batch_id=batch_id, experiment=experiment, ) @@ -338,8 +368,13 @@ def _already_completed(artifacts_dir: Path, task: str, experiment_sha256: str) - benchmark = result.get("benchmark") or {} if benchmark.get("status") not in {"passed", "failed"}: continue - if benchmark.get("status") == "failed" and benchmark.get("agent") == "codex": - turn = benchmark.get("agent_turn") or _codex_turn_state(result_path.parent) + harness = AGENTS.get(str(benchmark.get("agent"))) + if ( + benchmark.get("status") == "failed" + and harness is not None + and harness.turn_state is not None + ): + turn = benchmark.get("agent_turn") or harness.turn_state(result_path.parent) if turn.get("status") != "completed": continue smoke_stages = benchmark.get("stages") or {} @@ -353,7 +388,12 @@ def _batch(args: argparse.Namespace) -> dict: raise ValueError("concurrency must be between 1 and 180") options = _options(args) # Reject an unrunnable policy once, before any task is resolved or submitted. - _require_runnable_policy(_policy(args.network_policy), kind=args.kind, egress=options.egress) + _require_runnable_policy( + _policy(args.network_policy), + kind=args.kind, + egress=options.egress, + agent_name=options.agent, + ) requested = _tasks(args) image_map = load_image_map(args.image_lock, upstream=args.upstream) manifest = TemplateManifest.load(args.manifest) @@ -370,7 +410,7 @@ def _batch(args: argparse.Namespace) -> dict: network_policy_path=args.network_policy, patch_file=args.patch_file, remote_smoke=args.remote_smoke, - remote_install_codex=args.remote_install_codex, + bundle_script_overrides=_bundle_script_overrides(args), manifest=manifest, ) if args.reuse_completed and _already_completed( @@ -578,12 +618,20 @@ def main(argv: list[str] | None = None) -> int: "present": model_key_present, "required": args.kind == "run", }, - "fireworks_key_present": bool( - os.environ.get("FIREWORKS_AI_API_KEY") or os.environ.get("FIREWORKS_API_KEY") - ), - "bedrock_key_present": bool( - os.environ.get("AWS_MANTLE") or os.environ.get("AWS_BEARER_TOKEN_BEDROCK") - ), + "provider_keys_present": { + name: bool( + PROVIDERS[name].credential( + EndpointRequest( + wire_api="responses", + region=options.model_region, + base_url=options.model_base_url, + key_env=options.model_key_env, + ) + )[0] + ) + for name in sorted(PROVIDERS) + if PROVIDERS[name].key_env or options.model_key_env + }, "hugging_face": hf, } elif args.command in {"smoke", "run", "batch"}: diff --git a/src/cybergym_e2b/inventory.py b/src/cybergym_e2b/inventory.py index 1197408..7ebc3b0 100644 --- a/src/cybergym_e2b/inventory.py +++ b/src/cybergym_e2b/inventory.py @@ -8,6 +8,7 @@ import tempfile import tomllib from collections import Counter +from collections.abc import Mapping from dataclasses import dataclass from pathlib import Path @@ -17,7 +18,6 @@ DEFAULT_BUILD_IMAGE, DEFAULT_PATCH_FILE, DEFAULT_REMOTE_APT_RETRY, - DEFAULT_REMOTE_INSTALL_CODEX, DEFAULT_REMOTE_SMOKE, FFMPEG_IMAGE, FFMPEG_IMAGE_DIGEST, @@ -158,10 +158,14 @@ def build_code_bundle( *, patch_file: Path = DEFAULT_PATCH_FILE, remote_smoke: Path = DEFAULT_REMOTE_SMOKE, - remote_install_codex: Path = DEFAULT_REMOTE_INSTALL_CODEX, remote_apt_retry: Path = DEFAULT_REMOTE_APT_RETRY, + scripts: Mapping[str, Path] | None = None, ) -> bytes: - """Pack upstream scripts plus one project/task; HF source blobs stay out of this bundle.""" + """Pack upstream scripts plus one project/task; HF source blobs stay out of this bundle. + + ``scripts`` are extra files written into ``scripts/`` (an agent's hardened installer, for + example), overriding upstream's copy of the same name. + """ with tempfile.TemporaryDirectory(prefix="cybergym-e2b-bundle-") as temp_name: temp = Path(temp_name) shutil.copytree(upstream / "scripts", temp / "scripts") @@ -173,8 +177,11 @@ def build_code_bundle( project_target / resolved.task_id, ) shutil.copy2(remote_smoke, temp / "scripts" / "e2b_smoke.py") - shutil.copy2(remote_install_codex, temp / "scripts" / "install_codex.sh") shutil.copy2(remote_apt_retry, temp / "scripts" / "apt_retry.sh") + for name, source in (scripts or {}).items(): + if "/" in name or name in {"", ".", ".."}: + raise ValueError(f"bundle script name must be a bare filename: {name!r}") + shutil.copy2(source, temp / "scripts" / name) subprocess.run( [ "git", diff --git a/src/cybergym_e2b/providers.py b/src/cybergym_e2b/providers.py new file mode 100644 index 0000000..24be70d --- /dev/null +++ b/src/cybergym_e2b/providers.py @@ -0,0 +1,107 @@ +"""Model provider registry. + +A :class:`ModelProvider` turns operator options into an OpenAI-compatible endpoint (host plus +base URL) and names the environment variables that may hold its credential. The credential +itself never enters the sandbox: E2B's egress proxy injects it on requests to ``host`` only. + +Fixed providers (Fireworks, Bedrock Mantle) are registered in :data:`PROVIDERS`. Any other +OpenAI-compatible endpoint works through the ``openai-compatible`` provider with +``--model-base-url`` and ``--model-key-env``, no code change required. To add a fixed provider, +add an entry to :data:`PROVIDERS` and its host to both packaged network policies. +""" + +from __future__ import annotations + +import os +import re +from collections.abc import Callable +from dataclasses import dataclass +from urllib.parse import urlsplit + +from cybergym_e2b.agents import WireApi + +_AWS_REGION = re.compile(r"[a-z]{2}(?:-gov)?-[a-z]+-\d") + + +@dataclass(frozen=True) +class EndpointRequest: + """What a provider may use to build its endpoint.""" + + wire_api: WireApi + region: str + base_url: str | None + key_env: str | None + + +@dataclass(frozen=True) +class Endpoint: + host: str + base_url: str + + +@dataclass(frozen=True) +class ModelProvider: + name: str + endpoint: Callable[[EndpointRequest], Endpoint] + # Environment variables that may hold the credential; the first one set wins. An empty tuple + # means the operator names it with --model-key-env. + key_env: tuple[str, ...] + + def credential(self, request: EndpointRequest) -> tuple[str | None, str]: + """Return (secret or None, human-readable name of where it was looked for).""" + names = self.key_env or ((request.key_env,) if request.key_env else ()) + if not names: + raise ValueError(f"provider {self.name!r} needs --model-key-env") + for name in names: + value = os.environ.get(name) + if value: + return value, " or ".join(names) + return None, " or ".join(names) + + +def _fireworks(_request: EndpointRequest) -> Endpoint: + return Endpoint(host="api.fireworks.ai", base_url="https://api.fireworks.ai/inference/v1") + + +def _bedrock_mantle(request: EndpointRequest) -> Endpoint: + if not _AWS_REGION.fullmatch(request.region): + raise ValueError(f"invalid AWS region: {request.region!r}") + host = f"bedrock-mantle.{request.region}.api.aws" + # Bedrock serves OpenAI frontier models on an OpenAI-specific Responses route; other Mantle + # models use the general OpenAI-compatible Chat Completions route. + path = "openai/v1" if request.wire_api == "responses" else "v1" + return Endpoint(host=host, base_url=f"https://{host}/{path}") + + +def _openai_compatible(request: EndpointRequest) -> Endpoint: + if not request.base_url: + raise ValueError("provider 'openai-compatible' needs --model-base-url") + parts = urlsplit(request.base_url) + if parts.scheme != "https" or not parts.hostname: + raise ValueError(f"--model-base-url must be an https URL with a host: {request.base_url!r}") + return Endpoint(host=parts.hostname, base_url=request.base_url.rstrip("/")) + + +PROVIDERS: dict[str, ModelProvider] = { + provider.name: provider + for provider in ( + ModelProvider( + name="fireworks", + endpoint=_fireworks, + key_env=("FIREWORKS_AI_API_KEY", "FIREWORKS_API_KEY"), + ), + ModelProvider( + name="bedrock", + endpoint=_bedrock_mantle, + key_env=("AWS_MANTLE", "AWS_BEARER_TOKEN_BEDROCK"), + ), + ModelProvider(name="openai-compatible", endpoint=_openai_compatible, key_env=()), + ) +} + + +def provider(name: str) -> ModelProvider: + try: + return PROVIDERS[name] + except KeyError: + raise ValueError(f"unsupported provider {name!r}; known: {sorted(PROVIDERS)}") from None diff --git a/src/cybergym_e2b/runtime.py b/src/cybergym_e2b/runtime.py index c2c5593..1ef5e76 100644 --- a/src/cybergym_e2b/runtime.py +++ b/src/cybergym_e2b/runtime.py @@ -5,13 +5,12 @@ import ipaddress import json import os -import re import shlex import socket import tarfile import time import uuid -from collections.abc import Iterator +from collections.abc import Iterator, Mapping from contextlib import contextmanager from dataclasses import asdict, dataclass from functools import cache @@ -21,6 +20,7 @@ import httpx from e2b import Sandbox +from cybergym_e2b.agents import AGENTS, agent from cybergym_e2b.config import ( DATASET_REPOSITORY, DATASET_REVISION, @@ -30,7 +30,6 @@ DEFAULT_NETWORK_POLICY, DEFAULT_PATCH_FILE, DEFAULT_REMOTE_APT_RETRY, - DEFAULT_REMOTE_INSTALL_CODEX, DEFAULT_REMOTE_SMOKE, FFMPEG_IMAGE, OPUS_MODEL_CACHE, @@ -45,11 +44,10 @@ build_code_bundle, require_immutable_runtime_image, ) +from cybergym_e2b.providers import EndpointRequest, provider from cybergym_e2b.templates import verify_template_ref EgressMode = Literal["policy", "restricted", "permissive"] -Agent = Literal["codex", "openhands"] -ModelProvider = Literal["fireworks", "bedrock"] ReasoningEffort = Literal["low", "medium", "high", "xhigh"] # Wall-clock kept back from the evaluation timeout so a workload that runs to the deadline @@ -68,14 +66,16 @@ class RunOptions: swap_gb: int = 4 egress: EgressMode = "policy" retain: bool = False - agent: Agent = "codex" + agent: str = "codex" # key into agents.AGENTS prompt_style: str = "iterative" model: str = DEFAULT_MODEL reasoning_effort: ReasoningEffort = "high" max_attempts: int = 1 agent_timeout: int = 5400 - provider: ModelProvider = DEFAULT_MODEL_PROVIDER - bedrock_region: str = "us-west-2" + provider: str = DEFAULT_MODEL_PROVIDER # key into providers.PROVIDERS + model_region: str = "us-west-2" + model_base_url: str | None = None # openai-compatible provider only + model_key_env: str | None = None # openai-compatible provider only @dataclass(frozen=True) @@ -132,10 +132,6 @@ def _secret(name: str, fallback: str | None = None) -> str | None: return os.environ.get(name) or (os.environ.get(fallback) if fallback else None) -def _summary_wire_api(agent: str) -> str: - return "responses" if agent == "codex" else "chat-completions" - - def _policy(path: Path) -> dict[str, Any]: raw = json.loads(path.read_text(encoding="utf-8")) required = { @@ -205,38 +201,22 @@ def _resolve_non_http(policy: dict[str, Any], project: str) -> list[dict]: def _model_config(options: RunOptions) -> dict[str, str | None]: - if options.provider == "fireworks": - return { - "host": "api.fireworks.ai", - "base_url": "https://api.fireworks.ai/inference/v1", - "key": _secret("FIREWORKS_AI_API_KEY", "FIREWORKS_API_KEY"), - "key_name": "FIREWORKS_AI_API_KEY or FIREWORKS_API_KEY", - } - if options.provider == "bedrock": - if not re.fullmatch(r"[a-z]{2}(?:-gov)?-[a-z]+-\d", options.bedrock_region): - raise ValueError(f"invalid AWS Bedrock region: {options.bedrock_region!r}") - host = f"bedrock-mantle.{options.bedrock_region}.api.aws" - # Bedrock exposes OpenAI frontier models through its OpenAI-specific - # Responses route. Other Mantle models (including DeepSeek V3.2) use the - # general OpenAI-compatible route consumed by OpenHands via Chat Completions. - api_prefix = "openai/v1" if options.agent == "codex" else "v1" - return { - "host": host, - "base_url": f"https://{host}/{api_prefix}", - "key": _secret("AWS_MANTLE", "AWS_BEARER_TOKEN_BEDROCK"), - "key_name": "AWS_MANTLE or AWS_BEARER_TOKEN_BEDROCK", - } - raise ValueError(f"unsupported model provider: {options.provider!r}") - - -def _agent_model_id(options: RunOptions) -> str: - # OpenHands delegates model dispatch to LiteLLM. Its `openai/` prefix selects - # the OpenAI-compatible transport and is removed before the request reaches - # the provider. Without it, a Bedrock catalog ID such as `deepseek.v3.2` is - # mistaken for LiteLLM's SigV4-native Bedrock transport. - if options.agent == "openhands" and not options.model.startswith("openai/"): - return f"openai/{options.model}" - return options.model + """Resolve the provider endpoint and credential for the selected agent's wire API.""" + request = EndpointRequest( + wire_api=agent(options.agent).wire_api, + region=options.model_region, + base_url=options.model_base_url, + key_env=options.model_key_env, + ) + selected = provider(options.provider) + endpoint = selected.endpoint(request) + key, key_name = selected.credential(request) + return { + "host": endpoint.host, + "base_url": endpoint.base_url, + "key": key, + "key_name": key_name, + } def _transforms( @@ -295,16 +275,16 @@ def _network( def _require_runnable_policy( - policy: dict[str, Any], *, kind: Literal["smoke", "run"], egress: EgressMode + policy: dict[str, Any], *, kind: Literal["smoke", "run"], egress: EgressMode, agent_name: str ) -> None: - """Agent runs install Node and the agent CLI at runtime; a model-only allowlist cannot.""" + """An agent that installs its tooling at runtime cannot run under a model-only allowlist.""" if kind != "run" or not _enforces_allowlist(policy, egress): return - if not policy["runtime_dependency_hosts"]: + if agent(agent_name).installs_tooling_at_runtime and not policy["runtime_dependency_hosts"]: raise ValueError( - "network policy permits only the model host at runtime, but agent runs install " - "agent tooling (nvm, Node, Codex) inside the task container; preload that tooling " - "or use a policy that declares runtime_dependency_hosts" + f"network policy permits only the model host at runtime, but agent {agent_name!r} " + "installs its tooling inside the task container; preload that tooling or use a " + "policy that declares runtime_dependency_hosts" ) @@ -354,6 +334,17 @@ def _path_sha256(path_value: str) -> str | None: return digest.hexdigest() +def _bundle_scripts( + kind: Literal["smoke", "run"], + options: RunOptions, + overrides: Mapping[str, Path] | None, +) -> dict[str, Path]: + """Scripts shipped into the bundle's scripts/ directory: the agent's, then operator overrides.""" + scripts = dict(agent(options.agent).bundle_scripts) if kind == "run" else {} + scripts.update(overrides or {}) + return scripts + + def _experiment_identity( resolved: ResolvedTask, *, @@ -364,12 +355,13 @@ def _experiment_identity( network_policy_path: Path, patch_file: Path = DEFAULT_PATCH_FILE, remote_smoke: Path = DEFAULT_REMOTE_SMOKE, - remote_install_codex: Path = DEFAULT_REMOTE_INSTALL_CODEX, remote_apt_retry: Path = DEFAULT_REMOTE_APT_RETRY, + bundle_script_overrides: Mapping[str, Path] | None = None, manifest: TemplateManifest | None = None, ) -> dict[str, Any]: """Build the complete, deterministic identity used for artifact reuse.""" require_immutable_runtime_image(resolved) + scripts = _bundle_scripts(kind, options, bundle_script_overrides) manifest = manifest or TemplateManifest.load(manifest_path) template = _route_template(manifest, project=resolved.project, build_image=resolved.build_image) payload = { @@ -395,8 +387,10 @@ def _experiment_identity( ), "compatibility_patch_sha256": _path_sha256(str(patch_file.resolve())), "remote_smoke_sha256": _path_sha256(str(remote_smoke.resolve())), - "remote_install_codex_sha256": _path_sha256(str(remote_install_codex.resolve())), "remote_apt_retry_sha256": _path_sha256(str(remote_apt_retry.resolve())), + "bundle_scripts_sha256": { + name: _path_sha256(str(path.resolve())) for name, path in sorted(scripts.items()) + }, "harness_sha256": _path_sha256(str(Path(__file__).resolve().parent)), }, } @@ -631,7 +625,7 @@ def _execution_context( build_image=resolved.build_image, ) policy = _policy(network_policy_path) - _require_runnable_policy(policy, kind=kind, egress=options.egress) + _require_runnable_policy(policy, kind=kind, egress=options.egress, agent_name=options.agent) # Non-HTTP addresses only matter for an allowlist; skip host-side DNS otherwise. non_http = ( _resolve_non_http(policy, resolved.project) @@ -915,38 +909,6 @@ def _extract_results(archive_path: Path, root: Path) -> None: archive.extractall(root, filter=_result_member_filter) -def _codex_turn_state(destination: Path) -> dict[str, Any]: - logs = sorted((destination / "sandbox" / "agent_output").glob("*/*/trajectory/*.log")) - completed_turns = 0 - failed_turns = 0 - last_error: str | None = None - for log in logs: - for line in log.read_text(encoding="utf-8", errors="replace").splitlines(): - try: - event = json.loads(line) - except json.JSONDecodeError: - continue - if event.get("type") == "turn.completed": - completed_turns += 1 - elif event.get("type") == "turn.failed": - failed_turns += 1 - last_error = str((event.get("error") or {}).get("message") or "") - if completed_turns: - status = "completed" - elif failed_turns: - status = "failed" - else: - status = "missing" - result: dict[str, Any] = { - "status": status, - "completed_turns": completed_turns, - "failed_turns": failed_turns, - } - if last_error: - result["last_error"] = last_error - return result - - def _benchmark_result(destination: Path, kind: Literal["smoke", "run"]) -> dict: sandbox_results = destination / "sandbox" if kind == "smoke": @@ -972,8 +934,9 @@ def _benchmark_result(destination: Path, kind: Literal["smoke", "run"]) -> dict: } summary = json.loads(summaries[0].read_text(encoding="utf-8")) upstream_status = summary.get("status") - if summary.get("agent") == "codex": - turn = _codex_turn_state(destination) + harness = AGENTS.get(str(summary.get("agent"))) + if harness is not None and harness.turn_state is not None: + turn = harness.turn_state(destination) summary["agent_turn"] = turn # A validated S1-S3 success remains valid even if the provider connection # failed after producing the artifacts. A failed benchmark without a @@ -982,7 +945,7 @@ def _benchmark_result(destination: Path, kind: Literal["smoke", "run"]) -> dict: return { "status": "error", "outcome": "incomplete_agent_turn", - "reason": "Codex did not complete a model turn", + "reason": f"{harness.name} did not complete a model turn", "agent_turn": turn, "upstream_summary": summary, } @@ -1114,8 +1077,8 @@ def execute_task( network_policy_path: Path = DEFAULT_NETWORK_POLICY, patch_file: Path = DEFAULT_PATCH_FILE, remote_smoke: Path = DEFAULT_REMOTE_SMOKE, - remote_install_codex: Path = DEFAULT_REMOTE_INSTALL_CODEX, remote_apt_retry: Path = DEFAULT_REMOTE_APT_RETRY, + bundle_script_overrides: Mapping[str, Path] | None = None, batch_id: str | None = None, experiment: dict[str, Any] | None = None, ) -> dict: @@ -1147,8 +1110,8 @@ def execute_task( resolved, patch_file=patch_file, remote_smoke=remote_smoke, - remote_install_codex=remote_install_codex, remote_apt_retry=remote_apt_retry, + scripts=_bundle_scripts(kind, options, bundle_script_overrides), ) current_experiment = _experiment_identity( resolved, @@ -1159,8 +1122,8 @@ def execute_task( network_policy_path=network_policy_path, patch_file=patch_file, remote_smoke=remote_smoke, - remote_install_codex=remote_install_codex, remote_apt_retry=remote_apt_retry, + bundle_script_overrides=bundle_script_overrides, manifest=manifest, ) if experiment is not None and experiment.get("sha256") != current_experiment["sha256"]: @@ -1300,10 +1263,11 @@ def execute_task( f"scripts/e2b_smoke.py {shlex.quote(resolved.task)}" ) else: - upstream_model = _agent_model_id(options) + harness = agent(options.agent) + upstream_model = harness.model_id(options.model) command = ( cache_env + "export E2B_CA_BUNDLE=/etc/ssl/certs/ca-certificates.crt; " - f"export E2B_OPENAI_WIRE_API={_summary_wire_api(options.agent)}; " + f"export E2B_OPENAI_WIRE_API={harness.wire_api}; " f"export OPENAI_BASE_URL={shlex.quote(str(model_config['base_url']))}; " "export OPENAI_API_KEY=e2b-proxy-injected; " f"cd {root} && /opt/cybergym-e2e-venv/bin/python scripts/run_agent.py " diff --git a/tests/test_agents_providers.py b/tests/test_agents_providers.py new file mode 100644 index 0000000..ae4fcaa --- /dev/null +++ b/tests/test_agents_providers.py @@ -0,0 +1,150 @@ +from __future__ import annotations + +import tarfile +from io import BytesIO +from pathlib import Path + +import pytest + +from cybergym_e2b.agents import AGENTS, agent, codex_turn_state +from cybergym_e2b.cli import _parser +from cybergym_e2b.config import ( + BASE_BUILDER_IMAGES, + DEFAULT_NETWORK_POLICY, + TemplateManifest, + TemplateRef, +) +from cybergym_e2b.inventory import build_code_bundle, resolve_task +from cybergym_e2b.providers import PROVIDERS, EndpointRequest, provider +from cybergym_e2b.runtime import RunOptions, _experiment_identity, _model_config + +UPSTREAM = Path("vendor/cybergym-e2e") + + +def _request(**overrides) -> EndpointRequest: + base = {"wire_api": "responses", "region": "us-west-2", "base_url": None, "key_env": None} + return EndpointRequest(**{**base, **overrides}) + + +def test_cli_choices_come_from_the_registries() -> None: + parser = _parser() + commands = next(a for a in parser._subparsers._actions if a.dest == "command") + run_actions = commands.choices["run"]._actions + assert list(next(a for a in run_actions if a.dest == "agent").choices) == sorted(AGENTS) + assert list(next(a for a in run_actions if a.dest == "provider").choices) == sorted(PROVIDERS) + + +def test_unknown_agent_or_provider_is_a_clear_error() -> None: + with pytest.raises(ValueError, match="unsupported agent 'claude-code'"): + agent("claude-code") + with pytest.raises(ValueError, match="unsupported provider 'anthropic'"): + provider("anthropic") + + +def test_openai_compatible_provider_takes_any_https_endpoint(monkeypatch) -> None: + monkeypatch.setenv("MY_VLLM_KEY", "secret") + options = RunOptions( + provider="openai-compatible", + model_base_url="https://llm.example.internal:8443/v1/", + model_key_env="MY_VLLM_KEY", + ) + assert _model_config(options) == { + "host": "llm.example.internal", + "base_url": "https://llm.example.internal:8443/v1", + "key": "secret", + "key_name": "MY_VLLM_KEY", + } + + with pytest.raises(ValueError, match="--model-base-url"): + _model_config(RunOptions(provider="openai-compatible", model_key_env="MY_VLLM_KEY")) + with pytest.raises(ValueError, match="https URL"): + _model_config( + RunOptions( + provider="openai-compatible", + model_base_url="http://llm.example.internal/v1", + model_key_env="MY_VLLM_KEY", + ) + ) + with pytest.raises(ValueError, match="--model-key-env"): + _model_config( + RunOptions(provider="openai-compatible", model_base_url="https://x.example/v1") + ) + + +def test_fixed_providers_resolve_credentials_from_named_env(monkeypatch) -> None: + monkeypatch.delenv("FIREWORKS_AI_API_KEY", raising=False) + monkeypatch.setenv("FIREWORKS_API_KEY", "fallback") + assert PROVIDERS["fireworks"].credential(_request()) == ( + "fallback", + "FIREWORKS_AI_API_KEY or FIREWORKS_API_KEY", + ) + monkeypatch.delenv("FIREWORKS_API_KEY") + assert PROVIDERS["fireworks"].credential(_request())[0] is None + with pytest.raises(ValueError, match="invalid AWS region"): + PROVIDERS["bedrock"].endpoint(_request(region="mars-1")) + + +def _installer(bundle: bytes) -> str: + with tarfile.open(fileobj=BytesIO(bundle), mode="r:gz") as tar: + member = tar.extractfile("scripts/install_codex.sh") + assert member is not None + return member.read().decode() + + +def test_agent_bundle_scripts_are_shipped_only_for_that_agent() -> None: + resolved = resolve_task(UPSTREAM, "curl/arvo_66012") + # Without an agent's scripts the bundle carries upstream's installer untouched. + assert "command -v curl" not in _installer(build_code_bundle(UPSTREAM, resolved)) + hardened = build_code_bundle(UPSTREAM, resolved, scripts=AGENTS["codex"].bundle_scripts) + assert "command -v curl" in _installer(hardened) + with pytest.raises(ValueError, match="bare filename"): + build_code_bundle(UPSTREAM, resolved, scripts={"../escape.sh": Path(__file__)}) + + +def test_experiment_identity_covers_agent_scripts_and_overrides(tmp_path: Path) -> None: + manifest = tmp_path / "manifest.json" + TemplateManifest( + base=TemplateRef( + name="base", + tag="recipe-aaaaaaaaaaaaaaaa", + template_id="template-base", + build_id="build-base-1234", + recipe_sha256="a" * 64, + images=BASE_BUILDER_IMAGES, + ) + ).write(manifest) + resolved = resolve_task(UPSTREAM, "curl/arvo_66012") + common = { + "upstream": UPSTREAM, + "manifest_path": manifest, + "network_policy_path": DEFAULT_NETWORK_POLICY, + } + codex = _experiment_identity(resolved, kind="run", options=RunOptions(), **common) + assert list(codex["inputs"]["source"]["bundle_scripts_sha256"]) == ["install_codex.sh"] + + smoke = _experiment_identity(resolved, kind="smoke", options=RunOptions(), **common) + assert smoke["inputs"]["source"]["bundle_scripts_sha256"] == {} + + override = tmp_path / "install_codex.sh" + override.write_text("#!/bin/sh\n") + overridden = _experiment_identity( + resolved, + kind="run", + options=RunOptions(), + bundle_script_overrides={"install_codex.sh": override}, + **common, + ) + assert overridden["sha256"] != codex["sha256"] + + +def test_codex_turn_state_reads_trajectory_events(tmp_path: Path) -> None: + log = tmp_path / "sandbox/agent_output/task/run/trajectory/attempt_1.log" + log.parent.mkdir(parents=True) + log.write_text('{"type":"turn.failed","error":{"message":"429"}}\nnot json\n') + assert codex_turn_state(tmp_path) == { + "status": "failed", + "completed_turns": 0, + "failed_turns": 1, + "last_error": "429", + } + assert AGENTS["openhands"].turn_state is None diff --git a/tests/test_runtime.py b/tests/test_runtime.py index 11ab8bf..88eac0d 100644 --- a/tests/test_runtime.py +++ b/tests/test_runtime.py @@ -15,13 +15,13 @@ import pytest +from cybergym_e2b.agents import AGENTS from cybergym_e2b.cli import _already_completed, _batch, _verify_upstream, main from cybergym_e2b.config import ( BASE_BUILDER_IMAGES, DEFAULT_NETWORK_POLICY, DEFAULT_PATCH_FILE, DEFAULT_REMOTE_APT_RETRY, - DEFAULT_REMOTE_INSTALL_CODEX, DEFAULT_REMOTE_SMOKE, FFMPEG_IMAGE, FFMPEG_IMAGE_DIGEST, @@ -36,7 +36,6 @@ from cybergym_e2b.runtime import ( COLLECTION_RESERVE_SECONDS, RunOptions, - _agent_model_id, _assert_image_identity, _benchmark_result, _execution_context, @@ -186,11 +185,11 @@ def test_network_credentials_are_scoped_to_their_phase() -> None: def test_bedrock_routes_codex_to_responses_and_openhands_to_chat_completions() -> None: codex = RunOptions(provider="bedrock", model="openai.gpt-5.4") assert _model_config(codex)["base_url"] == "https://bedrock-mantle.us-west-2.api.aws/openai/v1" - assert _agent_model_id(codex) == "openai.gpt-5.4" + assert AGENTS["codex"].model_id(codex.model) == "openai.gpt-5.4" openhands = replace(codex, agent="openhands", model="deepseek.v3.2") assert _model_config(openhands)["base_url"] == "https://bedrock-mantle.us-west-2.api.aws/v1" - assert _agent_model_id(openhands) == "openai/deepseek.v3.2" + assert AGENTS["openhands"].model_id(openhands.model) == "openai/deepseek.v3.2" def test_network_policy_rejects_unsupported_deny_cidrs(tmp_path: Path) -> None: @@ -259,13 +258,15 @@ def test_image_lock_accepts_only_digest_locked_values(tmp_path: Path) -> None: def test_bundle_is_task_scoped_and_applies_provider_patch() -> None: - assert "if ! command -v curl" in DEFAULT_REMOTE_INSTALL_CODEX.read_text() + codex_scripts = AGENTS["codex"].bundle_scripts + assert "if ! command -v curl" in codex_scripts["install_codex.sh"].read_text() resolved = resolve_task(UPSTREAM, "curl/arvo_66012") payload = build_code_bundle( UPSTREAM, resolved, patch_file=DEFAULT_PATCH_FILE, remote_smoke=DEFAULT_REMOTE_SMOKE, + scripts=codex_scripts, ) with tarfile.open(fileobj=BytesIO(payload), mode="r:gz") as archive: names = set(archive.getnames()) @@ -662,7 +663,7 @@ def execute(_args, task: str, **_kwargs): network_policy=DEFAULT_NETWORK_POLICY, patch_file=DEFAULT_PATCH_FILE, remote_smoke=DEFAULT_REMOTE_SMOKE, - remote_install_codex=DEFAULT_REMOTE_SMOKE, + bundle_scripts=[], setup_timeout=1, evaluation_timeout=1, min_free_gb=1, @@ -676,7 +677,9 @@ def execute(_args, task: str, **_kwargs): max_attempts=1, agent_timeout=1, provider="bedrock", - bedrock_region="us-west-2", + model_region="us-west-2", + model_base_url=None, + model_key_env=None, ) summary = _batch(args) @@ -1088,17 +1091,17 @@ def test_packaged_policies_share_host_lists() -> None: def test_locked_policy_refuses_agent_runs_but_allows_smoke() -> None: locked = _policy(LOCKED_POLICY) - with pytest.raises(ValueError, match="agent tooling"): - _require_runnable_policy(locked, kind="run", egress="policy") - _require_runnable_policy(locked, kind="smoke", egress="policy") + with pytest.raises(ValueError, match="installs its tooling"): + _require_runnable_policy(locked, kind="run", egress="policy", agent_name="codex") + _require_runnable_policy(locked, kind="smoke", egress="policy", agent_name="codex") default = _policy(DEFAULT_NETWORK_POLICY) - _require_runnable_policy(default, kind="run", egress="policy") - _require_runnable_policy(default, kind="run", egress="restricted") + _require_runnable_policy(default, kind="run", egress="policy", agent_name="codex") + _require_runnable_policy(default, kind="run", egress="restricted", agent_name="codex") def test_cli_validates_asset_overrides_before_doing_work(tmp_path: Path, capsys) -> None: missing = tmp_path / "instal_codex.sh" - code = main(["preflight", "--remote-install-codex", str(missing)]) + code = main(["preflight", "--bundle-script", f"install_codex.sh={missing}"]) assert code == 1 error = json.loads(capsys.readouterr().err)["error"] assert error["type"] == "FileNotFoundError" @@ -1126,4 +1129,4 @@ def test_batch_refuses_unrunnable_policy_before_submitting_work( assert code == 1 error = json.loads(capsys.readouterr().err)["error"] assert error["type"] == "ValueError" - assert "agent tooling" in error["message"] + assert "installs its tooling" in error["message"] diff --git a/tests/test_templates.py b/tests/test_templates.py index 502aa9b..831dc19 100644 --- a/tests/test_templates.py +++ b/tests/test_templates.py @@ -224,7 +224,6 @@ def test_cli_defaults_resolve_packaged_runtime_assets_outside_checkout( assert args.network_policy.is_file() assert args.patch_file.is_file() assert args.remote_smoke.is_file() - assert args.remote_install_codex.is_file() assert Path("artifacts/templates/manifest.json") == DEFAULT_MANIFEST diff --git a/uv.lock b/uv.lock index 2e73021..28e01a0 100644 --- a/uv.lock +++ b/uv.lock @@ -278,7 +278,7 @@ wheels = [ [[package]] name = "cybergym-e2e-e2b" -version = "0.3.0" +version = "0.4.0" source = { editable = "." } dependencies = [ { name = "e2b" },