feat(data): validate rows against their environment server's task schema

environment_servers/legacy_agent/task_data.py

@@ -0,0 +1,11 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""Task fields of a ``legacy_agent`` relay: a single-agent run request for the agent's ``/run``.
+
+The agent's resources server, or the agent itself when it is self-contained, declares the other fields.
+"""
+
+from nemo_gym.task_data import SingleAgentTaskData as TaskData
+
+
+__all__ = ["TaskData"]

environment_servers/nemo_user_sim/task_data.py

@@ -0,0 +1,12 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""Task fields that the ``nemo_user_sim`` environment server reads: its resources server's ``TaskData``.
+
+The environment server runs the user and assistant agents itself and reads the same task fields its
+resources server declares, so it shares that schema instead of keeping a copy.
+"""
+
+from resources_servers.nemo_user_sim.task_data import TaskData
+
+
+__all__ = ["TaskData"]

environment_servers/single_agent_turn/task_data.py

@@ -0,0 +1,12 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""Task fields that ``single_agent_turn`` reads itself.
+
+It sends ``responses_create_params`` to its agent and passes every other field to its resources
+server, whose ``TaskData`` declares them. ``SingleAgentTurnTaskInput`` is the runtime model.
+"""
+
+from nemo_gym.task_data import SingleAgentTaskData as TaskData
+
+
+__all__ = ["TaskData"]

environment_servers/single_agent_turn_legacy/task_data.py

@@ -0,0 +1,8 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""Task fields that ``single_agent_turn_legacy`` reads itself: the same as ``single_agent_turn``."""
+
+from nemo_gym.task_data import SingleAgentTaskData as TaskData
+
+
+__all__ = ["TaskData"]

fern/versions/latest/pages/environment-tutorials/resources-server-apis.mdx

@@ -459,6 +459,24 @@ The source JSONL stays unchanged and can still be shared with flat-row recipes.
 Omit `taskset` to keep the current flat-row format. This option accepts flat input
 only; already materialized input should be passed directly with `--input`.
 
+Collation validates each source row against two `task_data.py` schemas. The first
+belongs to the Environment Server the dataset routes to, under
+`environment_servers/<implementation>/task_data.py`, and declares the fields that
+server reads itself. The second belongs to the resources server that server binds,
+and declares the remaining task fields. Each schema validates only the fields it
+declares and the fields neither declares, so a resources schema with
+`extra="forbid"` does not need to list the Environment Server's fields. Fields that
+neither schema declares are reported against the resources server's schema.
+
+The built-in single-agent Environment Servers declare `responses_create_params`,
+one agent's request to one model, with `SingleAgentTaskData` from
+`nemo_gym/task_data.py`. Datasets without `taskset` route by agent, so they are
+single-agent run requests and use the same schema. An Environment Server that runs
+several agents declares its own fields instead, and its rows need no
+`responses_create_params`. Rows that a preparation script already materialized
+name their own taskset, so collation checks them only against the resources
+server's schema, and their Environment Server validates them at dispatch.
+
 For example, add these keys to a config whose `swebench_pro_hermes` Environment
 Server binds the `swebench_pro_hermes_resources_server` instance:
 

nemo_gym/global_config.py

