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 96154dd506..c75d402f9d 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 @@ -20,6 +20,7 @@ from nemo_evaluator_sdk.agent_eval.runtimes.docker_sandbox import DockerSandboxAgentRuntime from nemo_evaluator_sdk.agent_eval.tasks import AgentEvalRunConfig, AgentEvalTask from nemo_evaluator_sdk.agent_eval.trials import AgentEvalTrial, AgentEvalTrialStatus, AgentOutput +from nemo_evaluator_sdk.agent_eval.workspace_seeds import SEED_FILES_INPUT_KEY, seed_workspace from nemo_evaluator_sdk.values.evidence import CandidateEvidence, EvidenceDescriptor DEFAULT_CODEX_TIMEOUT_S = 600 @@ -48,11 +49,6 @@ class EffectiveCodexRuntime(StrEnum): #: (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.""" @@ -113,8 +109,10 @@ async def _run_task(self, index: int, task: AgentEvalTask, config: AgentEvalRunC 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) + # just this task rather than aborting the whole run. Offload to a worker thread: seeding is + # synchronous (a handler may do blocking I/O, e.g. the plugin's fileset download), and this + # runs on the event loop shared by every concurrent task, so a blocking seed would stall them all. + seeded_files = await asyncio.to_thread(seed_workspace, workspace_dir, task.inputs.get(SEED_FILES_INPUT_KEY)) process = await self._process_factory( *command, stdin=subprocess.PIPE, @@ -412,27 +410,6 @@ def default_codex_prompt(task: AgentEvalTask) -> str: 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( task: AgentEvalTask, evidence_dir: Path, diff --git a/packages/nemo_evaluator_sdk/src/nemo_evaluator_sdk/agent_eval/workspace_seeds.py b/packages/nemo_evaluator_sdk/src/nemo_evaluator_sdk/agent_eval/workspace_seeds.py new file mode 100644 index 0000000000..93cc4d1e84 --- /dev/null +++ b/packages/nemo_evaluator_sdk/src/nemo_evaluator_sdk/agent_eval/workspace_seeds.py @@ -0,0 +1,171 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Workspace seed files for agent-eval tasks. + +A task can stage starter files into the agent's workspace before it runs, under +``inputs[SEED_FILES_INPUT_KEY]`` as a ``{relative_path: source}`` map. Each *source* is a +JSON-serializable value whose ``kind`` selects a registered :class:`SeedHandler` (a bare string is +sugar for inline text). + +The SDK ships handlers only for the kinds it can resolve with **no external dependency** — ``inline`` +and ``path``. Other kinds are contributed by consumers via :func:`register_seed_handler`; the SDK has +no knowledge of them (e.g. a platform ``fileset`` handler lives in the plugin and resolves against the +Files service at run time). An unregistered kind raises :class:`WorkspaceSeedError`. +""" + +from __future__ import annotations + +import base64 +from collections.abc import Mapping +from pathlib import Path +from typing import Any, Literal, Protocol, runtime_checkable + +from pydantic import BaseModel, ConfigDict, Field + +#: ``inputs`` key holding the ``{relative_path: seed}`` map of files to stage into the workspace. +SEED_FILES_INPUT_KEY = "files" + + +class WorkspaceSeedError(ValueError): + """A workspace seed could not be parsed, resolved, or written (bad value, unknown kind, ...). + + Subclasses ``ValueError`` so a runner's per-task error handling still catches it. + """ + + +class InlineSeed(BaseModel): + """File contents carried in the task itself. Portable to any runner.""" + + model_config = ConfigDict(extra="forbid") + + kind: Literal["inline"] = "inline" + content: str = Field(description="File contents; UTF-8 text, or base64-encoded bytes when encoding='base64'.") + encoding: Literal["text", "base64"] = "text" + + +class PathSeed(BaseModel): + """A file on the authoring host. Resolvable only where that path exists (local runs).""" + + model_config = ConfigDict(extra="forbid") + + kind: Literal["path"] = "path" + path: str = Field(description="Filesystem path on the authoring host (absolute or relative to the cwd).") + + +@runtime_checkable +class SeedHandler(Protocol): + """Parses + resolves one seed ``kind`` into the bytes to stage. + + Consumers implement this for kinds the SDK doesn't ship (e.g. a platform ``fileset`` handler) and + wire them in with :func:`register_seed_handler`. ``resolve`` runs at seeding time, so a handler + that needs external services (a client, credentials) acquires them there. + """ + + kind: str + + def parse(self, value: Mapping[str, Any]) -> BaseModel: + """Validate a raw seed mapping into this kind's typed model.""" + ... + + def resolve(self, seed: BaseModel) -> bytes: + """Resolve a parsed seed to the bytes to write into the workspace.""" + ... + + +_HANDLERS: dict[str, SeedHandler] = {} + + +def register_seed_handler(handler: SeedHandler) -> None: + """Register a :class:`SeedHandler` under its ``kind`` (replacing any handler already registered).""" + _HANDLERS[handler.kind] = handler + + +def _handler_for(kind: str) -> SeedHandler: + handler = _HANDLERS.get(kind) + if handler is None: + raise WorkspaceSeedError(f"no handler registered for seed kind {kind!r}") + return handler + + +class _InlineSeedHandler: + kind = "inline" + + def parse(self, value: Mapping[str, Any]) -> BaseModel: + return InlineSeed.model_validate(value) + + def resolve(self, seed: BaseModel) -> bytes: + assert isinstance(seed, InlineSeed) + if seed.encoding == "base64": + try: + return base64.b64decode(seed.content, validate=True) + except (ValueError, TypeError) as exc: + raise WorkspaceSeedError(f"inline seed is not valid base64: {exc}") from exc + return seed.content.encode("utf-8") + + +class _PathSeedHandler: + kind = "path" + + def parse(self, value: Mapping[str, Any]) -> BaseModel: + return PathSeed.model_validate(value) + + def resolve(self, seed: BaseModel) -> bytes: + assert isinstance(seed, PathSeed) + source = Path(seed.path).expanduser() + try: + return source.read_bytes() + except OSError as exc: + raise WorkspaceSeedError(f"path seed {seed.path!r} could not be read: {exc}") from exc + + +register_seed_handler(_InlineSeedHandler()) +register_seed_handler(_PathSeedHandler()) + + +def parse_seed(value: str | Mapping[str, Any]) -> BaseModel: + """Coerce a seed map value into its validated model. A bare string is inline UTF-8 text.""" + if isinstance(value, str): + return InlineSeed(content=value) + if not isinstance(value, Mapping): + raise WorkspaceSeedError(f"seed must be a string or mapping, got {type(value).__name__}") + kind = value.get("kind") + if not isinstance(kind, str): + raise WorkspaceSeedError("seed mapping is missing a string 'kind'") + handler = _handler_for(kind) + try: + return handler.parse(value) + except WorkspaceSeedError: + raise + except Exception as exc: # noqa: BLE001 - normalize a handler's validation error into our type + raise WorkspaceSeedError(f"invalid {kind!r} seed: {exc}") from exc + + +def _resolve_seed_bytes(seed: BaseModel) -> bytes: + """Resolve a parsed seed to bytes via its registered handler.""" + kind = getattr(seed, "kind", None) + if not isinstance(kind, str): + raise WorkspaceSeedError("parsed seed has no string 'kind'") + return _handler_for(kind).resolve(seed) + + +def seed_workspace(workspace_dir: str | Path, files: Mapping[str, Any] | None) -> list[str]: + """Write the ``files`` seed map into ``workspace_dir``; return the seeded relative paths. + + Each value is parsed into a seed model and resolved to bytes by its registered handler. Paths that + escape the workspace (absolute, or ``..`` traversal) are rejected so a task can only stage files + inside its own sandbox. ``None``/empty seeds nothing. + """ + if not isinstance(files, Mapping): + return [] + root = Path(workspace_dir).resolve() + written: list[str] = [] + for rel_path, value in files.items(): + target = (root / str(rel_path)).resolve() + if target != root and root not in target.parents: + raise WorkspaceSeedError(f"seed file path escapes the workspace: {rel_path!r}") + data = _resolve_seed_bytes(parse_seed(value)) + target.parent.mkdir(parents=True, exist_ok=True) + target.write_bytes(data) + written.append(str(rel_path)) + return written 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 7699910384..c78f50f7c9 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 @@ -2,13 +2,17 @@ # SPDX-License-Identifier: Apache-2.0 import json +import threading +from collections.abc import Mapping from pathlib import Path from typing import Any import pytest +from nemo_evaluator_sdk.agent_eval import workspace_seeds from nemo_evaluator_sdk.agent_eval.runtimes.codex import runtime as codex_runtime from nemo_evaluator_sdk.agent_eval.runtimes.docker_sandbox import DockerSandboxAgentRuntime from nemo_evaluator_sdk.agent_eval.tasks import AgentEvalTask +from pydantic import BaseModel # 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). @@ -314,6 +318,58 @@ async def fake_process_factory(*command: str, **kwargs: Any) -> FakeProcess: assert trials[0].metadata["seeded_files"] == ["buggy.py"] +@pytest.mark.asyncio +async def test_codex_cli_agent_runtime_seeds_off_the_event_loop_thread( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + # Seeding is synchronous and a handler may block (e.g. the plugin's fileset download). It runs on + # the event loop shared by every concurrent task, so it must be offloaded to a worker thread — + # otherwise one slow seed stalls the whole run. Register a probe handler that records the thread + # it resolves on and assert it is not the loop thread. + class _ProbeSeed(BaseModel): + kind: str = "thread_probe" + + class _ProbeHandler: + kind = "thread_probe" + resolved_on: int | None = None + + def parse(self, value: Mapping[str, Any]) -> BaseModel: + return _ProbeSeed() + + def resolve(self, seed: BaseModel) -> bytes: + _ProbeHandler.resolved_on = threading.get_ident() + return b"probe" + + monkeypatch.setitem(workspace_seeds._HANDLERS, "thread_probe", _ProbeHandler()) + + class FakeProcess: + returncode = 0 + + def __init__(self, command: tuple[str, ...]) -> None: + self.command = command + + async def communicate(self, input: bytes) -> tuple[bytes, bytes]: + 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", + process_factory=fake_process_factory, + ) + task = AgentEvalTask(id="probe", intent="probe", inputs={"files": {"p.txt": {"kind": "thread_probe"}}}) + + trials = await runtime.run_tasks([task]) + + assert trials[0].status == "completed" + assert _ProbeHandler.resolved_on is not None + assert _ProbeHandler.resolved_on != threading.get_ident() + + @pytest.mark.asyncio async def test_codex_cli_agent_runtime_rejects_seed_path_escaping_workspace( tmp_path: Path, monkeypatch: pytest.MonkeyPatch @@ -331,7 +387,7 @@ async def fake_process_factory(*command: str, **kwargs: Any) -> Any: # pragma: # 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["error_type"] == "WorkspaceSeedError" assert trials[0].metadata["agent_ok"] is False diff --git a/packages/nemo_evaluator_sdk/tests/agent_eval/test_workspace_seeds.py b/packages/nemo_evaluator_sdk/tests/agent_eval/test_workspace_seeds.py new file mode 100644 index 0000000000..4abcea7046 --- /dev/null +++ b/packages/nemo_evaluator_sdk/tests/agent_eval/test_workspace_seeds.py @@ -0,0 +1,113 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import base64 +from collections.abc import Mapping +from pathlib import Path +from typing import Any, Literal + +import pytest +from nemo_evaluator_sdk.agent_eval import workspace_seeds +from nemo_evaluator_sdk.agent_eval.workspace_seeds import ( + InlineSeed, + PathSeed, + WorkspaceSeedError, + parse_seed, + register_seed_handler, + seed_workspace, +) +from pydantic import BaseModel + + +def test_parse_bare_string_is_inline_text() -> None: + seed = parse_seed("hello") + assert isinstance(seed, InlineSeed) + assert seed.content == "hello" + assert seed.encoding == "text" + + +def test_parse_dispatches_on_registered_kind() -> None: + assert isinstance(parse_seed({"kind": "inline", "content": "x"}), InlineSeed) + assert isinstance(parse_seed({"kind": "path", "path": "/tmp/x"}), PathSeed) + + +def test_parse_rejects_unregistered_kind() -> None: + # The SDK ships only inline/path; a kind no consumer has registered is a generic unknown-kind error. + with pytest.raises(WorkspaceSeedError, match="no handler registered for seed kind 'url'"): + parse_seed({"kind": "url", "href": "http://x"}) + + +def test_seed_inline_text_and_bare_string(tmp_path: Path) -> None: + written = seed_workspace( + tmp_path, + {"a.txt": "bare", "b.txt": {"kind": "inline", "content": "typed"}}, + ) + assert sorted(written) == ["a.txt", "b.txt"] + assert (tmp_path / "a.txt").read_text() == "bare" + assert (tmp_path / "b.txt").read_text() == "typed" + + +def test_seed_inline_base64_writes_binary(tmp_path: Path) -> None: + payload = b"\x00\x01binary\xff" + seed_workspace( + tmp_path, {"blob.bin": {"kind": "inline", "content": base64.b64encode(payload).decode(), "encoding": "base64"}} + ) + assert (tmp_path / "blob.bin").read_bytes() == payload + + +def test_seed_bad_base64_raises(tmp_path: Path) -> None: + with pytest.raises(WorkspaceSeedError, match="base64"): + seed_workspace(tmp_path, {"x": {"kind": "inline", "content": "not base64!!", "encoding": "base64"}}) + + +def test_seed_path_reads_local_file(tmp_path: Path) -> None: + source = tmp_path / "src.py" + source.write_text("print('hi')\n") + dest_root = tmp_path / "ws" + seed_workspace(dest_root, {"nested/copied.py": {"kind": "path", "path": str(source)}}) + assert (dest_root / "nested" / "copied.py").read_text() == "print('hi')\n" + + +def test_seed_missing_path_raises(tmp_path: Path) -> None: + with pytest.raises(WorkspaceSeedError, match="could not be read"): + seed_workspace(tmp_path, {"x.py": {"kind": "path", "path": str(tmp_path / "nope.py")}}) + + +def test_seed_unregistered_kind_raises(tmp_path: Path) -> None: + # Seeding a kind the SDK doesn't ship (and nobody registered) is a generic unknown-kind error — + # the SDK carries no awareness of platform kinds like 'fileset' or where they resolve. + with pytest.raises(WorkspaceSeedError, match="no handler registered for seed kind 'url'"): + seed_workspace(tmp_path, {"data.csv": {"kind": "url", "href": "http://x/data.csv"}}) + + +def test_register_custom_handler_extends_kinds(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + # A consumer (e.g. the plugin) can register a new kind without the SDK knowing about it. + class _UpperSeed(BaseModel): + kind: Literal["upper"] = "upper" + content: str + + class _UpperHandler: + kind = "upper" + + def parse(self, value: Mapping[str, Any]) -> BaseModel: + return _UpperSeed.model_validate(value) + + def resolve(self, seed: BaseModel) -> bytes: + assert isinstance(seed, _UpperSeed) + return seed.content.upper().encode("utf-8") + + # Isolate the global registry to this test so the extra kind doesn't leak to others. + monkeypatch.setattr(workspace_seeds, "_HANDLERS", dict(workspace_seeds._HANDLERS)) + register_seed_handler(_UpperHandler()) + + seed_workspace(tmp_path, {"shout.txt": {"kind": "upper", "content": "hi"}}) + assert (tmp_path / "shout.txt").read_text() == "HI" + + +def test_seed_rejects_path_escaping_workspace(tmp_path: Path) -> None: + with pytest.raises(WorkspaceSeedError, match="escapes the workspace"): + seed_workspace(tmp_path / "ws", {"../escape.txt": "x"}) + + +def test_seed_none_is_noop(tmp_path: Path) -> None: + assert seed_workspace(tmp_path, None) == [] diff --git a/plugins/nemo-evaluator/src/nemo_evaluator/agent_seeds.py b/plugins/nemo-evaluator/src/nemo_evaluator/agent_seeds.py new file mode 100644 index 0000000000..0798a87e54 --- /dev/null +++ b/plugins/nemo-evaluator/src/nemo_evaluator/agent_seeds.py @@ -0,0 +1,75 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Platform ``fileset`` workspace-seed handler for agent-eval. + +The SDK's workspace-seed machinery ships only ``inline``/``path`` kinds and knows nothing about +filesets. This module registers a ``fileset`` :class:`~nemo_evaluator_sdk.agent_eval.workspace_seeds.SeedHandler` +as an import side effect, so that when the evaluator plugin runs an agent-eval job, tasks may stage a +file from a stored fileset. Resolution happens at seed time against the Files service, using the +running job's task SDK — the SDK layer never gains a files dependency or awareness of this kind. +""" + +from __future__ import annotations + +import tempfile +from collections.abc import Mapping +from typing import Any, Literal + +from nemo_evaluator.filesets import FilesetRef, download_dataset_sync +from nemo_evaluator_sdk.agent_eval.workspace_seeds import WorkspaceSeedError, register_seed_handler +from nemo_platform_plugin.sdk_provider import get_task_sdk +from pydantic import BaseModel, ConfigDict, Field, field_validator + +#: Service identity used to build the task SDK (matches ``tasks/agent_evaluate.py``). +_EVALUATOR_SERVICE = "evaluator" + + +class FilesetSeed(BaseModel): + """A reference to a stored fileset, resolved via the Files service at seed time.""" + + model_config = ConfigDict(extra="forbid") + + kind: Literal["fileset"] = "fileset" + ref: str = Field( + description="Fileset reference: 'workspace/name', or 'workspace/name#path' for a single file.", + ) + + @field_validator("ref") + @classmethod + def _validate_ref_shape(cls, value: str) -> str: + # Validate the 'workspace/name[#fragment]' shape only — no Files-service call. + base = value.split("#", 1)[0] + parts = [part for part in base.split("/") if part] + if len(parts) != 2: + raise ValueError(f"fileset ref must be 'workspace/name' or 'workspace/name#path', got {value!r}") + return value + + +class FilesetSeedHandler: + """Resolves a :class:`FilesetSeed` to bytes by downloading the referenced file via the Files service.""" + + kind = "fileset" + + def parse(self, value: Mapping[str, Any]) -> BaseModel: + return FilesetSeed.model_validate(value) + + def resolve(self, seed: BaseModel) -> bytes: + assert isinstance(seed, FilesetSeed) + # Acquire the client at resolve time from the running job's ambient identity. + sdk = get_task_sdk(_EVALUATOR_SERVICE) + with tempfile.TemporaryDirectory() as staging: + try: + downloaded = download_dataset_sync(sdk, FilesetRef(root=seed.ref), staging) + except Exception as exc: # noqa: BLE001 - surface any resolution failure as a seed error + raise WorkspaceSeedError(f"fileset seed {seed.ref!r} could not be resolved: {exc}") from exc + files = [downloaded] if downloaded.is_file() else sorted(p for p in downloaded.rglob("*") if p.is_file()) + if len(files) != 1: + raise WorkspaceSeedError( + f"fileset seed {seed.ref!r} resolved to {len(files)} files; a seed maps one path to one " + "file — reference a single file with a '#path' fragment." + ) + return files[0].read_bytes() + + +register_seed_handler(FilesetSeedHandler()) diff --git a/plugins/nemo-evaluator/src/nemo_evaluator/jobs/agent_evaluate.py b/plugins/nemo-evaluator/src/nemo_evaluator/jobs/agent_evaluate.py index b2c715234d..f332e3327b 100644 --- a/plugins/nemo-evaluator/src/nemo_evaluator/jobs/agent_evaluate.py +++ b/plugins/nemo-evaluator/src/nemo_evaluator/jobs/agent_evaluate.py @@ -23,6 +23,7 @@ from typing import Any, ClassVar from urllib.parse import urlsplit +import nemo_evaluator.agent_seeds # noqa: F401 - registers the platform 'fileset' workspace-seed handler from nemo_evaluator.api.schemas import MetricInline from nemo_evaluator.jobs.agent_compiler import compile_agent_eval_job from nemo_evaluator.jobs.agent_spec import ( diff --git a/plugins/nemo-evaluator/tests/test_agent_seeds.py b/plugins/nemo-evaluator/tests/test_agent_seeds.py new file mode 100644 index 0000000000..1b0891f848 --- /dev/null +++ b/plugins/nemo-evaluator/tests/test_agent_seeds.py @@ -0,0 +1,55 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Tests for the plugin's ``fileset`` workspace-seed handler (registered on import).""" + +from __future__ import annotations + +from pathlib import Path + +import pytest +from nemo_evaluator import agent_seeds +from nemo_evaluator.agent_seeds import FilesetSeed +from nemo_evaluator_sdk.agent_eval.workspace_seeds import WorkspaceSeedError, parse_seed, seed_workspace +from pydantic import ValidationError + + +def test_fileset_ref_shape_is_validated() -> None: + assert FilesetSeed(ref="ws/name").ref == "ws/name" + assert FilesetSeed(ref="ws/name#data.csv").ref == "ws/name#data.csv" + for bad in ("bad", "a/b/c", "/leading"): + with pytest.raises(ValidationError): + FilesetSeed(ref=bad) + + +def test_importing_plugin_registers_fileset_kind() -> None: + # Importing nemo_evaluator.agent_seeds is enough to teach the SDK registry the 'fileset' kind; + # the SDK itself has no awareness of it. + assert isinstance(parse_seed({"kind": "fileset", "ref": "ws/name"}), FilesetSeed) + + +def test_fileset_seed_resolves_single_file(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + def fake_download(sdk: object, ref: object, dest: str) -> Path: + target = Path(dest) / "data.csv" + target.write_bytes(b"col\n1\n") + return target + + monkeypatch.setattr(agent_seeds, "get_task_sdk", lambda service: object()) + monkeypatch.setattr(agent_seeds, "download_dataset_sync", fake_download) + + seed_workspace(tmp_path, {"seed/data.csv": {"kind": "fileset", "ref": "ws/name#data.csv"}}) + assert (tmp_path / "seed" / "data.csv").read_bytes() == b"col\n1\n" + + +def test_fileset_seed_rejects_multi_file_ref(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + def fake_download(sdk: object, ref: object, dest: str) -> Path: + root = Path(dest) + (root / "a.txt").write_text("a") + (root / "b.txt").write_text("b") + return root + + monkeypatch.setattr(agent_seeds, "get_task_sdk", lambda service: object()) + monkeypatch.setattr(agent_seeds, "download_dataset_sync", fake_download) + + with pytest.raises(WorkspaceSeedError, match="resolved to 2 files"): + seed_workspace(tmp_path, {"x": {"kind": "fileset", "ref": "ws/name"}}) 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 e7832a9d35..1443f7146a 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 @@ -20,6 +20,7 @@ from nemo_platform.beta.evaluator.agent_eval.runtimes.docker_sandbox import DockerSandboxAgentRuntime from nemo_platform.beta.evaluator.agent_eval.tasks import AgentEvalRunConfig, AgentEvalTask from nemo_platform.beta.evaluator.agent_eval.trials import AgentEvalTrial, AgentEvalTrialStatus, AgentOutput +from nemo_platform.beta.evaluator.agent_eval.workspace_seeds import SEED_FILES_INPUT_KEY, seed_workspace from nemo_platform.beta.evaluator.values.evidence import CandidateEvidence, EvidenceDescriptor DEFAULT_CODEX_TIMEOUT_S = 600 @@ -48,11 +49,6 @@ class EffectiveCodexRuntime(StrEnum): #: (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.""" @@ -113,8 +109,10 @@ async def _run_task(self, index: int, task: AgentEvalTask, config: AgentEvalRunC 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) + # just this task rather than aborting the whole run. Offload to a worker thread: seeding is + # synchronous (a handler may do blocking I/O, e.g. the plugin's fileset download), and this + # runs on the event loop shared by every concurrent task, so a blocking seed would stall them all. + seeded_files = await asyncio.to_thread(seed_workspace, workspace_dir, task.inputs.get(SEED_FILES_INPUT_KEY)) process = await self._process_factory( *command, stdin=subprocess.PIPE, @@ -412,27 +410,6 @@ def default_codex_prompt(task: AgentEvalTask) -> str: 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( task: AgentEvalTask, evidence_dir: Path, diff --git a/sdk/python/nemo-platform/src/nemo_platform/beta/evaluator/agent_eval/workspace_seeds.py b/sdk/python/nemo-platform/src/nemo_platform/beta/evaluator/agent_eval/workspace_seeds.py new file mode 100644 index 0000000000..93cc4d1e84 --- /dev/null +++ b/sdk/python/nemo-platform/src/nemo_platform/beta/evaluator/agent_eval/workspace_seeds.py @@ -0,0 +1,171 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Workspace seed files for agent-eval tasks. + +A task can stage starter files into the agent's workspace before it runs, under +``inputs[SEED_FILES_INPUT_KEY]`` as a ``{relative_path: source}`` map. Each *source* is a +JSON-serializable value whose ``kind`` selects a registered :class:`SeedHandler` (a bare string is +sugar for inline text). + +The SDK ships handlers only for the kinds it can resolve with **no external dependency** — ``inline`` +and ``path``. Other kinds are contributed by consumers via :func:`register_seed_handler`; the SDK has +no knowledge of them (e.g. a platform ``fileset`` handler lives in the plugin and resolves against the +Files service at run time). An unregistered kind raises :class:`WorkspaceSeedError`. +""" + +from __future__ import annotations + +import base64 +from collections.abc import Mapping +from pathlib import Path +from typing import Any, Literal, Protocol, runtime_checkable + +from pydantic import BaseModel, ConfigDict, Field + +#: ``inputs`` key holding the ``{relative_path: seed}`` map of files to stage into the workspace. +SEED_FILES_INPUT_KEY = "files" + + +class WorkspaceSeedError(ValueError): + """A workspace seed could not be parsed, resolved, or written (bad value, unknown kind, ...). + + Subclasses ``ValueError`` so a runner's per-task error handling still catches it. + """ + + +class InlineSeed(BaseModel): + """File contents carried in the task itself. Portable to any runner.""" + + model_config = ConfigDict(extra="forbid") + + kind: Literal["inline"] = "inline" + content: str = Field(description="File contents; UTF-8 text, or base64-encoded bytes when encoding='base64'.") + encoding: Literal["text", "base64"] = "text" + + +class PathSeed(BaseModel): + """A file on the authoring host. Resolvable only where that path exists (local runs).""" + + model_config = ConfigDict(extra="forbid") + + kind: Literal["path"] = "path" + path: str = Field(description="Filesystem path on the authoring host (absolute or relative to the cwd).") + + +@runtime_checkable +class SeedHandler(Protocol): + """Parses + resolves one seed ``kind`` into the bytes to stage. + + Consumers implement this for kinds the SDK doesn't ship (e.g. a platform ``fileset`` handler) and + wire them in with :func:`register_seed_handler`. ``resolve`` runs at seeding time, so a handler + that needs external services (a client, credentials) acquires them there. + """ + + kind: str + + def parse(self, value: Mapping[str, Any]) -> BaseModel: + """Validate a raw seed mapping into this kind's typed model.""" + ... + + def resolve(self, seed: BaseModel) -> bytes: + """Resolve a parsed seed to the bytes to write into the workspace.""" + ... + + +_HANDLERS: dict[str, SeedHandler] = {} + + +def register_seed_handler(handler: SeedHandler) -> None: + """Register a :class:`SeedHandler` under its ``kind`` (replacing any handler already registered).""" + _HANDLERS[handler.kind] = handler + + +def _handler_for(kind: str) -> SeedHandler: + handler = _HANDLERS.get(kind) + if handler is None: + raise WorkspaceSeedError(f"no handler registered for seed kind {kind!r}") + return handler + + +class _InlineSeedHandler: + kind = "inline" + + def parse(self, value: Mapping[str, Any]) -> BaseModel: + return InlineSeed.model_validate(value) + + def resolve(self, seed: BaseModel) -> bytes: + assert isinstance(seed, InlineSeed) + if seed.encoding == "base64": + try: + return base64.b64decode(seed.content, validate=True) + except (ValueError, TypeError) as exc: + raise WorkspaceSeedError(f"inline seed is not valid base64: {exc}") from exc + return seed.content.encode("utf-8") + + +class _PathSeedHandler: + kind = "path" + + def parse(self, value: Mapping[str, Any]) -> BaseModel: + return PathSeed.model_validate(value) + + def resolve(self, seed: BaseModel) -> bytes: + assert isinstance(seed, PathSeed) + source = Path(seed.path).expanduser() + try: + return source.read_bytes() + except OSError as exc: + raise WorkspaceSeedError(f"path seed {seed.path!r} could not be read: {exc}") from exc + + +register_seed_handler(_InlineSeedHandler()) +register_seed_handler(_PathSeedHandler()) + + +def parse_seed(value: str | Mapping[str, Any]) -> BaseModel: + """Coerce a seed map value into its validated model. A bare string is inline UTF-8 text.""" + if isinstance(value, str): + return InlineSeed(content=value) + if not isinstance(value, Mapping): + raise WorkspaceSeedError(f"seed must be a string or mapping, got {type(value).__name__}") + kind = value.get("kind") + if not isinstance(kind, str): + raise WorkspaceSeedError("seed mapping is missing a string 'kind'") + handler = _handler_for(kind) + try: + return handler.parse(value) + except WorkspaceSeedError: + raise + except Exception as exc: # noqa: BLE001 - normalize a handler's validation error into our type + raise WorkspaceSeedError(f"invalid {kind!r} seed: {exc}") from exc + + +def _resolve_seed_bytes(seed: BaseModel) -> bytes: + """Resolve a parsed seed to bytes via its registered handler.""" + kind = getattr(seed, "kind", None) + if not isinstance(kind, str): + raise WorkspaceSeedError("parsed seed has no string 'kind'") + return _handler_for(kind).resolve(seed) + + +def seed_workspace(workspace_dir: str | Path, files: Mapping[str, Any] | None) -> list[str]: + """Write the ``files`` seed map into ``workspace_dir``; return the seeded relative paths. + + Each value is parsed into a seed model and resolved to bytes by its registered handler. Paths that + escape the workspace (absolute, or ``..`` traversal) are rejected so a task can only stage files + inside its own sandbox. ``None``/empty seeds nothing. + """ + if not isinstance(files, Mapping): + return [] + root = Path(workspace_dir).resolve() + written: list[str] = [] + for rel_path, value in files.items(): + target = (root / str(rel_path)).resolve() + if target != root and root not in target.parents: + raise WorkspaceSeedError(f"seed file path escapes the workspace: {rel_path!r}") + data = _resolve_seed_bytes(parse_seed(value)) + target.parent.mkdir(parents=True, exist_ok=True) + target.write_bytes(data) + written.append(str(rel_path)) + return written