bug: never unpickle agent-writable files when loading experiment results (sfind-454ac387)

src/agentlab/benchmarks/gaia.py

@@ -21,6 +21,7 @@
 
 from agentlab.benchmarks.abstract_env import AbstractBenchmark, AbstractEnvArgs
 from agentlab.benchmarks.multitool_gym import MultiToolGym
+from agentlab.experiments.loop import CODE_SANDBOX_DIR_NAME
 
 logger = logging.getLogger(__name__)
 
@@ -96,7 +97,7 @@ def make_env(self, exp_dir: Path, action_mapping=None) -> GaiaGym:
 def init_code_sandbox(exp_dir: str) -> None:
     # Use a common code directory for all tasks in the experiment, which is mounted in the container
     root_exp_dir = Path(exp_dir).parent
-    code_path = os.path.join(root_exp_dir, "shared_code")
+    code_path = os.path.join(root_exp_dir, CODE_SANDBOX_DIR_NAME)
     os.makedirs(code_path, exist_ok=True)
     os.environ["COMPUTER_CONTAINER_NAME"] = CONTAINER_NAME
 

src/agentlab/experiments/loop.py

@@ -25,6 +25,8 @@
 from PIL import Image
 from tqdm import tqdm
 
+from agentlab.experiments import safe_pickle
+
 try:
     from agentlab.agents.tapeagent import TapeAgent, save_tape
 except ImportError:
@@ -35,6 +37,10 @@
 
 SEED_MAX = 2 ^ 32  # arbitrary max value (exclusive), seems large enough
 
+# Directory (directly under the study directory) mounted as the agent's code sandbox.
+# Its content is agent-controlled, so result discovery never descends into it.
+CODE_SANDBOX_DIR_NAME = "shared_code"
+
 
 @dataclass
 class EnvArgs(DataClassJsonMixin):
@@ -668,7 +674,7 @@ def __init__(self, exp_dir) -> None:
     def exp_args(self) -> ExpArgs:
         if self._exp_args is None:
             with open(self.exp_dir / "exp_args.pkl", "rb") as f:
-                self._exp_args = pickle.load(f)
+                self._exp_args = safe_pickle.load(f)
                 # in case experiments were moved
                 self._exp_args.exp_dir = self.exp_dir
         return self._exp_args
@@ -677,7 +683,7 @@ def get_step_info(self, step: int) -> StepInfo:
         """Load the step info from the file and return it."""
         if self._steps_info.get(step, None) is None:
             with gzip.open(self.exp_dir / f"step_{step}.pkl.gz", "rb") as f:
-                self._steps_info[step] = pickle.load(f)
+                self._steps_info[step] = safe_pickle.load(f)
             if self._steps_info[step].obs:
                 if "screenshot" not in self._steps_info[step].obs:
                     try:
@@ -699,7 +705,7 @@ def get_step_info(self, step: int) -> StepInfo:
                     and self._steps_info[step].obs["goal_object"] is None
                 ):
                     with gzip.open(self.exp_dir / "goal_object.pkl.gz", "rb") as f:
-                        goal_object = pickle.load(f)
+                        goal_object = safe_pickle.load(f)
                         self._steps_info[step].obs["goal_object"] = goal_object
 
         return self._steps_info[step]