@@ -1977,6 +1977,36 @@ def taskset_environment_server_name(global_config_dict: DictConfig, taskset: Opt
     return name if isinstance(name, str) else None
 
 
+@dataclass(frozen=True)
+class TasksetEnvironmentServer:
+    """The Environment Server a taskset routes to."""
+
+    name: str
+    implementation: str
+    config: DictConfig
+
+
+def taskset_environment_server(global_config_dict: DictConfig, taskset: str) -> TasksetEnvironmentServer:
+    """Resolve a taskset through ``environment_server_routes`` to its Environment Server instance.
+
+    Raises ``ConfigError`` when the taskset has no route or the route names no Environment Server.
+    """
+    routes = global_config_dict.get(ENVIRONMENT_SERVER_ROUTES_KEY_NAME)
+    environment_name = routes.get(taskset) if isinstance(routes, DictConfig) else None
+    if not isinstance(environment_name, str) or not environment_name:
+        raise ConfigError(f"No Environment Server route for taskset {taskset!r}; set environment_server_routes.")
+    environment = global_config_dict.get(environment_name)
+    servers = environment.get(ENVIRONMENT_SERVER_TYPE_KEY_NAME) if isinstance(environment, DictConfig) else None
+    if not isinstance(servers, DictConfig) or len(servers) != 1:
+        raise ConfigError(f"Taskset {taskset!r} route {environment_name!r} must name an Environment Server.")
+    implementation, environment_config = next(iter(servers.items()))
+    if not isinstance(environment_config, DictConfig):
+        raise ConfigError(f"Environment Server {environment_name!r} must have a configuration mapping.")
+    return TasksetEnvironmentServer(
+        name=environment_name, implementation=str(implementation), config=environment_config
+    )
+
+
 def resolve_dataset_agent(
     global_config_dict: DictConfig,
     declaring_instance_name: str,
@@ -2006,17 +2036,8 @@ def resolve_dataset_agent(
     is_agent = isinstance(block, DictConfig) and "responses_api_agents" in block
 
     if taskset is not None:
-        routes = global_config_dict.get(ENVIRONMENT_SERVER_ROUTES_KEY_NAME)
-        environment_name = routes.get(taskset) if isinstance(routes, DictConfig) else None
-        if not isinstance(environment_name, str) or not environment_name:
-            raise ConfigError(f"No Environment Server route for taskset {taskset!r}; set environment_server_routes.")
-        environment = global_config_dict.get(environment_name)
-        servers = environment.get(ENVIRONMENT_SERVER_TYPE_KEY_NAME) if isinstance(environment, DictConfig) else None
-        if not isinstance(servers, DictConfig) or len(servers) != 1:
-            raise ConfigError(f"Taskset {taskset!r} route {environment_name!r} must name an Environment Server.")
-        environment_config = next(iter(servers.values()))
-        if not isinstance(environment_config, DictConfig):
-            raise ConfigError(f"Environment Server {environment_name!r} must have a configuration mapping.")
+        environment = taskset_environment_server(global_config_dict, taskset)
+        environment_name, environment_config = environment.name, environment.config
         agent_names = []
         for agent_ref in environment_server_agent_refs(environment_config):
             name = agent_ref.get("name")

nemo_gym/task_data.py

@@ -8,6 +8,12 @@
 A self-contained agent (one that declares datasets but no ``resources_server`` reference) ships
 the same module under ``responses_api_agents/<implementation>/`` and owns its rows' schema; an
 agent WITH a reference uses that resources server's schema instead.
+An environment server ships the same module under ``environment_servers/<implementation>/`` to
+declare the task fields it reads itself. Rows are validated against the environment server they
+route to, composed with the bound resources server's schema: each schema validates the fields it
+declares plus the fields neither declares. Single-agent environment servers declare
+``responses_create_params``, one agent's request to one model, through ``SingleAgentTaskData``.
+Rows routed by agent rather than by taskset are single-agent run requests and use it too.
 It describes the task-owned fields of that server's dataset rows, written FLAT in the planned
 end-state shape: the fields as they will appear inside the unified ``task_data`` row key after
 the row-format migration. Framework-owned keys (see ``RESERVED_ROW_KEYS``) are never part of
@@ -45,9 +51,9 @@
 import sys
 from dataclasses import dataclass, field
 from pathlib import Path
-from typing import Any, Dict, Iterable, List, Mapping, Optional
+from typing import Any, Dict, Iterable, List, Mapping, Optional, Union, get_args
 
-from pydantic import BaseModel, TypeAdapter
+from pydantic import BaseModel, ConfigDict, Field, TypeAdapter
 
 from nemo_gym.episode_types import is_materialized_task_row
 
@@ -56,12 +62,13 @@
 TASK_DATA_EXPORT_NAME = "TaskData"
 
 # Framework-owned top-level row keys. Everything else in a row is task-owned and is the subject of
-# the server's TaskData schema. (The row-format migration will extend this set with reserved keys such as
+# the environment server's and resources server's TaskData schemas. ``responses_create_params`` is
+# not one of them: it is a task field of single-agent environment servers (``SingleAgentTaskData``).
+# (The row-format migration will extend this set with reserved keys such as
 # ``task_id``/``subset_for_metrics``/``provenance``; until then some servers legitimately use those
 # names as task fields, so they stay task-owned here.)
 RESERVED_ROW_KEYS = frozenset(
     {
-        "responses_create_params",
         "agent_ref",
         "task_source",
         "_ng_task_index",
@@ -70,6 +77,32 @@
 )
 
 
+class SingleAgentResponsesCreateParams(BaseModel):
+    """The part of a Responses API request that a dataset row supplies.
+
+    ``input`` may be missing from a source row whose dataset declares a ``prompt_config``, which
+    fills it during collation. The environment server validates the full request at runtime; this
+    keeps schemas free of server dependencies.
+    """
+
+    model_config = ConfigDict(extra="allow")
+
+    input: Optional[Union[str, List[Any]]] = None
+
+
+class SingleAgentTaskData(BaseModel):
+    """Task fields that every single-agent environment server reads: one agent's request to one model.
+
+    The bound resources server's ``TaskData`` declares the remaining fields.
+    """
+
+    model_config = ConfigDict(extra="allow")
+
+    responses_create_params: SingleAgentResponsesCreateParams = Field(
+        json_schema_extra={"consumed_by": ["prompt"]},
+    )
+
+
 class TaskDataSchemaError(Exception):
     """A server's ``task_data.py`` exists but does not satisfy the protocol."""
 
@@ -197,9 +230,8 @@ def summary(self) -> str:
         return "\n".join(parts)
 
 
-def legacy_metadata_fields(adapter: TypeAdapter) -> frozenset:
-    """Schema fields annotated ``legacy_location: verifier_metadata`` (today's wire reads them there)."""
-    from typing import get_args
+def _schema_models(adapter: TypeAdapter) -> List[type[BaseModel]]:
+    """The models a schema is built from: the model itself, or each member of a union."""
 
     def models_of(tp, out):
         if isinstance(tp, type) and issubclass(tp, BaseModel):
@@ -209,25 +241,94 @@ def models_of(tp, out):
             models_of(arg, out)
         return out
 
+    return models_of(getattr(adapter, "_type", None), [])
+
+
+def declared_fields(adapter: TypeAdapter) -> frozenset:
+    """The field names, and their aliases, that any model of a schema declares."""
+    names = set()
+    for model in _schema_models(adapter):
+        for field_name, info in model.model_fields.items():
+            names.add(field_name)
+            if isinstance(info.alias, str):
+                names.add(info.alias)
+    return frozenset(names)
+
+
+def legacy_metadata_fields(adapter: TypeAdapter) -> frozenset:
+    """Schema fields annotated ``legacy_location: verifier_metadata`` (today's wire reads them there)."""
     names = set()
-    for model in models_of(getattr(adapter, "_type", None), []):
+    for model in _schema_models(adapter):
         for field_name, info in model.model_fields.items():
             extra = info.json_schema_extra
             if isinstance(extra, dict) and extra.get("legacy_location") == LEGACY_METADATA_KEY:
                 names.add(field_name)
     return frozenset(names)
 
 
+SINGLE_AGENT_TASK_DATA = TypeAdapter(SingleAgentTaskData)
+
+
 class TaskDataValidator:
-    """Validates dataset rows against a server's ``TaskData`` schema, accumulating a report."""
+    """Validates dataset rows against an environment server's and a resources server's ``TaskData``.
+
+    ``environment_adapter`` is the schema of the environment server the rows route to, such as
+    ``SINGLE_AGENT_TASK_DATA`` for rows routed by agent. Each schema validates the fields it
+    declares plus the fields neither declares, so a resources schema with ``extra="forbid"`` does not
+    see the environment server's fields. Undeclared fields are reported against the resources server's
+    schema. Either adapter may be ``None`` when its server ships no schema. The environment schema does
+    not validate materialized rows (see ``validate_row``).
+    """
 
-    def __init__(self, server_name: str, adapter: TypeAdapter, dataset_fpath: str):
+    def __init__(
+        self,
+        server_name: str,
+        adapter: Optional[TypeAdapter],
+        dataset_fpath: str,
+        *,
+        environment_adapter: Optional[TypeAdapter] = None,
+    ):
         self._adapter = adapter
-        self._legacy_fields = legacy_metadata_fields(adapter)
+        self._environment_adapter = environment_adapter
+        self._fields = declared_fields(adapter) if adapter is not None else frozenset()
+        self._environment_fields = (
+            declared_fields(environment_adapter) if environment_adapter is not None else frozenset()
+        )
+        self._legacy_fields = legacy_metadata_fields(adapter) if adapter is not None else frozenset()
         self.report = TaskDataValidationReport(server_name=server_name, dataset_fpath=dataset_fpath)
 
+    def _validate_fields(self, subject: Dict[str, Any], *, environment_applies: bool) -> tuple[List[str], set]:
+        """Validate one row's task fields against each schema; return its errors and undeclared fields."""
+        errors: List[str] = []
+        undeclared: set = set()
+        # The environment server's fields stay out of the resources schema's subject either way, so
+        # they are never reported as fields the resources server does not declare.
+        environment_adapter = self._environment_adapter if environment_applies else None
+        for adapter, own, other, reports_unknown in (
+            (environment_adapter, self._environment_fields, self._fields, False),
+            (self._adapter, self._fields, self._environment_fields, True),
+        ):
+            if adapter is None:
+                continue
+            fields = {key: value for key, value in subject.items() if key in own or key not in other}
+            try:
+                validated = adapter.validate_python(fields)
+            except Exception as e:
+                errors.append(str(e).strip().replace("\n", "; "))
+                continue
+            # Pydantic returns the concrete model (the selected union member for union schemas), and
+            # with extra="allow" it stores undeclared inputs on __pydantic_extra__ — so unknown-field
+            # reporting works uniformly for plain models and discriminated unions.
+            if reports_unknown and isinstance(validated, BaseModel):
+                undeclared = set(getattr(validated, "__pydantic_extra__", None) or {})
+        return errors, undeclared
+
     def validate_row(self, row_index: int, row: Dict[str, Any]) -> None:
         self.report.rows += 1
+        # A materialized row already names its taskset, and the environment server that taskset routes
+        # to validates its input at dispatch. The environment schema does not validate it here, so such a
+        # row in a dataset routed by agent is not required to be a single-agent run request.
+        environment_applies = not is_materialized_task_row(row)
         # New materialized tasks keep the flat source fields; historical single-agent
         # tasks use task_input.task_data. Normalize both through the existing schema path.
         task_input = row.get("task_input")
@@ -270,27 +371,31 @@ def validate_row(self, row_index: int, row: Dict[str, Any]) -> None:
         subject, conflicts = normalize_task_fields(row)
         for key in conflicts:
             self.report.conflicting_keys[key] = self.report.conflicting_keys.get(key, 0) + 1
-        try:
-            validated = self._adapter.validate_python(subject)
-        except Exception as e:
+        errors, undeclared = self._validate_fields(subject, environment_applies=environment_applies)
+        if errors:
             self.report.error_rows += 1
             if len(self.report.errors) < TaskDataValidationReport.MAX_RECORDED_ERRORS:
-                first_line = str(e).strip().replace("\n", "; ")
-                self.report.errors.append(f"{row_index}: {first_line[:400]}")
+                self.report.errors.append(f"{row_index}: {'; '.join(error[:400] for error in errors)}")
             return
-        # Pydantic returns the concrete model (the selected union member for union schemas), and
-        # with extra="allow" it stores undeclared inputs on __pydantic_extra__ — so unknown-field
-        # reporting works uniformly for plain models and discriminated unions.
-        if isinstance(validated, BaseModel):
-            for key in getattr(validated, "__pydantic_extra__", None) or {}:
-                self.report.unknown_keys[key] = self.report.unknown_keys.get(key, 0) + 1
+        for key in undeclared:
+            self.report.unknown_keys[key] = self.report.unknown_keys.get(key, 0) + 1
 
 
 def validate_jsonl_rows(
-    server_name: str, adapter: TypeAdapter, dataset_fpath: str, lines: Iterable[str]
+    server_name: str,
+    adapter: Optional[TypeAdapter],
+    dataset_fpath: str,
+    lines: Iterable[str],
+    *,
+    environment_adapter: Optional[TypeAdapter] = None,
 ) -> TaskDataValidationReport:
     """Validate an iterable of JSONL lines against one schema; entry point for whole-file validation."""
-    validator = TaskDataValidator(server_name=server_name, adapter=adapter, dataset_fpath=dataset_fpath)
+    validator = TaskDataValidator(
+        server_name=server_name,
+        adapter=adapter,
+        dataset_fpath=dataset_fpath,
+        environment_adapter=environment_adapter,
+    )
     for i, line in enumerate(lines):
         line = line.strip()
         if not line: