feat(checkpoint): protocol for harnesses that borrow a sandbox

environment_servers/single_agent_turn/README.md

@@ -31,6 +31,8 @@ environment_servers:
 
 `direct_http` grants the Agent Server scoped HTTP access to the Resources Server. `mcp` requires the Resources seed response to provide HTTP MCP connection metadata. Selecting neither transport sends `tool_accesses=[]` when the agent session is seeded.
 
-With partial-rollout checkpointing on, an episode that grants `mcp` access restarts from its input after a crash instead of continuing from a checkpoint. Its agent holds an MCP session token that names the episode's attempt, and the Resources Server refuses that token once a restore replaces the attempt.
+With partial-rollout checkpointing on, the agent of an episode that grants `mcp` access holds an MCP session token that names the episode's attempt, and the Resources Server refuses that token once a restore replaces the attempt. Such an episode continues from a checkpoint after a crash only if its agent borrows the sandbox the Resources seed lent (`sandbox_access`) and re-binds restored sessions to new tool grants (`rebinds_tool_accesses`, as the OpenCode agent does). Otherwise it restarts from its input.
+
+A restore rebuilds a borrowed sandbox under a new ID. When the replacement attempt continues an episode whose agent had not closed yet, it first seeds the Resources and agent sessions again: the Resources Server returns the restored sandbox, and the agent re-binds its restored session to that sandbox and to the replacement attempt's tool grants, such as a new MCP session token. Seeding again is idempotent, so a checkpoint does not wait for it, and a crash during it seeds again after the next restore. An episode without a borrowed sandbox seeds nothing again.
 
 The server supports `num_workers > 1`. Episode state and cleanup remain local to the worker handling that request. `max_concurrent_episodes` and queue admission are also enforced independently by each worker, so the deployment-wide concurrency limit is approximately `num_workers * max_concurrent_episodes`.

environment_servers/single_agent_turn/app.py

@@ -3,6 +3,8 @@
 
 """Resources-backed single-agent environment server."""
 
+import logging
+from collections.abc import Callable
 from typing import Any, Literal
 from uuid import uuid4
 from warnings import warn
@@ -54,6 +56,7 @@
     BaseEnvironmentServer,
     BaseEnvironmentServerConfig,
     CleanupContext,
+    CleanupHandle,
     DirectHTTPToolAccess,
     MCPStreamableHTTPConnection,
     MCPToolAccess,
@@ -68,6 +71,7 @@
     dependency_failure,
     handled_episode_failure,
     is_handled_episode_failure,
+    seed_rebinds_tools,
     seed_restarts,
     seed_verify_step_type,
 )
@@ -82,6 +86,9 @@
 from nemo_gym.utils import is_nemo_gym_fastapi_entrypoint
 
 
+LOGGER = logging.getLogger(__name__)
+
+
 class SessionHandles(BaseModel):
     """The sessions an episode owns: enough to reuse them after a restore and to close them."""
 
@@ -92,6 +99,9 @@ class SessionHandles(BaseModel):
     # Whether a checkpoint may replay the resources server's /verify, as its seed reply reported.
     resources_verify_step_type: StepType = "wait"
     observation: str | None = None
+    # Whether the agent borrows a sandbox the resources server lent in its seed reply.
+    # A restore rebuilds that sandbox under a new id, so the replacement attempt seeds both sessions again.
+    borrowed_sandbox: bool = False
 
     @classmethod
     def new(cls) -> "SessionHandles":
@@ -197,10 +207,17 @@ async def close_agent() -> None:
             await self.checkpoint_boundary(request, {"next": "seed", "handles": handles.model_dump()})
             # Seeds are idempotent for a caller-assigned session ID, so a checkpoint need not wait for them.
             async with self.checkpoint_step(request, "replay"):
-                seed = await self._seed_resources(request, handles)
-                agent_cleanup = cleanup.register_cleanup("agent session", close_agent)
-                await self._seed_agent(request, handles, seed)
+                agent_cleanup = await self.seed_episode_sessions(
+                    request,
+                    handles,
+                    register_agent_cleanup=lambda: cleanup.register_cleanup("agent session", close_agent),
+                )
             stage = "invoke_agent"
+        elif continuation and stage in ("invoke_agent", "close_agent"):
+            # Restored before the agent's work ended: the sandbox the agent borrowed was rebuilt under a new id.
+            # Re-seeding is idempotent, so a checkpoint need not wait for it, and a crash during it re-seeds again.
+            async with self.checkpoint_step(request, "replay"):
+                await self.reseed_after_restore(request, handles)
 
         if stage == "invoke_agent":
             await self.checkpoint_boundary(request, {"next": "invoke_agent", "handles": handles.model_dump()})
