Skip to content

Commit fdaeae2

Browse files
committed
feat(agents): sandbox-bound CLI agents (Codex, Claude Code, custom) with capture + SWE-bench/Terminal-Bench grading
Adds a reusable layer for running CLI coding agents *inside* a Gym sandbox and collecting RL-ready trajectories, built on the sandbox provider + adapter framework. Benchmark-agnostic: only task metadata + the grader change. SandboxCliAgent (nemo_gym/sandbox_cli_agent.py) - Owns the whole per-rollout lifecycle: start a sandbox, stand up a per-rollout capture proxy, install + run the CLI in-box, collect the git patch, assemble the trajectory, and grade. Subclasses supply only a thin seam (build_launch + parse_stdout); adding an agent is a small subclass or YAML. Agents - codex_swe_agent: Codex (Responses wire) as a thin subclass. - claude_code_swe_agent: Claude Code (Messages wire) via the translate_anthropic interceptor so it runs against any OpenAI-compatible backend; reuses claude_code_agent's stream-json parser. - custom_agent: manifest-driven (YAML-only) onboarding. Adapters (sandbox-bound, not model-bound) - capture interceptor + CaptureStore: durable, session-keyed JSONL of model exchanges (token-ids when the backend is a Gym model server); trajectory assembly for chat + responses wires. choose_trajectory prefers a well-formed (paired tool calls) trajectory, so codex's streamed-response capture and claude's duplicate-call redundancy both fall back to the CLI's clean stdout. - translate_anthropic interceptor: Anthropic Messages <-> OpenAI Chat. In-box grading (the only benchmark-specific seam) - SWE-bench: the official swebench harness (make_test_spec eval_script + get_eval_report), like mini_swe_agent_2. - Terminal-Bench / Harbor: per-task public docker_image (auto_mirrored), stage the task's tests, run test.sh, read the verifier reward file. - Fallback: a lightweight pytest membership grader for a custom eval_command. Runnable under ng_run / ng_collect_rollouts; per-agent + per-task examples and a Terminal-Bench example row are in the agent READMEs/data. Agents install nemo-gym[sandbox,sandbox-ecs] (editable) + swebench. Validated on real ECS Fargate: - SWE-bench Verified: golden -> 1.0, no-patch -> 0.0; Codex + Claude Code each solve real instances; 10 diverse repos via ng_run + ng_collect_rollouts. - Terminal-Bench 2.0 (Harbor): across 5 tasks x Codex + Claude Code — multiple solves (reward 1.0 via the verifier reward file), grader discriminates, and trajectories are healthy (paired tool calls) for both agents. Unit-tested: grader, image-tag mapping, SWE-bench instance reconstruction, Harbor grading, trajectory selection, capture assembly, Anthropic<->OpenAI translation. Note: stacked on the sandbox provider PR (feat/ecs-fargate-sandbox) and the adapter framework PR (feat/adapter-base); review/merge after them. Signed-off-by: Michal Bien <mbien@nvidia.com>
1 parent 73051c7 commit fdaeae2

40 files changed

Lines changed: 3979 additions & 0 deletions

nemo_gym/adapters/capture_store.py

Lines changed: 276 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,276 @@
1+
# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2+
# SPDX-License-Identifier: Apache-2.0
3+
#
4+
# Licensed under the Apache License, Version 2.0 (the "License");
5+
# you may not use this file except in compliance with the License.
6+
# You may obtain a copy of the License at
7+
#
8+
# http://www.apache.org/licenses/LICENSE-2.0
9+
#
10+
# Unless required by applicable law or agreed to in writing, software
11+
# distributed under the License is distributed on an "AS IS" BASIS,
12+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
# See the License for the specific language governing permissions and
14+
# limitations under the License.
15+
"""Session-keyed capture store and trajectory assembly.
16+
17+
The sandbox-bound model proxy records every model exchange (request + response,
18+
including token-ids when the policy is the Gym model server) into a per-session
19+
JSONL file. Writing one file per session — keyed by the rollout's identity, i.e.
20+
"one box = one session" — means a sandbox reaped mid-run cannot lose turns that
21+
were already streamed out (durable gather). After the run the harness reads the
22+
exchanges back and assembles a NeMoGym trajectory (for eval) carrying per-turn
23+
token-ids (for RL).
24+
"""
25+
26+
from __future__ import annotations
27+
28+
import json
29+
import os
30+
import threading
31+
from pathlib import Path
32+
from typing import Any
33+
34+
from nemo_gym.openai_utils import (
35+
NeMoGymFunctionCallOutput,
36+
NeMoGymResponseFunctionToolCall,
37+
NeMoGymResponseOutputMessageForTraining,
38+
NeMoGymResponseOutputText,
39+
)
40+
41+
42+
# Per-message token-id / logprob fields a Gym model server returns when
43+
# ``return_token_id_information`` is enabled. We look for them on either the
44+
# choice or the message object, since servers differ on placement.
45+
_TOKEN_FIELDS = ("prompt_token_ids", "generation_token_ids", "generation_log_probs")
46+
47+
48+
def _sanitize(session_id: str) -> str:
49+
cleaned = "".join(c for c in session_id if c.isalnum() or c in ("-", "_", "."))
50+
return cleaned or "session"
51+
52+
53+
class CaptureStore:
54+
"""Append-only, session-keyed JSONL sink for model exchanges."""
55+
56+
def __init__(self, root: str | Path) -> None:
57+
self._root = Path(root)
58+
self._root.mkdir(parents=True, exist_ok=True)
59+
self._lock = threading.Lock()
60+
61+
@property
62+
def root(self) -> Path:
63+
return self._root
64+
65+
def path_for(self, session_id: str) -> Path:
66+
return self._root / f"{_sanitize(session_id)}.capture.jsonl"
67+
68+
def record(self, session_id: str, exchange: dict[str, Any]) -> None:
69+
"""Append one exchange and fsync, so a killed box can't lose it."""
70+
line = json.dumps(exchange, default=str, ensure_ascii=False)
71+
path = self.path_for(session_id)
72+
with self._lock:
73+
with path.open("a", encoding="utf-8") as handle:
74+
handle.write(line + "\n")
75+
handle.flush()
76+
os.fsync(handle.fileno())
77+
78+
def read(self, session_id: str) -> list[dict[str, Any]]:
79+
path = self.path_for(session_id)
80+
if not path.exists():
81+
return []
82+
exchanges: list[dict[str, Any]] = []
83+
# Stream line-by-line: a capture can be hundreds of MB (token-ids/logprobs), so avoid
84+
# read_text().splitlines() which would hold the whole file as a string + a list of lines.
85+
with path.open("r", encoding="utf-8") as handle:
86+
for line in handle:
87+
stripped = line.strip()
88+
if not stripped:
89+
continue
90+
try:
91+
exchanges.append(json.loads(stripped))
92+
except json.JSONDecodeError:
93+
continue
94+
return exchanges
95+
96+
97+
def _token_fields(obj: Any) -> dict[str, list]:
98+
if not isinstance(obj, dict):
99+
return {}
100+
return {key: obj[key] for key in _TOKEN_FIELDS if isinstance(obj.get(key), list)}
101+
102+
103+
def _response_message(response: dict[str, Any]) -> dict[str, Any]:
104+
"""First-choice assistant message, with token-ids hoisted onto it."""
105+
if not isinstance(response, dict):
106+
return {}
107+
choices = response.get("choices") or []
108+
if not choices or not isinstance(choices[0], dict):
109+
return {}
110+
choice = choices[0]
111+
message = choice.get("message") or choice.get("delta") or {}
112+
if not isinstance(message, dict):
113+
return {}
114+
merged = dict(message)
115+
merged.update(_token_fields(choice))
116+
merged.update(_token_fields(message))
117+
return merged
118+
119+
120+
def _content_text(content: Any) -> str:
121+
if isinstance(content, str):
122+
return content
123+
if isinstance(content, list):
124+
return "".join(
125+
(part.get("text", "") if isinstance(part, dict) else getattr(part, "text", "")) for part in content
126+
)
127+
return "" if content is None else str(content)
128+
129+
130+
def _assistant_message(index: int, text: str, token_fields: dict[str, Any]) -> NeMoGymResponseOutputMessageForTraining:
131+
"""A training assistant message, carrying token-ids/logprobs when the policy returned them."""
132+
return NeMoGymResponseOutputMessageForTraining(
133+
id=f"msg-{index}",
134+
content=[NeMoGymResponseOutputText(type="output_text", text=text, annotations=[])],
135+
role="assistant",
136+
status="completed",
137+
type="message",
138+
prompt_token_ids=token_fields.get("prompt_token_ids") or [],
139+
generation_token_ids=token_fields.get("generation_token_ids") or [],
140+
generation_log_probs=token_fields.get("generation_log_probs") or [],
141+
)
142+
143+
144+
def assemble_trajectory(exchanges: list[dict[str, Any]], wire: str = "chat") -> list[Any]:
145+
"""Reconstruct ordered NeMoGym output items from captured model exchanges.
146+
147+
``wire`` selects how the captured request/response bodies are shaped:
148+
``"chat"`` (OpenAI Chat Completions / the Anthropic-translated path) or
149+
``"responses"`` (OpenAI Responses API, e.g. codex). Assistant messages carry
150+
token-ids/logprobs when the policy returned them.
151+
"""
152+
if wire == "responses":
153+
return _assemble_responses(exchanges)
154+
return _assemble_chat(exchanges)
155+
156+
157+
def _assemble_chat(exchanges: list[dict[str, Any]]) -> list[Any]:
158+
"""Chat-Completions wire: a turn's tool *results* arrive as ``role: tool``
159+
messages in the *next* request, emitted just before the following assistant
160+
message to yield the natural ``assistant -> tool_output -> assistant`` order."""
161+
output: list[Any] = []
162+
seen_tool_call_ids: set[str] = set()
163+
assistant_index = 0
164+
165+
for exchange in exchanges:
166+
request = exchange.get("request") or {}
167+
response = exchange.get("response") or {}
168+
169+
for message in request.get("messages") or []:
170+
if not isinstance(message, dict) or message.get("role") != "tool":
171+
continue
172+
call_id = message.get("tool_call_id") or ""
173+
if call_id and call_id in seen_tool_call_ids:
174+
continue
175+
if call_id:
176+
seen_tool_call_ids.add(call_id)
177+
output.append(
178+
NeMoGymFunctionCallOutput(
179+
type="function_call_output",
180+
call_id=call_id,
181+
output=_content_text(message.get("content")),
182+
status="completed",
183+
)
184+
)
185+
186+
message = _response_message(response)
187+
if not message:
188+
continue
189+
output.append(
190+
_assistant_message(assistant_index, _content_text(message.get("content")), _token_fields(message))
191+
)
192+
assistant_index += 1
193+
194+
for tool_call in message.get("tool_calls") or []:
195+
function = tool_call.get("function") if isinstance(tool_call, dict) else None
196+
if not function:
197+
continue
198+
output.append(
199+
NeMoGymResponseFunctionToolCall(
200+
arguments=function.get("arguments", ""),
201+
call_id=tool_call.get("id", ""),
202+
name=function.get("name", ""),
203+
type="function_call",
204+
id=tool_call.get("id"),
205+
status="completed",
206+
)
207+
)
208+
209+
return output
210+
211+
212+
def _assemble_responses(exchanges: list[dict[str, Any]]) -> list[Any]:
213+
"""Responses-API wire: each response carries an ``output`` list of items
214+
(``message`` / ``reasoning`` / ``function_call``); token-ids ride on the
215+
assistant ``message`` item when the policy is a Gym model. A turn's tool
216+
*results* arrive as ``function_call_output`` items in the *next* request's
217+
``input``, so we emit those just before that turn's response — yielding the
218+
natural ``assistant -> function_call -> function_call_output -> assistant``
219+
interleave (mirrors the chat-wire assembler)."""
220+
output: list[Any] = []
221+
assistant_index = 0
222+
seen_output_call_ids: set[str] = set()
223+
for exchange in exchanges:
224+
request = exchange.get("request") or {}
225+
request_input = request.get("input")
226+
if isinstance(request_input, list):
227+
for item in request_input:
228+
if not isinstance(item, dict) or item.get("type") != "function_call_output":
229+
continue
230+
call_id = item.get("call_id") or item.get("id") or ""
231+
if call_id and call_id in seen_output_call_ids:
232+
continue
233+
if call_id:
234+
seen_output_call_ids.add(call_id)
235+
output.append(
236+
NeMoGymFunctionCallOutput(
237+
type="function_call_output",
238+
call_id=call_id,
239+
output=_content_text(item.get("output")),
240+
status="completed",
241+
)
242+
)
243+
244+
response = exchange.get("response") or {}
245+
for item in response.get("output") or []:
246+
if not isinstance(item, dict):
247+
continue
248+
kind = item.get("type")
249+
if kind == "message" and item.get("role") == "assistant":
250+
text = "".join(
251+
block.get("text", "")
252+
for block in (item.get("content") or [])
253+
if isinstance(block, dict) and block.get("type") == "output_text"
254+
)
255+
output.append(_assistant_message(assistant_index, text, _token_fields(item)))
256+
assistant_index += 1
257+
elif kind in ("function_call", "tool_call"):
258+
output.append(
259+
NeMoGymResponseFunctionToolCall(
260+
arguments=item.get("arguments", "") or "",
261+
call_id=item.get("call_id") or item.get("id") or "",
262+
name=item.get("name", ""),
263+
type="function_call",
264+
id=item.get("id"),
265+
status="completed",
266+
)
267+
)
268+
return output
269+
270+
271+
def has_token_ids(output_items: list[Any]) -> bool:
272+
"""True if any assistant item carries generation token-ids (RL-usable)."""
273+
for item in output_items:
274+
if getattr(item, "generation_token_ids", None):
275+
return True
276+
return False
Lines changed: 101 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,101 @@
1+
# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2+
# SPDX-License-Identifier: Apache-2.0
3+
#
4+
# Licensed under the Apache License, Version 2.0 (the "License");
5+
# you may not use this file except in compliance with the License.
6+
# You may obtain a copy of the License at
7+
#
8+
# http://www.apache.org/licenses/LICENSE-2.0
9+
#
10+
# Unless required by applicable law or agreed to in writing, software
11+
# distributed under the License is distributed on an "AS IS" BASIS,
12+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
# See the License for the specific language governing permissions and
14+
# limitations under the License.
15+
"""Capture interceptor — records model exchanges (with token-ids) per session.
16+
17+
The model-bound ``endpoint``/``logging`` interceptors forward and log; neither
18+
persists what would be needed to rebuild a trajectory. This interceptor is meant
19+
to run inside a *sandbox-bound* proxy that the harness starts per rollout: it
20+
tags every model call with that rollout's ``session_id`` and writes the full
21+
request + response to a durable :class:`CaptureStore`.
22+
23+
It also optionally injects request-level fields (``inject_extra_body``) so the
24+
policy returns the extra information the trajectory needs — e.g. flags that turn
25+
on token-id reporting, or ``chat_template_kwargs`` that keep reasoning history
26+
intact for on-policy RL.
27+
"""
28+
29+
from __future__ import annotations
30+
31+
import logging
32+
import time
33+
from typing import Any
34+
35+
from nemo_gym.adapters.capture_store import CaptureStore
36+
from nemo_gym.adapters.types import (
37+
AdapterRequest,
38+
AdapterResponse,
39+
RequestInterceptor,
40+
ResponseInterceptor,
41+
)
42+
43+
44+
logger = logging.getLogger(__name__)
45+
46+
# Stashed on the per-request context so the response phase can pair the
47+
# captured response with the request body that produced it.
48+
_REQUEST_KEY = "_capture_request_body"
49+
_SESSION_KEY = "session_id"
50+
51+
52+
class Interceptor(RequestInterceptor, ResponseInterceptor):
53+
"""Record each model exchange to a session-keyed store.
54+
55+
Best-effort: a capture failure logs and is swallowed by the pipeline rather
56+
than breaking the rollout (the agent's model call still succeeds).
57+
"""
58+
59+
best_effort = True
60+
61+
def __init__(
62+
self,
63+
*,
64+
store_dir: str,
65+
session_id: str | None = None,
66+
inject_extra_body: dict[str, Any] | None = None,
67+
upstream_api_key: str | None = None,
68+
) -> None:
69+
self._store = CaptureStore(store_dir)
70+
self._session_id = session_id
71+
self._inject_extra_body = inject_extra_body or {}
72+
# When set, the real upstream key is stamped onto the *forwarded* request
73+
# so the in-box agent only ever holds a dummy. It is applied to the
74+
# headers (which are NOT persisted), never to the recorded body.
75+
self._upstream_api_key = upstream_api_key
76+
77+
def _session(self, ctx_extra: dict[str, Any]) -> str:
78+
return self._session_id or ctx_extra.get(_SESSION_KEY) or "session"
79+
80+
async def intercept_request(self, req: AdapterRequest) -> AdapterRequest:
81+
if isinstance(req.body, dict):
82+
for key, value in self._inject_extra_body.items():
83+
req.body.setdefault(key, value)
84+
req.ctx.extra[_REQUEST_KEY] = req.body
85+
if self._upstream_api_key:
86+
req.headers["Authorization"] = f"Bearer {self._upstream_api_key}"
87+
return req
88+
89+
async def intercept_response(self, resp: AdapterResponse) -> AdapterResponse:
90+
session_id = self._session(resp.ctx.extra)
91+
exchange = {
92+
"request_id": resp.ctx.request_id,
93+
"session_id": session_id,
94+
"ts": time.time(),
95+
"status": resp.status_code,
96+
"latency_ms": resp.latency_ms,
97+
"request": resp.ctx.extra.get(_REQUEST_KEY),
98+
"response": resp.body if isinstance(resp.body, dict) else None,
99+
}
100+
self._store.record(session_id, exchange)
101+
return resp

0 commit comments

Comments
 (0)