feat(swe): serve typed resources sessions from SWEmer Ultra, SWEmer OML, and SWE internal v1

resources_servers/r2e_gym/app.py

@@ -47,7 +47,8 @@
     BaseSeedSessionResponse,
     BaseVerifyRequest,
     BaseVerifyResponse,
-    SimpleResourcesServer,
+    ResourcesSeedSessionRequest,
+    ResourcesSeedSessionResponse,
 )
 from nemo_gym.global_config import get_global_config_dict
 from nemo_gym.sandbox import AsyncSandbox, SandboxResources, SandboxSpec
@@ -72,6 +73,7 @@
     capture_model_patch,
     prepare_git_for_commits,
 )
+from resources_servers.swebench.sandbox_sessions import SandboxSessionResourcesServer
 
 
 class R2EGymResourcesServerConfig(BaseResourcesServerConfig):
@@ -140,7 +142,7 @@ class R2EGymVerifyResponse(BaseVerifyResponse):
     model_patch_sha256: str = ""
 
 
-class R2EGymResourcesServer(SimpleResourcesServer):
+class R2EGymResourcesServer(SandboxSessionResourcesServer):
     ray_enabled = False
     config: R2EGymResourcesServerConfig
 
@@ -238,28 +240,48 @@ async def _extract_model_patch(self, session_id: str, workdir: str) -> PatchCapt
                 drop_sections=drop_patch_sections,
             )
         finally:
-            await self._stop_sandbox(original_sandbox)
+            await self._release_task_sandbox(session_id, original_sandbox)
 
-    async def seed_session(self, request: Request, body: R2EGymSeedSessionRequest) -> R2EGymSeedSessionResponse:
-        """Start the instance's image so an agent can work in it, minus everything that gives the task away."""
+    async def seed_session(
+        self, request: Request, body: R2EGymSeedSessionRequest | ResourcesSeedSessionRequest
+    ) -> R2EGymSeedSessionResponse | ResourcesSeedSessionResponse:
+        """Start the instance's image so an agent can work in it, minus everything that gives the task away.
+
+        An Environment Server seeds a typed session and gets the sandbox back as ``sandbox_access``; an
+        agent's ``/run`` seeds with the row and gets the sandbox handle.
+        """
+        if isinstance(body, ResourcesSeedSessionRequest):
+            return await self.seed_task_sandbox_session(request, body, R2EGymInstanceRequest)
         session_id = request.session[SESSION_ID_KEY]
         await self._stop_sandbox(self._session_id_to_sandbox.pop(session_id, None))
-        self._session_id_to_pristine_untracked.pop(session_id, None)
-        sandbox = await self._create_sandbox(body)
         try:
-            await self._hide_tests_from_agent(sandbox, body.instance_id)
-            if self.config.apply_anti_cheating:
-                await apply_anti_cheat_setup(sandbox, REPO_DIRECTORY, body.instance_id, "r2e_gym")
-            # The anti-cheat scrub leaves no committer identity, so the agent's `git commit` would fail.
-            await prepare_git_for_commits(sandbox, REPO_DIRECTORY, "r2e_gym")
-            self._session_id_to_pristine_untracked[session_id] = await self._pristine_untracked_files(
-                sandbox, REPO_DIRECTORY
-            )
+            await self._start_task_sandbox(session_id, body)
         except BaseException:
-            await self._stop_sandbox(sandbox)
+            # A sandbox that leaked hidden tests must not be handed out or left running.
+            await self._stop_sandbox(self._session_id_to_sandbox.pop(session_id, None))
             raise
+        return R2EGymSeedSessionResponse(
+            sandbox_handle=str(self._session_id_to_sandbox[session_id]._handle.sandbox_id), workdir=REPO_DIRECTORY
+        )
+
+    async def _start_task_sandbox(self, session_id: str, body: R2EGymInstanceRequest) -> str:
+        """Start the task sandbox for ``session_id`` and return the directory the agent works in."""
+        self._forget_task_sandbox_state(session_id)
+        sandbox = await self._create_sandbox(body)
+        # Own the sandbox before preparing it, so a failed seed can still stop it.
         self._session_id_to_sandbox[session_id] = sandbox
