Fix load_from_state to work with subclass constructors

Moved load_from_state to SweBenchEnv and HarborEnv since they need
different constructor arguments (tasks list). Base class now provides
_restore_from_snapshot helper that subclasses call after init.

- Add load_from_state implementation to SweBenchEnv
- Add load_from_state implementation to HarborEnv
- Refactor base class to use _restore_from_snapshot helper
- Add test for load_from_state functionality

Author: sarvanithin

src/ares/environments/base.py

@@ -565,55 +565,27 @@ async def export_state(
 
         return snap
 
-    @classmethod
-    async def load_from_state(
-        cls,
-        snapshot_path: snapshot.EnvironmentSnapshot | pathlib.Path,
-        *,
-        container_factory: containers.ContainerFactory | None = None,
-        code_agent_factory: code_agent_base.CodeAgentFactory | None = None,
-    ) -> "CodeBaseEnv":
-        """Restore environment from snapshot.
+    async def _restore_from_snapshot(self, snap: snapshot.EnvironmentSnapshot) -> None:
+        """Internal helper to restore state from snapshot.
 
-        Args:
-            snapshot_path: EnvironmentSnapshot or path to snapshot.json
-            container_factory: Override factory (uses snapshot metadata if None)
-            code_agent_factory: Override factory (uses default if None)
+        Called by subclass load_from_state implementations after creating the environment.
 
-        Returns:
-            Restored environment (NOT active, use async with)
+        Args:
+            snap: The snapshot to restore from
         """
-        # Load snapshot if path provided
-        if isinstance(snapshot_path, pathlib.Path):
-            snap = snapshot.EnvironmentSnapshot.load_from_file(snapshot_path)
-        else:
-            snap = snapshot_path
-
-        # Deserialize task
-        task = cls._deserialize_task(snap.task_data, snap.task_type)
-
-        # Create environment instance
-        container_factory = container_factory or cls._default_container_factory(snap)
-        code_agent_factory = code_agent_factory or mini_swe_agent.MiniSWECodeAgent
+        # Restore container
+        self._container = await self._restore_container(snap)
 
-        # Note: This creates a base CodeBaseEnv which is abstract
-        # In practice, this should be called on SweBenchEnv or HarborEnv subclasses
-        env = cls(
-            container_factory=container_factory,
-            code_agent_factory=code_agent_factory,
-            step_limit=snap.step_limit,
-        )
+        # Deserialize and set task
+        task = self._deserialize_task(snap.task_data, snap.task_type)
+        self._current_task = task
 
-        # Restore container
-        env._container = await env._restore_container(snap)
-        env._current_task = task
-        env._step_count = snap.step_count
-        env._requires_reset = snap.requires_reset
+        # Restore state
+        self._step_count = snap.step_count
+        self._requires_reset = snap.requires_reset
 
         # Store messages for later restoration
-        env._saved_agent_messages = snap.agent_messages
-
-        return env
+        self._saved_agent_messages = snap.agent_messages
 
     @abc.abstractmethod
     def _serialize_task(self, task: TaskType) -> dict:

src/ares/environments/harbor_env.py

@@ -179,3 +179,51 @@ def _deserialize_task(cls, task_data: dict, task_type: str) -> harbor_task.Task:
         """Deserialize Harbor task (reload from directory)."""
         del task_type  # Unused - validated by caller
         return harbor_task.Task(task_dir=pathlib.Path(task_data["task_dir"]))
+
+    @classmethod
+    async def load_from_state(
+        cls,
+        snapshot_path: "base.snapshot.EnvironmentSnapshot | pathlib.Path",
+        *,
+        container_factory: containers.ContainerFactory | None = None,
+        code_agent_factory: code_agent_base.CodeAgentFactory | None = None,
+        tracker: stat_tracker.StatTracker | None = None,
+    ) -> "HarborEnv":
+        """Restore HarborEnv from snapshot.
+
+        Args:
+            snapshot_path: EnvironmentSnapshot or path to snapshot.json
+            container_factory: Override factory (uses snapshot metadata if None)
+            code_agent_factory: Override factory (uses default if None)
+            tracker: Optional stat tracker
+
+        Returns:
+            Restored environment (NOT active, use async with)
+        """
+        from ares.environments import snapshot as snapshot_module
+
+        # Load snapshot if path provided
+        if isinstance(snapshot_path, pathlib.Path):
+            snap = snapshot_module.EnvironmentSnapshot.load_from_file(snapshot_path)
+        else:
+            snap = snapshot_path
+
+        # Deserialize task
+        task = cls._deserialize_task(snap.task_data, snap.task_type)
+
+        # Create environment instance with tasks argument
+        container_factory = container_factory or base.ares_daytona.DaytonaContainer
+        code_agent_factory = code_agent_factory or mini_swe_agent.MiniSWECodeAgent
+
+        env = cls(
+            tasks=[task],
+            container_factory=container_factory,
+            code_agent_factory=code_agent_factory,
+            step_limit=snap.step_limit,
+            tracker=tracker,
+        )
+
+        # Restore state using base helper
+        await env._restore_from_snapshot(snap)
+
+        return env

src/ares/environments/snapshot_test.py

@@ -334,3 +334,65 @@ async def test_export_state_raises_if_no_task():
 
     with tempfile.TemporaryDirectory() as tmp_dir, pytest.raises(RuntimeError, match="No current task set"):
         await env.export_state(pathlib.Path(tmp_dir))
+
+
+@pytest.mark.asyncio
+async def test_load_from_state_creates_valid_env(tmp_path: pathlib.Path):
+    """Test load_from_state creates a properly initialized environment."""
+
+    # Create a mock container with download_dir and upload_dir support
+    class MockContainerWithDirOps(mock_container.MockContainer):
+        def __init__(self):
+            super().__init__()
+            self.resources = None
+            self.image = "python:3.12"  # Add image attribute for snapshot
+
+        async def download_dir(self, remote_path: str, local_path: pathlib.Path):
+            del remote_path  # Unused in mock
+            local_path.parent.mkdir(parents=True, exist_ok=True)
+            local_path.write_text("mock tarball")
+
+        async def upload_dir(self, local_path: pathlib.Path, remote_path: str):
+            """Mock upload_dir for container restoration."""
+            del local_path, remote_path  # Unused in mock
+
+    # Create and export state
+    env = swebench_env.SweBenchEnv(
+        tasks=[_MOCK_SWEBENCH_TASK],
+        container_factory=mock_container.MockContainerFactory,
+        step_limit=42,
+    )
+
+    container = MockContainerWithDirOps()
+    await container.start()
+    env._container = container
+    env._current_task = _MOCK_SWEBENCH_TASK
+    env._step_count = 7
+    env._requires_reset = False
+    env._code_agent_task = None
+
+    snap = await env.export_state(tmp_path, snapshot_id="test-load")
+
+    # Load from snapshot
+    class MockContainerFactory:
+        @classmethod
+        def from_image(cls, *, image: str, name: str | None = None, resources=None):
+            del image, name, resources
+            return MockContainerWithDirOps()
+
+        @classmethod
+        def from_dockerfile(cls, *, dockerfile_path, name: str | None = None, resources=None):
+            del dockerfile_path, name, resources
+            return MockContainerWithDirOps()
+
+    restored_env = await swebench_env.SweBenchEnv.load_from_state(snap, container_factory=MockContainerFactory)
+
+    # Verify restoration
+    assert restored_env._step_count == 7
+    assert restored_env._step_limit == 42
+    assert restored_env._requires_reset is False
+    assert restored_env._current_task.instance_id == _MOCK_SWEBENCH_TASK.instance_id
+    assert restored_env._container is not None
+
+    # Cleanup
+    await restored_env.close()

src/ares/environments/swebench_env.py

@@ -10,6 +10,7 @@
 import functools
 import json
 import logging
+import pathlib
 import random
 import time
 from typing import Any, Literal, cast
@@ -266,3 +267,53 @@ def _deserialize_task(cls, task_data: dict, task_type: str) -> SwebenchTask:
         del task_type  # Unused - validated by caller
         # The field validators will convert JSON strings to lists
         return SwebenchTask.model_validate(task_data)
+
+    @classmethod
+    async def load_from_state(
+        cls,
+        snapshot_path: "base.snapshot.EnvironmentSnapshot | pathlib.Path",
+        *,
+        container_factory: containers.ContainerFactory | None = None,
+        code_agent_factory: code_agent_base.CodeAgentFactory | None = None,
+        tracker: stat_tracker.StatTracker | None = None,
+    ) -> "SweBenchEnv":
+        """Restore SweBenchEnv from snapshot.
+
+        Args:
+            snapshot_path: EnvironmentSnapshot or path to snapshot.json
+            container_factory: Override factory (uses snapshot metadata if None)
+            code_agent_factory: Override factory (uses default if None)
+            tracker: Optional stat tracker
+
+        Returns:
+            Restored environment (NOT active, use async with)
+        """
+        import pathlib as pathlib_module
+
+        from ares.environments import snapshot as snapshot_module
+
+        # Load snapshot if path provided
+        if isinstance(snapshot_path, pathlib_module.Path):
+            snap = snapshot_module.EnvironmentSnapshot.load_from_file(snapshot_path)
+        else:
+            snap = snapshot_path
+
+        # Deserialize task
+        task = cls._deserialize_task(snap.task_data, snap.task_type)
+
+        # Create environment instance with tasks argument
+        container_factory = container_factory or base.ares_daytona.DaytonaContainer
+        code_agent_factory = code_agent_factory or mini_swe_agent.MiniSWECodeAgent
+
+        env = cls(
+            tasks=[task],
+            container_factory=container_factory,
+            code_agent_factory=code_agent_factory,
+            step_limit=snap.step_limit,
+            tracker=tracker,
+        )
+
+        # Restore state using base helper
+        await env._restore_from_snapshot(snap)
+
+        return env