diff --git a/README.md b/README.md index cd289d31e5..0c3108a763 100644 --- a/README.md +++ b/README.md @@ -293,6 +293,8 @@ The Dataset column links to publicly available datasets (e.g., on HuggingFace). | Rdkit Chemistry | knowledge | [DEPRECATED — use litmus_agent] Molecular chemistry question answering: calculate properties of SMILES. Includes a mix of tool-use (python + rdkit) and no-tool-use questions. Superseded by the domain-agnostic litmus_agent server, which generalizes this server's scoring path. | Improve molecular reasoning and SMILES parsing. | ✓ | - | TBD | rdkit_chemistry.yaml | - | | Reasoning Gym | knowledge | Claude Code agent harness for reasoning gym tasks | Evaluate model capabilities in the Claude Code agent harness | ✓ | - | Creative Commons Attribution 4.0 International | reasoning_gym_claude_code_agent.yaml | Nemotron-RL-ReasoningGym-v1 | | Reasoning Gym | knowledge | Claude Code agent harness for reasoning gym tasks, via a Gym model server's /v1/messages | Showcase Claude Code running against any Gym model backend | - | - | - | reasoning_gym_claude_code_agent_model_server.yaml | - | +| Reasoning Gym | knowledge | Codex agent harness for reasoning gym tasks | Evaluate model capabilities in the Codex agent harness | - | - | - | reasoning_gym_codex_agent.yaml | - | +| Reasoning Gym | knowledge | Codex agent harness for reasoning gym tasks, via a Gym model server's /v1/responses | Showcase Codex running against any Gym model backend | - | - | - | reasoning_gym_codex_agent_model_server.yaml | - | | Reasoning Gym | knowledge | LangGraph orchestrator agent compatible with resource servers that do not use tools; enables diverse agent training data and test time scaling vs a simple agent, extensible to use tools or other agent architectures | Iterative test time scaling for improved performance in reasoning tasks | ✓ | - | Apache 2.0 | orchestrator_agent.yaml | - | | Reasoning Gym | knowledge | LangGraph parallel thinking agent compatible with resource servers that do not use tools; enables diverse agent training data and test time scaling vs a simple agent, extensible to use tools or other agent architectures | Iterative test time scaling for improved performance in reasoning tasks | ✓ | - | Apache 2.0 | parallel_thinking_agent.yaml | - | | Reasoning Gym | knowledge | LangGraph reflection agent compatible with resource servers that do not use tools; provides iterative reflection for diverse agent training data and test time scaling, extensible to use tools or other agent architectures | Iterative test time scaling for improved performance in reasoning tasks | ✓ | - | Apache 2.0 | reflection_agent.yaml | - | diff --git a/fern/versions/latest/pages/agent-server/agent-skills.mdx b/fern/versions/latest/pages/agent-server/agent-skills.mdx index 01a34fc7d9..8431643002 100644 --- a/fern/versions/latest/pages/agent-server/agent-skills.mdx +++ b/fern/versions/latest/pages/agent-server/agent-skills.mdx @@ -84,7 +84,7 @@ To support skills in a new agent, implement the three-step adapter. NeMo Gym pro | Agent | Discovery location | |---|---| | Claude Code | `CLAUDE_CONFIG_DIR/skills/` | - | Codex CLI | the Codex config home (e.g. `$CODEX_HOME`) | + | Codex CLI | `CODEX_HOME/skills/` | | MCP-backed agent | configure the MCP skill server to serve from the directory | Prefer staging into a **per-request** ephemeral location so concurrent requests with different skills do not contaminate one another, and so nothing leaks between rollouts. @@ -94,7 +94,7 @@ To support skills in a new agent, implement the three-step adapter. NeMo Gym pro Note that enabling discovery is usually all-or-nothing: with Claude Code, dropping `--bare` so skills load also re-enables *all* other auto-discovery — hooks, plugins, MCP servers, memory, and `CLAUDE.md` — not just skills. If you are measuring a skill's isolated impact, be aware that a baseline (`--bare`, no skills) versus a skills run may differ by more than the skill content alone. -The reference implementation is the [Claude Code agent](https://github.com/NVIDIA-NeMo/Gym/tree/main/responses_api_agents/claude_code_agent). It stages skills into a fresh per-request `CLAUDE_CONFIG_DIR/skills/` and drops `--bare` when skills are active. +The reference implementation is the [Claude Code agent](https://github.com/NVIDIA-NeMo/Gym/tree/main/responses_api_agents/claude_code_agent). It stages skills into a fresh per-request `CLAUDE_CONFIG_DIR/skills/` and drops `--bare` when skills are active. The [Codex agent](https://github.com/NVIDIA-NeMo/Gym/tree/main/responses_api_agents/codex_agent) follows the same pattern, staging into a fresh per-request `CODEX_HOME/skills/` (Codex discovers them natively — no discovery flag needed). ## What NeMo Gym does not do diff --git a/nemo_gym/base_responses_api_model.py b/nemo_gym/base_responses_api_model.py index 402d28a437..7d29549ad9 100644 --- a/nemo_gym/base_responses_api_model.py +++ b/nemo_gym/base_responses_api_model.py @@ -42,8 +42,9 @@ import orjson from fastapi import Body, FastAPI, Request +from fastapi.exceptions import RequestValidationError from fastapi.responses import StreamingResponse -from pydantic import BaseModel, Field, model_validator +from pydantic import BaseModel, Field, ValidationError, model_validator from nemo_gym.anthropic_converter import AnthropicConverter from nemo_gym.config_types import ROLLOUT_PATH_PREFIX, ModelServerRef @@ -58,6 +59,12 @@ NeMoGymResponse, NeMoGymResponseCreateParamsNonStreaming, ) +from nemo_gym.responses_streaming import ( + sanitize_streaming_responses_body, + synthesize_responses_failure_sse, + synthesize_responses_sse, + validate_streaming_responses_params, +) from nemo_gym.server_utils import ( BaseRunServerInstanceConfig, BaseServer, @@ -90,7 +97,7 @@ def setup_webserver(self) -> FastAPI: app.post("/v1/chat/completions")(self.chat_completions) - app.post("/v1/responses")(self.responses) + app.post("/v1/responses")(self.responses_dispatch) # Every Gym model server speaks the Anthropic Messages API by default, mapping # Messages <-> Responses around its own responses() implementation. This lets blackbox @@ -110,6 +117,46 @@ async def chat_completions( async def responses(self, body: NeMoGymResponseCreateParamsNonStreaming = Body()) -> NeMoGymResponse: pass + async def responses_dispatch(self, request: Request, body: dict = Body()): + """Default ``/v1/responses`` entrypoint shared by every Gym model server. + + A plain JSON request validates strictly against + ``NeMoGymResponseCreateParamsNonStreaming`` and delegates to this server's own + ``responses()``, preserving the historical non-streaming behavior. When the client + requests ``stream: true`` (blackbox Responses-over-SSE harnesses like the Codex CLI + always do), the request is first sanitized from the streaming wire dialect (extra + bookkeeping fields, ``namespace`` tool specs — see ``nemo_gym.responses_streaming``), + delegated to the same ``responses()``, and the complete response is re-emitted as a + synthesized Responses SSE event stream. A ``responses()`` failure on this path is turned + into a terminal ``response.failed`` event rather than an HTTP 500 (bad-request validation + still fails eagerly, before the stream is committed). + """ + if not body.get("stream"): + params = _validate_responses_params(body) + return await self._invoke_responses(request, params) + + cleaned, ns_map = sanitize_streaming_responses_body(body) + try: + params = validate_streaming_responses_params(cleaned) + except ValidationError as exc: + raise RequestValidationError([{**error, "loc": ("body", *error["loc"])} for error in exc.errors()]) + + try: + response = await self._invoke_responses(request, params) + response_json = response.model_dump(mode="json") if isinstance(response, BaseModel) else dict(response) + except Exception as exc: + # The streaming contract is already the response's shape, so a backend failure must be a + # terminal response.failed event, not an HTTP 500 the client would see as a broken stream. + logger.exception("responses() failed while serving a streaming /v1/responses request") + return StreamingResponse( + synthesize_responses_failure_sse(str(exc)), + media_type="text/event-stream", + ) + return StreamingResponse( + synthesize_responses_sse(response_json, ns_map), + media_type="text/event-stream", + ) + async def messages(self, request: Request, body: dict = Body()): """Default Anthropic Messages <-> Responses mapping shared by every Gym model server. @@ -141,6 +188,14 @@ async def _invoke_responses( return await self.responses(body=params) +def _validate_responses_params(body: dict) -> NeMoGymResponseCreateParamsNonStreaming: + """Validate a /v1/responses body dict, surfacing failures as FastAPI's standard 422.""" + try: + return NeMoGymResponseCreateParamsNonStreaming.model_validate(body) + except ValidationError as exc: + raise RequestValidationError([{**error, "loc": ("body", *error["loc"])} for error in exc.errors()]) + + # --- Capture configuration + rollout-keyed storage --- diff --git a/nemo_gym/responses_converter.py b/nemo_gym/responses_converter.py index 638690177d..a515e5157e 100644 --- a/nemo_gym/responses_converter.py +++ b/nemo_gym/responses_converter.py @@ -60,6 +60,13 @@ ) +def _message_content_to_text(content: Any) -> str: + """Plain text of a chat message content (a string, or a list of text parts).""" + if isinstance(content, str): + return content + return "".join(part.get("text", "") for part in content or [] if isinstance(part, dict)) + + class ResponsesConverterState(BaseModel): return_token_id_information: bool @@ -166,6 +173,22 @@ def responses_to_chat_completion_create_params( state.flush_assistant() + # The Responses API inserts `instructions` as a system message at the start of the model's + # context. Chat Completions has no such parameter, so map it explicitly — otherwise it is + # silently dropped when the remaining params are passed through (extra fields are ignored). + # The leading run of system/developer messages is folded into the same single system + # message: chat backends commonly admit only one system message, at position 0 (harnesses + # like the Codex CLI send instructions plus a leading developer message). + instructions = responses_create_params.pop("instructions", None) + if instructions: + leading_parts = [instructions] + while state.messages and state.messages[0]["role"] in ("system", "developer"): + leading_parts.append(_message_content_to_text(state.messages.pop(0)["content"])) + state.messages.insert( + 0, + NeMoGymChatCompletionSystemMessageParam(content="\n\n".join(leading_parts), role="system"), + ) + model = responses_create_params.pop("model", None) if model is not None: responses_create_params["model"] = model diff --git a/nemo_gym/responses_streaming.py b/nemo_gym/responses_streaming.py new file mode 100644 index 0000000000..0f9c1d9e09 --- /dev/null +++ b/nemo_gym/responses_streaming.py @@ -0,0 +1,282 @@ +# 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. +"""Streaming Responses-dialect support shared by every Gym model server. + +Blackbox harnesses that speak the OpenAI Responses API over SSE (e.g. the Codex CLI) send +requests the strict ``NeMoGymResponseCreateParamsNonStreaming`` model would reject: a +``stream: true`` flag, client-bookkeeping fields (``client_metadata``, ``prompt_cache_key``), +and ``namespace`` tool specs — functions grouped under a namespace that the model calls back +with separate ``namespace`` + ``name`` fields on the ``function_call`` item. Backends behind a +Gym model server (chat-completions endpoints, vLLM) only understand flat function tools, so +this module provides: + +- the request-side sanitizer that flattens namespace tools into plain functions (joined as + ``__``), rewrites replayed namespaced calls in the input history the same + way, and drops the fields the params model does not know; +- the response-side synthesizer that re-emits a complete ``NeMoGymResponse`` as the minimal + Responses SSE event sequence streaming clients require (``response.created`` -> + ``response.output_item.done`` per output item -> ``response.completed``), splitting flattened + function-call names back into ``namespace`` + ``name`` on the way out. +""" + +import json +import logging +from copy import deepcopy +from typing import Any, Iterator, Optional +from uuid import uuid4 + +from openai.types.responses.response_create_params import ToolParam +from pydantic import TypeAdapter, ValidationError + +from nemo_gym.openai_utils import NeMoGymResponseCreateParamsNonStreaming, NeMoGymResponseInputItem + + +LOG = logging.getLogger(__name__) + +NAMESPACE_TOOL_DELIMITER = "__" + +_PARAM_FIELDS = frozenset(NeMoGymResponseCreateParamsNonStreaming.model_fields) +_TOOL_ADAPTER = TypeAdapter(ToolParam) +_INPUT_ITEM_ADAPTER = TypeAdapter(NeMoGymResponseInputItem) + +# ns_map: flattened function name -> (namespace, name) for the round trip back to the client. +NamespaceMap = dict[str, tuple[str, str]] + + +def flatten_namespace_tools(tools: Any) -> tuple[list[Any], NamespaceMap]: + """Flatten ``namespace`` tool specs into plain function tools with joined names.""" + flat: list[Any] = [] + ns_map: NamespaceMap = {} + for tool in tools or []: + if not (isinstance(tool, dict) and tool.get("type") == "namespace"): + flat.append(tool) + continue + namespace = str(tool.get("name") or "") + for sub in tool.get("tools") or []: + if not isinstance(sub, dict) or not sub.get("name"): + continue + joined = f"{namespace}{NAMESPACE_TOOL_DELIMITER}{sub['name']}" + flat.append({**sub, "type": "function", "name": joined}) + ns_map[joined] = (namespace, str(sub["name"])) + return flat, ns_map + + +def _tool_valid(tool: Any) -> bool: + """Whether a tool spec validates against the params model's tool union; failures are logged.""" + try: + _TOOL_ADAPTER.validate_python(tool) + return True + except ValidationError: + LOG.warning( + "Dropping unsupported tool spec of type %r from a streaming /v1/responses request.", + tool.get("type") if isinstance(tool, dict) else type(tool).__name__, + ) + return False + + +def _is_system_like_message(item: Any) -> bool: + return ( + isinstance(item, dict) + and item.get("type") in (None, "message") + and item.get("role") in ("system", "developer") + ) + + +def _input_message_text(item: dict[str, Any]) -> str: + content = item.get("content") + if isinstance(content, str): + return content + parts = [] + for part in content or []: + if isinstance(part, dict) and part.get("type") in ("input_text", "output_text", "text"): + parts.append(part.get("text") or "") + return "".join(parts) + + +def sanitize_streaming_responses_body(body: dict[str, Any]) -> tuple[dict[str, Any], NamespaceMap]: + """Map a streaming-dialect request body onto the strict non-streaming params shape. + + Returns the cleaned body dict (ready for ``NeMoGymResponseCreateParamsNonStreaming`` + validation) and the namespace map needed to restore namespaced call names in the + synthesized SSE response. Tool entries that still fail per-entry validation after + flattening are dropped with a warning rather than failing the whole request, since a + harness's exotic tool is better lost than the rollout. + """ + body = deepcopy(body) + # The params model only admits `stream: false` (Gym responses are non-streaming internally); + # the streaming envelope is synthesized by the caller, so the flag is dropped here. + body.pop("stream", None) + + ns_map: NamespaceMap = {} + if "tools" in body: + tools, ns_map = flatten_namespace_tools(body.get("tools")) + body["tools"] = [tool for tool in tools if _tool_valid(tool)] + + # Replayed history: a namespaced call the client echoes back must match the flattened tool + # name the model actually saw, so the conversation stays self-consistent for the backend. + input_items = body.get("input") + if isinstance(input_items, list): + kept_items = [] + carrier_tools: list[Any] = [] + for item in input_items: + if isinstance(item, dict) and item.get("type") == "function_call" and item.get("namespace"): + item["name"] = f"{item.pop('namespace')}{NAMESPACE_TOOL_DELIMITER}{item.get('name')}" + # Codex's code mode ships tools inside an `additional_tools` input item instead of the + # `tools` param. Hoist what a Gym backend can express (plain and namespaced function + # tools) into `tools`; the rest (e.g. the freeform JS `exec` tool) has no function-call + # representation and is dropped with a warning. + if isinstance(item, dict) and item.get("type") == "additional_tools": + carrier_tools.extend(item.get("tools") or []) + continue + # Streaming harnesses interleave item kinds the Gym input union has no representation + # for. Drop those individually -- a lost carrier item is recoverable, a 422 on the + # whole request kills the rollout. + try: + _INPUT_ITEM_ADAPTER.validate_python(item) + except ValidationError: + LOG.warning( + "Dropping unsupported input item of type %r from a streaming /v1/responses request.", + item.get("type") if isinstance(item, dict) else type(item).__name__, + ) + continue + kept_items.append(item) + + if carrier_tools: + flat, carrier_map = flatten_namespace_tools(carrier_tools) + ns_map.update(carrier_map) + usable = [t for t in flat if isinstance(t, dict) and t.get("type") == "function" and _tool_valid(t)] + skipped = [t.get("name") for t in flat if not (isinstance(t, dict) and t.get("type") == "function")] + if skipped: + LOG.warning( + "Dropping non-function tools %s from an additional_tools item; a Gym backend can only " + "express function calls. Configure the harness with an explicit unknown model name so " + "it advertises classic function tools instead.", + skipped, + ) + if usable: + body["tools"] = [*(body.get("tools") or []), *usable] + + # Hoist the leading run of system/developer messages into `instructions` (prepended by any + # instructions already present). Streaming harnesses may open with several developer + # messages and no instructions (Codex does, for newer model families); downstream + # Responses -> Chat conversion renders `instructions` as the single leading system message + # strict chat backends require. + leading_parts = [body.get("instructions") or ""] + while kept_items and _is_system_like_message(kept_items[0]): + leading_parts.append(_input_message_text(kept_items.pop(0))) + hoisted = "\n\n".join(part for part in leading_parts if part) + if hoisted: + body["instructions"] = hoisted + + body["input"] = kept_items + + dropped = sorted(set(body) - _PARAM_FIELDS) + if dropped: + LOG.debug("Dropping unsupported fields from a streaming /v1/responses request: %s", dropped) + return {key: value for key, value in body.items() if key in _PARAM_FIELDS}, ns_map + + +def _delete_loc(body: Any, loc: tuple) -> bool: + """Delete the value at a pydantic error loc from a nested dict/list structure. + + Returns False when the loc cannot be walked literally (e.g. it contains a union-arm label + rather than a real key), in which case nothing is deleted. + """ + node = body + for part in loc[:-1]: + if isinstance(node, dict) and part in node: + node = node[part] + elif isinstance(node, list) and isinstance(part, int) and part < len(node): + node = node[part] + else: + return False + last = loc[-1] + if isinstance(node, dict) and last in node: + del node[last] + return True + return False + + +def validate_streaming_responses_params(body: dict[str, Any]) -> NeMoGymResponseCreateParamsNonStreaming: + """Validate a sanitized streaming-dialect body, pruning fields newer than the params model. + + Harness wire formats evolve faster than the pinned OpenAI SDK types (e.g. Codex sending + ``reasoning.context``, which the SDK's ``Reasoning`` model forbids). Any field pydantic flags + as ``extra_forbidden`` is removed and validation retried, so only errors that cannot be fixed + by dropping an unknown field surface to the client. + """ + body = deepcopy(body) + while True: + try: + return NeMoGymResponseCreateParamsNonStreaming.model_validate(body) + except ValidationError as exc: + removed = False + for error in exc.errors(): + if error["type"] == "extra_forbidden" and _delete_loc(body, tuple(error["loc"])): + LOG.warning( + "Dropping unsupported field %s from a streaming /v1/responses request.", + ".".join(str(part) for part in error["loc"]), + ) + removed = True + if not removed: + raise + + +def _sse_event(payload: dict[str, Any]) -> str: + return f"event: {payload['type']}\ndata: {json.dumps(payload)}\n\n" + + +def synthesize_responses_sse(response_json: dict[str, Any], ns_map: Optional[NamespaceMap] = None) -> Iterator[str]: + """Re-emit a complete Responses API response object as an SSE event stream. + + Streaming clients build their view of the turn from ``response.output_item.done`` events and + treat ``response.completed`` (which carries the response id and usage) as the terminal event, + so those two are the required minimum; ``response.created`` is included for clients that wait + for an acknowledgement before reading items. + """ + output_items = [] + for item in response_json.get("output") or []: + if ns_map and isinstance(item, dict) and item.get("type") == "function_call" and item.get("name") in ns_map: + namespace, name = ns_map[item["name"]] + item = {**item, "namespace": namespace, "name": name} + output_items.append(item) + + yield _sse_event( + {"type": "response.created", "response": {**response_json, "status": "in_progress", "output": []}} + ) + for index, item in enumerate(output_items): + yield _sse_event({"type": "response.output_item.done", "output_index": index, "item": item}) + yield _sse_event({"type": "response.completed", "response": {**response_json, "output": output_items}}) + + +def synthesize_responses_failure_sse(message: str, *, code: str = "server_error") -> Iterator[str]: + """Emit a terminal Responses SSE stream for a backend failure. + + Once the streaming contract is committed (HTTP 200 + ``text/event-stream``), a ``responses()`` + error can no longer surface as an HTTP 500. Streaming clients (e.g. Codex) expect a terminal + ``response.failed`` event; without one they see a truncated stream and cannot tell an + application error from a transport failure. Emitting ``response.failed`` lets the client report + a clean turn failure, and lets model-call capture classify it as an upstream error (its + terminal-SSE table maps ``response.failed`` to an error) rather than a stream truncation. + """ + response = { + "id": f"resp_{uuid4().hex}", + "object": "response", + "status": "failed", + "output": [], + "error": {"code": code, "message": message}, + } + yield _sse_event({"type": "response.created", "response": {**response, "status": "in_progress"}}) + yield _sse_event({"type": "response.failed", "response": response}) diff --git a/resources_servers/reasoning_gym/configs/reasoning_gym_codex_agent.yaml b/resources_servers/reasoning_gym/configs/reasoning_gym_codex_agent.yaml new file mode 100644 index 0000000000..9881450aaa --- /dev/null +++ b/resources_servers/reasoning_gym/configs/reasoning_gym_codex_agent.yaml @@ -0,0 +1,33 @@ +reasoning_gym: + resources_servers: + reasoning_gym: + entrypoint: app.py + domain: knowledge + verified: false + description: Codex agent harness for reasoning gym tasks + value: Evaluate model capabilities in the Codex agent harness + +reasoning_gym_codex_agent: + responses_api_agents: + codex_agent: + entrypoint: app.py + resources_server: + type: resources_servers + name: reasoning_gym + concurrency: 32 + model: ${openai_model_name} + openai_api_key: ${openai_api_key} + # Must serve the OpenAI Responses API over SSE (include /v1; Codex appends /responses). + # null -> https://api.openai.com/v1. Chat-completions-only endpoints need the + # reasoning_gym_codex_agent_model_server config instead (Gym adapts the dialect). + openai_base_url: ${openai_base_url} + sandbox_mode: danger-full-access + timeout: 600 + codex_version: 0.144.4 + system_prompt: | + You are a precise reasoning assistant. You can run shell commands to use Python for calculations. + For every problem: think step by step, use code to verify when helpful, and state your final answer clearly. + datasets: + - name: example + type: example + jsonl_fpath: resources_servers/reasoning_gym/data/example.jsonl diff --git a/resources_servers/reasoning_gym/configs/reasoning_gym_codex_agent_model_server.yaml b/resources_servers/reasoning_gym/configs/reasoning_gym_codex_agent_model_server.yaml new file mode 100644 index 0000000000..8907f5f36d --- /dev/null +++ b/resources_servers/reasoning_gym/configs/reasoning_gym_codex_agent_model_server.yaml @@ -0,0 +1,48 @@ +# Showcase: codex_agent against a NeMo Gym model server's /v1/responses (not OpenAI). +# +# Every Gym model server serves the streaming Responses dialect Codex speaks (request sanitizing + +# synthesized SSE on SimpleResponsesAPIModel). The agent's `model_server` ref resolves the Codex +# model provider's base_url to the model server. Needs NO openai_* env vars. +# +# Compose with any model server as `policy_model`, e.g. a vLLM OpenAI-compatible endpoint: +# gym env start \ +# --resources-server reasoning_gym/reasoning_gym_codex_agent_model_server \ +# --model-type vllm_model +# gym eval run --no-serve \ +# --agent reasoning_gym_codex_agent_model_server \ +# --input resources_servers/reasoning_gym/data/example.jsonl \ +# --output codex_via_model_server_rollout.jsonl --limit 1 + +reasoning_gym: + resources_servers: + reasoning_gym: + entrypoint: app.py + domain: knowledge + verified: false + description: Codex agent harness for reasoning gym tasks, via a Gym model server's /v1/responses + value: Showcase Codex running against any Gym model backend + +reasoning_gym_codex_agent_model_server: + responses_api_agents: + codex_agent: + entrypoint: app.py + resources_server: + type: resources_servers + name: reasoning_gym + model_server: + type: responses_api_models + name: policy_model + concurrency: 32 + model: ${policy_model_name} + openai_api_key: EMPTY # pragma: allowlist secret + openai_base_url: null + sandbox_mode: danger-full-access + timeout: 600 + codex_version: 0.144.4 + system_prompt: | + You are a precise reasoning assistant. You can run shell commands to use Python for calculations. + For every problem: think step by step, use code to verify when helpful, and state your final answer clearly. + datasets: + - name: example + type: example + jsonl_fpath: resources_servers/reasoning_gym/data/example.jsonl diff --git a/responses_api_agents/codex_agent/README.md b/responses_api_agents/codex_agent/README.md new file mode 100644 index 0000000000..6ccf227178 --- /dev/null +++ b/responses_api_agents/codex_agent/README.md @@ -0,0 +1,155 @@ +# Codex Agent + +Runs the OpenAI Codex CLI (`codex exec`) as a NeMo Gym agent server. + +## Quick start + +### env.yaml + +For the OpenAI API: + +```yaml +openai_api_key: sk-... +``` + +For any endpoint that serves the OpenAI Responses API over SSE: + +```yaml +openai_api_key: EMPTY +``` + +and set `openai_base_url` (must include `/v1`; Codex appends `/responses` itself). + +### Launch + +For a quick eval against OpenAI (or any Responses endpoint set via `openai_base_url`), pass the resources server config, which includes the agent server config: + +```bash +gym env start --resources-server reasoning_gym/reasoning_gym_codex_agent +``` + +#### Against a Gym model server + +Every Gym model server serves the streaming Responses dialect Codex speaks on `POST /v1/responses` (`SimpleResponsesAPIModel` sanitizes the request and synthesizes the SSE stream), so Codex can run against any backend Gym serves — vLLM, OpenAI, an inference provider. Set the agent's `model_server` ref to that server (it takes precedence over `openai_base_url`); the harness resolves the provider `base_url` to it. + +`reasoning_gym_codex_agent_model_server.yaml` wires the agent's `model_server` ref to `policy_model`. Compose it with any model server (here a vLLM serving `policy_model`): + +```bash +gym env start \ + --resources-server reasoning_gym/reasoning_gym_codex_agent_model_server \ + --model-type vllm_model +``` + +This path needs only the model server's `policy_base_url`, `policy_api_key`, and `policy_model_name` (in `env.yaml` or as `+` overrides) — no `openai_*` vars. + +### Run the agent + +```bash +gym eval run --no-serve \ + --agent reasoning_gym_codex_agent \ + --input resources_servers/reasoning_gym/data/example.jsonl \ + --output codex_rollout.jsonl \ + --limit 1 +``` + +For the model-server config above, use `--agent reasoning_gym_codex_agent_model_server`. + +### Smoke test + +Check the streaming `/v1/responses` dialect and the real-CLI seam without a full rollout. Launch a model server, then take its URL from the `gym env start` log (`'url': 'http://127.0.0.1:'`): + +```bash +gym env start --model-type vllm_model \ + +policy_base_url=https://integrate.api.nvidia.com/v1 \ + '+policy_api_key=${oc.env:NVIDIA_API_KEY}' +policy_model_name=meta/llama-3.1-8b-instruct + +# 1. the endpoint speaks the streaming Responses dialect: +curl -N $URL/v1/responses -H 'content-type: application/json' \ + -d '{"model":"x","stream":true,"input":[{"type":"message","role":"user","content":[{"type":"input_text","text":"2+2?"}]}]}' + +# 2. the real Codex CLI runs against it: +mkdir -p /tmp/codex_home && cat > /tmp/codex_home/config.toml <__` functions on the way in and splits the names back on the way out, so third-party models can call them. + +## Skills evaluation + +Skills are evaluated as a run-level variable, not a dataset field — point `skills.path` at a directory of [Agent Skills standard](https://agentskills.io/specification) skill directories on `gym eval run`, and the agent stages them into each request's `CODEX_HOME/skills/`, where Codex's native skill discovery picks them up: + +```bash +gym eval run --agent reasoning_gym_codex_agent \ + --input resources_servers/reasoning_gym/data/example.jsonl \ + --output rollouts_variant_a.jsonl \ + +skills.path=skills/variant_a/ +``` + +Each rollout result is stamped with a `skills_ref` for provenance and grouping during reward profiling, exactly as for the Claude Code agent (see its README for the full workflow). + +## Limitations + +- Eval only for now. Token IDs and logprobs are not wired up yet. +- Token counts come from Codex's own usage reporting (`turn.completed`). +- `turns_used` counts assistant messages right now, not tool calls. +- Codex has no `--max-turns` equivalent; runaway rollouts are bounded by `timeout`. +- Multi-turn dataset inputs are collapsed to a single prompt: only the first `system` message (as `developer_instructions`) and the last `user` message are passed to `codex exec`; any earlier user/assistant/tool turns in `responses_create_params.input` are dropped. This matches the Claude Code agent and is fine for single-turn datasets like reasoning_gym, but datasets that encode prior conversation turns in `input` will not see that history. diff --git a/responses_api_agents/codex_agent/__init__.py b/responses_api_agents/codex_agent/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/responses_api_agents/codex_agent/app.py b/responses_api_agents/codex_agent/app.py new file mode 100644 index 0000000000..e5eaec65bd --- /dev/null +++ b/responses_api_agents/codex_agent/app.py @@ -0,0 +1,686 @@ +# 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. + +import asyncio +import json +import logging +import os +import re +import shutil +import signal +import subprocess +import tempfile +from asyncio import Semaphore +from copy import deepcopy +from pathlib import Path +from time import time +from typing import Any, Literal, Optional +from uuid import uuid4 + +from fastapi import Request +from pydantic import ConfigDict, Field + +from nemo_gym.base_resources_server import NEMO_GYM_MCP_METADATA_KEY, BaseRunRequest, BaseVerifyResponse +from nemo_gym.base_responses_api_agent import BaseResponsesAPIAgentConfig, Body, SimpleResponsesAPIAgent +from nemo_gym.config_types import ModelServerRef, ResourcesServerRef +from nemo_gym.global_config import SKILLS_REF_KEY_NAME, get_first_server_config_dict +from nemo_gym.openai_utils import ( + NeMoGymEasyInputMessage, + NeMoGymFunctionCallOutput, + NeMoGymResponse, + NeMoGymResponseCreateParamsNonStreaming, + NeMoGymResponseFunctionToolCall, + NeMoGymResponseInputTokensDetails, + NeMoGymResponseOutputMessage, + NeMoGymResponseOutputText, + NeMoGymResponseOutputTokensDetails, + NeMoGymResponseUsage, +) +from nemo_gym.server_utils import get_response_json, raise_for_status +from nemo_gym.skills import stage_skills +from responses_api_agents.codex_agent.setup_codex import ensure_codex + + +LOG = logging.getLogger(__name__) + + +def _toml_key(key: str) -> str: + if re.fullmatch(r"[A-Za-z0-9_-]+", key): + return key + return json.dumps(key) + + +def _toml_value(value: Any) -> str: + if isinstance(value, bool): + return "true" if value else "false" + if isinstance(value, (int, float)): + return str(value) + if isinstance(value, str): + # JSON string escaping is a valid TOML basic string. + return json.dumps(value) + if isinstance(value, list): + return "[" + ", ".join(_toml_value(v) for v in value) + "]" + raise TypeError(f"Unsupported TOML value type: {type(value).__name__}") + + +def toml_dumps(data: dict[str, Any], _prefix: str = "") -> str: + """Serialize a nested dict of scalars/lists/dicts to TOML (the subset Codex config uses).""" + lines: list[str] = [] + tables: list[tuple[str, dict]] = [] + for key, value in data.items(): + if isinstance(value, dict): + tables.append((key, value)) + else: + lines.append(f"{_toml_key(key)} = {_toml_value(value)}") + chunks = ["\n".join(lines)] if lines else [] + for key, value in tables: + full_key = f"{_prefix}.{_toml_key(key)}" if _prefix else _toml_key(key) + body = toml_dumps(value, full_key) + chunks.append(f"[{full_key}]" + (f"\n{body}" if body else "")) + return "\n\n".join(chunks) + + +def _deep_merge(base: dict[str, Any], override: dict[str, Any]) -> dict[str, Any]: + merged = dict(base) + for key, value in override.items(): + if isinstance(value, dict) and isinstance(merged.get(key), dict): + merged[key] = _deep_merge(merged[key], value) + else: + merged[key] = value + return merged + + +def _mcp_result_text(item: dict[str, Any]) -> str: + if item.get("error"): + return f"error: {item['error']}" + result = item.get("result") + if isinstance(result, dict): + content = result.get("content") + if isinstance(content, list): + texts = [b.get("text", "") for b in content if isinstance(b, dict) and b.get("type") == "text"] + if any(texts): + return "".join(texts) + return json.dumps(result) + return "" if result is None else str(result) + + +def parse_exec_jsonl(stdout: str) -> tuple[list[Any], dict]: + """Convert ``codex exec --json`` JSONL stdout into (output_items, metadata). + + Codex emits ``item.completed`` events for each unit of work (assistant messages, reasoning, + shell commands, MCP tool calls, file changes, ...) and a terminal ``turn.completed`` carrying + token usage summed over every model call in the turn. Tool-shaped items are mapped to a + ``function_call`` + ``function_call_output`` pair so verifiers see one uniform trajectory + shape across agent harnesses; reasoning is buffered and prepended to the next assistant + message inside ```` tags (mirroring the Claude Code agent). + """ + output_items: list[Any] = [] + buffered_think: Optional[str] = None + metadata: dict[str, Any] = {"input_tokens": 0, "output_tokens": 0, "cached_input_tokens": 0, "reasoning_tokens": 0} + errors: list[str] = [] + + def _add_tool_pair(item: dict[str, Any], name: str, arguments: dict[str, Any], output: str) -> None: + call_id = str(item.get("id") or f"call-{uuid4().hex[:8]}") + status = "completed" if item.get("status") != "failed" else "incomplete" + output_items.append( + NeMoGymResponseFunctionToolCall( + arguments=json.dumps(arguments), + call_id=call_id, + name=name, + type="function_call", + id=call_id, + status=status, + ) + ) + output_items.append( + NeMoGymFunctionCallOutput(type="function_call_output", call_id=call_id, output=output, status="completed") + ) + + for line in stdout.splitlines(): + line = line.strip() + if not line: + continue + try: + event = json.loads(line) + except json.JSONDecodeError: + continue + if not isinstance(event, dict): + continue + + etype = event.get("type") + + if etype == "turn.completed": + usage = event.get("usage") or {} + metadata["input_tokens"] += int(usage.get("input_tokens") or 0) + metadata["output_tokens"] += int(usage.get("output_tokens") or 0) + metadata["cached_input_tokens"] += int(usage.get("cached_input_tokens") or 0) + metadata["reasoning_tokens"] += int(usage.get("reasoning_output_tokens") or 0) + continue + + if etype == "turn.failed": + message = (event.get("error") or {}).get("message") or "turn failed" + errors.append(message) + continue + + if etype != "item.completed": + continue + + item = event.get("item") + if not isinstance(item, dict): + continue + itype = item.get("type") + + if itype == "agent_message": + text = item.get("text") or "" + if buffered_think: + text = f"\n{buffered_think}\n\n\n{text}" + buffered_think = None + output_items.append( + NeMoGymResponseOutputMessage( + id=str(item.get("id") or f"msg-{len(output_items)}"), + content=[NeMoGymResponseOutputText(type="output_text", text=text, annotations=[])], + role="assistant", + status="completed", + type="message", + ) + ) + elif itype == "reasoning": + think = item.get("text") or "" + if think: + buffered_think = (buffered_think + "\n" + think) if buffered_think else think + elif itype == "command_execution": + output = item.get("aggregated_output") or "" + exit_code = item.get("exit_code") + if exit_code not in (None, 0): + output = f"{output}\n[exit code: {exit_code}]" + _add_tool_pair(item, "exec_command", {"cmd": item.get("command") or ""}, output) + elif itype == "mcp_tool_call": + _add_tool_pair(item, str(item.get("tool") or ""), item.get("arguments") or {}, _mcp_result_text(item)) + elif itype == "file_change": + _add_tool_pair(item, "apply_patch", {"changes": item.get("changes")}, item.get("status") or "completed") + elif itype == "web_search": + _add_tool_pair(item, "web_search", {"query": item.get("query") or ""}, "") + elif itype == "todo_list": + _add_tool_pair(item, "update_plan", {"items": item.get("items") or []}, "") + elif itype == "error": + errors.append(item.get("message") or "unknown error") + + # Some backends route the final answer through the reasoning channel (e.g. a vLLM reasoning + # parser labeling the closing message as reasoning). If the run ends on buffered reasoning with + # no assistant message after it, surface it as a think-tagged message rather than dropping it. + if buffered_think: + output_items.append( + NeMoGymResponseOutputMessage( + id=f"msg-{len(output_items)}", + content=[ + NeMoGymResponseOutputText( + type="output_text", text=f"\n{buffered_think}\n", annotations=[] + ) + ], + role="assistant", + status="completed", + type="message", + ) + ) + + if errors: + metadata["errors"] = errors + return output_items, metadata + + +def _kill_process_group(proc: Any) -> None: + """Kill the codex subprocess and every child in its process group. + + Killing only the direct child leaves the npm shim's vendored-binary child alive, holding the + stdout pipe open — the post-kill ``communicate()`` would then block until the orphan exits. + """ + try: + os.killpg(os.getpgid(proc.pid), signal.SIGKILL) + except Exception: + proc.kill() + + +def _extract_instruction(body_input) -> tuple[str, Optional[str]]: + """Return (user_message, system_message) from a responses body input list.""" + items = list(body_input) + system_message: Optional[str] = None + + if items: + first = items[0] + role = getattr(first, "role", None) or (first.get("role") if isinstance(first, dict) else None) + if role == "system": + content = getattr(first, "content", None) or (first.get("content") if isinstance(first, dict) else None) + if isinstance(content, list): + content = "".join( + (p.get("text", "") if isinstance(p, dict) else getattr(p, "text", "")) for p in content + ) + system_message = content or "" + items = items[1:] + + user_message = "" + for item in reversed(items): + role = getattr(item, "role", None) or (item.get("role") if isinstance(item, dict) else None) + if role == "user": + content = getattr(item, "content", None) or (item.get("content") if isinstance(item, dict) else None) + if isinstance(content, list): + content = "".join( + (p.get("text", "") if isinstance(p, dict) else getattr(p, "text", "")) for p in content + ) + user_message = content or "" + break + + return user_message, system_message + + +class CodexAgentConfig(BaseResponsesAPIAgentConfig): + resources_server: ResourcesServerRef + # When model_server is set, the Codex model provider's base_url is resolved from the Gym model + # server's URL (every Gym model server speaks the streaming Responses dialect on /v1/responses). + # When None, openai_base_url is used directly (default: the real OpenAI API). + model_server: Optional[ModelServerRef] = None + concurrency: int = 32 + # None -> omit `model` from the generated config and use the Codex CLI's own default. Gym model + # servers substitute their configured model anyway; set explicitly for direct endpoints. + model: Optional[str] = None + openai_api_key: str = "" # pragma: allowlist secret + openai_base_url: Optional[str] = None + sandbox_mode: Literal["read-only", "workspace-write", "danger-full-access"] = "danger-full-access" + timeout: int = 600 + system_prompt: Optional[str] = None + reasoning_effort: Optional[str] = None + # Required: every config pins an explicit npm version so auto-install is reproducible and cannot + # silently drift as new Codex releases land. Version bumps are then explicit, tested changes. + codex_version: str + # Working root handed to `codex exec --cd`. None -> a fresh temp dir per request, removed + # afterwards, so rollouts cannot see each other's files. + cwd: Optional[str] = None + # Provider stream idle timeout. Gym model servers emit the synthesized SSE only once the full + # response is computed, so the idle budget must cover an entire generation; None -> timeout * 1000. + stream_idle_timeout_ms: Optional[int] = None + # Extra config.toml content deep-merged over the generated base config (mcp_servers, features, + # tools, model_verbosity, ...). Per-rollout Gym MCP entries take precedence on name collisions. + extra_config: dict[str, Any] = Field(default_factory=dict) + + +class CodexAgentRunRequest(BaseRunRequest): + model_config = ConfigDict(extra="allow") + + +class CodexAgentVerifyResponse(BaseVerifyResponse): + model_config = ConfigDict(extra="allow") + turns_used: int = 0 + finished_naturally: bool = False + + +class CodexAgent(SimpleResponsesAPIAgent): + config: CodexAgentConfig + sem: Semaphore = None + model_config = ConfigDict(arbitrary_types_allowed=True) + + def model_post_init(self, __context: Any) -> None: + self.sem = Semaphore(self.config.concurrency) + ensure_codex(self.config.codex_version) + try: + ver = subprocess.run(["codex", "--version"], capture_output=True, text=True, timeout=10).stdout.strip() + LOG.warning("codex version: %s", ver or "(unknown)") + except Exception as exc: + LOG.warning("could not determine codex version: %s", exc) + + def _resolve_call_base_url(self, rollout_id: Optional[str]) -> str: + """Provider base_url for the CLI's model calls (Codex appends ``/responses`` to it). + + A Gym model server gets the per-rollout capture prefix plus the ``/v1`` suffix; a direct + endpoint (``model_server`` unset) is used verbatim and never prefixed — it has no + prefix-stripping middleware, so a prefix would 404 every call. + """ + if self.config.model_server: + return self.resolve_model_base_url(self.config.model_server.name, rollout_id) + # Mirrors claude_code_agent's null anthropic_base_url: null means the provider's real API. + return self.config.openai_base_url or "https://api.openai.com/v1" + + def _effective_model(self) -> Optional[str]: + """The model name written into the generated config (and reported on the response). + + An explicit (even unknown) model name keeps Codex from applying model-family feature + gating: for models it recognizes, Codex may switch tools into code-mode carriers that + models served through a Gym model server cannot drive. Gym model servers substitute their + own configured model anyway, so a placeholder never reaches the backend. Returns None only + for a direct endpoint with no configured model, where Codex uses its own default. + """ + if self.config.model: + return self.config.model + if self.config.model_server: + return "gym-policy-model" + return None + + def _build_config( + self, + base_url: str, + developer_instructions: Optional[str] = None, + mcp_servers: Optional[dict[str, Any]] = None, + ) -> dict[str, Any]: + """Assemble the per-run CODEX_HOME/config.toml content. + + The base config pins a Gym-owned model provider (bypassing Codex's login flow), disables + everything that would make a rollout depend on ambient host state or phone home (analytics, + update checks, on-disk history), and turns off the tools a Gym-served model cannot execute + (server-side web search, multi-agent). ``extra_config`` is deep-merged on top; the + per-rollout Gym MCP entries are applied last so they win name collisions. + """ + config: dict[str, Any] = { + "model_provider": "gym", + "approval_policy": "never", + "sandbox_mode": self.config.sandbox_mode, + "web_search": "disabled", + "check_for_update_on_startup": False, + "analytics": {"enabled": False}, + "history": {"persistence": "none"}, + # multi_agent and code_mode add tool shapes (namespace fan-out, custom JS-exec tools) + # that models served through a Gym model server cannot execute or express. + "features": {"multi_agent": False, "code_mode": False}, + "model_providers": { + "gym": { + "name": "gym", + "base_url": base_url, + # A custom provider reads its API key only from the env var named here; the + # agent sets it on the codex subprocess from `openai_api_key` (see _run_codex), + # so no `codex login` is ever needed. + "env_key": "OPENAI_API_KEY", + "wire_api": "responses", + "stream_idle_timeout_ms": self.config.stream_idle_timeout_ms or self.config.timeout * 1000, + } + }, + } + model = self._effective_model() + if model: + config["model"] = model + if self.config.reasoning_effort: + config["model_reasoning_effort"] = self.config.reasoning_effort + if developer_instructions: + config["developer_instructions"] = developer_instructions + if self.config.extra_config: + config = _deep_merge(config, deepcopy(self.config.extra_config)) + if mcp_servers: + config["mcp_servers"] = {**config.get("mcp_servers", {}), **mcp_servers} + return config + + def _setup_codex_home(self, config: dict[str, Any], skills_path: Optional[str] = None) -> Path: + """Create a per-run CODEX_HOME and stage config.toml (and optionally skills) into it. + + The directory lives for the duration of a single ``_run_codex`` call. When ``skills_path`` + is provided, the directory of skills is copied into ``/skills/`` where Codex's native + skill discovery picks them up. Each request gets its own ephemeral copy, so concurrent + requests with different skills do not contaminate one another. If setup fails partway + (e.g. a bad ``skills_path``), the partially-created dir is removed before re-raising. + """ + codex_home = Path.home() / ".codex_agent" / uuid4().hex + codex_home.mkdir(parents=True) + try: + (codex_home / "config.toml").write_text(toml_dumps(config)) + if skills_path: + stage_skills(skills_path, codex_home / "skills") + except Exception: + shutil.rmtree(codex_home, ignore_errors=True) + raise + return codex_home + + def _build_command(self, instruction: str, cwd: str) -> list[str]: + """Construct the ``codex exec`` argv. + + ``--json`` emits machine-readable JSONL events; ``--ephemeral`` skips session persistence; + ``--skip-git-repo-check`` allows running in the per-rollout scratch dir. Sandboxing and + approvals are pinned in the generated config.toml (``approval_policy = "never"``), not argv. + The ``--`` separator keeps prompts from being parsed as flags or subcommands. + """ + return [ + "codex", + "exec", + "--json", + "--ephemeral", + "--skip-git-repo-check", + "--cd", + cwd, + "--", + instruction, + ] + + async def _run_codex( + self, + instruction: str, + system_prompt: Optional[str] = None, + mcp_servers: Optional[dict[str, Any]] = None, + skills_path: Optional[str] = None, + rollout_id: Optional[str] = None, + ) -> tuple[str, str]: + """Run ``codex exec --json`` and return (stdout, model_name). + + When ``rollout_id`` is set and a model server is configured, the per-rollout capture prefix + is applied to the provider base_url so the CLI's streaming /v1/responses calls correlate to + this rollout. + """ + base_url = self._resolve_call_base_url(rollout_id) + # Report the name the config actually pins (so response.model matches what Codex was told); + # falls back to a sentinel only for a direct endpoint that lets Codex pick its own default. + model = self._effective_model() or "codex-default" + + config = self._build_config(base_url, developer_instructions=system_prompt, mcp_servers=mcp_servers) + + codex_home: Optional[Path] = None + scratch_cwd: Optional[str] = None + try: + # Inside the try so a bad skills_path (raising in stage_skills) still cleans up the + # partially-created home in the finally rather than leaking it per failing request. + codex_home = self._setup_codex_home(config, skills_path=skills_path) + cwd = self.config.cwd + if cwd is None: + cwd = scratch_cwd = tempfile.mkdtemp(prefix="nemo_gym_codex_ws_") + + env = { + **os.environ, + "CODEX_HOME": str(codex_home), + # The provider's `env_key` in the generated config.toml; always set from config so + # a key inherited from the server's environment can never leak into a rollout. + "OPENAI_API_KEY": self.config.openai_api_key or "local", # pragma: allowlist secret + } + + proc = await asyncio.create_subprocess_exec( + *self._build_command(instruction, cwd), + stdin=asyncio.subprocess.DEVNULL, # codex appends piped stdin to the prompt and blocks on it + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, + env=env, + # Own process group: `codex` on PATH is an npm shim whose child (the vendored + # binary) must die with it, or it keeps the stdout pipe open past the kill below. + start_new_session=True, + ) + try: + stdout, stderr = await asyncio.wait_for(proc.communicate(), timeout=self.config.timeout) + except asyncio.TimeoutError: + _kill_process_group(proc) + await proc.communicate() + LOG.warning("codex timed out after %ds", self.config.timeout) + return "", model + + if proc.returncode not in (0, None): + LOG.warning("codex exited %d: %s", proc.returncode, stderr.decode(errors="replace")[:500]) + + LOG.debug("codex stdout (%d chars): %s", len(stdout), stdout[:2000].decode(errors="replace")) + return stdout.decode(errors="replace"), model + finally: + if codex_home is not None: + shutil.rmtree(codex_home, ignore_errors=True) + if scratch_cwd is not None: + shutil.rmtree(scratch_cwd, ignore_errors=True) + + def _resources_server_base_url(self) -> str: + cfg = get_first_server_config_dict( + self.server_client.global_config_dict, + self.config.resources_server.name, + ) + return self.server_client._build_server_base_url(cfg) + + def _rollout_mcp_servers(self, seed_response_json: dict[str, Any]) -> Optional[dict[str, Any]]: + """Per-rollout ``mcp_servers`` config.toml entries from /seed_session MCP metadata. + + Codex reaches Gym MCP tools over streamable HTTP; the per-rollout session token rides on a + custom header via ``http_headers``. + """ + metadata = seed_response_json.get(NEMO_GYM_MCP_METADATA_KEY) + if not isinstance(metadata, dict): + return None + + server_name = str(metadata.get("server_name") or self.config.resources_server.name) + url_path = str(metadata.get("url_path") or "/mcp") + entry: dict[str, Any] = { + "url": f"{self._resources_server_base_url().rstrip('/')}/{url_path.lstrip('/')}", + } + headers = metadata.get("headers") + if isinstance(headers, dict) and headers: + entry["http_headers"] = {str(key): str(value) for key, value in headers.items()} + else: + LOG.warning( + "MCP seed metadata for %r has no headers; the tool endpoint will be called without a " + "session token and will reject the calls.", + server_name, + ) + return {server_name: entry} + + async def _create_response( + self, + body: NeMoGymResponseCreateParamsNonStreaming, + mcp_servers: Optional[dict[str, Any]] = None, + skills_path: Optional[str] = None, + rollout_id: Optional[str] = None, + ) -> NeMoGymResponse: + body = body.model_copy(deep=True) + if isinstance(body.input, str): + body.input = [NeMoGymEasyInputMessage(role="user", content=body.input)] + + user_message, input_system = _extract_instruction(body.input) + system_parts = [p for p in [self.config.system_prompt, input_system] if p] + system_prompt = "\n\n".join(system_parts) if system_parts else None + + stdout, model_name = await self._run_codex( + user_message, + system_prompt=system_prompt, + mcp_servers=mcp_servers, + skills_path=skills_path, + rollout_id=rollout_id, + ) + output_items, usage = parse_exec_jsonl(stdout) + + if usage.get("errors"): + LOG.warning("codex reported errors: %s", usage["errors"]) + + if not any( + getattr(item, "type", None) == "message" and getattr(item, "role", None) == "assistant" + for item in output_items + ): + LOG.warning("codex produced no assistant message; padding empty output") + output_items.append( + NeMoGymResponseOutputMessage( + id=f"msg_{uuid4().hex}", + content=[NeMoGymResponseOutputText(text="", annotations=[])], + role="assistant", + status="completed", + type="message", + ) + ) + + input_tokens = usage.get("input_tokens", 0) + output_tokens = usage.get("output_tokens", 0) + + return NeMoGymResponse( + id=f"resp_{uuid4().hex}", + created_at=int(time()), + model=model_name, + object="response", + output=output_items, + tool_choice=body.tool_choice, + tools=body.tools, + parallel_tool_calls=body.parallel_tool_calls, + usage=NeMoGymResponseUsage( + input_tokens=input_tokens, + input_tokens_details=NeMoGymResponseInputTokensDetails( + cached_tokens=usage.get("cached_input_tokens", 0) + ), + output_tokens=output_tokens, + output_tokens_details=NeMoGymResponseOutputTokensDetails( + reasoning_tokens=usage.get("reasoning_tokens", 0) + ), + total_tokens=input_tokens + output_tokens, + ), + ) + + async def responses( + self, + request: Request, + body: NeMoGymResponseCreateParamsNonStreaming = Body(), + ) -> NeMoGymResponse: + return await self._create_response(body) + + async def run(self, request: Request, body: CodexAgentRunRequest) -> CodexAgentVerifyResponse: + async with self.sem: + cookies = request.cookies + + seed_resp = await self.server_client.post( + server_name=self.config.resources_server.name, + url_path="/seed_session", + json=body.model_dump(), + cookies=cookies, + ) + await raise_for_status(seed_resp) + cookies = seed_resp.cookies + seed_resp_json = await get_response_json(seed_resp) + + # The run-level skills_ref (stamped by rollout collection) rides on the request body + # (extra="allow"). Pass its path straight into _create_response so the CLI invocation + # can stage the skills into its per-request CODEX_HOME. + skills_path = ((body.model_extra or {}).get(SKILLS_REF_KEY_NAME) or {}).get("path") + rollout_id = self.rollout_id_from_run(body) + + agent_resp = await self._create_response( + body.responses_create_params, + mcp_servers=self._rollout_mcp_servers(seed_resp_json), + skills_path=skills_path, + rollout_id=rollout_id, + ) + agent_resp_json = agent_resp.model_dump(mode="json") + + verify_resp = await self.server_client.post( + server_name=self.config.resources_server.name, + url_path="/verify", + json=body.model_dump() | {"response": agent_resp_json}, + cookies=cookies, + ) + await raise_for_status(verify_resp) + verify_json = await get_response_json(verify_resp) + + gym_resp = NeMoGymResponse.model_validate(agent_resp_json) + turns = sum( + 1 + for item in gym_resp.output + if getattr(item, "type", None) == "message" and getattr(item, "role", None) == "assistant" + ) + last = gym_resp.output[-1] if gym_resp.output else None + naturally = getattr(last, "type", None) == "message" and getattr(last, "role", None) == "assistant" + + return CodexAgentVerifyResponse.model_validate( + verify_json | {"turns_used": turns, "finished_naturally": naturally} + ) + + +if __name__ == "__main__": + CodexAgent.run_webserver() diff --git a/responses_api_agents/codex_agent/configs/codex_agent.yaml b/responses_api_agents/codex_agent/configs/codex_agent.yaml new file mode 100644 index 0000000000..829ceeeb0d --- /dev/null +++ b/responses_api_agents/codex_agent/configs/codex_agent.yaml @@ -0,0 +1,19 @@ +codex_agent: + responses_api_agents: + codex_agent: + entrypoint: app.py + resources_server: + type: resources_servers + name: ??? + concurrency: 32 + model: null # null -> the Codex CLI's own default model; set explicitly for direct endpoints + openai_api_key: ${openai_api_key} + openai_base_url: null # null -> https://api.openai.com/v1 (must include /v1; Codex appends /responses) + sandbox_mode: danger-full-access # read-only | workspace-write | danger-full-access + timeout: 600 + system_prompt: null # inserted as a developer message via `developer_instructions` + reasoning_effort: null # passed as `model_reasoning_effort` (e.g. low, medium, high) + codex_version: 0.144.4 # required: npm version pinned on auto-install for reproducibility + cwd: null # working root for `codex exec --cd`; null -> fresh temp dir per request + stream_idle_timeout_ms: null # provider stream idle budget; null -> timeout * 1000 + extra_config: {} # extra config.toml content deep-merged over the generated base config diff --git a/responses_api_agents/codex_agent/requirements.txt b/responses_api_agents/codex_agent/requirements.txt new file mode 100644 index 0000000000..00ed83213e --- /dev/null +++ b/responses_api_agents/codex_agent/requirements.txt @@ -0,0 +1 @@ +-e nemo-gym[dev] @ ../../ diff --git a/responses_api_agents/codex_agent/setup_codex.py b/responses_api_agents/codex_agent/setup_codex.py new file mode 100644 index 0000000000..36a5384e9b --- /dev/null +++ b/responses_api_agents/codex_agent/setup_codex.py @@ -0,0 +1,103 @@ +# 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. + +import logging +import os +import shutil +import subprocess +import tarfile +import urllib.request +from pathlib import Path + + +LOG = logging.getLogger(__name__) + +_CODEX_PKG = "@openai/codex" +_NODE_VERSION = "22.15.0" +_NODE_DIST_URL = f"https://nodejs.org/dist/v{_NODE_VERSION}/node-v{_NODE_VERSION}-linux-x64.tar.xz" +_LOCAL_PREFIX = Path(__file__).parent / ".codex_node" + + +def _npm_install(npm_bin: str, version: str | None) -> None: + pkg = f"{_CODEX_PKG}@{version}" if version else f"{_CODEX_PKG}@latest" + subprocess.run([npm_bin, "install", "-g", pkg], check=True) + + +def _install_node_locally() -> Path: + node_bin = _LOCAL_PREFIX / "bin" / "node" + if node_bin.is_file(): + return _LOCAL_PREFIX / "bin" + + _LOCAL_PREFIX.mkdir(parents=True, exist_ok=True) + tarball = _LOCAL_PREFIX / "node.tar.xz" + + LOG.info("downloading Node.js %s", _NODE_VERSION) + urllib.request.urlretrieve(_NODE_DIST_URL, tarball) # noqa: S310 + + with tarfile.open(tarball, "r:xz") as tf: + tf.extractall(_LOCAL_PREFIX, filter="data") + + nested = next(p for p in _LOCAL_PREFIX.iterdir() if p.is_dir() and p.name.startswith("node-")) + for item in nested.iterdir(): + item.rename(_LOCAL_PREFIX / item.name) + nested.rmdir() + tarball.unlink(missing_ok=True) + return _LOCAL_PREFIX / "bin" + + +def ensure_codex(version: str | None = None) -> None: + """Ensure ``codex`` is on PATH, installing it if necessary.""" + if shutil.which("codex"): + return + + # Check ~/.local/bin + local_bin = Path.home() / ".local" / "bin" + if (local_bin / "codex").is_file(): + os.environ["PATH"] = str(local_bin) + os.pathsep + os.environ.get("PATH", "") + return + + npm = shutil.which("npm") + if npm: + LOG.info("installing codex via system npm (%s)", npm) + _npm_install(npm, version) + else: + LOG.info("npm not found; installing local Node.js") + bin_dir = _install_node_locally() + os.environ["PATH"] = str(bin_dir) + os.pathsep + os.environ.get("PATH", "") + npm = shutil.which("npm") + if not npm: + raise RuntimeError(f"npm not found after local Node.js install in {bin_dir}") + _npm_install(npm, version) + + # npm install -g may put the binary in a prefix not yet on PATH. `npm bin -g` was removed in + # npm >= 9 (Node 22.15 ships npm 10), so resolve the global prefix and append its bin dir. + if not shutil.which("codex"): + npm_prefix = subprocess.run( + [shutil.which("npm") or "npm", "prefix", "-g"], + capture_output=True, + text=True, + ).stdout.strip() + npm_bin_dir = str(Path(npm_prefix) / "bin") if npm_prefix else "" + if npm_bin_dir and Path(npm_bin_dir).is_dir(): + os.environ["PATH"] = npm_bin_dir + os.pathsep + os.environ.get("PATH", "") + + # Also check ~/.local/bin after install + if not shutil.which("codex") and (local_bin / "codex").is_file(): + os.environ["PATH"] = str(local_bin) + os.pathsep + os.environ.get("PATH", "") + + if not shutil.which("codex"): + raise RuntimeError("codex install appeared to succeed but 'codex' is still not on PATH") + + LOG.info("codex is ready at %s", shutil.which("codex")) diff --git a/responses_api_agents/codex_agent/tests/__init__.py b/responses_api_agents/codex_agent/tests/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/responses_api_agents/codex_agent/tests/test_app.py b/responses_api_agents/codex_agent/tests/test_app.py new file mode 100644 index 0000000000..78bdf18d99 --- /dev/null +++ b/responses_api_agents/codex_agent/tests/test_app.py @@ -0,0 +1,857 @@ +# 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. + +import asyncio +import json +import tomllib +from pathlib import Path +from unittest.mock import AsyncMock, MagicMock, patch + +import orjson +import pytest +import yaml +from fastapi import Request +from pydantic import ValidationError + +from nemo_gym.global_config import SKILLS_REF_KEY_NAME +from nemo_gym.openai_utils import ( + NeMoGymEasyInputMessage, + NeMoGymFunctionCallOutput, + NeMoGymResponseCreateParamsNonStreaming, + NeMoGymResponseFunctionToolCall, + NeMoGymResponseOutputMessage, +) +from nemo_gym.server_utils import ServerClient +from responses_api_agents.codex_agent.app import ( + CodexAgent, + CodexAgentConfig, + CodexAgentRunRequest, + ModelServerRef, + ResourcesServerRef, + _extract_instruction, + parse_exec_jsonl, + toml_dumps, +) + + +def _write_skill_dir(root: Path, name: str = "cot_enhanced") -> Path: + skills_dir = root / "variant_a" + skill = skills_dir / name + skill.mkdir(parents=True) + (skill / "SKILL.md").write_text(f"---\nname: {name}\ndescription: A skill.\n---\n# Body\n") + return skills_dir + + +def _config(**kwargs) -> CodexAgentConfig: + kwargs.setdefault("resources_server", ResourcesServerRef(type="resources_servers", name="")) + kwargs.setdefault("codex_version", "0.144.4") + return CodexAgentConfig( + host="0.0.0.0", + port=8080, + entrypoint="", + name="", + **kwargs, + ) + + +def _make_agent(**kwargs) -> CodexAgent: + # Patch only the external side effect (codex install/version check) so the real + # model_post_init still runs — it initializes the semaphore. + with patch("responses_api_agents.codex_agent.app.ensure_codex"): + return CodexAgent(config=_config(**kwargs), server_client=MagicMock(spec=ServerClient)) + + +def _event(type_: str, **kwargs) -> str: + return json.dumps({"type": type_, **kwargs}) + + +def _item_completed(item: dict) -> str: + return _event("item.completed", item=item) + + +class FakeAioHTTPResponse: + ok = True + + def __init__(self, payload: dict, cookies: dict | None = None): + self.payload = payload + self.cookies = cookies or {} + + async def read(self) -> bytes: + return json.dumps(self.payload).encode() + + +class TestSanity: + def test_config_defaults(self) -> None: + cfg = _config() + assert cfg.concurrency == 32 + assert cfg.timeout == 600 + assert cfg.model is None + assert cfg.sandbox_mode == "danger-full-access" + assert cfg.cwd is None + assert cfg.extra_config == {} + + def test_semaphore_initialized(self) -> None: + agent = _make_agent(concurrency=4) + assert agent.sem._value == 4 + + def test_codex_version_is_required(self) -> None: + # Pinning is mandatory so auto-install cannot silently drift; omitting it is a config error. + with pytest.raises(ValidationError): + CodexAgentConfig( + host="0.0.0.0", + port=8080, + entrypoint="", + name="", + resources_server=ResourcesServerRef(type="resources_servers", name=""), + ) + + +class TestTomlDumps: + def test_round_trips_via_tomllib(self) -> None: + data = { + "model": "gpt-5", + "check_for_update_on_startup": False, + "analytics": {"enabled": False}, + "model_providers": {"gym": {"name": "gym", "base_url": "http://x/v1", "stream_idle_timeout_ms": 600000}}, + "mcp_servers": { + "weather": { + "url": "http://h:1/mcp", + "http_headers": {"X-NeMo-Gym-Session-Token": "tok"}, + } + }, + } + parsed = tomllib.loads(toml_dumps(data)) + assert parsed == data + + def test_quotes_non_bare_keys_and_escapes_strings(self) -> None: + parsed = tomllib.loads(toml_dumps({"a b": 'quo"te\nnl', "list": ["x", "y"]})) + assert parsed == {"a b": 'quo"te\nnl', "list": ["x", "y"]} + + +class TestBuildCommand: + def test_command_shape(self) -> None: + agent = _make_agent() + cmd = agent._build_command("do the thing", "/work/dir") + assert cmd[:5] == ["codex", "exec", "--json", "--ephemeral", "--skip-git-repo-check"] + assert cmd[cmd.index("--cd") + 1] == "/work/dir" + # instruction is the final positional after the `--` separator + assert cmd[-2:] == ["--", "do the thing"] + + +class TestBuildConfig: + def test_base_config_isolated_and_pinned_to_gym_provider(self) -> None: + agent = _make_agent(timeout=30) + config = agent._build_config("http://model:9000/v1") + assert config["model_provider"] == "gym" + assert config["approval_policy"] == "never" + assert config["sandbox_mode"] == "danger-full-access" + assert config["web_search"] == "disabled" + assert config["check_for_update_on_startup"] is False + assert config["analytics"] == {"enabled": False} + assert config["history"] == {"persistence": "none"} + assert config["features"] == {"multi_agent": False, "code_mode": False} + provider = config["model_providers"]["gym"] + assert provider["base_url"] == "http://model:9000/v1" + assert provider["env_key"] == "OPENAI_API_KEY" + assert provider["wire_api"] == "responses" + # idle budget defaults to the whole-run timeout (Gym servers stream only at completion) + assert provider["stream_idle_timeout_ms"] == 30_000 + assert "model" not in config + assert "developer_instructions" not in config + + def test_optional_knobs_threaded_through(self) -> None: + agent = _make_agent(model="gpt-5-codex", reasoning_effort="high", stream_idle_timeout_ms=42) + config = agent._build_config("http://x/v1", developer_instructions="be terse") + assert config["model"] == "gpt-5-codex" + assert config["model_reasoning_effort"] == "high" + assert config["developer_instructions"] == "be terse" + assert config["model_providers"]["gym"]["stream_idle_timeout_ms"] == 42 + + def test_model_server_without_model_pins_placeholder(self) -> None: + # With a model server and no explicit model, config pins a placeholder to avoid Codex's + # model-family code-mode gating, and the effective name is reported consistently. + agent = _make_agent(model_server=ModelServerRef(type="responses_api_models", name="policy_model")) + config = agent._build_config("http://x/v1") + assert config["model"] == "gym-policy-model" + assert agent._effective_model() == "gym-policy-model" + + def test_direct_endpoint_without_model_omits_model(self) -> None: + # A direct endpoint with no model lets Codex pick its own default: no model key in config. + agent = _make_agent() + config = agent._build_config("http://x/v1") + assert "model" not in config + assert agent._effective_model() is None + + def test_extra_config_deep_merged(self) -> None: + agent = _make_agent( + extra_config={ + "features": {"web_search": True}, + "model_verbosity": "low", + "mcp_servers": {"static": {"command": "server"}}, + } + ) + config = agent._build_config("http://x/v1") + # deep merge preserves base keys next to user keys + assert config["features"] == {"multi_agent": False, "code_mode": False, "web_search": True} + assert config["model_verbosity"] == "low" + assert config["mcp_servers"] == {"static": {"command": "server"}} + + def test_rollout_mcp_servers_win_name_collisions(self) -> None: + agent = _make_agent(extra_config={"mcp_servers": {"weather": {"command": "stale"}}}) + config = agent._build_config("http://x/v1", mcp_servers={"weather": {"url": "http://h:1/mcp"}}) + assert config["mcp_servers"]["weather"] == {"url": "http://h:1/mcp"} + + def test_extra_config_not_mutated_across_calls(self) -> None: + agent = _make_agent(extra_config={"mcp_servers": {"static": {"command": "server"}}}) + agent._build_config("http://x/v1", mcp_servers={"dynamic": {"url": "http://h:1/mcp"}}) + assert agent.config.extra_config == {"mcp_servers": {"static": {"command": "server"}}} + + +class TestSetupCodexHome: + def test_creates_home_with_config_toml(self, tmp_path: Path) -> None: + agent = _make_agent() + with patch("responses_api_agents.codex_agent.app.Path.home", return_value=tmp_path): + codex_home = agent._setup_codex_home(agent._build_config("http://x/v1")) + try: + parsed = tomllib.loads((codex_home / "config.toml").read_text()) + assert parsed["model_provider"] == "gym" + assert parsed["history"]["persistence"] == "none" + finally: + import shutil as _shutil + + _shutil.rmtree(codex_home, ignore_errors=True) + + def test_stages_skills_into_home(self, tmp_path: Path) -> None: + skills_dir = _write_skill_dir(tmp_path) + home = tmp_path / "home" + home.mkdir() + agent = _make_agent() + with patch("responses_api_agents.codex_agent.app.Path.home", return_value=home): + codex_home = agent._setup_codex_home(agent._build_config("http://x/v1"), skills_path=str(skills_dir)) + try: + assert (codex_home / "skills" / "cot_enhanced" / "SKILL.md").is_file() + finally: + import shutil as _shutil + + _shutil.rmtree(codex_home, ignore_errors=True) + + +class _FakeHttpResp: + def __init__(self, payload: dict) -> None: + self._payload = payload + self.cookies: dict = {} + self.ok = True + + async def read(self) -> bytes: + return orjson.dumps(self._payload) + + +def _gym_response(text: str = "done") -> dict: + return { + "id": "resp_x", + "created_at": 0.0, + "model": "codex-default", + "object": "response", + "output": [ + { + "id": "msg_x", + "content": [{"annotations": [], "text": text, "type": "output_text"}], + "role": "assistant", + "status": "completed", + "type": "message", + } + ], + "parallel_tool_calls": True, + "tool_choice": "auto", + "tools": [], + "usage": { + "input_tokens": 1, + "input_tokens_details": {"cached_tokens": 0}, + "output_tokens": 1, + "output_tokens_details": {"reasoning_tokens": 0}, + "total_tokens": 2, + }, + } + + +class TestRunForwardsSkillsPath: + """run() reads skills_ref off the request's model_extra (extra='allow') and forwards its path + directly to _create_response/_run_codex.""" + + def _seed_and_verify_post(self): + async def _post(server_name, url_path, json=None, cookies=None, **kw): + if url_path == "/verify": + return _FakeHttpResp( + {"responses_create_params": {"input": []}, "response": _gym_response(), "reward": 1.0} + ) + return _FakeHttpResp({}) + + return AsyncMock(side_effect=_post) + + def _run(self, agent: CodexAgent, body: CodexAgentRunRequest, run_codex: AsyncMock): + agent.server_client.post = self._seed_and_verify_post() + req = MagicMock() + req.cookies = {} + # Stub the CLI invocation; _create_response still runs for real, so we exercise the full + # run() -> _create_response -> _run_codex argument threading. + with patch.object(CodexAgent, "_run_codex", run_codex): + return asyncio.run(agent.run(req, body)) + + def test_skills_ref_path_forwarded(self) -> None: + agent = _make_agent() + run_codex = AsyncMock(return_value=("", "codex-default")) + body = CodexAgentRunRequest.model_validate( + { + "responses_create_params": {"input": []}, + SKILLS_REF_KEY_NAME: {"path": "skills/variant_a/", "hash": "abc123", "skills": []}, + } + ) + + self._run(agent, body, run_codex) + + assert run_codex.call_args.kwargs["skills_path"] == "skills/variant_a/" + + def test_no_skills_ref_forwards_none(self) -> None: + agent = _make_agent() + run_codex = AsyncMock(return_value=("", "codex-default")) + body = CodexAgentRunRequest.model_validate({"responses_create_params": {"input": []}}) + + self._run(agent, body, run_codex) + + assert run_codex.call_args.kwargs["skills_path"] is None + + +class TestRunCodex: + def test_wires_command_env_and_cleans_up(self, tmp_path: Path) -> None: + agent = _make_agent(openai_api_key="sk-test", system_prompt=None) # pragma: allowlist secret + captured: dict = {} + + class FakeProc: + returncode = 0 + + async def communicate(self): + return ( + b'{"type":"turn.completed","usage":{"input_tokens":3,"output_tokens":4}}\n', + b"", + ) + + async def fake_exec(*cmd, **kwargs): + env = kwargs["env"] + codex_home = env["CODEX_HOME"] + captured["cmd"] = list(cmd) + captured["codex_home"] = codex_home + captured["api_key"] = env["OPENAI_API_KEY"] + captured["stdin"] = kwargs.get("stdin") + captured["start_new_session"] = kwargs.get("start_new_session") + # the staged home + config must exist while the subprocess runs + captured["config_during_run"] = tomllib.loads((Path(codex_home) / "config.toml").read_text()) + captured["cwd_exists_during_run"] = Path(cmd[cmd.index("--cd") + 1]).is_dir() + captured["scratch_cwd"] = cmd[cmd.index("--cd") + 1] + return FakeProc() + + with ( + patch("responses_api_agents.codex_agent.app.Path.home", return_value=tmp_path), + patch("responses_api_agents.codex_agent.app.asyncio.create_subprocess_exec", fake_exec), + ): + stdout, model = asyncio.run(agent._run_codex("hello", system_prompt="be terse")) + + assert captured["cmd"][0] == "codex" + assert captured["cmd"][-1] == "hello" + assert captured["api_key"] == "sk-test" # pragma: allowlist secret + assert captured["stdin"] == asyncio.subprocess.DEVNULL + # own process group, so a timeout kill reaps the npm shim's vendored-binary child too + assert captured["start_new_session"] is True + assert captured["config_during_run"]["developer_instructions"] == "be terse" + assert captured["cwd_exists_during_run"] is True + # per-run home and scratch cwd are removed after the run (no leakage between rollouts) + assert not Path(captured["codex_home"]).exists() + assert not Path(captured["scratch_cwd"]).exists() + assert "turn.completed" in stdout + assert model == "codex-default" + + def test_explicit_cwd_is_used_and_kept(self, tmp_path: Path) -> None: + workdir = tmp_path / "work" + workdir.mkdir() + agent = _make_agent(cwd=str(workdir)) + + class FakeProc: + returncode = 0 + + async def communicate(self): + return b"", b"" + + captured: dict = {} + + async def fake_exec(*cmd, **kwargs): + captured["cwd"] = cmd[cmd.index("--cd") + 1] + return FakeProc() + + with ( + patch("responses_api_agents.codex_agent.app.Path.home", return_value=tmp_path), + patch("responses_api_agents.codex_agent.app.asyncio.create_subprocess_exec", fake_exec), + ): + asyncio.run(agent._run_codex("hello")) + + assert captured["cwd"] == str(workdir) + assert workdir.is_dir() # a user-provided cwd is never removed + + def test_bad_skills_path_does_not_leak_codex_home(self, tmp_path: Path) -> None: + # stage_skills raises for a missing skills dir; the partially-created home must + # still be cleaned up (setup happens inside the try whose finally rmtree's it). + home = tmp_path / "home" + home.mkdir() + agent = _make_agent() + + with patch("responses_api_agents.codex_agent.app.Path.home", return_value=home): + with pytest.raises(ValueError): + asyncio.run(agent._run_codex("hello", skills_path=str(tmp_path / "does_not_exist"))) + + leaked = home / ".codex_agent" + assert not leaked.exists() or not any(leaked.iterdir()) + + def test_timeout_returns_empty(self, tmp_path: Path) -> None: + agent = _make_agent(timeout=1) + killed = {"called": False} + + class SlowProc: + returncode = None + + def kill(self): + killed["called"] = True + + async def communicate(self): + return b"", b"" + + async def fake_exec(*cmd, **kwargs): + return SlowProc() + + async def fake_wait_for(coro, timeout): + coro.close() # avoid un-awaited coroutine warning + raise asyncio.TimeoutError + + with ( + patch("responses_api_agents.codex_agent.app.Path.home", return_value=tmp_path), + patch("responses_api_agents.codex_agent.app.asyncio.create_subprocess_exec", fake_exec), + patch("responses_api_agents.codex_agent.app.asyncio.wait_for", fake_wait_for), + ): + stdout, model = asyncio.run(agent._run_codex("hello")) + + assert stdout == "" + assert killed["called"] is True + assert model == "codex-default" + + +class TestRolloutMCPServers: + def _agent_with_resources_server(self, **kwargs) -> CodexAgent: + agent = _make_agent( + resources_server=ResourcesServerRef(type="resources_servers", name="example_mcp_weather"), **kwargs + ) + agent.server_client.global_config_dict = { + "example_mcp_weather": { + "resources_servers": { + "example_mcp_weather": { + "host": "127.0.0.1", + "port": 8123, + } + } + } + } + agent.server_client._build_server_base_url.side_effect = lambda cfg: f"http://{cfg['host']}:{cfg['port']}" + return agent + + def test_no_metadata_returns_none(self) -> None: + agent = self._agent_with_resources_server() + assert agent._rollout_mcp_servers({}) is None + + def test_builds_streamable_http_entry_with_session_header(self) -> None: + agent = self._agent_with_resources_server() + servers = agent._rollout_mcp_servers( + { + "mcp": { + "server_name": "example_mcp_weather", + "url_path": "/mcp", + "headers": {"X-NeMo-Gym-Session-Token": "secret-token"}, + } + } + ) + assert servers == { + "example_mcp_weather": { + "url": "http://127.0.0.1:8123/mcp", + "http_headers": {"X-NeMo-Gym-Session-Token": "secret-token"}, + } + } + + def test_run_writes_mcp_entry_into_config(self) -> None: + agent = self._agent_with_resources_server() + + async def fake_post(server_name, url_path, json=None, cookies=None): + if url_path == "/seed_session": + return FakeAioHTTPResponse( + { + "mcp": { + "server_name": "example_mcp_weather", + "url_path": "/mcp", + "headers": {"X-NeMo-Gym-Session-Token": "tok"}, + } + }, + cookies={"session": "abc"}, + ) + if url_path == "/verify": + return FakeAioHTTPResponse(json | {"reward": 1.0}) + raise AssertionError(f"unexpected post: {server_name} {url_path}") + + captured: dict = {} + + async def fake_run_codex(instruction, system_prompt=None, mcp_servers=None, **kwargs): + captured["instruction"] = instruction + captured["mcp_servers"] = mcp_servers + captured["config"] = agent._build_config("http://x/v1", mcp_servers=mcp_servers) + return _item_completed( + {"id": "item_1", "type": "agent_message", "text": "The weather in Paris is sunny and 72 F."} + ), "codex-default" + + agent.server_client.post.side_effect = fake_post + object.__setattr__(agent, "_run_codex", fake_run_codex) + request = MagicMock(spec=Request) + request.cookies = {} + body = CodexAgentRunRequest( + responses_create_params=NeMoGymResponseCreateParamsNonStreaming(input="use the weather tool"), + expected_city="Paris", + ) + + result = asyncio.run(agent.run(request, body)) + + assert result.reward == 1.0 + assert captured["instruction"] == "use the weather tool" + server = captured["config"]["mcp_servers"]["example_mcp_weather"] + assert server["url"] == "http://127.0.0.1:8123/mcp" + assert server["http_headers"]["X-NeMo-Gym-Session-Token"] == "tok" + + def test_run_threads_session_cookie_seed_to_verify(self) -> None: + agent = self._agent_with_resources_server() + captured: dict = {} + + async def fake_post(server_name, url_path, json=None, cookies=None): + if url_path == "/seed_session": + # the resources server sets a session cookie on the seed response + return FakeAioHTTPResponse( + { + "mcp": { + "server_name": "example_mcp_weather", + "url_path": "/mcp", + "headers": {"X-NeMo-Gym-Session-Token": "tok"}, + } + }, + cookies={"session": "sess-cookie"}, + ) + if url_path == "/verify": + captured["verify_cookies"] = cookies + return FakeAioHTTPResponse(json | {"reward": 1.0}) + raise AssertionError(f"unexpected post: {server_name} {url_path}") + + async def fake_run_codex(instruction, system_prompt=None, mcp_servers=None, **kwargs): + captured["mcp_servers"] = mcp_servers + return _item_completed({"id": "item_1", "type": "agent_message", "text": "ok"}), "codex-default" + + agent.server_client.post.side_effect = fake_post + object.__setattr__(agent, "_run_codex", fake_run_codex) + request = MagicMock(spec=Request) + request.cookies = {} + body = CodexAgentRunRequest( + responses_create_params=NeMoGymResponseCreateParamsNonStreaming(input="use the weather tool"), + verifier_metadata={"expected_city": "Paris"}, + ) + + asyncio.run(agent.run(request, body)) + + # the cookie set on /seed_session is threaded into the /verify call (same rollout session), + # and the per-rollout token from seed metadata reaches the generated MCP config. + assert captured["verify_cookies"] == {"session": "sess-cookie"} + assert captured["mcp_servers"]["example_mcp_weather"]["http_headers"]["X-NeMo-Gym-Session-Token"] == "tok" + + +class TestRolloutCorrelation: + """The CLI streams /v1/responses, so correlation rides on the provider base_url path prefix.""" + + def _fake_proc(self): + class FakeProc: + returncode = 0 + + async def communicate(self): + return b'{"type":"turn.completed","usage":{"input_tokens":1,"output_tokens":1}}\n', b"" + + return FakeProc() + + def _run_and_capture_base_url(self, agent, tmp_path: Path, **run_kwargs) -> str: + captured: dict = {} + + async def fake_exec(*cmd, **kwargs): + config = tomllib.loads((Path(kwargs["env"]["CODEX_HOME"]) / "config.toml").read_text()) + captured["base_url"] = config["model_providers"]["gym"]["base_url"] + return self._fake_proc() + + def fake_resolve(name, rollout_id=None): + prefix = f"/ng-rollout/{rollout_id}" if rollout_id else "" + return f"http://model-server:9000{prefix}/v1" + + with ( + patch("responses_api_agents.codex_agent.app.Path.home", return_value=tmp_path), + patch.object(type(agent), "resolve_model_base_url", side_effect=fake_resolve), + patch("responses_api_agents.codex_agent.app.asyncio.create_subprocess_exec", fake_exec), + ): + asyncio.run(agent._run_codex("hi", **run_kwargs)) + return captured["base_url"] + + def test_base_url_correlation(self, tmp_path: Path) -> None: + agent = _make_agent(model_server=ModelServerRef(type="responses_api_models", name="policy_model")) + base_url = self._run_and_capture_base_url(agent, tmp_path, rollout_id="task3-roll1") + # Codex appends /responses -> server strips /ng-rollout/ and keys capture by it. + assert base_url == "http://model-server:9000/ng-rollout/task3-roll1/v1" + + # Direct endpoint (no model server): never prefixed -- it has no stripping middleware, + # so a prefix would 404 every /v1/responses call. + direct = _make_agent(openai_base_url="https://api.openai.com/v1") + assert direct._resolve_call_base_url("t3-r1") == "https://api.openai.com/v1" + + def test_defaults_to_openai_when_nothing_configured(self) -> None: + agent = _make_agent() + assert agent._resolve_call_base_url(None) == "https://api.openai.com/v1" + + +class TestExtractInstruction: + def test_user_only(self) -> None: + items = [NeMoGymEasyInputMessage(role="user", content="hello")] + user, system = _extract_instruction(items) + assert user == "hello" + assert system is None + + def test_system_plus_user(self) -> None: + items = [ + NeMoGymEasyInputMessage(role="system", content="be concise"), + NeMoGymEasyInputMessage(role="user", content="hi"), + ] + user, system = _extract_instruction(items) + assert user == "hi" + assert system == "be concise" + + def test_empty(self) -> None: + user, system = _extract_instruction([]) + assert user == "" + assert system is None + + +class TestParseExecJsonl: + def test_empty(self) -> None: + items, usage = parse_exec_jsonl("") + assert items == [] + assert usage["input_tokens"] == 0 + assert usage["output_tokens"] == 0 + + def test_agent_message(self) -> None: + line = _item_completed({"id": "item_1", "type": "agent_message", "text": "hello"}) + items, _ = parse_exec_jsonl(line) + assert len(items) == 1 + assert isinstance(items[0], NeMoGymResponseOutputMessage) + assert items[0].content[0].text == "hello" + + def test_reasoning_prepended_to_next_message(self) -> None: + lines = "\n".join( + [ + _item_completed({"id": "item_1", "type": "reasoning", "text": "let me reason"}), + _item_completed({"id": "item_2", "type": "agent_message", "text": "answer"}), + ] + ) + items, _ = parse_exec_jsonl(lines) + assert len(items) == 1 + text = items[0].content[0].text + assert "\nlet me reason\n" in text + assert "answer" in text + + def test_trailing_reasoning_surfaced_as_think_message(self) -> None: + # Some backends route the final answer through the reasoning channel (vLLM reasoning + # parsers); a run ending on reasoning must not lose it. + line = _item_completed({"id": "item_1", "type": "reasoning", "text": "the answer is 391"}) + items, _ = parse_exec_jsonl(line) + assert len(items) == 1 + assert isinstance(items[0], NeMoGymResponseOutputMessage) + assert items[0].content[0].text == "\nthe answer is 391\n" + + def test_reasoning_cleared_after_message(self) -> None: + lines = "\n".join( + [ + _item_completed({"id": "item_1", "type": "reasoning", "text": "think"}), + _item_completed({"id": "item_2", "type": "agent_message", "text": "msg1"}), + _item_completed({"id": "item_3", "type": "agent_message", "text": "msg2"}), + ] + ) + items, _ = parse_exec_jsonl(lines) + assert len(items) == 2 + assert "" in items[0].content[0].text + assert "" not in items[1].content[0].text + + def test_command_execution_maps_to_call_and_output(self) -> None: + line = _item_completed( + { + "id": "item_1", + "type": "command_execution", + "command": "/bin/bash -lc ls", + "aggregated_output": "file.txt\n", + "exit_code": 0, + "status": "completed", + } + ) + items, _ = parse_exec_jsonl(line) + assert len(items) == 2 + assert isinstance(items[0], NeMoGymResponseFunctionToolCall) + assert items[0].name == "exec_command" + assert json.loads(items[0].arguments) == {"cmd": "/bin/bash -lc ls"} + assert isinstance(items[1], NeMoGymFunctionCallOutput) + assert "file.txt" in items[1].output + assert items[0].call_id == items[1].call_id == "item_1" + + def test_command_execution_nonzero_exit_annotated(self) -> None: + line = _item_completed( + { + "id": "item_1", + "type": "command_execution", + "command": "false", + "aggregated_output": "", + "exit_code": 1, + "status": "failed", + } + ) + items, _ = parse_exec_jsonl(line) + assert "[exit code: 1]" in items[1].output + + def test_mcp_tool_call_maps_result_text(self) -> None: + line = _item_completed( + { + "id": "item_1", + "type": "mcp_tool_call", + "server": "gymweather", + "tool": "get_weather", + "arguments": {"city": "Paris"}, + "result": {"content": [{"type": "text", "text": "sunny, 72F"}]}, + "error": None, + "status": "completed", + } + ) + items, _ = parse_exec_jsonl(line) + assert items[0].name == "get_weather" + assert json.loads(items[0].arguments) == {"city": "Paris"} + assert items[1].output == "sunny, 72F" + + def test_mcp_tool_call_error_surfaced(self) -> None: + line = _item_completed( + { + "id": "item_1", + "type": "mcp_tool_call", + "server": "s", + "tool": "t", + "arguments": {}, + "result": None, + "error": "boom", + "status": "failed", + } + ) + items, _ = parse_exec_jsonl(line) + assert items[1].output == "error: boom" + + def test_message_then_command(self) -> None: + lines = "\n".join( + [ + _item_completed({"id": "item_1", "type": "agent_message", "text": "running ls"}), + _item_completed( + { + "id": "item_2", + "type": "command_execution", + "command": "ls", + "aggregated_output": "x", + "exit_code": 0, + "status": "completed", + } + ), + ] + ) + items, _ = parse_exec_jsonl(lines) + assert [type(i).__name__ for i in items] == [ + "NeMoGymResponseOutputMessage", + "NeMoGymResponseFunctionToolCall", + "NeMoGymFunctionCallOutput", + ] + + def test_malformed_lines_skipped(self) -> None: + good = _item_completed({"id": "item_1", "type": "agent_message", "text": "ok"}) + items, _ = parse_exec_jsonl(f"not-json\n{good}\n{{bad") + assert len(items) == 1 + + def test_turn_completed_accumulates_usage(self) -> None: + line = _event( + "turn.completed", + usage={"input_tokens": 100, "cached_input_tokens": 40, "output_tokens": 50, "reasoning_output_tokens": 5}, + ) + _, usage = parse_exec_jsonl(line) + assert usage["input_tokens"] == 100 + assert usage["cached_input_tokens"] == 40 + assert usage["output_tokens"] == 50 + assert usage["reasoning_tokens"] == 5 + + def test_errors_collected(self) -> None: + lines = "\n".join( + [ + _item_completed({"id": "item_0", "type": "error", "message": "model metadata missing"}), + _event("turn.failed", error={"message": "stream disconnected"}), + ] + ) + items, usage = parse_exec_jsonl(lines) + assert items == [] + assert usage["errors"] == ["model metadata missing", "stream disconnected"] + + def test_item_started_events_ignored(self) -> None: + lines = "\n".join( + [ + _event( + "item.started", + item={"id": "item_1", "type": "command_execution", "command": "ls", "status": "in_progress"}, + ), + _item_completed( + { + "id": "item_1", + "type": "command_execution", + "command": "ls", + "aggregated_output": "x", + "exit_code": 0, + "status": "completed", + } + ), + ] + ) + items, _ = parse_exec_jsonl(lines) + assert len(items) == 2 # one call + one output, not doubled + + +class TestConfigYaml: + def test_module_parses(self) -> None: + app_path = Path(__file__).resolve().parent.parent / "app.py" + compile(app_path.read_text(), str(app_path), "exec") + + def test_config_yaml_parses(self) -> None: + cfg_path = Path(__file__).resolve().parent.parent / "configs" / "codex_agent.yaml" + data = yaml.safe_load(cfg_path.read_text()) + assert "codex_agent" in data + inner = data["codex_agent"]["responses_api_agents"]["codex_agent"] + assert inner["entrypoint"] == "app.py" + assert inner["concurrency"] == 32 + assert inner["sandbox_mode"] == "danger-full-access" diff --git a/tests/unit_tests/test_responses_api_model_streaming.py b/tests/unit_tests/test_responses_api_model_streaming.py new file mode 100644 index 0000000000..fbb25e7583 --- /dev/null +++ b/tests/unit_tests/test_responses_api_model_streaming.py @@ -0,0 +1,446 @@ +# 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. +"""Tests for the streaming Responses dialect on ``SimpleResponsesAPIModel``. + +Every Gym model server's ``/v1/responses`` accepts the wire dialect streaming harnesses +(e.g. the Codex CLI) speak: ``stream: true`` plus extra bookkeeping fields and ``namespace`` +tool specs. The request is sanitized onto the strict params model and the complete response +is re-emitted as a synthesized SSE event stream. Non-streaming requests keep the historical +strict-validation behavior. +""" + +import json +from time import time +from unittest.mock import MagicMock +from uuid import uuid4 + +import pytest +from fastapi import Body, Request +from fastapi.testclient import TestClient + +from nemo_gym.base_responses_api_model import BaseResponsesAPIModelConfig, SimpleResponsesAPIModel +from nemo_gym.openai_utils import ( + NeMoGymChatCompletion, + NeMoGymChatCompletionCreateParamsNonStreaming, + NeMoGymResponse, + NeMoGymResponseCreateParamsNonStreaming, +) +from nemo_gym.responses_streaming import ( + flatten_namespace_tools, + sanitize_streaming_responses_body, + synthesize_responses_failure_sse, + synthesize_responses_sse, + validate_streaming_responses_params, +) +from nemo_gym.server_utils import ServerClient + + +def _build_response(output: list) -> NeMoGymResponse: + return NeMoGymResponse( + id=f"resp_{uuid4().hex}", + created_at=int(time()), + model="downstream-model", + object="response", + output=output, + tool_choice="auto", + parallel_tool_calls=True, + tools=[], + usage={ + "input_tokens": 7, + "input_tokens_details": {"cached_tokens": 0}, + "output_tokens": 3, + "output_tokens_details": {"reasoning_tokens": 0}, + "total_tokens": 10, + }, + ) + + +def _message_item(text: str) -> dict: + return { + "type": "message", + "id": f"msg_{uuid4().hex}", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": text, "annotations": []}], + } + + +def _function_call_item(name: str) -> dict: + return { + "type": "function_call", + "id": "fc_1", + "call_id": "call_1", + "name": name, + "arguments": "{}", + "status": "completed", + } + + +NAMESPACE_TOOL = { + "type": "namespace", + "name": "mcp__weather", + "description": "Tools in the mcp__weather namespace.", + "tools": [ + { + "type": "function", + "name": "get_weather", + "description": "Get the weather.", + "strict": False, + "parameters": {"type": "object", "properties": {"city": {"type": "string"}}}, + } + ], +} + + +class TestSanitizeStreamingBody: + def test_drops_unknown_top_level_fields(self) -> None: + cleaned, _ = sanitize_streaming_responses_body( + {"input": [], "stream": True, "client_metadata": {"x": 1}, "prompt_cache_key": "abc", "store": False} + ) + assert set(cleaned) == {"input", "store"} + # the cleaned body validates against the strict params model + NeMoGymResponseCreateParamsNonStreaming.model_validate(cleaned) + + def test_flattens_namespace_tools(self) -> None: + flat, ns_map = flatten_namespace_tools([NAMESPACE_TOOL]) + assert len(flat) == 1 + assert flat[0]["type"] == "function" + assert flat[0]["name"] == "mcp__weather__get_weather" + assert ns_map == {"mcp__weather__get_weather": ("mcp__weather", "get_weather")} + + def test_sanitize_keeps_function_tools_and_flattens_namespaces(self) -> None: + function_tool = { + "type": "function", + "name": "exec_command", + "description": "Run a command.", + "strict": False, + "parameters": {"type": "object", "properties": {}}, + } + cleaned, ns_map = sanitize_streaming_responses_body( + {"input": [], "stream": True, "tools": [function_tool, NAMESPACE_TOOL]} + ) + names = [t["name"] for t in cleaned["tools"]] + assert names == ["exec_command", "mcp__weather__get_weather"] + assert "mcp__weather__get_weather" in ns_map + params = NeMoGymResponseCreateParamsNonStreaming.model_validate(cleaned) + assert len(params.tools) == 2 + + def test_drops_unsupported_tool_specs(self) -> None: + cleaned, _ = sanitize_streaming_responses_body( + {"input": [], "stream": True, "tools": [{"type": "totally_unknown_tool_kind", "config": 1}]} + ) + assert cleaned["tools"] == [] + + def test_rewrites_namespaced_calls_in_input_history(self) -> None: + cleaned, _ = sanitize_streaming_responses_body( + { + "stream": True, + "input": [ + { + "type": "function_call", + "namespace": "mcp__weather", + "name": "get_weather", + "arguments": "{}", + "call_id": "call_1", + }, + {"type": "function_call_output", "call_id": "call_1", "output": "sunny"}, + ], + } + ) + call = cleaned["input"][0] + assert call["name"] == "mcp__weather__get_weather" + assert "namespace" not in call + NeMoGymResponseCreateParamsNonStreaming.model_validate(cleaned) + + def test_drops_unsupported_input_items(self) -> None: + # Codex's code_mode interleaves an `additional_tools` carrier item into the input history; + # the Gym input union has no representation for it, so it is dropped item-by-item. + cleaned, _ = sanitize_streaming_responses_body( + { + "stream": True, + "input": [ + {"type": "additional_tools", "role": "developer", "tools": [{"type": "custom", "name": "exec"}]}, + {"type": "message", "role": "user", "content": [{"type": "input_text", "text": "hi"}]}, + ], + } + ) + assert [i["type"] for i in cleaned["input"]] == ["message"] + NeMoGymResponseCreateParamsNonStreaming.model_validate(cleaned) + + def test_additional_tools_carrier_functions_hoisted_into_tools(self) -> None: + # Codex's code mode ships tools inside an `additional_tools` input item; plain and + # namespaced function tools are hoisted into `tools`, non-function tools are dropped. + cleaned, ns_map = sanitize_streaming_responses_body( + { + "stream": True, + "input": [ + { + "type": "additional_tools", + "role": "developer", + "tools": [ + {"type": "custom", "name": "exec", "description": "JS orchestrator", "format": {}}, + { + "type": "function", + "name": "wait", + "description": "Wait.", + "strict": False, + "parameters": {"type": "object", "properties": {}}, + }, + NAMESPACE_TOOL, + ], + }, + {"type": "message", "role": "user", "content": [{"type": "input_text", "text": "hi"}]}, + ], + } + ) + names = [t["name"] for t in cleaned["tools"]] + assert "wait" in names + assert "mcp__weather__get_weather" in names + assert "exec" not in names + assert [i.get("type") for i in cleaned["input"]] == ["message"] + assert "mcp__weather__get_weather" in ns_map + NeMoGymResponseCreateParamsNonStreaming.model_validate(cleaned) + + def test_hoists_leading_developer_messages_into_instructions(self) -> None: + # Codex may open with several developer messages and no `instructions`; strict chat + # backends admit a single leading system message, so they are hoisted into instructions. + cleaned, _ = sanitize_streaming_responses_body( + { + "stream": True, + "input": [ + { + "type": "message", + "role": "developer", + "content": [{"type": "input_text", "text": "You are Codex."}], + }, + { + "type": "message", + "role": "developer", + "content": [{"type": "input_text", "text": ""}], + }, + {"type": "message", "role": "user", "content": [{"type": "input_text", "text": "hi"}]}, + ], + } + ) + assert cleaned["instructions"] == "You are Codex.\n\n" + assert [i["role"] for i in cleaned["input"]] == ["user"] + + def test_hoisting_prepends_existing_instructions(self) -> None: + cleaned, _ = sanitize_streaming_responses_body( + { + "stream": True, + "instructions": "base instructions", + "input": [ + {"type": "message", "role": "developer", "content": "perms"}, + {"type": "message", "role": "user", "content": "hi"}, + ], + } + ) + assert cleaned["instructions"] == "base instructions\n\nperms" + + def test_mid_conversation_developer_messages_not_hoisted(self) -> None: + cleaned, _ = sanitize_streaming_responses_body( + { + "stream": True, + "input": [ + {"type": "message", "role": "user", "content": "hi"}, + {"type": "message", "role": "developer", "content": "mid-run steer"}, + ], + } + ) + assert "instructions" not in cleaned + assert [i["role"] for i in cleaned["input"]] == ["user", "developer"] + + def test_does_not_mutate_caller_body(self) -> None: + body = {"input": [], "stream": True, "tools": [NAMESPACE_TOOL]} + sanitize_streaming_responses_body(body) + assert body["tools"] == [NAMESPACE_TOOL] + assert body["stream"] is True + + +class TestValidateStreamingParams: + def test_prunes_nested_extra_fields(self) -> None: + # Codex sends `reasoning.context`, which the pinned SDK's Reasoning model forbids. + params = validate_streaming_responses_params( + {"input": [], "reasoning": {"effort": "medium", "context": "all_turns"}} + ) + assert params.reasoning == {"effort": "medium"} + + def test_unfixable_errors_still_raise(self) -> None: + import pydantic + + with pytest.raises(pydantic.ValidationError): + validate_streaming_responses_params({"input": [], "temperature": "not-a-number"}) + + +class TestSynthesizeSSE: + def _events(self, sse_text: str) -> list[dict]: + events = [] + for block in sse_text.split("\n\n"): + for line in block.splitlines(): + if line.startswith("data: "): + events.append(json.loads(line[len("data: ") :])) + return events + + def test_event_sequence(self) -> None: + response = _build_response([_message_item("hello")]).model_dump(mode="json") + events = self._events("".join(synthesize_responses_sse(response))) + assert [e["type"] for e in events] == ["response.created", "response.output_item.done", "response.completed"] + assert events[0]["response"]["status"] == "in_progress" + assert events[0]["response"]["output"] == [] + assert events[1]["output_index"] == 0 + assert events[1]["item"]["content"][0]["text"] == "hello" + completed = events[2]["response"] + assert completed["id"] == response["id"] + assert completed["usage"]["input_tokens"] == 7 + assert len(completed["output"]) == 1 + + def test_namespaced_call_names_restored(self) -> None: + response = _build_response([_function_call_item("mcp__weather__get_weather")]).model_dump(mode="json") + ns_map = {"mcp__weather__get_weather": ("mcp__weather", "get_weather")} + events = self._events("".join(synthesize_responses_sse(response, ns_map))) + item = events[1]["item"] + assert item["namespace"] == "mcp__weather" + assert item["name"] == "get_weather" + # the terminal envelope carries the same rewritten item + assert events[2]["response"]["output"][0]["name"] == "get_weather" + + def test_unmapped_calls_left_alone(self) -> None: + response = _build_response([_function_call_item("exec_command")]).model_dump(mode="json") + events = self._events("".join(synthesize_responses_sse(response, {"other__tool": ("other", "tool")}))) + assert events[1]["item"]["name"] == "exec_command" + assert "namespace" not in events[1]["item"] + + def test_failure_stream_is_terminal_response_failed(self) -> None: + events = self._events("".join(synthesize_responses_failure_sse("boom", code="server_error"))) + assert [e["type"] for e in events] == ["response.created", "response.failed"] + failed = events[-1]["response"] + assert failed["status"] == "failed" + assert failed["error"] == {"code": "server_error", "message": "boom"} + assert failed["output"] == [] + + +class _EchoModel(SimpleResponsesAPIModel): + """Fake model server capturing the params its responses() receives.""" + + config: BaseResponsesAPIModelConfig + last_params: object = None + model_config = {"arbitrary_types_allowed": True} + + async def responses(self, body: NeMoGymResponseCreateParamsNonStreaming = Body()) -> NeMoGymResponse: + object.__setattr__(self, "last_params", body) + output = [_message_item("hi")] + if body.tools: + output.insert(0, _function_call_item(body.tools[0].get("name", ""))) + return _build_response(output) + + async def chat_completions( + self, body: NeMoGymChatCompletionCreateParamsNonStreaming = Body() + ) -> NeMoGymChatCompletion: + raise NotImplementedError + + +class _RequestAwareEchoModel(_EchoModel): + saw_request: bool = False + + async def responses( + self, request: Request, body: NeMoGymResponseCreateParamsNonStreaming = Body() + ) -> NeMoGymResponse: + object.__setattr__(self, "saw_request", isinstance(request, Request)) + return await super().responses(body) + + +class _FailingModel(_EchoModel): + async def responses(self, body: NeMoGymResponseCreateParamsNonStreaming = Body()) -> NeMoGymResponse: + raise RuntimeError("backend exploded") + + +def _client(model_cls) -> tuple[TestClient, SimpleResponsesAPIModel]: + server = model_cls( + config=BaseResponsesAPIModelConfig(host="0.0.0.0", port=8099, entrypoint="", name=""), + server_client=MagicMock(spec=ServerClient, global_config_dict={}), + ) + return TestClient(server.setup_webserver()), server + + +class TestResponsesDispatchRoute: + def test_non_streaming_request_returns_plain_json(self) -> None: + client, server = _client(_EchoModel) + resp = client.post("/v1/responses", json={"input": [{"role": "user", "content": "hi"}]}) + assert resp.status_code == 200 + assert resp.headers["content-type"].startswith("application/json") + assert resp.json()["output"][-1]["content"][0]["text"] == "hi" + assert server.last_params.input[0].content == "hi" + + def test_non_streaming_request_still_validates_strictly(self) -> None: + client, _ = _client(_EchoModel) + resp = client.post("/v1/responses", json={"input": [], "client_metadata": {"x": 1}}) + assert resp.status_code == 422 + assert resp.json()["detail"][0]["loc"][0] == "body" + + def test_streaming_request_returns_synthesized_sse(self) -> None: + client, server = _client(_EchoModel) + resp = client.post( + "/v1/responses", + json={ + "stream": True, + "client_metadata": {"cli": "codex"}, + "prompt_cache_key": "abc", + "input": [{"role": "user", "content": "hi"}], + "tools": [NAMESPACE_TOOL], + }, + ) + assert resp.status_code == 200 + assert resp.headers["content-type"].startswith("text/event-stream") + assert "event: response.completed" in resp.text + # the server saw sanitized params: flattened tools, no bookkeeping fields + assert server.last_params.tools[0]["name"] == "mcp__weather__get_weather" + # and the synthesized items restore the namespaced call shape + done_events = [line for line in resp.text.splitlines() if '"response.output_item.done"' in line] + first_item = json.loads(done_events[0][len("data: ") :])["item"] + assert first_item["namespace"] == "mcp__weather" + assert first_item["name"] == "get_weather" + + def test_streaming_request_aware_signature(self) -> None: + client, server = _client(_RequestAwareEchoModel) + resp = client.post("/v1/responses", json={"stream": True, "input": [{"role": "user", "content": "hi"}]}) + assert resp.status_code == 200 + assert server.saw_request is True + + def test_streaming_invalid_body_is_422(self) -> None: + client, _ = _client(_EchoModel) + resp = client.post("/v1/responses", json={"stream": True}) # no input at all + assert resp.status_code == 422 + + def test_streaming_backend_error_yields_response_failed(self) -> None: + # A responses() failure after the streaming contract is committed becomes a terminal + # response.failed event (HTTP 200 SSE), not a broken-stream HTTP 500. + client, _ = _client(_FailingModel) + resp = client.post("/v1/responses", json={"stream": True, "input": [{"role": "user", "content": "hi"}]}) + assert resp.status_code == 200 + assert resp.headers["content-type"].startswith("text/event-stream") + assert "event: response.failed" in resp.text + assert "event: response.completed" not in resp.text + failed = [line for line in resp.text.splitlines() if line.startswith("data: ") and "response.failed" in line] + payload = json.loads(failed[0][len("data: ") :]) + assert payload["response"]["status"] == "failed" + assert "backend exploded" in payload["response"]["error"]["message"] + + def test_non_streaming_backend_error_still_raises(self) -> None: + # Without the streaming contract, a backend failure is a normal exception (HTTP 500), not a + # synthesized response.failed — only the stream path swallows it into a terminal event. + client, _ = _client(_FailingModel) + with pytest.raises(RuntimeError, match="backend exploded"): + client.post("/v1/responses", json={"input": [{"role": "user", "content": "hi"}]}) diff --git a/tests/unit_tests/test_responses_converter.py b/tests/unit_tests/test_responses_converter.py index 817212dd9e..400b1a4ba0 100644 --- a/tests/unit_tests/test_responses_converter.py +++ b/tests/unit_tests/test_responses_converter.py @@ -168,6 +168,44 @@ def test_responses_to_chat_completion_all_message_roles(converter: ResponsesConv assert params.messages[-1]["content"] == "assistant content" +def test_responses_to_chat_completion_instructions_become_leading_system_message(converter: ResponsesConverter): + params = converter.responses_to_chat_completion_create_params( + NeMoGymResponseCreateParamsNonStreaming( + instructions="you are a coding agent", + input=[NeMoGymEasyInputMessage(role="user", content="usr", type="message")], + ) + ) + # instructions are inserted before any input-derived messages (Responses API semantics) + assert params.messages[0] == {"role": "system", "content": "you are a coding agent"} + assert [m["role"] for m in params.messages] == ["system", "user"] + + +def test_responses_to_chat_completion_instructions_fold_leading_system_and_developer(converter: ResponsesConverter): + params = converter.responses_to_chat_completion_create_params( + NeMoGymResponseCreateParamsNonStreaming( + instructions="you are a coding agent", + input=[ + NeMoGymEasyInputMessage(role="system", content="sys", type="message"), + NeMoGymEasyInputMessage(role="developer", content="dev", type="message"), + NeMoGymEasyInputMessage(role="user", content="usr", type="message"), + ], + ) + ) + # chat backends commonly admit a single system message at position 0, so the leading run of + # system/developer messages is folded into the instructions message + assert params.messages[0] == {"role": "system", "content": "you are a coding agent\n\nsys\n\ndev"} + assert [m["role"] for m in params.messages] == ["system", "user"] + + +def test_responses_to_chat_completion_no_instructions_adds_no_message(converter: ResponsesConverter): + params = converter.responses_to_chat_completion_create_params( + NeMoGymResponseCreateParamsNonStreaming( + input=[NeMoGymEasyInputMessage(role="user", content="usr", type="message")] + ) + ) + assert [m["role"] for m in params.messages] == ["user"] + + def test_responses_to_chat_completion_input_image_part(converter: ResponsesConverter): params = converter.responses_to_chat_completion_create_params( NeMoGymResponseCreateParamsNonStreaming(