-        return R2EGymSeedSessionResponse(sandbox_handle=str(sandbox._handle.sandbox_id), workdir=REPO_DIRECTORY)
+        await self._hide_tests_from_agent(sandbox, body.instance_id)
+        if self.config.apply_anti_cheating:
+            await apply_anti_cheat_setup(sandbox, REPO_DIRECTORY, body.instance_id, "r2e_gym")
+        # The anti-cheat scrub leaves no committer identity, so the agent's `git commit` would fail.
+        await prepare_git_for_commits(sandbox, REPO_DIRECTORY, "r2e_gym")
+        self._session_id_to_pristine_untracked[session_id] = await self._pristine_untracked_files(
+            sandbox, REPO_DIRECTORY
+        )
+        return REPO_DIRECTORY
+
+    def _forget_task_sandbox_state(self, session_id: str) -> None:
+        self._session_id_to_pristine_untracked.pop(session_id, None)
 
     def _response(self, body: R2EGymVerifyRequest, **fields: Any) -> R2EGymVerifyResponse:
         # Spread the request: BaseVerifyResponse extends BaseVerifyRequest, so responses_create_params
@@ -277,6 +299,7 @@ async def verify(self, request: Request, body: R2EGymVerifyRequest) -> R2EGymVer
         if self.config.is_verifying_golden_patch:
             capture = PatchCapture.static(body.patch, mode, "golden")
         else:
+            self._claim_task_sandbox(session_id)
             try:
                 capture = await self._extract_model_patch(session_id, REPO_DIRECTORY)
             except Exception as exc:

resources_servers/r2e_gym/tests/test_app.py