@@ -352,6 +369,40 @@ def dumped_trajectory() -> dict | None:
             ),
         )
 
+    async def seed_episode_sessions(
+        self,
+        request: SingleAgentTurnRequest,
+        handles: "SessionHandles",
+        *,
+        register_agent_cleanup: Callable[[], CleanupHandle] | None = None,
+    ) -> CleanupHandle | None:
+        """Seed the episode's resources session, then its agent session with what the resources seed lent.
+
+        ``register_agent_cleanup`` runs between the two seeds, so a lost agent seed reply is still closed;
+        its result is returned.
+        An episode whose agent holds resources MCP grants restarts after a crash
+        unless the agent borrows a sandbox the resources seed lent and re-binds them.
+        """
+        seed = await self._seed_resources(request, handles)
+        registered = register_agent_cleanup() if register_agent_cleanup is not None else None
+        await self._seed_agent(request, handles, seed)
+        return registered
+
+    async def reseed_after_restore(self, request: SingleAgentTurnRequest, handles: "SessionHandles") -> None:
+        """Seed both sessions again for a replacement attempt that continues a checkpoint before the agent closed.
+
+        The resources server returns the sandbox its restore forked, without setting the task up again;
+        the agent re-binds its restored session to that sandbox and to this attempt's tool grants.
+        Without a borrowed sandbox nothing changed hands, and nothing is seeded.
+        """
+        if not handles.borrowed_sandbox:
+            return
+        LOGGER.info(
+            "Re-seeding the sessions of %s after a restore: the agent's sandbox was rebuilt",
+            request.episode_id.capture_key,
+        )
+        await self.seed_episode_sessions(request, handles)
+
     async def _seed_resources(
         self, request: SingleAgentTurnRequest, handles: "SessionHandles"
     ) -> ResourcesSeedSessionResponse:
@@ -377,6 +428,7 @@ async def _seed_resources(
             seed = ResourcesSeedSessionResponse.model_validate(await get_response_json(seed_http_response))
             if seed.resources_session_id != handles.resources_session_id:
                 raise ValueError("Resources seed returned a different resources_session_id")
+            handles.borrowed_sandbox = seed.sandbox_access is not None
         except Exception as error:
             raise dependency_failure(
                 error,
@@ -420,6 +472,13 @@ async def _seed_agent(
                     failure_reason=f"Unsupported resources MCP transport: {seed.resources_tools.transport}",
                     terminal=True,
                 )
+            if seed.sandbox_access is None:
+                # The MCP session token names this attempt, and only a restore's re-seed hands the agent the token
+                # of the attempt that replaces it. A restore re-seeds only an episode whose agent borrows a sandbox,
+                # so without one the agent would keep a token the Resources Server refuses, whatever it declares:
+                # the episode restarts from its input after a crash.
+                # Marked before the agent seed, so no checkpoint exports the agent session it creates.
+                await self.checkpoint_restart(request)
             url_path = seed.resources_tools.url_path.lstrip("/")
             tool_accesses.append(
                 MCPToolAccess(
@@ -431,10 +490,6 @@ async def _seed_agent(
                     ),
                 )
             )
-            # The MCP session token names this attempt, and the Resources Server refuses it once a restore replaces it.
-            # An agent continued from a checkpoint would keep that token and see every tool call fail.
-            # So the episode restarts from its input after a crash, and its new seed mints a token for the new attempt.
-            await self.checkpoint_restart(request)
 
         try:
             agent_create_http_response = await self.server_client.post(
@@ -455,7 +510,15 @@ async def _seed_agent(
             if agent_session.agent_session_id != handles.agent_session_id:
                 raise ValueError("Agent seed returned a different agent_session_id")
             handles.agent_cookies = _cookies(agent_create_http_response)
-            if seed_restarts(agent_create_http_response.headers):
+            # The agent cannot capture its session, or holds an MCP session token it cannot exchange for the one a
+            # restore's re-seed carries: the Resources Server refuses the token once a restore replaces this attempt,
+            # so the episode restarts from its input after a crash.
+            # Such an agent reports the restart on its own seed reply, so no checkpoint exports its session;
+            # the MCP check covers an agent whose own check misses that case.
+            if seed_restarts(agent_create_http_response.headers) or (
+                "mcp" in self.config.resources_tool_transports
+                and not seed_rebinds_tools(agent_create_http_response.headers)
+            ):
                 await self.checkpoint_restart(request)
         except Exception as error:
             raise dependency_failure(