From 6c82897aa33496a69c93e8c539a75defc8367d57 Mon Sep 17 00:00:00 2001 From: Atharva Sehgal Date: Thu, 9 Apr 2026 04:28:48 +0000 Subject: [PATCH 1/4] Make docker_build_env.sh editable and persist build_env_sh --- src/datasmith/agents/sandbox.py | 38 +++++++++++-- src/datasmith/agents/synthesizer.py | 9 ++- .../agents/templates/sandbox_verify.py | 55 ++++++++++++++++++- 3 files changed, 92 insertions(+), 10 deletions(-) diff --git a/src/datasmith/agents/sandbox.py b/src/datasmith/agents/sandbox.py index db96f2d4..f615256d 100644 --- a/src/datasmith/agents/sandbox.py +++ b/src/datasmith/agents/sandbox.py @@ -11,6 +11,7 @@ import hashlib import json +import os import shutil import subprocess import sys @@ -35,7 +36,6 @@ _IMMUTABLE_FILES = ( "Dockerfile.pr", "docker_build_base.sh", - "docker_build_env.sh", "docker_build_final.sh", "profile.sh", "run-tests.sh", @@ -54,14 +54,34 @@ def _compute_immutable_hashes(task_dir: Path) -> dict[str, str]: return hashes +def _read_env_payload_override(task_dir: Path) -> str | None: + """Read and validate ``env_payload_override.json`` from *task_dir*. + + Returns the raw JSON string if the file exists and contains a valid + JSON list, otherwise ``None``. + """ + override_file = task_dir / "env_payload_override.json" + if not override_file.exists(): + return None + try: + raw = override_file.read_text() + parsed = json.loads(raw) + if isinstance(parsed, list): + return raw + logger.warning("env_payload_override.json is not a JSON list, ignoring") + except (json.JSONDecodeError, Exception): + logger.warning("Failed to parse env_payload_override.json, ignoring") + return None + + @dataclass class SandboxConfig: """Configuration for the Codex sandbox runner.""" - timeout_s: int = 3600 + timeout_s: int = int(os.environ.get("SYNTHESIS_TIMEOUT_S", "14400")) """Total wall-clock timeout for the codex session (seconds).""" - codex_timeout_s: int = 3600 + codex_timeout_s: int = int(os.environ.get("SYNTHESIS_TIMEOUT_S", "14400")) """Timeout passed to subprocess.run for the codex process (seconds).""" @@ -78,6 +98,7 @@ class SandboxResult: agent_name: str = "" files_changed: list[str] = field(default_factory=list) resource_metrics: dict = field(default_factory=dict) + env_payload_override: str | None = None class SandboxRunner: @@ -296,7 +317,7 @@ def _extract_results(self, workspace: Path, codex_result: AgentResult, agent_nam success = success_file.exists() - # Read back only the two agent-editable scripts (the rest are templates) + # Read back the agent-editable scripts (the rest are templates) docker_context: DockerContext | None = None try: pkg_sh = ( @@ -305,10 +326,16 @@ def _extract_results(self, workspace: Path, codex_result: AgentResult, agent_nam run_sh = ( (task_dir / "docker_build_run.sh").read_text() if (task_dir / "docker_build_run.sh").exists() else "" ) - docker_context = DockerContext(build_pkg_sh=pkg_sh, build_run_sh=run_sh) + env_sh = ( + (task_dir / "docker_build_env.sh").read_text() if (task_dir / "docker_build_env.sh").exists() else "" + ) + docker_context = DockerContext(build_pkg_sh=pkg_sh, build_run_sh=run_sh, build_env_sh=env_sh) except Exception: logger.warning("Failed to read Docker context from workspace") + # Read env_payload override if the agent wrote one + env_payload_override = _read_env_payload_override(task_dir) + # Read failure.json if present failure_json: dict | None = None if failure_file.exists(): @@ -347,6 +374,7 @@ def _extract_results(self, workspace: Path, codex_result: AgentResult, agent_nam agent_name=agent_name, files_changed=codex_result.files_changed, resource_metrics=resource_metrics, + env_payload_override=env_payload_override if success else None, ) diff --git a/src/datasmith/agents/synthesizer.py b/src/datasmith/agents/synthesizer.py index 311e5dea..6c4f07e9 100644 --- a/src/datasmith/agents/synthesizer.py +++ b/src/datasmith/agents/synthesizer.py @@ -343,11 +343,13 @@ def _save_context( issue_number: int, ctx: DockerContext, resource_metrics: dict | None = None, + env_payload_override: str | None = None, ) -> None: """Persist the agent-edited scripts to the ``candidate_containers`` table. - Only ``build_pkg_sh`` and ``build_run_sh`` are saved — the other - fields come from templates and don't need to be persisted. + Saves ``build_pkg_sh``, ``build_run_sh``, and ``build_env_sh``. + When the agent also modified the env payload, ``env_payload_override`` + is persisted to the ``env_payload`` column. """ if not sha: return @@ -360,9 +362,12 @@ def _save_context( "issue_number": issue_number, "build_pkg_sh": ctx.build_pkg_sh, "build_run_sh": ctx.build_run_sh, + "build_env_sh": ctx.build_env_sh, } if resource_metrics: row["resource_metrics"] = resource_metrics + if env_payload_override: + row["env_payload"] = env_payload_override client.table("candidate_containers").upsert(row).execute() logger.info("Saved context for %s/%s@%s", owner, repo, sha[:12]) except Exception: diff --git a/src/datasmith/agents/templates/sandbox_verify.py b/src/datasmith/agents/templates/sandbox_verify.py index 9eaa8546..fdd07c09 100644 --- a/src/datasmith/agents/templates/sandbox_verify.py +++ b/src/datasmith/agents/templates/sandbox_verify.py @@ -34,7 +34,6 @@ _IMMUTABLE_FILES = ( "Dockerfile.pr", "docker_build_base.sh", - "docker_build_env.sh", "docker_build_final.sh", "profile.sh", "run-tests.sh", @@ -173,6 +172,31 @@ def __init__(self, message: str, stdout: str, stderr: str, rc: int) -> None: self.rc = rc +def _parse_failed_stage(build_log: str) -> str: + """Identify which Dockerfile stage (env/pkg/run) failed from Docker build output. + + Scans for BuildKit stage markers like ``[env 2/2]`` or legacy markers like + ``FROM env AS pkg``. Returns the name of the last stage seen before the + error, or ``"build"`` as a fallback. + """ + # BuildKit format: #N [stage_name step/total] ... + buildkit_re = re.compile(r"\[(\w+)\s+\d+/\d+\]") + # Legacy format: Step N/M : FROM x AS stage + legacy_re = re.compile(r"Step \d+/\d+\s*:\s*FROM\s+\S+\s+AS\s+(\w+)", re.IGNORECASE) + + last_stage = "" + for line in build_log.splitlines(): + m = buildkit_re.search(line) + if m: + last_stage = m.group(1) + continue + m = legacy_re.search(line) + if m: + last_stage = m.group(1) + + return last_stage if last_stage else "build" + + _MEM_UNITS = {"B": 1, "KIB": 1024, "MIB": 1024**2, "GIB": 1024**3, "TIB": 1024**4} _MEM_RE = re.compile(r"([\d.]+)\s*((?:[KMGT]i)?B)", re.IGNORECASE) @@ -380,13 +404,38 @@ def verify(task_dir: Path) -> bool: _write_failure(task_dir, "parse", stderr="Task.sha is None", metrics=metrics) return False + # Check for env_payload override written by the agent + override_file = task_dir / "env_payload_override.json" + if override_file.exists(): + try: + raw = override_file.read_text() + parsed = json.loads(raw) + if isinstance(parsed, list): + print(f"Using env_payload_override.json ({len(parsed)} packages)") + task = Task( + owner=task.owner, + repo=task.repo, + sha=task.sha, + commit_date=task.commit_date, + env_payload=json.dumps(parsed), + python_version=task.python_version, + tag=task.tag, + benchmarks=task.benchmarks, + repo_image=task.repo_image, + ) + else: + print("WARNING: env_payload_override.json is not a JSON list, ignoring") + except (json.JSONDecodeError, Exception) as e: + print(f"WARNING: Failed to parse env_payload_override.json: {e}") + # Build try: tag = build_image(docker, task_dir, task, target="run", metrics=metrics) print(f"Build succeeded: {tag}") except BuildError as e: - print(f"Build failed: {e}") - _write_failure(task_dir, "build", stdout=e.stdout, stderr=e.stderr, rc=e.rc, metrics=metrics) + stage = _parse_failed_stage(e.stdout) + print(f"Build failed at stage '{stage}': {e}") + _write_failure(task_dir, stage, stdout=e.stdout, stderr=e.stderr, rc=e.rc, metrics=metrics) return False except Exception as e: print(f"Build failed: {str(e)[:200]}") From 546a7368eec2958f0f875a2ab7bbd7d520cb4aa1 Mon Sep 17 00:00:00 2001 From: Atharva Sehgal Date: Thu, 9 Apr 2026 04:28:54 +0000 Subject: [PATCH 2/4] Update AGENTS.md.j2 for env editability, payload override, and stage detection --- src/datasmith/agents/templates/AGENTS.md.j2 | 29 ++++++++++++++++----- 1 file changed, 22 insertions(+), 7 deletions(-) diff --git a/src/datasmith/agents/templates/AGENTS.md.j2 b/src/datasmith/agents/templates/AGENTS.md.j2 index 78c7188a..bb45d455 100644 --- a/src/datasmith/agents/templates/AGENTS.md.j2 +++ b/src/datasmith/agents/templates/AGENTS.md.j2 @@ -35,24 +35,31 @@ Follow this iterative cycle: 1. Run `python3 sandbox_verify.py` (A good timeout is 60 minutes, but adjust as needed) 2. If it fails, read `task/failure.json` for error details -3. Edit `task/docker_build_pkg.sh` and/or `task/docker_build_run.sh` +3. Edit `task/docker_build_pkg.sh`, `task/docker_build_run.sh`, and/or `task/docker_build_env.sh` 4. Re-run `python3 sandbox_verify.py` 5. Repeat until `task/verification_success.json` is created -**IMPORTANT**: Only modify `task/docker_build_pkg.sh` and `task/docker_build_run.sh`. +**IMPORTANT**: Only modify `task/docker_build_pkg.sh`, `task/docker_build_run.sh`, and `task/docker_build_env.sh`. Do **not** modify the Dockerfile, `task.txt`, or any other scripts. +You may also write `task/env_payload_override.json` — a JSON list of pinned package strings +(e.g. `["numpy==1.21.0", "scipy==1.7.0"]`) — to replace the default env payload when the +original package versions fail to build. + ## File Constraints ### Files you CAN edit - `task/docker_build_pkg.sh` — Primary build script. Install the package, add build deps. - `task/docker_build_run.sh` — Runtime setup. Add test/benchmark deps, repo-specific config. +- `task/docker_build_env.sh` — Python environment setup. Add build prerequisites needed before env payload installation (e.g. Cython, distutils). + +### Files you CAN create +- `task/env_payload_override.json` — JSON list of pinned package strings to replace the default env payload. Use this when the original package versions are fundamentally broken (e.g. no wheels available, incompatible with the Python version). ### Files you MUST NOT edit - `task/Dockerfile.pr` — Multi-stage build definition (base and repo stages are pre-built) - `task/task.txt` — Task configuration (owner, repo, sha, deps) - `task/docker_build_base.sh` — Base system setup (template) -- `task/docker_build_env.sh` — Python environment creation (template) - `task/docker_build_final.sh` — Final image setup (template) - `task/profile.sh` — ASV benchmark runner - `task/run-tests.sh` — Test runner @@ -64,7 +71,7 @@ When verification fails, `task/failure.json` contains: ```json { - "stage": "build|tests", + "stage": "env|pkg|run|tests|build", "return_code": 1, "stderr": "Full error output from the failed stage", "stdout": "Full standard output from the failed stage", @@ -72,8 +79,11 @@ When verification fails, `task/failure.json` contains: } ``` -- **build** failures: Usually missing dependencies or compilation errors → fix in `docker_build_pkg.sh` +- **env** failures: Env payload packages fail to install (missing build tools, broken wheels) → fix in `docker_build_env.sh` or write `env_payload_override.json` +- **pkg** failures: Package installation errors or missing build deps → fix in `docker_build_pkg.sh` +- **run** failures: Runtime setup issues → fix in `docker_build_run.sh` - **tests** failures: Missing test deps, import errors, or benchmark setup issues → fix in `docker_build_run.sh` or `docker_build_pkg.sh` +- **build** (generic): Stage could not be determined from logs — check stdout/stderr for details ## Common Fixes for docker_build_pkg.sh @@ -132,13 +142,15 @@ These are set by earlier build stages and available in both scripts: ## Docker Build Stages The Dockerfile.pr has 4 stages: env → pkg → run → final (base and repo images are pre-built). -Only the `pkg` and `run` stages use your editable scripts. +The `env`, `pkg`, and `run` stages use your editable scripts. +- **env** stage: Runs `docker_build_env.sh` to set up the Python environment and install pinned deps - **pkg** stage: Runs `docker_build_pkg.sh` to install the package - **run** stage: Runs `docker_build_run.sh` to prepare for tests/benchmarks +Changes to `docker_build_env.sh` rebuild from the env stage onward (slowest). Changes to `docker_build_pkg.sh` rebuild from the pkg stage onward (fast iteration). -Changes to `docker_build_run.sh` rebuild only the run stage. +Changes to `docker_build_run.sh` rebuild only the run stage (fastest). ## Tips @@ -148,3 +160,6 @@ Changes to `docker_build_run.sh` rebuild only the run stage. - If the build times out, look for ways to simplify or skip expensive steps - Docker layer caching means earlier stages are cached — only your changed stage rebuilds - A lot of bookkeeping code is used in downstream stages; avoid removing code unless you understand its purpose. +- **Fixing env payload packages**: If packages in the env payload fail to build (e.g. no wheels, missing build deps), you have two options: + 1. **Edit `docker_build_env.sh`** to install build prerequisites (Cython, compiler toolchains, etc.) *before* the payload is installed. This is preferred when the package versions are correct but just need build tools. + 2. **Write `task/env_payload_override.json`** with a corrected JSON list of pinned packages (e.g. `["numpy==1.21.0", "h5py==3.1.0"]`). `sandbox_verify.py` will use this instead of the original payload. Use this when a package version is fundamentally broken (no source dist, incompatible with the Python version, yanked). Keep changes minimal — stay as close as possible to the originals. From 242e07396f40689a7b59047da5370f3b5fc81d6a Mon Sep 17 00:00:00 2001 From: Atharva Sehgal Date: Thu, 9 Apr 2026 04:29:03 +0000 Subject: [PATCH 3/4] Add tests for env editability, payload override, and build stage detection --- tests/agents/test_sandbox.py | 369 ++++++++++++++++++++++++++++++- tests/agents/test_synthesizer.py | 69 ++++++ 2 files changed, 437 insertions(+), 1 deletion(-) diff --git a/tests/agents/test_sandbox.py b/tests/agents/test_sandbox.py index dc6f0915..b5f5313a 100644 --- a/tests/agents/test_sandbox.py +++ b/tests/agents/test_sandbox.py @@ -39,6 +39,12 @@ def test_defaults(self) -> None: assert r.duration_s == 0.0 assert r.agent_output == "" assert r.resource_metrics == {} + assert r.env_payload_override is None + + def test_with_env_payload_override(self) -> None: + override = '["numpy==1.21.0"]' + r = SandboxResult(success=True, env_payload_override=override) + assert r.env_payload_override == override class TestGenerateTaskTxt: @@ -155,7 +161,6 @@ def test_immutable_hashes_written(self, tmp_path: Path) -> None: for fname in ( "Dockerfile.pr", "docker_build_base.sh", - "docker_build_env.sh", "docker_build_final.sh", "profile.sh", "run-tests.sh", @@ -169,6 +174,7 @@ def test_immutable_hashes_written(self, tmp_path: Path) -> None: # Editable files should NOT be in hashes assert "docker_build_pkg.sh" not in hashes assert "docker_build_run.sh" not in hashes + assert "docker_build_env.sh" not in hashes def test_all_files_from_templates(self, tmp_path: Path) -> None: runner = SandboxRunner() @@ -460,3 +466,364 @@ def fake_launch(workspace: Path) -> tuple[str, MagicMock]: assert result.success is False assert result.failure_json is not None assert result.failure_json["stage"] == "build" + + +# ── docker_build_env.sh editability ──────────────────────────────── + + +class TestEnvShEditability: + """docker_build_env.sh should be editable by the agent, not immutable.""" + + def test_env_sh_not_in_immutable_files(self) -> None: + """docker_build_env.sh must NOT appear in _IMMUTABLE_FILES.""" + from datasmith.agents.sandbox import _IMMUTABLE_FILES + + assert "docker_build_env.sh" not in _IMMUTABLE_FILES + + def test_env_sh_not_in_verify_immutable_files(self) -> None: + """docker_build_env.sh must NOT appear in sandbox_verify.py's _IMMUTABLE_FILES.""" + verify_src = ( + Path(__file__).parents[1] / ".." / "src" / "datasmith" / "agents" / "templates" / "sandbox_verify.py" + ) + content = verify_src.resolve().read_text() + # Parse the _IMMUTABLE_FILES tuple from the source + import ast + + tree = ast.parse(content) + for node in ast.walk(tree): + if isinstance(node, ast.Assign): + for target in node.targets: + if isinstance(target, ast.Name) and target.id == "_IMMUTABLE_FILES": + immutable = ast.literal_eval(node.value) + assert "docker_build_env.sh" not in immutable, ( + "docker_build_env.sh should not be in sandbox_verify.py _IMMUTABLE_FILES" + ) + return + raise AssertionError("Could not find _IMMUTABLE_FILES in sandbox_verify.py") + + def test_env_sh_not_in_immutable_hashes(self, tmp_path: Path) -> None: + """_prepare_workspace should NOT hash docker_build_env.sh (it's editable).""" + runner = SandboxRunner() + runner._prepare_workspace( + workspace=tmp_path, + owner="o", + repo="r", + sha="s", + repo_image="formulacode/o-r:latest", + env_payload="", + python_version="3.10", + pr_context="", + ) + + hashes = json.loads((tmp_path / ".immutable_hashes.json").read_text()) + assert "docker_build_env.sh" not in hashes + # Editable files should not be hashed + assert "docker_build_pkg.sh" not in hashes + assert "docker_build_run.sh" not in hashes + + def test_integrity_allows_env_sh_edits(self, tmp_path: Path) -> None: + """Editing docker_build_env.sh must NOT trigger an integrity violation.""" + task_dir = tmp_path / "task" + task_dir.mkdir() + + # Write immutable files and the editable env script + (task_dir / "Dockerfile.pr").write_text("FROM base") + (task_dir / "docker_build_env.sh").write_text("#!/bin/bash\noriginal env") + hashes = _compute_immutable_hashes(task_dir) + (tmp_path / ".immutable_hashes.json").write_text(json.dumps(hashes)) + + # Agent edits env script and succeeds + (task_dir / "docker_build_env.sh").write_text("#!/bin/bash\npip install Cython\noriginal env") + (task_dir / "docker_build_pkg.sh").write_text("#!/bin/bash\necho pkg") + (task_dir / "docker_build_run.sh").write_text("#!/bin/bash\necho run") + (task_dir / "verification_success.json").write_text('{"local_image": "test:latest"}') + + runner = SandboxRunner() + codex_result = MagicMock() + codex_result.output = "done" + codex_result.raw_output = "" + codex_result.files_changed = ["task/docker_build_env.sh"] + + result = runner._extract_results(tmp_path, codex_result) + + assert result.success is True + assert result.docker_context is not None + + def test_extract_results_reads_build_env_sh(self, tmp_path: Path) -> None: + """_extract_results must read back docker_build_env.sh into DockerContext.build_env_sh.""" + task_dir = tmp_path / "task" + task_dir.mkdir() + + env_content = "#!/bin/bash\npip install Cython\n# modified env" + (task_dir / "docker_build_env.sh").write_text(env_content) + (task_dir / "docker_build_pkg.sh").write_text("#!/bin/bash\necho pkg") + (task_dir / "docker_build_run.sh").write_text("#!/bin/bash\necho run") + (task_dir / "verification_success.json").write_text('{"local_image": "test:latest"}') + + runner = SandboxRunner() + codex_result = MagicMock() + codex_result.output = "done" + codex_result.raw_output = "" + codex_result.files_changed = [] + + result = runner._extract_results(tmp_path, codex_result) + + assert result.success is True + assert result.docker_context is not None + assert result.docker_context.build_env_sh == env_content + + def test_extract_results_env_sh_empty_when_missing(self, tmp_path: Path) -> None: + """If docker_build_env.sh doesn't exist, build_env_sh should be empty.""" + task_dir = tmp_path / "task" + task_dir.mkdir() + + (task_dir / "docker_build_pkg.sh").write_text("#!/bin/bash\necho pkg") + (task_dir / "docker_build_run.sh").write_text("#!/bin/bash\necho run") + (task_dir / "verification_success.json").write_text('{"local_image": "test:latest"}') + + runner = SandboxRunner() + codex_result = MagicMock() + codex_result.output = "done" + codex_result.raw_output = "" + codex_result.files_changed = [] + + result = runner._extract_results(tmp_path, codex_result) + + assert result.success is True + assert result.docker_context is not None + assert result.docker_context.build_env_sh == "" + + def test_failed_result_does_not_return_context(self, tmp_path: Path) -> None: + """On failure, docker_context should still be None even with env edits.""" + task_dir = tmp_path / "task" + task_dir.mkdir() + + (task_dir / "docker_build_env.sh").write_text("#!/bin/bash\nmodified") + (task_dir / "docker_build_pkg.sh").write_text("#!/bin/bash\necho pkg") + failure = {"stage": "build", "return_code": 1, "error_message": "err"} + (task_dir / "failure.json").write_text(json.dumps(failure)) + + runner = SandboxRunner() + codex_result = MagicMock() + codex_result.output = "" + codex_result.error = "" + codex_result.raw_output = "" + codex_result.files_changed = [] + + result = runner._extract_results(tmp_path, codex_result) + + assert result.success is False + assert result.docker_context is None + + +# ── env_payload override ─────────────────────────────────────────── + + +class TestEnvPayloadOverride: + """Agent can write env_payload_override.json to modify the package list.""" + + def test_extract_results_reads_payload_override(self, tmp_path: Path) -> None: + """_extract_results reads env_payload_override.json when present.""" + task_dir = tmp_path / "task" + task_dir.mkdir() + + override = '["numpy==1.21.0", "scipy==1.7.0"]' + (task_dir / "env_payload_override.json").write_text(override) + (task_dir / "docker_build_pkg.sh").write_text("#!/bin/bash") + (task_dir / "docker_build_run.sh").write_text("#!/bin/bash") + (task_dir / "verification_success.json").write_text('{"local_image": "test:latest"}') + + runner = SandboxRunner() + codex_result = MagicMock() + codex_result.output = "done" + codex_result.raw_output = "" + codex_result.files_changed = [] + + result = runner._extract_results(tmp_path, codex_result) + + assert result.success is True + assert result.env_payload_override == override + + def test_extract_results_no_override_returns_none(self, tmp_path: Path) -> None: + """When no override file exists, env_payload_override should be None.""" + task_dir = tmp_path / "task" + task_dir.mkdir() + + (task_dir / "docker_build_pkg.sh").write_text("#!/bin/bash") + (task_dir / "docker_build_run.sh").write_text("#!/bin/bash") + (task_dir / "verification_success.json").write_text('{"local_image": "test:latest"}') + + runner = SandboxRunner() + codex_result = MagicMock() + codex_result.output = "done" + codex_result.raw_output = "" + codex_result.files_changed = [] + + result = runner._extract_results(tmp_path, codex_result) + + assert result.success is True + assert result.env_payload_override is None + + def test_override_not_returned_on_failure(self, tmp_path: Path) -> None: + """On failure, env_payload_override should be None.""" + task_dir = tmp_path / "task" + task_dir.mkdir() + + (task_dir / "env_payload_override.json").write_text('["numpy==1.21.0"]') + (task_dir / "docker_build_pkg.sh").write_text("#!/bin/bash") + failure = {"stage": "build", "return_code": 1, "error_message": "err"} + (task_dir / "failure.json").write_text(json.dumps(failure)) + + runner = SandboxRunner() + codex_result = MagicMock() + codex_result.output = "" + codex_result.error = "" + codex_result.raw_output = "" + codex_result.files_changed = [] + + result = runner._extract_results(tmp_path, codex_result) + + assert result.success is False + assert result.env_payload_override is None + + def test_override_must_be_valid_json_list(self, tmp_path: Path) -> None: + """Malformed env_payload_override.json should be treated as no override.""" + task_dir = tmp_path / "task" + task_dir.mkdir() + + (task_dir / "env_payload_override.json").write_text("not valid json") + (task_dir / "docker_build_pkg.sh").write_text("#!/bin/bash") + (task_dir / "docker_build_run.sh").write_text("#!/bin/bash") + (task_dir / "verification_success.json").write_text('{"local_image": "test:latest"}') + + runner = SandboxRunner() + codex_result = MagicMock() + codex_result.output = "done" + codex_result.raw_output = "" + codex_result.files_changed = [] + + result = runner._extract_results(tmp_path, codex_result) + + assert result.success is True + assert result.env_payload_override is None + + def test_override_rejects_non_list(self, tmp_path: Path) -> None: + """env_payload_override.json must be a JSON list, not an object or string.""" + task_dir = tmp_path / "task" + task_dir.mkdir() + + (task_dir / "env_payload_override.json").write_text('{"numpy": "1.21.0"}') + (task_dir / "docker_build_pkg.sh").write_text("#!/bin/bash") + (task_dir / "docker_build_run.sh").write_text("#!/bin/bash") + (task_dir / "verification_success.json").write_text('{"local_image": "test:latest"}') + + runner = SandboxRunner() + codex_result = MagicMock() + codex_result.output = "done" + codex_result.raw_output = "" + codex_result.files_changed = [] + + result = runner._extract_results(tmp_path, codex_result) + + assert result.success is True + assert result.env_payload_override is None + + +# ── Build stage detection ────────────────────────────────────────── + + +class TestBuildStageDetection: + """sandbox_verify.py should identify which Dockerfile stage (env/pkg/run) failed.""" + + def test_parse_failed_stage_env(self) -> None: + """Detect failure in the 'env' stage from Docker build logs.""" + from datasmith.agents.templates.sandbox_verify import _parse_failed_stage + + # Typical buildkit log: the last stage marker before the error + logs = ( + "#5 [env 1/2] RUN git checkout abc123\n" + "#5 DONE 0.5s\n" + "#6 [env 2/2] RUN chmod +x ... && /workspace/repo/docker_build_env.sh\n" + "#6 ERROR: process returned non-zero exit code: 1\n" + ) + assert _parse_failed_stage(logs) == "env" + + def test_parse_failed_stage_pkg(self) -> None: + """Detect failure in the 'pkg' stage from Docker build logs.""" + from datasmith.agents.templates.sandbox_verify import _parse_failed_stage + + logs = ( + "#5 [env 2/2] RUN chmod +x ... && /workspace/repo/docker_build_env.sh\n" + "#5 DONE 30.0s\n" + "#6 [pkg 1/2] COPY docker_build_pkg.sh /workspace/repo/docker_build_pkg.sh\n" + "#6 DONE 0.1s\n" + "#7 [pkg 2/2] RUN chmod +x ... && /workspace/repo/docker_build_pkg.sh\n" + "#7 ERROR: process returned non-zero exit code: 1\n" + ) + assert _parse_failed_stage(logs) == "pkg" + + def test_parse_failed_stage_run(self) -> None: + """Detect failure in the 'run' stage from Docker build logs.""" + from datasmith.agents.templates.sandbox_verify import _parse_failed_stage + + logs = ( + "#8 [pkg 2/2] RUN chmod +x ... && /workspace/repo/docker_build_pkg.sh\n" + "#8 DONE 15.0s\n" + "#9 [run 1/2] COPY docker_build_run.sh /docker_build_run.sh\n" + "#9 DONE 0.1s\n" + "#10 [run 2/2] RUN chmod +x /docker_build_run.sh && /docker_build_run.sh\n" + "#10 ERROR: process returned non-zero exit code: 2\n" + ) + assert _parse_failed_stage(logs) == "run" + + def test_parse_failed_stage_no_markers(self) -> None: + """When no stage markers are found, return 'build' as fallback.""" + from datasmith.agents.templates.sandbox_verify import _parse_failed_stage + + logs = "Some random error output\nwithout any stage markers\n" + assert _parse_failed_stage(logs) == "build" + + def test_parse_failed_stage_empty(self) -> None: + """Empty log returns 'build' fallback.""" + from datasmith.agents.templates.sandbox_verify import _parse_failed_stage + + assert _parse_failed_stage("") == "build" + + def test_parse_failed_stage_legacy_docker(self) -> None: + """Legacy (non-buildkit) Docker output with 'Step N/M : FROM x AS stage'.""" + from datasmith.agents.templates.sandbox_verify import _parse_failed_stage + + logs = ( + "Step 3/10 : FROM env AS pkg\n" + "Step 4/10 : COPY docker_build_pkg.sh /workspace/repo/docker_build_pkg.sh\n" + "Step 5/10 : RUN chmod +x ... && /workspace/repo/docker_build_pkg.sh\n" + "The command '/bin/sh -c chmod +x ...' returned a non-zero code: 1\n" + ) + assert _parse_failed_stage(logs) == "pkg" + + def test_write_failure_uses_detected_stage(self, tmp_path: Path) -> None: + """_write_failure should write the specific stage name, not generic 'build'.""" + from datasmith.agents.templates.sandbox_verify import _write_failure + + _write_failure(tmp_path, "env", stdout="log output", stderr="error", rc=1) + + failure = json.loads((tmp_path / "failure.json").read_text()) + assert failure["stage"] == "env" + assert failure["return_code"] == 1 + + def test_build_error_carries_stage_through_verify(self, tmp_path: Path) -> None: + """When build_image raises BuildError, verify() should detect and write the correct stage.""" + from datasmith.agents.templates.sandbox_verify import _parse_failed_stage, _write_failure + + # Simulate: BuildError with env-stage logs + build_stdout = ( + "#5 [env 2/2] RUN chmod +x ... && /workspace/repo/docker_build_env.sh\n" + "#5 ERROR: process returned non-zero exit code: 1\n" + ) + stage = _parse_failed_stage(build_stdout) + assert stage == "env" + + # Write failure with detected stage + _write_failure(tmp_path, stage, stdout=build_stdout, stderr="", rc=1) + failure = json.loads((tmp_path / "failure.json").read_text()) + assert failure["stage"] == "env" diff --git a/tests/agents/test_synthesizer.py b/tests/agents/test_synthesizer.py index 62d4ea4a..399093f0 100644 --- a/tests/agents/test_synthesizer.py +++ b/tests/agents/test_synthesizer.py @@ -399,6 +399,75 @@ def test_log_attempt_includes_resource_metrics(self, mock_get_client: MagicMock) assert row["resource_metrics"] == metrics +class TestSaveContextEnvAndPayload: + """_save_context should persist build_env_sh and env_payload_override.""" + + @patch("datasmith.agents.synthesizer.get_client") + def test_save_context_persists_build_env_sh(self, mock_get_client: MagicMock) -> None: + """When DockerContext has build_env_sh, it should be included in the upsert row.""" + mock_client = MagicMock() + mock_get_client.return_value = mock_client + + ctx = DockerContext( + build_pkg_sh="#!/bin/bash\npkg", + build_run_sh="#!/bin/bash\nrun", + build_env_sh="#!/bin/bash\npip install Cython\n# modified env", + ) + + synth = Synthesizer() + synth._save_context("owner", "repo", "abc123", 42, ctx) + + upsert_call = mock_client.table.return_value.upsert + upsert_call.assert_called_once() + row = upsert_call.call_args[0][0] + assert row["build_env_sh"] == "#!/bin/bash\npip install Cython\n# modified env" + assert row["build_pkg_sh"] == "#!/bin/bash\npkg" + assert row["build_run_sh"] == "#!/bin/bash\nrun" + + @patch("datasmith.agents.synthesizer.get_client") + def test_save_context_empty_env_sh(self, mock_get_client: MagicMock) -> None: + """When build_env_sh is empty, it should still be included as empty string.""" + mock_client = MagicMock() + mock_get_client.return_value = mock_client + + ctx = DockerContext(build_pkg_sh="pkg", build_run_sh="run") + + synth = Synthesizer() + synth._save_context("owner", "repo", "abc123", 42, ctx) + + row = mock_client.table.return_value.upsert.call_args[0][0] + assert row["build_env_sh"] == "" + + @patch("datasmith.agents.synthesizer.get_client") + def test_save_context_persists_env_payload_override(self, mock_get_client: MagicMock) -> None: + """When env_payload_override is provided, it should be persisted.""" + mock_client = MagicMock() + mock_get_client.return_value = mock_client + + ctx = DockerContext(build_pkg_sh="pkg", build_run_sh="run") + override = '["numpy==1.21.0", "scipy==1.7.0"]' + + synth = Synthesizer() + synth._save_context("owner", "repo", "abc123", 42, ctx, env_payload_override=override) + + row = mock_client.table.return_value.upsert.call_args[0][0] + assert row["env_payload"] == override + + @patch("datasmith.agents.synthesizer.get_client") + def test_save_context_no_payload_override(self, mock_get_client: MagicMock) -> None: + """When no env_payload_override is given, env_payload should not be in the row.""" + mock_client = MagicMock() + mock_get_client.return_value = mock_client + + ctx = DockerContext(build_pkg_sh="pkg", build_run_sh="run") + + synth = Synthesizer() + synth._save_context("owner", "repo", "abc123", 42, ctx) + + row = mock_client.table.return_value.upsert.call_args[0][0] + assert "env_payload" not in row + + class TestFormatPriorAttempts: def test_formats_failed_attempts(self) -> None: ctx = DockerContext(build_pkg_sh="#!/bin/bash\npkg", build_run_sh="#!/bin/bash\nrun") From 22dcfdfe81ecf5a782fa50b9dfda70a85035efd0 Mon Sep 17 00:00:00 2001 From: Atharva Sehgal Date: Thu, 9 Apr 2026 04:36:17 +0000 Subject: [PATCH 4/4] Fix test_defaults to match updated SandboxConfig timeout (14400s) --- tests/agents/test_sandbox.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/agents/test_sandbox.py b/tests/agents/test_sandbox.py index b5f5313a..307c452c 100644 --- a/tests/agents/test_sandbox.py +++ b/tests/agents/test_sandbox.py @@ -21,8 +21,8 @@ class TestSandboxConfig: def test_defaults(self) -> None: cfg = SandboxConfig() - assert cfg.timeout_s == 3600 - assert cfg.codex_timeout_s == 3600 + assert cfg.timeout_s == 14400 + assert cfg.codex_timeout_s == 14400 def test_custom(self) -> None: cfg = SandboxConfig(timeout_s=600, codex_timeout_s=300)