@@ -486,7 +486,7 @@ def test_seed_session_hides_tests_before_scrubbing_history_and_handing_over(self
         assert (
             seed.index("_hide_tests_from_agent(")
             < seed.index("apply_anti_cheat_setup(")
-            < seed.index("self._session_id_to_sandbox[session_id] = sandbox")
+            < seed.index("return REPO_DIRECTORY")
         )
         assert "_create_sandbox(body)" in seed and "files=" not in seed, "the agent sandbox gets no row files"
 

resources_servers/r2e_gym/tests/test_sessions.py

@@ -0,0 +1,149 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Environment Server sessions and the legacy seed, through the server's HTTP routes with a fake sandbox."""
+
+from unittest.mock import AsyncMock, MagicMock
+
+from fastapi.testclient import TestClient
+from pytest import MonkeyPatch
+
+from nemo_gym.base_resources_server import ResourcesSeedSessionRequest
+from nemo_gym.episode_types import EpisodeId, TaskId
+from nemo_gym.sandbox import SandboxExecResult, SandboxHandle
+from nemo_gym.server_utils import ServerClient
+from nemo_gym.testing.session_conformance import check_resources_session_contract
+from resources_servers.r2e_gym import app as r2e_gym_app
+from resources_servers.r2e_gym.app import R2EGymResourcesServer, R2EGymResourcesServerConfig
+from resources_servers.r2e_gym.verification import VerificationResult
+from resources_servers.swebench.patch_capture import PatchCapture
+
+
+_ROW = {
+    "instance_id": "namanjain12__pillow_final__6e97dd5c",
+    "repo_name": "pillow",
+    "image_name": "namanjain12/pillow_final:6e97dd5cecc074a0d8622e4aa5658b4263324760",
+    "expected_output_json": '{"test_a": "PASSED"}',
+}
+_RESPONSE = {
+    "output": [],
+    "id": "",
+    "created_at": 0,
+    "model": "",
+    "object": "response",
+    "parallel_tool_calls": False,
+    "tool_choice": "auto",
+    "tools": [],
+}
+_VERIFY = _ROW | {"responses_create_params": {"input": "Fix it"}, "response": _RESPONSE}
+
+
+def _server(monkeypatch: MonkeyPatch) -> tuple[R2EGymResourcesServer, list[MagicMock]]:
+    """A server whose sandboxes are fakes, with the agent's patch capture and the grading run stubbed."""
+    sandboxes: list[MagicMock] = []
+
+    async def create_sandbox(self, body, files=None) -> MagicMock:
+        sandbox = MagicMock()
+        sandbox._handle = SandboxHandle(sandbox_id=f"sb-{len(sandboxes) + 1}", provider_name="test", raw=None)
+        sandbox.exec = AsyncMock(return_value=SandboxExecResult(return_code=0, stdout="", stderr=""))
+        sandbox.serialize = AsyncMock(return_value={"sandbox_id": sandbox._handle.sandbox_id})
+        sandbox.stop = AsyncMock()
+        sandboxes.append(sandbox)
+        return sandbox
+
+    monkeypatch.setattr(R2EGymResourcesServer, "_create_sandbox", create_sandbox)
+    monkeypatch.setattr(r2e_gym_app, "prepare_git_for_commits", AsyncMock())
+    monkeypatch.setattr(
+        r2e_gym_app,
+        "capture_model_patch",
+        AsyncMock(return_value=PatchCapture.static("diff --git a/x b/x\n", "worktree", "worktree")),
+    )
+    monkeypatch.setattr(
+        r2e_gym_app, "run_verification", AsyncMock(return_value=VerificationResult(True, True, True, {}, "ok"))
+    )
+    config = R2EGymResourcesServerConfig(
+        host="0.0.0.0",
+        port=8080,
+        entrypoint="",
+        name="r2e_gym",
+        sandbox_provider="docker",
+        sandbox_config={},
+        apply_anti_cheating=False,
+    )
+    return R2EGymResourcesServer(config=config, server_client=MagicMock(spec=ServerClient)), sandboxes
+
+
+def _seed(session_id: str = "resources-session") -> ResourcesSeedSessionRequest:
+    return ResourcesSeedSessionRequest(
+        resources_session_id=session_id,
+        episode_id=EpisodeId(rollout_id="rollout"),
+        task_id=TaskId(taskset="r2e_gym", task_id=_ROW["instance_id"]),
+        task_data=_ROW | {"responses_create_params": {"input": "Fix it"}},
+    )
+
+
+def test_typed_episode_hands_over_the_task_sandbox_grades_it_and_closes(monkeypatch: MonkeyPatch) -> None:
+    server, sandboxes = _server(monkeypatch)
+    seed = _seed().model_dump(mode="json")
+    close = {"resources_session_id": "resources-session", "episode_id": {"rollout_id": "rollout"}}
+
+    with TestClient(server.setup_webserver(), raise_server_exceptions=False) as client:
+        seeded = client.post("/seed_session", json=seed)
+        assert seeded.status_code == 200, seeded.text
+        assert seeded.json()["sandbox_access"] == {
+            "connection": {"kind": "direct", "provider_config_ref": "docker", "descriptor": {"sandbox_id": "sb-1"}},
+            "workdir": "/testbed",
+        }
+
+        verified = client.post("/verify", json=_VERIFY)
+        assert verified.status_code == 200, verified.text
+        assert verified.json()["reward"] == 1.0
+        assert verified.json()["model_patch"] == "diff --git a/x b/x\n"
+        sandboxes[0].stop.assert_awaited_once()
+        assert client.post("/verify", json=_VERIFY).status_code == 500, "a repeated verify must not grade again"
+
+        assert client.post("/close_session", json=close).status_code == 200
+        assert client.post("/seed_session", json=seed).status_code == 500
+
+    sandboxes[0].stop.assert_awaited_once()
+    assert server._session_id_to_sandbox == {} and server._session_id_to_pristine_untracked == {}
+
+
+def test_legacy_seed_still_returns_the_sandbox_handle(monkeypatch: MonkeyPatch) -> None:
+    server, sandboxes = _server(monkeypatch)
+    with TestClient(server.setup_webserver(), raise_server_exceptions=False) as client:
+        seeded = client.post("/seed_session", json=_ROW)
+        assert seeded.status_code == 200, seeded.text
+        assert seeded.json() == {"sandbox_handle": "sb-1", "workdir": "/testbed"}
+        verified = client.post("/verify", json=_VERIFY)
+        assert verified.status_code == 200, verified.text
+        assert verified.json()["reward"] == 1.0
+    sandboxes[0].stop.assert_awaited_once()
+
+
+def test_session_conformance(monkeypatch: MonkeyPatch) -> None:
+    server, _ = _server(monkeypatch)
+    check_resources_session_contract(server.setup_webserver(), _seed("conformance-session"), keeps_state=True)
+
+
+def test_a_failed_test_scrub_refuses_the_seed_and_stops_the_sandbox(monkeypatch: MonkeyPatch) -> None:
+    server, sandboxes = _server(monkeypatch)
+    monkeypatch.setattr(
+        R2EGymResourcesServer, "_hide_tests_from_agent", AsyncMock(side_effect=RuntimeError("scrub failed"))
+    )
+    with TestClient(server.setup_webserver(), raise_server_exceptions=False) as client:
+        assert client.post("/seed_session", json=_seed().model_dump(mode="json")).status_code == 500
+        assert client.post("/seed_session", json=_ROW).status_code == 500
+    assert [sandbox.stop.await_count for sandbox in sandboxes] == [1, 1]
+    assert server._session_id_to_sandbox == {}

resources_servers/swe_gym/app.py

@@ -46,7 +46,8 @@
     BaseSeedSessionResponse,
     BaseVerifyRequest,
     BaseVerifyResponse,
-    SimpleResourcesServer,
+    ResourcesSeedSessionRequest,
+    ResourcesSeedSessionResponse,
 )
 from nemo_gym.global_config import get_global_config_dict
 from nemo_gym.sandbox import AsyncSandbox, SandboxResources, SandboxSpec
