feat: add resources session lifecycle API

nemo_gym/base_resources_server.py

@@ -14,10 +14,10 @@
 # limitations under the License.
 from abc import abstractmethod
 from enum import Enum
-from typing import TYPE_CHECKING, Any, ClassVar, Optional
+from typing import TYPE_CHECKING, Any, ClassVar, Generic, Optional, TypeVar
 
 from fastapi import FastAPI
-from pydantic import BaseModel, ConfigDict, Field, field_validator
+from pydantic import BaseModel, ConfigDict, Field, JsonValue, field_validator
 
 
 if TYPE_CHECKING:
@@ -26,6 +26,7 @@
     from nemo_gym.mcp_auto_exposure import MCPTool
 
 from nemo_gym.config_types import AggregateMetrics, AggregateMetricsRequest
+from nemo_gym.episode_types import EpisodeId, TaskId
 from nemo_gym.failure_kinds import validate_failure_kind
 from nemo_gym.judge import judge_failsafe
 from nemo_gym.openai_utils import (
@@ -34,6 +35,7 @@
 )
 from nemo_gym.reward_profile import AggregateMetricsMixin, compute_aggregate_metrics
 from nemo_gym.rollout_correlation import RolloutContextMiddleware
+from nemo_gym.sandbox.access import SandboxAccess
 from nemo_gym.server_utils import BaseRunServerInstanceConfig, BaseServer, SimpleServer
 from nemo_gym.telemetry.endpoints import traced_verify_endpoint
 
@@ -67,7 +69,7 @@ def normalize_tool_name(name: str, server_name: Optional[str] = None) -> str:
 
 
 # Tool names that would collide with the resources server's own endpoints if advertised over MCP.
-RESERVED_MCP_TOOL_NAMES = frozenset({"verify", "seed_session", "aggregate_metrics", "mcp"})
+RESERVED_MCP_TOOL_NAMES = frozenset({"verify", "seed_session", "close_session", "aggregate_metrics", "mcp"})
 
 
 class ReverifyMode(str, Enum):
@@ -179,6 +181,62 @@ class MCPServerMetadata(BaseModel):
     headers: dict[str, str]
 
 
+class ResourcesSeedSessionRequest(BaseModel):
+    """Initialize resources-server state for one episode."""
+
+    model_config = ConfigDict(extra="forbid")
+
+    episode_id: EpisodeId
+    task_id: TaskId
+    task_data: dict[str, JsonValue]
+
+
+class ResourcesSeedSessionResponse(BaseModel):
+    """Return resources state and optional agent access."""
+
+    model_config = ConfigDict(extra="forbid")
+
+    resources_session_id: str
+    resources_tools: MCPServerMetadata | None = None
+    sandbox_access: SandboxAccess | None = None
+
+
+VerificationInputT = TypeVar("VerificationInputT")
+
+
+class ResourcesVerifyRequest(BaseModel, Generic[VerificationInputT]):
+    """Carry typed environment output to a resources server."""
+
+    model_config = ConfigDict(extra="forbid")
+
+    episode_id: EpisodeId
+    task_id: TaskId
+    verification_input: VerificationInputT
+
+
+class ResourcesVerifyResponse(BaseVerifyResponse):
+    """Preserve environment-specific verification fields across the server boundary."""
+
+    model_config = ConfigDict(extra="allow")
+
+
+class ResourcesCloseSessionRequest(BaseModel):
+    """Close resources-server state."""
+
+    model_config = ConfigDict(extra="forbid")
+
+    resources_session_id: str
+    episode_id: EpisodeId
+
+
+class ResourcesCloseSessionResponse(BaseModel):
+    """Confirm resources-server state was closed."""
+
+    model_config = ConfigDict(extra="forbid")
+
+    resources_session_id: str
+
+
 class SimpleResourcesServer(BaseResourcesServer, AggregateMetricsMixin, SimpleServer):
     config: BaseResourcesServerConfig
 

tests/unit_tests/test_resources_sessions.py

@@ -0,0 +1,33 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from nemo_gym.base_resources_server import BaseSeedSessionResponse, BaseVerifyResponse, ResourcesVerifyResponse
+from nemo_gym.openai_utils import NeMoGymResponse
+
+
+def test_empty_seed_response_preserves_legacy_wire_shape() -> None:
+    assert BaseSeedSessionResponse().model_dump() == {}
+
+
+def test_base_verify_response_preserves_legacy_extra_field_behavior() -> None:
+    response = BaseVerifyResponse.model_validate(
+        {
+            "responses_create_params": {"input": "hi"},
+            "response": NeMoGymResponse.model_construct(id="response", output=[]),
+            "reward": 0.5,
+            "request_only_field": "ignored",
+        }
+    )
+    assert "request_only_field" not in response.model_dump()
+
+
+def test_resources_verify_response_preserves_benchmark_fields() -> None:
+    response = ResourcesVerifyResponse.model_validate(
+        {
+            "responses_create_params": {"input": "hi"},
+            "response": NeMoGymResponse.model_construct(id="response", output=[]),
+            "reward": 0.5,
+            "benchmark_diagnostic": "kept",
+        }
+    )
+    assert response.model_dump()["benchmark_diagnostic"] == "kept"