@@ -867,15 +873,14 @@ def yield_all_exp_results(
     if not isinstance(savedir_base, list):
         savedir_base = [savedir_base]
 
-    exp_args_paths = []
-    for exp_dir in savedir_base:
-        exp_args_paths.extend(list(Path(exp_dir).glob("**/exp_args.pkl")))
+    exp_dirs = []
+    for base_dir in savedir_base:
+        exp_dirs.extend(_find_exp_dirs(base_dir))
 
     if progress_fn is not None:
-        exp_args_paths = progress_fn(exp_args_paths, desc="Searching experiments directories.")
+        exp_dirs = progress_fn(exp_dirs, desc="Searching experiments directories.")
 
-    for exp_args_path in exp_args_paths:
-        exp_dir = exp_args_path.parent
+    for exp_dir in exp_dirs:
         if not load_hidden:
             if exp_dir.name.startswith("_") or exp_dir.name.startswith("."):
                 continue
@@ -885,6 +890,29 @@ def yield_all_exp_results(
             yield ExpResult(exp_dir)
 
 
+def _find_exp_dirs(base_dir: str | Path) -> list[Path]:
+    """Find experiment directories (those containing an `exp_args.pkl`) under base_dir.
+
+    Symlinks are not followed, the code sandbox directory is skipped and experiment
+    directories are not searched further: they may contain agent-written files (e.g. the
+    `code` symlink to the sandbox), and AgentLab never nests experiments inside one another.
+
+    Args:
+        base_dir: Root directory to search.
+
+    Returns:
+        list[Path]: The experiment directories found under base_dir.
+    """
+    exp_dirs = []
+    for dirpath, dirnames, filenames in os.walk(base_dir, followlinks=False):
+        if "exp_args.pkl" in filenames:
+            exp_dirs.append(Path(dirpath))
+            dirnames.clear()
+        else:
+            dirnames[:] = sorted(d for d in dirnames if d != CODE_SANDBOX_DIR_NAME)
+    return exp_dirs
+
+
 class DataclassJSONEncoder(json.JSONEncoder):
     def default(self, obj):
         if is_dataclass(obj):

src/agentlab/experiments/safe_pickle.py

@@ -0,0 +1,154 @@
+"""Restricted unpickling for experiment result files.
+
+Result directories can contain files written by code the agent controls (e.g. a code
+sandbox mounted inside the study directory), so `exp_args.pkl`, `step_*.pkl.gz` and
+`goal_object.pkl.gz` are loaded with an allowlist instead of a bare `pickle.load`.
+
+Only classes are accepted, never arbitrary callables (apart from a few reconstruction
+helpers used by the standard library and numpy):
+- any class from a trusted package (agentlab, browsergym, bgym, plus the comma-separated
+  module roots in the `AGENTLAB_TRUSTED_PICKLE_MODULES` environment variable),
+- subclasses of those classes defined in modules that are already imported (e.g. custom
+  agent args declared in the launching script or a test module); nothing is imported for them,
+- data classes (dataclasses, enums, pydantic models, omegaconf nodes) from dependencies
+  whose objects end up in agent configs and LLM traces,
+- plain value types from the standard library and numpy.
+"""
+
+import builtins
+import collections
+import copyreg
+import dataclasses
+import datetime
+import decimal
+import enum
+import os
+import pathlib
+import pickle
+import sys
+import typing
+import uuid
+
+import numpy as np
+
+TRUSTED_MODULES_ENV_VAR = "AGENTLAB_TRUSTED_PICKLE_MODULES"
+
+_TRUSTED_ROOTS = frozenset({"agentlab", "browsergym", "bgym"})
+_DATA_CLASS_ROOTS = frozenset(
+    {"tapeagents", "pydantic", "pydantic_core", "omegaconf", "openai", "anthropic"}
+)
+
+_SAFE_GLOBALS = {
+    *(
+        getattr(builtins, name)
+        for name in (
+            "object", "bool", "int", "float", "complex", "str", "bytes", "bytearray",
+            "list", "tuple", "dict", "set", "frozenset", "slice", "range",
+        )
+    ),
+    datetime.datetime, datetime.date, datetime.time, datetime.timedelta, datetime.timezone,
+    pathlib.Path, pathlib.PosixPath, pathlib.WindowsPath,
+    pathlib.PurePath, pathlib.PurePosixPath, pathlib.PureWindowsPath,
+    collections.OrderedDict, collections.defaultdict, collections.deque, collections.Counter,
+    decimal.Decimal,
+    uuid.UUID,
+    copyreg._reconstructor,
+    np.ndarray,
+    typing.Any,
+}  # fmt: skip
+
+_SAFE_FUNCTIONS = {
+    ("_codecs", "encode"),
+    ("numpy.core.multiarray", "_reconstruct"),
+    ("numpy._core.multiarray", "_reconstruct"),
+    ("numpy.core.multiarray", "scalar"),
+    ("numpy._core.multiarray", "scalar"),
+    ("numpy.core.numeric", "_frombuffer"),
+    ("numpy._core.numeric", "_frombuffer"),
+}
+
+_SAFE_MODULE_ROOTS = frozenset(
+    {"builtins", "copyreg", "_codecs", "collections", "datetime", "decimal", "pathlib"}
+    | {"uuid", "numpy", "PIL", "typing"}
+)
+
+
+class UnsafePickleError(pickle.UnpicklingError):
+    """Raised when a pickle references a global that is not on the allowlist."""
+
+
+def _root(module: str) -> str:
+    return module.split(".", 1)[0]
+
+
+def _trusted_roots() -> frozenset:
+    extra = os.environ.get(TRUSTED_MODULES_ENV_VAR, "")
+    return _TRUSTED_ROOTS | {m.strip() for m in extra.split(",") if m.strip()}
+
+
+def _is_data_class(obj) -> bool:
+    if dataclasses.is_dataclass(obj) or issubclass(obj, enum.Enum):
+        return True
+    pydantic = sys.modules.get("pydantic")
+    if pydantic is not None and issubclass(obj, pydantic.BaseModel):
+        return True
+    omegaconf = sys.modules.get("omegaconf")
+    return omegaconf is not None and issubclass(obj, omegaconf.Node)
+
+
+def _is_allowed(obj, trusted_roots) -> bool:
+    if isinstance(obj, type):
+        if obj in _SAFE_GLOBALS:
+            return True
+        root = _root(getattr(obj, "__module__", None) or "")
+        if root in trusted_roots:
+            return True
+        if root in _DATA_CLASS_ROOTS:
+            return _is_data_class(obj)
+        if root == "numpy":
+            return issubclass(obj, (np.dtype, np.generic))
+        if root == "PIL":
+            from PIL import Image
+
+            return obj is Image.Image
+        return False
+    try:
+        return obj in _SAFE_GLOBALS
+    except TypeError:  # unhashable
+        return False
+
+
+def _imported_trusted_subclass(module: str, name: str, trusted_roots) -> type | None:
+    obj = sys.modules.get(module)
+    for attr in name.split("."):
+        obj = getattr(obj, attr, None)
+    if isinstance(obj, type) and any(
+        _root(getattr(base, "__module__", None) or "") in trusted_roots for base in obj.__mro__[1:]
+    ):
+        return obj
+    return None
+
+
+class RestrictedUnpickler(pickle.Unpickler):
+    def find_class(self, module, name):
+        trusted_roots = _trusted_roots()
+        allowed_roots = trusted_roots | _DATA_CLASS_ROOTS | _SAFE_MODULE_ROOTS
+        if _root(module) in allowed_roots:
+            if (module, name) in _SAFE_FUNCTIONS:
+                return super().find_class(module, name)
+            obj = super().find_class(module, name)
+            if _is_allowed(obj, trusted_roots):
+                return obj
+        obj = _imported_trusted_subclass(module, name, trusted_roots)
+        if obj is not None:
+            return obj
+        raise UnsafePickleError(
+            f"Refusing to unpickle global '{module}.{name}' from an untrusted result file. "
+            f"If this class belongs to your own agent or benchmark package, add its top-level "
+            f"module to the {TRUSTED_MODULES_ENV_VAR} environment variable (comma-separated)."
+        )
+
+
+def load(file):
+    """Drop-in replacement for `pickle.load` restricted to the allowlist above."""
+    return RestrictedUnpickler(file).load()