@@ -72,6 +73,7 @@
     capture_model_patch,
     prepare_git_for_commits,
 )
+from resources_servers.swebench.sandbox_sessions import SandboxSessionResourcesServer
 
 
 # The image's default PATH points at the conda *base* interpreter (Python 3.11 without the repo's
@@ -156,7 +158,7 @@ class SWEGymVerifyResponse(BaseVerifyResponse):
     model_patch_sha256: str = ""
 
 
-class SWEGymResourcesServer(SimpleResourcesServer):
+class SWEGymResourcesServer(SandboxSessionResourcesServer):
     ray_enabled = False
     config: SWEGymResourcesServerConfig
 
@@ -246,28 +248,47 @@ async def _extract_model_patch(self, session_id: str, workdir: str, base_commit:
                 drop_sections=drop_patch_sections,
             )
         finally:
-            await self._stop_sandbox(original_sandbox)
+            await self._release_task_sandbox(session_id, original_sandbox)
 
-    async def seed_session(self, request: Request, body: SWEGymSeedSessionRequest) -> SWEGymSeedSessionResponse:
+    async def seed_session(
+        self, request: Request, body: SWEGymSeedSessionRequest | ResourcesSeedSessionRequest
+    ) -> SWEGymSeedSessionResponse | ResourcesSeedSessionResponse:
         """Start the instance's image so an agent can work in it.
 
+        An Environment Server seeds a typed session and gets the sandbox back as ``sandbox_access``; an
+        agent's ``/run`` seeds with the row and gets the sandbox handle.
+        """
+        if isinstance(body, ResourcesSeedSessionRequest):
+            return await self.seed_task_sandbox_session(request, body, SWEGymInstanceRequest)
+        session_id = request.session[SESSION_ID_KEY]
+        await self._stop_sandbox(self._session_id_to_sandbox.pop(session_id, None))
+        await self._start_task_sandbox(session_id, body)
+        return SWEGymSeedSessionResponse(
+            sandbox_handle=str(self._session_id_to_sandbox[session_id]._handle.sandbox_id), workdir=REPO_DIRECTORY
+        )
+
+    async def _start_task_sandbox(self, session_id: str, body: SWEGymInstanceRequest) -> str:
+        """Start the task sandbox for ``session_id`` and return the directory the agent works in.
+
         The sandbox gets no files from the row. The git scrub drops every ref but HEAD and prunes the
         objects behind them: SWE-bench images clone the whole upstream history, so without it the fix
         commit and every later release tag sit in ``.git`` within a ``git log --all`` of the agent.
         """
-        session_id = request.session[SESSION_ID_KEY]
-        await self._stop_sandbox(self._session_id_to_sandbox.pop(session_id, None))
-        self._session_id_to_pristine_untracked.pop(session_id, None)
+        self._forget_task_sandbox_state(session_id)
         sandbox = await self._create_sandbox(body)
+        # Own the sandbox before preparing it, so a failed seed can still stop it.
+        self._session_id_to_sandbox[session_id] = sandbox
         if self.config.apply_anti_cheating:
             await apply_anti_cheat_setup(sandbox, REPO_DIRECTORY, body.instance_id, "swe_gym")
         # The anti-cheat scrub leaves no committer identity, so the agent's `git commit` would fail.
         await prepare_git_for_commits(sandbox, REPO_DIRECTORY, "swe_gym")
         self._session_id_to_pristine_untracked[session_id] = await self._pristine_untracked_files(
             sandbox, REPO_DIRECTORY
         )
-        self._session_id_to_sandbox[session_id] = sandbox
-        return SWEGymSeedSessionResponse(sandbox_handle=str(sandbox._handle.sandbox_id), workdir=REPO_DIRECTORY)
+        return REPO_DIRECTORY
+
+    def _forget_task_sandbox_state(self, session_id: str) -> None:
+        self._session_id_to_pristine_untracked.pop(session_id, None)
 
     def _response(self, body: SWEGymVerifyRequest, **fields: Any) -> SWEGymVerifyResponse:
         # Spread the request: BaseVerifyResponse extends BaseVerifyRequest, so responses_create_params
@@ -290,6 +311,7 @@ async def verify(self, request: Request, body: SWEGymVerifyRequest) -> SWEGymVer
         if self.config.is_verifying_golden_patch:
             capture = PatchCapture.static(body.patch, mode, "golden")
         else:
+            self._claim_task_sandbox(session_id)
             try:
                 capture = await self._extract_model_patch(session_id, REPO_DIRECTORY, body.base_commit)
             except Exception as exc: