From 0928ad64f32c60af0ebbe0a4fceb8f8da2c3b89e Mon Sep 17 00:00:00 2001 From: Sandy Chapman Date: Fri, 3 Jul 2026 10:55:18 -0300 Subject: [PATCH] refactor(evaluator-sdk): make the Codex runtime a general coding-agent runtime MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The shipped CodexCliAgentRuntime was ProfBench-shaped: a hardcoded "return only the final answer text" prompt, no way to seed a workspace, and no success signal for the standard AgentPhaseSuccessMetric. That made it unusable for general coding tasks (write docs, write tests, fix a bug) and leaked a benchmark's opinions into src/. Generalize the runtime: - Neutral, injectable prompt: default_codex_prompt presents the task and invites workspace edits; both runtimes take a prompt_builder to override framing. - Workspace seeding: inputs["files"] = {path: contents} is staged into the agent's workspace before it runs (paths escaping the workspace are rejected), so a task can hand the agent starter code. - Stamp agent_ok from the exit code so AgentPhaseSuccessMetric works over Codex trials. Keep runtime selection generic and in the SDK (resolve_codex_runtime + the RuntimeChoice/EffectiveCodexRuntime enums); move only ProfBench's own policy — its candidate prompt and the candidate+judge score_source labels — into the ProfBench example, inlined at the point of use in runner.py. Backward compatible: the plugin's CodexRunnerTarget path is unchanged (the new constructor args are optional) and now gets the neutral prompt for free. Signed-off-by: Sandy Chapman --- .../examples/profbench/runner.py | 32 +++- .../agent_eval/runtimes/codex/runtime.py | 119 +++++++++--- .../tests/agent_eval/test_codex_runtime.py | 169 +++++++++++++----- .../agent_eval/runtimes/codex/runtime.py | 119 +++++++++--- 4 files changed, 338 insertions(+), 101 deletions(-) diff --git a/packages/nemo_evaluator_sdk/examples/profbench/runner.py b/packages/nemo_evaluator_sdk/examples/profbench/runner.py index 1eba8276a7..d8435f32f7 100644 --- a/packages/nemo_evaluator_sdk/examples/profbench/runner.py +++ b/packages/nemo_evaluator_sdk/examples/profbench/runner.py @@ -27,9 +27,9 @@ EffectiveCodexRuntime, RuntimeChoice, print_codex_agent_models, - resolve_codex_target, + resolve_codex_runtime, ) -from nemo_evaluator_sdk.agent_eval.tasks import AgentEvalRunConfig +from nemo_evaluator_sdk.agent_eval.tasks import AgentEvalRunConfig, AgentEvalTask from nemo_evaluator_sdk.agent_eval.trials import AgentEvalTarget, AgentEvalTrial from nemo_evaluator_sdk.values import InferenceParams, Model, RunConfigOnlineModel, SecretRef @@ -241,6 +241,28 @@ def _judge_model() -> Model: ) +# ProfBench-specific Codex policy. Runtime *selection* is generic (resolve_codex_runtime, in the SDK); +# ProfBench owns only how a task is framed for the agent and how the resulting score is labeled. +# ProfBench runs Codex as a *candidate* whose single text answer a live judge grades, so the prompt +# forbids tool chatter and the score_source records the candidate+judge topology. +PROFBENCH_SCORE_SOURCE = { + EffectiveCodexRuntime.LOCAL_CLI: "codex_cli_candidate_and_live_judge", + EffectiveCodexRuntime.DOCKER_CLI: "codex_docker_cli_candidate_and_live_judge", + EffectiveCodexRuntime.DOCKER_SANDBOX: "docker_sandbox_candidate_and_live_judge", +} + + +def profbench_codex_prompt(task: AgentEvalTask) -> str: + """Frame a task as a ProfBench candidate: return only the final answer text, no tooling chatter.""" + return ( + "Answer the ProfBench task below. Return only the final answer text; do not include " + "analysis, markdown fences, tool logs, or commentary.\n\n" + f"Task id: {task.id}\n" + f"Intent: {task.intent}\n" + f"Inputs: {task.inputs}\n" + ) + + def _live_candidate_target( *, agent: AgentChoice, @@ -265,13 +287,15 @@ def _live_candidate_target( None, ) if agent == AgentChoice.CODEX: - target, score_source, effective_runtime = resolve_codex_target( + # SDK picks the runtime; ProfBench supplies the candidate prompt and the score_source label. + target, effective_runtime = resolve_codex_runtime( runtime=runtime, model=agent_model, output_dir=output_dir, env=env, + prompt_builder=profbench_codex_prompt, ) - return target, None, score_source, effective_runtime + return target, None, PROFBENCH_SCORE_SOURCE[effective_runtime], effective_runtime raise ValueError(f"unsupported ProfBench agent {agent!r}") diff --git a/packages/nemo_evaluator_sdk/src/nemo_evaluator_sdk/agent_eval/runtimes/codex/runtime.py b/packages/nemo_evaluator_sdk/src/nemo_evaluator_sdk/agent_eval/runtimes/codex/runtime.py index 19bb55fb54..96154dd506 100644 --- a/packages/nemo_evaluator_sdk/src/nemo_evaluator_sdk/agent_eval/runtimes/codex/runtime.py +++ b/packages/nemo_evaluator_sdk/src/nemo_evaluator_sdk/agent_eval/runtimes/codex/runtime.py @@ -30,16 +30,30 @@ class RuntimeChoice(StrEnum): + """Which Codex execution mode the caller wants.""" + DOCKER = "docker" LOCAL = "local" class EffectiveCodexRuntime(StrEnum): + """The concrete runtime chosen for a :class:`RuntimeChoice` + environment.""" + DOCKER_SANDBOX = "docker_sandbox" DOCKER_CLI = "docker_cli" LOCAL_CLI = "local_cli" +#: Builds the prompt handed to Codex on stdin for a task. Swap it to change how a task is framed +#: (e.g. a benchmark-specific preamble); the default presents the task and invites workspace edits. +CodexPromptBuilder = Callable[[AgentEvalTask], str] + +#: Optional ``inputs`` key holding a ``{relative_path: contents}`` map of files to seed into the +#: agent's workspace before it runs — how a task hands the agent starter code (a buggy file to fix, +#: a module to test, a project skeleton). Excluded from the default prompt body (listed by name). +SEED_FILES_INPUT_KEY = "files" + + class CodexCliAgentRuntime: """AgentTaskRunner that uses the locally installed Codex CLI credentials.""" @@ -50,6 +64,7 @@ def __init__( work_root: str | Path | None = None, codex_bin: str = "codex", timeout_s: int = DEFAULT_CODEX_TIMEOUT_S, + prompt_builder: CodexPromptBuilder | None = None, process_factory: ProcessFactory | None = None, runtime_name: str = "codex_cli", ) -> None: @@ -57,6 +72,7 @@ def __init__( self._work_root = Path(work_root).expanduser() if work_root is not None else None self._codex_bin = codex_bin self._timeout_s = timeout_s + self._prompt_builder = prompt_builder or default_codex_prompt self._process_factory = process_factory or asyncio.create_subprocess_exec self._runtime_name = runtime_name @@ -83,7 +99,7 @@ async def _run_task(self, index: int, task: AgentEvalTask, config: AgentEvalRunC evidence_dir.mkdir(parents=True, exist_ok=True) workspace_dir.mkdir(parents=True, exist_ok=True) - prompt = _codex_prompt(task) + prompt = self._prompt_builder(task) prompt_path = evidence_dir / "prompt.txt" task_path = evidence_dir / "task.json" stdout_path = evidence_dir / "stdout.jsonl" @@ -96,6 +112,9 @@ async def _run_task(self, index: int, task: AgentEvalTask, config: AgentEvalRunC command = self._command(workspace_dir=workspace_dir, final_output_path=final_output_path) process: Any | None = None try: + # Seed inside the guarded block so a bad seed (e.g. a path escaping the workspace) fails + # just this task rather than aborting the whole run. + seeded_files = _seed_workspace(workspace_dir, task) process = await self._process_factory( *command, stdin=subprocess.PIPE, @@ -158,6 +177,8 @@ async def _run_task(self, index: int, task: AgentEvalTask, config: AgentEvalRunC "runtime": self._runtime_name, "agent": "codex", "agent_model": self._model, + "agent_ok": True, + "seeded_files": seeded_files, "generated": True, }, ) @@ -204,12 +225,14 @@ def __init__( codex_package: str = DEFAULT_CODEX_DOCKER_CLI_PACKAGE, auth_path: str | Path | None = None, timeout_s: int = DEFAULT_CODEX_TIMEOUT_S, + prompt_builder: CodexPromptBuilder | None = None, process_factory: ProcessFactory | None = None, ) -> None: super().__init__( model=model, work_root=work_root, timeout_s=timeout_s, + prompt_builder=prompt_builder, process_factory=process_factory, runtime_name="codex_docker_cli", ) @@ -280,33 +303,43 @@ def _command(self, *, workspace_dir: Path, final_output_path: Path) -> list[str] ] -def resolve_codex_target( +def resolve_codex_runtime( *, runtime: RuntimeChoice, model: str | None, output_dir: Path, env: Mapping[str, str] = os.environ, -) -> tuple[CodexCliAgentRuntime | CodexDockerCliAgentRuntime | DockerSandboxAgentRuntime, str, EffectiveCodexRuntime]: - """Resolve a Codex-backed agent-eval target for ProfBench-style candidate runs.""" + prompt_builder: CodexPromptBuilder | None = None, +) -> tuple[CodexCliAgentRuntime | CodexDockerCliAgentRuntime | DockerSandboxAgentRuntime, EffectiveCodexRuntime]: + """Pick and construct a Codex runtime for a run-mode + environment. + + ``local`` runs the on-PATH Codex CLI. ``docker`` prefers the OpenAI-Agents ``DockerSandbox`` when + ``OPENAI_API_KEY`` is an OpenAI platform secret (``sk-...``) and otherwise falls back to the + containerized Codex CLI (which mounts ``~/.codex/auth.json``). ``prompt_builder`` is threaded into + the CLI runtimes; the sandbox runtime does its own prompting. Returns the runtime plus the + :class:`EffectiveCodexRuntime` actually chosen so callers can label/report it. + """ effective_runtime = _resolve_codex_runtime(runtime, env) if effective_runtime == EffectiveCodexRuntime.LOCAL_CLI: return ( - CodexCliAgentRuntime(model=model, work_root=output_dir / "evidence" / "codex"), - "codex_cli_candidate_and_live_judge", + CodexCliAgentRuntime( + model=model, + work_root=output_dir / "evidence" / "codex", + prompt_builder=prompt_builder, + ), effective_runtime, ) if effective_runtime == EffectiveCodexRuntime.DOCKER_CLI: return ( - CodexDockerCliAgentRuntime(model=model, work_root=output_dir / "evidence" / "codex-docker"), - "codex_docker_cli_candidate_and_live_judge", + CodexDockerCliAgentRuntime( + model=model, + work_root=output_dir / "evidence" / "codex-docker", + prompt_builder=prompt_builder, + ), effective_runtime, ) if effective_runtime == EffectiveCodexRuntime.DOCKER_SANDBOX: - return ( - DockerSandboxAgentRuntime(model=model or DEFAULT_CODEX_DOCKER_MODEL), - "docker_sandbox_candidate_and_live_judge", - effective_runtime, - ) + return DockerSandboxAgentRuntime(model=model or DEFAULT_CODEX_DOCKER_MODEL), effective_runtime raise ValueError(f"unsupported Codex runtime {runtime!r}") @@ -320,6 +353,10 @@ def _resolve_codex_runtime(runtime: RuntimeChoice, env: Mapping[str, str] = os.e raise ValueError(f"unsupported Codex runtime {runtime!r}") +def _openai_sdk_secret_key_is_set(env: Mapping[str, str] = os.environ) -> bool: + return env.get("OPENAI_API_KEY", "").strip().startswith("sk-") + + def list_codex_agent_models(*, codex_bin: str = "codex") -> list[dict[str, Any]]: """Return visible Codex model descriptors from the local Codex CLI.""" if shutil.which(codex_bin) is None: @@ -351,14 +388,49 @@ def print_codex_agent_models(*, codex_bin: str = "codex") -> None: print(slug) -def _codex_prompt(task: AgentEvalTask) -> str: - return ( - "Answer the ProfBench task below. Return only the final answer text; do not include " - "analysis, markdown fences, tool logs, or commentary.\n\n" - f"Task id: {task.id}\n" - f"Intent: {task.intent}\n" - f"Inputs: {task.inputs}\n" - ) +def default_codex_prompt(task: AgentEvalTask) -> str: + """Frame a task for Codex as an agent that works in its current directory. + + Task-agnostic: it states the intent and inputs and invites the agent to read/create/edit files, + rather than constraining the answer to a single text reply. Seed files (``inputs[SEED_FILES_INPUT_KEY]``) + are listed by name instead of dumped inline — the agent finds them already in its workspace. Pass a + custom :data:`CodexPromptBuilder` to the runtime to override this framing for a specific benchmark. + """ + body_inputs = {key: value for key, value in task.inputs.items() if key != SEED_FILES_INPUT_KEY} + lines = [f"Task id: {task.id}", f"Intent: {task.intent}"] + if body_inputs: + lines += ["", "Inputs:", json.dumps(body_inputs, indent=2, default=str)] + seeded = task.inputs.get(SEED_FILES_INPUT_KEY) + if isinstance(seeded, Mapping) and seeded: + lines += ["", "These files are already in your working directory:"] + lines += [f" - {path}" for path in seeded] + lines += [ + "", + "Complete the task by working in your current directory. You may read, create, and edit files " + "as needed. When you are done, briefly summarize what you changed.", + ] + return "\n".join(lines) + "\n" + + +def _seed_workspace(workspace_dir: Path, task: AgentEvalTask) -> list[str]: + """Write any ``inputs[SEED_FILES_INPUT_KEY]`` files into the workspace before the agent runs. + + Returns the seeded relative paths (for trial metadata). Paths that escape the workspace (absolute + or ``..`` traversal) are rejected so a task can only stage files inside its own sandbox. + """ + seeds = task.inputs.get(SEED_FILES_INPUT_KEY) + if not isinstance(seeds, Mapping): + return [] + workspace_root = workspace_dir.resolve() + written: list[str] = [] + for rel_path, contents in seeds.items(): + target = (workspace_root / str(rel_path)).resolve() + if target != workspace_root and workspace_root not in target.parents: + raise ValueError(f"seed file path escapes the workspace: {rel_path!r}") + target.parent.mkdir(parents=True, exist_ok=True) + target.write_text(contents if isinstance(contents, str) else str(contents), encoding="utf-8") + written.append(str(rel_path)) + return written def _failed_codex_trial( @@ -384,6 +456,7 @@ def _failed_codex_trial( metadata={ "runtime": runtime_name, "agent": "codex", + "agent_ok": False, "error_type": exc.__class__.__name__, "error": str(exc), }, @@ -408,7 +481,3 @@ def _decode_process_output(value: bytes | str | None) -> str: def _safe_path_name(value: str) -> str: return "".join(char if char.isalnum() or char in "._-" else "-" for char in value).strip(".-")[:120] - - -def _openai_sdk_secret_key_is_set(env: Mapping[str, str] = os.environ) -> bool: - return env.get("OPENAI_API_KEY", "").strip().startswith("sk-") diff --git a/packages/nemo_evaluator_sdk/tests/agent_eval/test_codex_runtime.py b/packages/nemo_evaluator_sdk/tests/agent_eval/test_codex_runtime.py index 625aac98ea..7699910384 100644 --- a/packages/nemo_evaluator_sdk/tests/agent_eval/test_codex_runtime.py +++ b/packages/nemo_evaluator_sdk/tests/agent_eval/test_codex_runtime.py @@ -10,76 +10,56 @@ from nemo_evaluator_sdk.agent_eval.runtimes.docker_sandbox import DockerSandboxAgentRuntime from nemo_evaluator_sdk.agent_eval.tasks import AgentEvalTask +# The runtime *selection* (local vs docker-cli vs docker-sandbox) is generic and lives here; only the +# ProfBench ``score_source`` labels + candidate prompt live in the example (test_profbench_codex_target). -def test_resolve_codex_target_selects_local_cli_runtime(tmp_path: Path) -> None: - target, score_source, effective_runtime = codex_runtime.resolve_codex_target( + +def _prompt_builder(task: AgentEvalTask) -> str: + return f"do: {task.id}\n" + + +def test_resolve_codex_runtime_local_cli_threads_prompt_builder(tmp_path: Path) -> None: + target, effective = codex_runtime.resolve_codex_runtime( runtime=codex_runtime.RuntimeChoice.LOCAL, model="gpt-5", - output_dir=tmp_path / "live-candidate", + output_dir=tmp_path / "run", env={"OPENAI_API_KEY": "sk-test-key"}, + prompt_builder=_prompt_builder, ) assert isinstance(target, codex_runtime.CodexCliAgentRuntime) assert target._model == "gpt-5" - assert target._work_root == tmp_path / "live-candidate" / "evidence" / "codex" - assert score_source == "codex_cli_candidate_and_live_judge" - assert effective_runtime == codex_runtime.EffectiveCodexRuntime.LOCAL_CLI + assert target._work_root == tmp_path / "run" / "evidence" / "codex" + assert target._prompt_builder is _prompt_builder + assert effective == codex_runtime.EffectiveCodexRuntime.LOCAL_CLI -def test_resolve_codex_target_uses_sdk_docker_when_openai_api_key_is_set(tmp_path: Path) -> None: - target, score_source, effective_runtime = codex_runtime.resolve_codex_target( +def test_resolve_codex_runtime_docker_uses_sandbox_for_openai_secret_key(tmp_path: Path) -> None: + target, effective = codex_runtime.resolve_codex_runtime( runtime=codex_runtime.RuntimeChoice.DOCKER, model=None, - output_dir=tmp_path / "live-candidate", + output_dir=tmp_path / "run", env={"OPENAI_API_KEY": "sk-test-key"}, ) assert isinstance(target, DockerSandboxAgentRuntime) assert target._model == codex_runtime.DEFAULT_CODEX_DOCKER_MODEL - assert score_source == "docker_sandbox_candidate_and_live_judge" - assert effective_runtime == codex_runtime.EffectiveCodexRuntime.DOCKER_SANDBOX - - -def test_resolve_codex_target_uses_sdk_docker_agent_model(tmp_path: Path) -> None: - target, score_source, effective_runtime = codex_runtime.resolve_codex_target( - runtime=codex_runtime.RuntimeChoice.DOCKER, - model="gpt-5.4", - output_dir=tmp_path / "live-candidate", - env={"OPENAI_API_KEY": "sk-test-key"}, - ) - - assert isinstance(target, DockerSandboxAgentRuntime) - assert target._model == "gpt-5.4" - assert score_source == "docker_sandbox_candidate_and_live_judge" - assert effective_runtime == codex_runtime.EffectiveCodexRuntime.DOCKER_SANDBOX + assert effective == codex_runtime.EffectiveCodexRuntime.DOCKER_SANDBOX -def test_resolve_codex_target_falls_back_to_docker_cli_without_openai_api_key(tmp_path: Path) -> None: - target, score_source, effective_runtime = codex_runtime.resolve_codex_target( +def test_resolve_codex_runtime_docker_falls_back_to_cli_without_sdk_key(tmp_path: Path) -> None: + target, effective = codex_runtime.resolve_codex_runtime( runtime=codex_runtime.RuntimeChoice.DOCKER, model="gpt-5.4", - output_dir=tmp_path / "live-candidate", - env={}, - ) - - assert isinstance(target, codex_runtime.CodexDockerCliAgentRuntime) - assert target._model == "gpt-5.4" - assert target._work_root == tmp_path / "live-candidate" / "evidence" / "codex-docker" - assert score_source == "codex_docker_cli_candidate_and_live_judge" - assert effective_runtime == codex_runtime.EffectiveCodexRuntime.DOCKER_CLI - - -def test_resolve_codex_target_falls_back_to_docker_cli_with_oauth_token(tmp_path: Path) -> None: - target, score_source, effective_runtime = codex_runtime.resolve_codex_target( - runtime=codex_runtime.RuntimeChoice.DOCKER, - model="gpt-5.4", - output_dir=tmp_path / "live-candidate", + output_dir=tmp_path / "run", env={"OPENAI_API_KEY": "oauth-token"}, + prompt_builder=_prompt_builder, ) assert isinstance(target, codex_runtime.CodexDockerCliAgentRuntime) - assert score_source == "codex_docker_cli_candidate_and_live_judge" - assert effective_runtime == codex_runtime.EffectiveCodexRuntime.DOCKER_CLI + assert target._work_root == tmp_path / "run" / "evidence" / "codex-docker" + assert target._prompt_builder is _prompt_builder + assert effective == codex_runtime.EffectiveCodexRuntime.DOCKER_CLI def test_list_codex_agent_models_prints_visible_models( @@ -127,7 +107,9 @@ def __init__(self, command: tuple[str, ...]) -> None: self.command = command async def communicate(self, input: bytes) -> tuple[bytes, bytes]: - assert b"Answer the ProfBench task below" in input + # Default prompt is task-agnostic: it states the task and invites workspace edits. + assert b"Task id: task/1" in input + assert b"Intent:" in input final_output_path = Path(self.command[self.command.index("--output-last-message") + 1]) final_output_path.write_text("codex answer", encoding="utf-8") return b'{"type":"event"}\n', b"" @@ -179,7 +161,7 @@ def __init__(self, command: tuple[str, ...]) -> None: self.command = command async def communicate(self, input: bytes) -> tuple[bytes, bytes]: - assert b"Answer the ProfBench task below" in input + assert b"Task id: task/1" in input evidence_mount = self.command[self.command.index(f"{auth_path.resolve()}:/root/.codex/auth.json:ro") + 4] evidence_dir = Path(evidence_mount.split(":/evidence", maxsplit=1)[0]) (evidence_dir / "final_output.txt").write_text("docker codex answer", encoding="utf-8") @@ -289,3 +271,96 @@ async def fake_process_factory(*command: str, **kwargs: Any) -> FakeProcess: assert trials[0].output.output_text == "stdout fallback\n" final_output = tmp_path / "codex" / "000000-task-2" / "final_output.txt" assert final_output.read_text(encoding="utf-8") == "stdout fallback\n" + + +@pytest.mark.asyncio +async def test_codex_cli_agent_runtime_seeds_workspace_and_stamps_agent_ok( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + class FakeProcess: + returncode = 0 + + def __init__(self, command: tuple[str, ...]) -> None: + self.command = command + + async def communicate(self, input: bytes) -> tuple[bytes, bytes]: + # Seed files are staged before the agent runs and listed (by name) in the prompt. + workspace_dir = Path(self.command[self.command.index("--cd") + 1]) + assert (workspace_dir / "buggy.py").read_text(encoding="utf-8") == "def add(a, b)\n return a + b\n" + assert b"buggy.py" in input + final_output_path = Path(self.command[self.command.index("--output-last-message") + 1]) + final_output_path.write_text("fixed it", encoding="utf-8") + return b"", b"" + + async def fake_process_factory(*command: str, **kwargs: Any) -> FakeProcess: + return FakeProcess(command) + + monkeypatch.setattr(codex_runtime.shutil, "which", lambda value: f"/bin/{value}") + runtime = codex_runtime.CodexCliAgentRuntime( + work_root=tmp_path / "codex", + process_factory=fake_process_factory, + ) + task = AgentEvalTask( + id="fix-bug", + intent="Fix the syntax error.", + inputs={"files": {"buggy.py": "def add(a, b)\n return a + b\n"}}, + ) + + trials = await runtime.run_tasks([task]) + + assert trials[0].status == "completed" + # agent_ok is stamped so AgentPhaseSuccessMetric works over Codex trials. + assert trials[0].metadata["agent_ok"] is True + assert trials[0].metadata["seeded_files"] == ["buggy.py"] + + +@pytest.mark.asyncio +async def test_codex_cli_agent_runtime_rejects_seed_path_escaping_workspace( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + async def fake_process_factory(*command: str, **kwargs: Any) -> Any: # pragma: no cover - never reached + raise AssertionError("agent should not run when seeding fails") + + monkeypatch.setattr(codex_runtime.shutil, "which", lambda value: f"/bin/{value}") + runtime = codex_runtime.CodexCliAgentRuntime( + work_root=tmp_path / "codex", + process_factory=fake_process_factory, + ) + task = AgentEvalTask(id="evil", intent="escape", inputs={"files": {"../escape.txt": "x"}}) + + # A traversal path is surfaced as a failed trial (the exception is caught per-task). + trials = await runtime.run_tasks([task]) + assert trials[0].status == "failed" + assert trials[0].metadata["error_type"] == "ValueError" + assert trials[0].metadata["agent_ok"] is False + + +@pytest.mark.asyncio +async def test_codex_cli_agent_runtime_uses_injected_prompt_builder( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + class FakeProcess: + returncode = 0 + + def __init__(self, command: tuple[str, ...]) -> None: + self.command = command + + async def communicate(self, input: bytes) -> tuple[bytes, bytes]: + assert input == b"CUSTOM: fix-bug\n" + final_output_path = Path(self.command[self.command.index("--output-last-message") + 1]) + final_output_path.write_text("ok", encoding="utf-8") + return b"", b"" + + async def fake_process_factory(*command: str, **kwargs: Any) -> FakeProcess: + return FakeProcess(command) + + monkeypatch.setattr(codex_runtime.shutil, "which", lambda value: f"/bin/{value}") + runtime = codex_runtime.CodexCliAgentRuntime( + work_root=tmp_path / "codex", + prompt_builder=lambda task: f"CUSTOM: {task.id}\n", + process_factory=fake_process_factory, + ) + task = AgentEvalTask(id="fix-bug", intent="Fix.", inputs={}) + + trials = await runtime.run_tasks([task]) + assert trials[0].status == "completed" diff --git a/sdk/python/nemo-platform/src/nemo_platform/beta/evaluator/agent_eval/runtimes/codex/runtime.py b/sdk/python/nemo-platform/src/nemo_platform/beta/evaluator/agent_eval/runtimes/codex/runtime.py index b9e929d3da..e7832a9d35 100644 --- a/sdk/python/nemo-platform/src/nemo_platform/beta/evaluator/agent_eval/runtimes/codex/runtime.py +++ b/sdk/python/nemo-platform/src/nemo_platform/beta/evaluator/agent_eval/runtimes/codex/runtime.py @@ -30,16 +30,30 @@ class RuntimeChoice(StrEnum): + """Which Codex execution mode the caller wants.""" + DOCKER = "docker" LOCAL = "local" class EffectiveCodexRuntime(StrEnum): + """The concrete runtime chosen for a :class:`RuntimeChoice` + environment.""" + DOCKER_SANDBOX = "docker_sandbox" DOCKER_CLI = "docker_cli" LOCAL_CLI = "local_cli" +#: Builds the prompt handed to Codex on stdin for a task. Swap it to change how a task is framed +#: (e.g. a benchmark-specific preamble); the default presents the task and invites workspace edits. +CodexPromptBuilder = Callable[[AgentEvalTask], str] + +#: Optional ``inputs`` key holding a ``{relative_path: contents}`` map of files to seed into the +#: agent's workspace before it runs — how a task hands the agent starter code (a buggy file to fix, +#: a module to test, a project skeleton). Excluded from the default prompt body (listed by name). +SEED_FILES_INPUT_KEY = "files" + + class CodexCliAgentRuntime: """AgentTaskRunner that uses the locally installed Codex CLI credentials.""" @@ -50,6 +64,7 @@ def __init__( work_root: str | Path | None = None, codex_bin: str = "codex", timeout_s: int = DEFAULT_CODEX_TIMEOUT_S, + prompt_builder: CodexPromptBuilder | None = None, process_factory: ProcessFactory | None = None, runtime_name: str = "codex_cli", ) -> None: @@ -57,6 +72,7 @@ def __init__( self._work_root = Path(work_root).expanduser() if work_root is not None else None self._codex_bin = codex_bin self._timeout_s = timeout_s + self._prompt_builder = prompt_builder or default_codex_prompt self._process_factory = process_factory or asyncio.create_subprocess_exec self._runtime_name = runtime_name @@ -83,7 +99,7 @@ async def _run_task(self, index: int, task: AgentEvalTask, config: AgentEvalRunC evidence_dir.mkdir(parents=True, exist_ok=True) workspace_dir.mkdir(parents=True, exist_ok=True) - prompt = _codex_prompt(task) + prompt = self._prompt_builder(task) prompt_path = evidence_dir / "prompt.txt" task_path = evidence_dir / "task.json" stdout_path = evidence_dir / "stdout.jsonl" @@ -96,6 +112,9 @@ async def _run_task(self, index: int, task: AgentEvalTask, config: AgentEvalRunC command = self._command(workspace_dir=workspace_dir, final_output_path=final_output_path) process: Any | None = None try: + # Seed inside the guarded block so a bad seed (e.g. a path escaping the workspace) fails + # just this task rather than aborting the whole run. + seeded_files = _seed_workspace(workspace_dir, task) process = await self._process_factory( *command, stdin=subprocess.PIPE, @@ -158,6 +177,8 @@ async def _run_task(self, index: int, task: AgentEvalTask, config: AgentEvalRunC "runtime": self._runtime_name, "agent": "codex", "agent_model": self._model, + "agent_ok": True, + "seeded_files": seeded_files, "generated": True, }, ) @@ -204,12 +225,14 @@ def __init__( codex_package: str = DEFAULT_CODEX_DOCKER_CLI_PACKAGE, auth_path: str | Path | None = None, timeout_s: int = DEFAULT_CODEX_TIMEOUT_S, + prompt_builder: CodexPromptBuilder | None = None, process_factory: ProcessFactory | None = None, ) -> None: super().__init__( model=model, work_root=work_root, timeout_s=timeout_s, + prompt_builder=prompt_builder, process_factory=process_factory, runtime_name="codex_docker_cli", ) @@ -280,33 +303,43 @@ def _command(self, *, workspace_dir: Path, final_output_path: Path) -> list[str] ] -def resolve_codex_target( +def resolve_codex_runtime( *, runtime: RuntimeChoice, model: str | None, output_dir: Path, env: Mapping[str, str] = os.environ, -) -> tuple[CodexCliAgentRuntime | CodexDockerCliAgentRuntime | DockerSandboxAgentRuntime, str, EffectiveCodexRuntime]: - """Resolve a Codex-backed agent-eval target for ProfBench-style candidate runs.""" + prompt_builder: CodexPromptBuilder | None = None, +) -> tuple[CodexCliAgentRuntime | CodexDockerCliAgentRuntime | DockerSandboxAgentRuntime, EffectiveCodexRuntime]: + """Pick and construct a Codex runtime for a run-mode + environment. + + ``local`` runs the on-PATH Codex CLI. ``docker`` prefers the OpenAI-Agents ``DockerSandbox`` when + ``OPENAI_API_KEY`` is an OpenAI platform secret (``sk-...``) and otherwise falls back to the + containerized Codex CLI (which mounts ``~/.codex/auth.json``). ``prompt_builder`` is threaded into + the CLI runtimes; the sandbox runtime does its own prompting. Returns the runtime plus the + :class:`EffectiveCodexRuntime` actually chosen so callers can label/report it. + """ effective_runtime = _resolve_codex_runtime(runtime, env) if effective_runtime == EffectiveCodexRuntime.LOCAL_CLI: return ( - CodexCliAgentRuntime(model=model, work_root=output_dir / "evidence" / "codex"), - "codex_cli_candidate_and_live_judge", + CodexCliAgentRuntime( + model=model, + work_root=output_dir / "evidence" / "codex", + prompt_builder=prompt_builder, + ), effective_runtime, ) if effective_runtime == EffectiveCodexRuntime.DOCKER_CLI: return ( - CodexDockerCliAgentRuntime(model=model, work_root=output_dir / "evidence" / "codex-docker"), - "codex_docker_cli_candidate_and_live_judge", + CodexDockerCliAgentRuntime( + model=model, + work_root=output_dir / "evidence" / "codex-docker", + prompt_builder=prompt_builder, + ), effective_runtime, ) if effective_runtime == EffectiveCodexRuntime.DOCKER_SANDBOX: - return ( - DockerSandboxAgentRuntime(model=model or DEFAULT_CODEX_DOCKER_MODEL), - "docker_sandbox_candidate_and_live_judge", - effective_runtime, - ) + return DockerSandboxAgentRuntime(model=model or DEFAULT_CODEX_DOCKER_MODEL), effective_runtime raise ValueError(f"unsupported Codex runtime {runtime!r}") @@ -320,6 +353,10 @@ def _resolve_codex_runtime(runtime: RuntimeChoice, env: Mapping[str, str] = os.e raise ValueError(f"unsupported Codex runtime {runtime!r}") +def _openai_sdk_secret_key_is_set(env: Mapping[str, str] = os.environ) -> bool: + return env.get("OPENAI_API_KEY", "").strip().startswith("sk-") + + def list_codex_agent_models(*, codex_bin: str = "codex") -> list[dict[str, Any]]: """Return visible Codex model descriptors from the local Codex CLI.""" if shutil.which(codex_bin) is None: @@ -351,14 +388,49 @@ def print_codex_agent_models(*, codex_bin: str = "codex") -> None: print(slug) -def _codex_prompt(task: AgentEvalTask) -> str: - return ( - "Answer the ProfBench task below. Return only the final answer text; do not include " - "analysis, markdown fences, tool logs, or commentary.\n\n" - f"Task id: {task.id}\n" - f"Intent: {task.intent}\n" - f"Inputs: {task.inputs}\n" - ) +def default_codex_prompt(task: AgentEvalTask) -> str: + """Frame a task for Codex as an agent that works in its current directory. + + Task-agnostic: it states the intent and inputs and invites the agent to read/create/edit files, + rather than constraining the answer to a single text reply. Seed files (``inputs[SEED_FILES_INPUT_KEY]``) + are listed by name instead of dumped inline — the agent finds them already in its workspace. Pass a + custom :data:`CodexPromptBuilder` to the runtime to override this framing for a specific benchmark. + """ + body_inputs = {key: value for key, value in task.inputs.items() if key != SEED_FILES_INPUT_KEY} + lines = [f"Task id: {task.id}", f"Intent: {task.intent}"] + if body_inputs: + lines += ["", "Inputs:", json.dumps(body_inputs, indent=2, default=str)] + seeded = task.inputs.get(SEED_FILES_INPUT_KEY) + if isinstance(seeded, Mapping) and seeded: + lines += ["", "These files are already in your working directory:"] + lines += [f" - {path}" for path in seeded] + lines += [ + "", + "Complete the task by working in your current directory. You may read, create, and edit files " + "as needed. When you are done, briefly summarize what you changed.", + ] + return "\n".join(lines) + "\n" + + +def _seed_workspace(workspace_dir: Path, task: AgentEvalTask) -> list[str]: + """Write any ``inputs[SEED_FILES_INPUT_KEY]`` files into the workspace before the agent runs. + + Returns the seeded relative paths (for trial metadata). Paths that escape the workspace (absolute + or ``..`` traversal) are rejected so a task can only stage files inside its own sandbox. + """ + seeds = task.inputs.get(SEED_FILES_INPUT_KEY) + if not isinstance(seeds, Mapping): + return [] + workspace_root = workspace_dir.resolve() + written: list[str] = [] + for rel_path, contents in seeds.items(): + target = (workspace_root / str(rel_path)).resolve() + if target != workspace_root and workspace_root not in target.parents: + raise ValueError(f"seed file path escapes the workspace: {rel_path!r}") + target.parent.mkdir(parents=True, exist_ok=True) + target.write_text(contents if isinstance(contents, str) else str(contents), encoding="utf-8") + written.append(str(rel_path)) + return written def _failed_codex_trial( @@ -384,6 +456,7 @@ def _failed_codex_trial( metadata={ "runtime": runtime_name, "agent": "codex", + "agent_ok": False, "error_type": exc.__class__.__name__, "error": str(exc), }, @@ -408,7 +481,3 @@ def _decode_process_output(value: bytes | str | None) -> str: def _safe_path_name(value: str) -> str: return "".join(char if char.isalnum() or char in "._-" else "-" for char in value).strip(".-")[:120] - - -def _openai_sdk_secret_key_is_set(env: Mapping[str, str] = os.environ) -> bool: - return env.get("OPENAI_API_KEY", "").strip().startswith("sk-")