From 8186d8fbec18b578e9c75a577f2aac32112caa2c Mon Sep 17 00:00:00 2001 From: adongwanai Date: Thu, 8 Oct 2026 14:47:37 +0800 Subject: [PATCH] fix(envs): confine benchmark task output paths --- skillopt/envs/docvqa/rollout.py | 5 +- .../envs/livemathematicianbench/rollout.py | 4 +- skillopt/envs/officeqa/rollout.py | 6 +- skillopt/envs/searchqa/rollout.py | 3 +- skillopt/envs/spreadsheetbench/rollout.py | 53 +------- skillopt/envs/task_output.py | 54 ++++++++ skillopt/optimizer/slow_update.py | 16 ++- tests/test_benchmark_task_output_paths.py | 128 ++++++++++++++++++ 8 files changed, 214 insertions(+), 55 deletions(-) create mode 100644 skillopt/envs/task_output.py create mode 100644 tests/test_benchmark_task_output_paths.py diff --git a/skillopt/envs/docvqa/rollout.py b/skillopt/envs/docvqa/rollout.py index 6396163fd..e41a73c45 100644 --- a/skillopt/envs/docvqa/rollout.py +++ b/skillopt/envs/docvqa/rollout.py @@ -6,6 +6,7 @@ from concurrent.futures import FIRST_COMPLETED, ThreadPoolExecutor, wait from skillopt.envs.docvqa.evaluator import evaluate +from skillopt.envs.task_output import confined_task_output_dir from skillopt.model import chat_target_messages, get_target_backend, is_target_exec_backend from skillopt.model.codex_harness import prepare_workspace, render_skill_md, run_target_exec from skillopt.prompts import load_prompt @@ -139,6 +140,7 @@ def process_one( diagnostic_instruction: str = "", ) -> dict: item_id = str(item["id"]) + pred_dir = confined_task_output_dir(out_root, item_id) result = { "id": item_id, "question": item["question"], @@ -170,7 +172,7 @@ def process_one( ] for turn in range(max_turns): response, _raw, system_prompt, user_text = _run_codex_once( - pred_dir=os.path.join(out_root, "predictions", item_id), + pred_dir=pred_dir, item=item, skill_content=skill_content, model=_llm.TARGET_DEPLOYMENT, @@ -229,7 +231,6 @@ def process_one( result["agent_ok"] = True result["n_turns"] = len(conversation) - 1 - pred_dir = os.path.join(out_root, "predictions", item_id) os.makedirs(pred_dir, exist_ok=True) with open(os.path.join(pred_dir, "target_system_prompt.txt"), "w", encoding="utf-8") as f: f.write(system_prompt) diff --git a/skillopt/envs/livemathematicianbench/rollout.py b/skillopt/envs/livemathematicianbench/rollout.py index 01de404f8..50a9b15a9 100644 --- a/skillopt/envs/livemathematicianbench/rollout.py +++ b/skillopt/envs/livemathematicianbench/rollout.py @@ -7,10 +7,12 @@ from concurrent.futures import FIRST_COMPLETED, ThreadPoolExecutor, wait from skillopt.envs.livemathematicianbench.evaluator import evaluate +from skillopt.envs.task_output import confined_task_output_dir from skillopt.model import chat_target, get_target_backend, is_target_exec_backend from skillopt.model.codex_harness import prepare_workspace, render_skill_md, run_target_exec from skillopt.prompts import load_prompt + def _build_system(skill_content: str) -> str: if skill_content.strip(): skill_section = f"## Skill\n{skill_content.strip()}\n\n" @@ -123,6 +125,7 @@ def process_one( max_completion_tokens: int = 16384, ) -> dict: item_id = str(item["id"]) + pred_dir = confined_task_output_dir(out_root, item_id, map_unsafe=True) result = { "id": item_id, "question": item["question"], @@ -141,7 +144,6 @@ def process_one( } try: - pred_dir = os.path.join(out_root, "predictions", item_id) os.makedirs(pred_dir, exist_ok=True) llm_timeout = int(exec_timeout) if exec_timeout and int(exec_timeout) > 0 else None diff --git a/skillopt/envs/officeqa/rollout.py b/skillopt/envs/officeqa/rollout.py index 01afe8b51..7d4c8e045 100644 --- a/skillopt/envs/officeqa/rollout.py +++ b/skillopt/envs/officeqa/rollout.py @@ -1,8 +1,10 @@ from __future__ import annotations + import json import os import re from concurrent.futures import ThreadPoolExecutor, as_completed + from skillopt.envs.officeqa.evaluator import evaluate from skillopt.envs.officeqa.tool_runtime import ( build_oracle_parsed_pages_context, @@ -11,9 +13,11 @@ resolve_docs_roots, run_tool, ) +from skillopt.envs.task_output import confined_task_output_dir from skillopt.model import chat_target_messages, get_target_backend, is_target_exec_backend from skillopt.model.codex_harness import prepare_workspace, render_skill_md, run_target_exec from skillopt.prompts import load_prompt + _TOOL_SCHEMAS = [ { "type": "function", @@ -527,7 +531,7 @@ def process_one( diagnostic_instruction: str = "", ) -> dict: item_id = str(item["id"]) - pred_dir = os.path.join(out_root, "predictions", item_id) + pred_dir = confined_task_output_dir(out_root, item_id) os.makedirs(pred_dir, exist_ok=True) normalized_search_mode = _normalize_search_mode(search_mode) docs_roots: list[str] = [] diff --git a/skillopt/envs/searchqa/rollout.py b/skillopt/envs/searchqa/rollout.py index 83165b5e5..7f96393d0 100644 --- a/skillopt/envs/searchqa/rollout.py +++ b/skillopt/envs/searchqa/rollout.py @@ -17,6 +17,7 @@ from concurrent.futures import FIRST_COMPLETED, ThreadPoolExecutor, wait from skillopt.envs.searchqa.evaluator import evaluate +from skillopt.envs.task_output import confined_task_output_dir from skillopt.model import chat_target, is_target_exec_backend from skillopt.model.codex_harness import prepare_workspace, render_skill_md, run_target_exec from skillopt.prompts import load_prompt @@ -180,6 +181,7 @@ def process_one( Result with ``hard`` (EM as int), ``soft`` (F1), etc. """ item_id = str(item["id"]) + pred_dir = confined_task_output_dir(out_root, item_id) question = item["question"] context = item.get("context", "") gold_answers = item.get("answers", []) @@ -201,7 +203,6 @@ def process_one( } try: - pred_dir = os.path.join(out_root, "predictions", item_id) os.makedirs(pred_dir, exist_ok=True) if is_target_exec_backend(): diff --git a/skillopt/envs/spreadsheetbench/rollout.py b/skillopt/envs/spreadsheetbench/rollout.py index a1709dd24..bfdc85813 100644 --- a/skillopt/envs/spreadsheetbench/rollout.py +++ b/skillopt/envs/spreadsheetbench/rollout.py @@ -12,7 +12,6 @@ import glob as _glob import json import os -import re import shutil import tempfile import time @@ -31,6 +30,10 @@ evaluate, _generate_cell_names, _compare_cell_value, ) from skillopt.envs.spreadsheetbench.executor import run_generated_code +from skillopt.envs.task_output import ( + confined_task_output_dir as _confined_task_out_dir, + is_safe_task_id as _is_safe_task_id, +) # ── Data loading ───────────────────────────────────────────────────────────── @@ -208,54 +211,6 @@ def _auto_verify_output( return report -# ── Task identifier → confined output directory ────────────────────────────── - -# A task id becomes ONE path segment under ``/predictions``. Keep the -# accepted alphabet deliberately narrow: released dataset ids are alphanumeric -# with ``-``/``_`` (e.g. "1-1", "80-42"), and anything that is not a plausible -# single segment must not be trusted with a path. -_SAFE_TASK_ID = re.compile(r"\A[A-Za-z0-9][A-Za-z0-9._-]{0,127}\Z") - - -def _is_safe_task_id(task_id: str) -> bool: - """True when ``task_id`` is safe to use as one path segment. - - The alphabet excludes path separators and a leading dot, so ``..`` cannot - match the pattern at all. The substring check on top is belt-and-braces: - it also rejects traversal-shaped values such as ``a..b``, which no released - dataset id uses. - """ - tid = str(task_id) - return bool(_SAFE_TASK_ID.match(tid)) and ".." not in tid - - -def _confined_task_out_dir(out_root: str, task_id: str) -> str: - """Return ``/predictions/``, confined to that root. - - Raises ``ValueError`` for an identifier that cannot be a single safe path - segment, and for a destination that resolves outside ``predictions`` (a - pre-existing symlink, say). Callers validate before any filesystem access - so a hostile id never reaches the destination write, the agent, or the - code-execution path. - """ - tid = str(task_id) - if not _is_safe_task_id(tid): - raise ValueError(f"unsafe spreadsheet task id: {tid!r}") - predictions = os.path.join(os.path.abspath(out_root), "predictions") - dest = os.path.join(predictions, tid) - try: - contained = ( - os.path.commonpath([os.path.realpath(predictions), os.path.realpath(dest)]) - == os.path.realpath(predictions) - ) - except ValueError: - # Different drives / mixed absolute-relative shapes cannot be compared. - contained = False - if not contained: - raise ValueError(f"spreadsheet destination escapes out_root: {dest!r}") - return dest - - # ── Per-task worker ────────────────────────────────────────────────────────── diff --git a/skillopt/envs/task_output.py b/skillopt/envs/task_output.py new file mode 100644 index 000000000..439485bd9 --- /dev/null +++ b/skillopt/envs/task_output.py @@ -0,0 +1,54 @@ +from __future__ import annotations + +import hashlib +import os +import re + +_SAFE_TASK_ID = re.compile(r"\A[A-Za-z0-9][A-Za-z0-9._-]{0,127}\Z") + + +def is_safe_task_id(task_id: str) -> bool: + """Return whether a task id is safe as one portable path segment.""" + value = str(task_id) + return bool(_SAFE_TASK_ID.match(value)) and ".." not in value + + +def task_output_segment(task_id: str, *, map_unsafe: bool = False) -> str: + """Return the stable directory segment for a task id.""" + value = str(task_id) + if is_safe_task_id(value): + return value + if not map_unsafe: + raise ValueError(f"unsafe task id: {value!r}") + digest = hashlib.sha256(value.encode("utf-8")).hexdigest() + return f"_mapped-{digest}" + + +def confined_task_output_dir( + out_root: str, + task_id: str, + *, + map_unsafe: bool = False, +) -> str: + """Return a task directory confined below ``out_root/predictions``.""" + predictions = os.path.join(os.path.abspath(out_root), "predictions") + destination = os.path.join( + predictions, + task_output_segment(task_id, map_unsafe=map_unsafe), + ) + if os.path.commonpath( + [os.path.realpath(predictions), os.path.realpath(destination)] + ) != os.path.realpath(predictions): + raise ValueError(f"task output escapes predictions: {destination!r}") + return destination + + +def confined_legacy_task_output_dir(out_root: str, task_id: str) -> str: + """Return an existing raw-id task directory without allowing an escape.""" + predictions = os.path.join(os.path.abspath(out_root), "predictions") + destination = os.path.join(predictions, str(task_id)) + if os.path.commonpath( + [os.path.realpath(predictions), os.path.realpath(destination)] + ) != os.path.realpath(predictions): + raise ValueError(f"legacy task output escapes predictions: {destination!r}") + return destination diff --git a/skillopt/optimizer/slow_update.py b/skillopt/optimizer/slow_update.py index 9f8dd5260..ba05b7507 100644 --- a/skillopt/optimizer/slow_update.py +++ b/skillopt/optimizer/slow_update.py @@ -23,6 +23,10 @@ import traceback from typing import Any +from skillopt.envs.task_output import ( + confined_legacy_task_output_dir, + task_output_segment, +) from skillopt.model import chat_optimizer from skillopt.prompts import load_prompt from skillopt.utils import extract_json @@ -111,7 +115,17 @@ def _clip_text(value, limit: int | None = None) -> str: def _read_trajectory(rollout_dir: str, task_id: str) -> str: """Read and format a single trajectory from a rollout directory.""" - conv_path = os.path.join(rollout_dir, "predictions", task_id, "conversation.json") + segment = task_output_segment(task_id, map_unsafe=True) + conv_path = os.path.join(rollout_dir, "predictions", segment, "conversation.json") + if not os.path.exists(conv_path) and segment != task_id: + try: + legacy_dir = confined_legacy_task_output_dir(rollout_dir, task_id) + except ValueError: + pass + else: + legacy_path = os.path.join(legacy_dir, "conversation.json") + if os.path.exists(legacy_path): + conv_path = legacy_path if not os.path.exists(conv_path): return "(trajectory not available)" try: diff --git a/tests/test_benchmark_task_output_paths.py b/tests/test_benchmark_task_output_paths.py new file mode 100644 index 000000000..786e038c4 --- /dev/null +++ b/tests/test_benchmark_task_output_paths.py @@ -0,0 +1,128 @@ +from __future__ import annotations + +import json +import os + +import pytest + +from skillopt.envs.task_output import ( + confined_task_output_dir, + task_output_segment, +) +from skillopt.optimizer.slow_update import _read_trajectory + + +@pytest.mark.parametrize( + "module_name", + [ + "skillopt.envs.docvqa.rollout", + "skillopt.envs.searchqa.rollout", + "skillopt.envs.officeqa.rollout", + ], +) +def test_rollouts_reject_unconfined_task_ids_before_running(monkeypatch, tmp_path, module_name): + module = __import__(module_name, fromlist=["process_one"]) + reached = [] + + def _bomb(*args, **kwargs): + reached.append((args, kwargs)) + raise AssertionError("model reached for an unsafe task id") + + for name in ( + "chat_target", + "chat_target_messages", + "_run_offline_no_tools_process", + ): + if hasattr(module, name): + monkeypatch.setattr(module, name, _bomb) + + with pytest.raises(ValueError, match="unsafe task id"): + module.process_one({"id": "../../escape"}, str(tmp_path / "out"), "") + + assert reached == [] + assert not (tmp_path / "escape").exists() + assert not (tmp_path / "out").exists() + + +def test_live_math_id_uses_portable_directory_and_slow_update_reads_it(tmp_path): + task_id = "202602:12" + segment = task_output_segment(task_id, map_unsafe=True) + assert segment.startswith("_mapped-") + assert ":" not in segment + + pred_dir = confined_task_output_dir( + str(tmp_path), + task_id, + map_unsafe=True, + ) + os.makedirs(pred_dir) + with open(os.path.join(pred_dir, "conversation.json"), "w", encoding="utf-8") as f: + json.dump([{"type": "message", "content": "mapped trajectory"}], f) + + assert "mapped trajectory" in _read_trajectory(str(tmp_path), task_id) + + +def test_live_math_rollout_writes_and_reads_mapped_task_directory(monkeypatch, tmp_path): + from skillopt.envs.livemathematicianbench import rollout + + monkeypatch.setattr(rollout, "is_target_exec_backend", lambda: False) + monkeypatch.setattr( + rollout, + "chat_target", + lambda **_kwargs: ("A", None), + ) + task_id = "202602:12" + item = { + "id": task_id, + "question": "Pick A.", + "choices": [ + {"label": "A", "text": "A"}, + {"label": "B", "text": "B"}, + ], + "correct_choice": {"label": "A", "text": "A"}, + } + + result = rollout.process_one(item, str(tmp_path), "") + + segment = task_output_segment(task_id, map_unsafe=True) + assert result["hard"] == 1 + assert (tmp_path / "predictions" / segment / "conversation.json").is_file() + assert not (tmp_path / "predictions" / task_id).exists() + assert "A" in _read_trajectory(str(tmp_path), task_id) + + +def test_slow_update_reads_legacy_confined_live_math_directory(tmp_path): + task_id = "202602:12" + legacy_dir = tmp_path / "predictions" / task_id + legacy_dir.mkdir(parents=True) + (legacy_dir / "conversation.json").write_text( + json.dumps([{"type": "message", "content": "legacy trajectory"}]), + encoding="utf-8", + ) + + assert "legacy trajectory" in _read_trajectory(str(tmp_path), task_id) + + +def test_slow_update_never_reads_escaping_legacy_directory(tmp_path): + outside = tmp_path.parent / "escape" / "conversation.json" + outside.parent.mkdir(exist_ok=True) + outside.write_text( + json.dumps([{"type": "message", "content": "outside trajectory"}]), + encoding="utf-8", + ) + + assert _read_trajectory(str(tmp_path), "../../escape") == "(trajectory not available)" + + +def test_confined_task_output_dir_rejects_symlink_escape(tmp_path): + predictions = tmp_path / "out" / "predictions" + predictions.mkdir(parents=True) + outside = tmp_path / "outside" + outside.mkdir() + try: + os.symlink(outside, predictions / "task-1", target_is_directory=True) + except (OSError, NotImplementedError): + pytest.skip("directory symlinks are not available on this host") + + with pytest.raises(ValueError, match="escapes predictions"): + confined_task_output_dir(str(tmp_path / "out"), "task-1")