diff --git a/nemo_gym/base_resources_server.py b/nemo_gym/base_resources_server.py index 35f37a06675..3ae483b19c0 100644 --- a/nemo_gym/base_resources_server.py +++ b/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 diff --git a/tests/unit_tests/test_resources_sessions.py b/tests/unit_tests/test_resources_sessions.py new file mode 100644 index 00000000000..ec304d8a162 --- /dev/null +++ b/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"